Pyfunc mlflow loading Databricks Model failing
- Mark as New
- Bookmark
- Subscribe
- Mute
- Subscribe to RSS Feed
- Permalink
- Report Inappropriate Content
10-01-2024 12:34 PM
I am running this example notebook https://docs.databricks.com/_extras/notebooks/source/machine-learning/large-language-models/pyspark-...
I am running most as if except switching embeddings and llm to databricks from OpenAI.
On the last cell I am using this.
model_uri = f"runs:/{ mlflow_run.info.run_id }/retrieval_qa_chain"
loaded_pyfunc_model = mlflow.pyfunc.load_model(model_uri)
langchain_input = {"query": "Who is Harrison Schmitt"}
loaded_pyfunc_model.predict([langchain_input])
And I get the following error
ValueError: Must specify a chain Type in config
File <command-1908280184094495>, line 5
1 model_uri = f"runs:/{ mlflow_run.info.run_id }/retrieval_qa_chain"
----> 5 loaded_pyfunc_model = mlflow.pyfunc.load_model(model_uri)
6 langchain_input = {"query": "Who is Harrison Schmitt"}
7 loaded_pyfunc_model.predict([langchain_input])
.....
File /local_disk0/.ephemeral_nfs/envs/pythonEnv-0c4b76b4-2202-4daf-abbd-9d314dbf7542/lib/python3.11/site-packages/langchain/chains/loading.py:664, in load_chain_from_config(config, **kwargs) 662 """Load chain from Config Dict.""" 663 if "_type" not in config: --> 664 raise ValueError("Must specify a chain Type in config") 665 config_type = config.pop("_type") 667 if config_type not in type_to_loader_dict:
https://mlflow.org/docs/latest/python_api/mlflow.pyfunc.html
The documentation on mlflow pyfunc does not give anything else. Every other cell seems to work until this very last cell.
- Mark as New
- Bookmark
- Subscribe
- Mute
- Subscribe to RSS Feed
- Permalink
- Report Inappropriate Content
10-01-2024 12:47 PM
with mlflow.start_run() as mlflow_run:
logged_model = mlflow.langchain.log_model(
retrieval_qa,
"retrieval_qa_chain",
loader_fn=load_retriever,
persist_dir=persist_directory,
model_config={"chain_type": "stuff"}
)
logging model as above
- Mark as New
- Bookmark
- Subscribe
- Mute
- Subscribe to RSS Feed
- Permalink
- Report Inappropriate Content
10-09-2024 03:30 AM
Any updates? I'm facing the same problem. Thanks!
- Mark as New
- Bookmark
- Subscribe
- Mute
- Subscribe to RSS Feed
- Permalink
- Report Inappropriate Content
10-16-2024 08:46 AM
We abandoned this. Too many compatibility errors and we stuck with the notebook level. Good luck!

