Pyfunc mlflow loading Databricks Model failing

[swollazn]

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.

[swollazn]

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

[stefa]

Any updates? I'm facing the same problem. Thanks!

[swollazn]

We abandoned this. Too many compatibility errors and we stuck with the notebook level. Good luck!