fail to log model by mlflow

[jinholee-makinarocks]

System Info

python version : 3.9.16
mlflow: 2.4.2
transformers[torch] version : 4.31.0
optimum : 1.11.2

Who can help?

@philschmid , @JingyaHuang , @ec

Information

• [ ] The official example scripts
• [X] My own modified scripts

Tasks

• [X] An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• [ ] My own task or dataset (give details below)

Reproduction (minimal, reproducible, runnable)

Try to run following code:

from transformers import AutoTokenizer
from optimum.onnxruntime import ORTModelForCausalLM
from optimum.pipelines import pipeline
import mlflow

tokenizer = AutoTokenizer.from_pretrained("optimum/gpt2")
model = ORTModelForCausalLM.from_pretrained("optimum/gpt2")
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

mlflow.transformers.log_model(transformers_model=generator, artifact_path="my_pipeline")

You can see error log as:

<stdin>:1: FutureWarning: The 'transformers' MLflow Models integration is known to be compatible with the following package version ranges: ``4.25.1`` -  ``4.29.2``. MLflow Models integrations with transformers may not succeed when used with package versions outside of this range.
/Users/mark/.pyenv/versions/3.9.15/lib/python3.9/site-packages/mlflow/models/model.py:562: FutureWarning: The 'transformers' MLflow Models integration is known to be compatible with the following package version ranges: ``4.25.1`` -  ``4.29.2``. MLflow Models integrations with transformers may not succeed when used with package versions outside of this range.
  flavor.save_model(path=local_path, mlflow_model=mlflow_model, **kwargs)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/mark/.pyenv/versions/3.9.15/lib/python3.9/site-packages/mlflow/utils/docstring_utils.py", line 235, in version_func
    return func(*args, **kwargs)
  File "/Users/mark/.pyenv/versions/3.9.15/lib/python3.9/site-packages/mlflow/transformers.py", line 728, in log_model
    return Model.log(
  File "/Users/mark/.pyenv/versions/3.9.15/lib/python3.9/site-packages/mlflow/models/model.py", line 562, in log
    flavor.save_model(path=local_path, mlflow_model=mlflow_model, **kwargs)
  File "/Users/mark/.pyenv/versions/3.9.15/lib/python3.9/site-packages/mlflow/utils/docstring_utils.py", line 235, in version_func
    return func(*args, **kwargs)
  File "/Users/mark/.pyenv/versions/3.9.15/lib/python3.9/site-packages/mlflow/transformers.py", line 390, in save_model
    _validate_transformers_model_dict(transformers_model)
  File "/Users/mark/.pyenv/versions/3.9.15/lib/python3.9/site-packages/mlflow/transformers.py", line 176, in _validate_transformers_model_dict
    raise MlflowException(
mlflow.exceptions.MlflowException: The submitted model type ORTModelForCausalLM does not inherit from a transformers pre-trained model. It is missing the attribute 'name_or_path'. Please verify that the model is a supported transformers model.

Expected behavior

Since it's a transformers model, I think it should be logged normally through the transformers flavor of mlflow. when I try to run with normally a pipeline of the transformer, it works well.

from transformers import pipeline
qa_pipe = pipeline("question-answering", "csarron/mobilebert-uncased-squad-v2")
mlflow.transformers.log_model(transformers_model=qa_pipe, artifact_path="my_pipeline")

[fxmarty]

Thank you for the report @jinholee-makinarocks . ONNX Runtime-based models indeed do not inherit from transformers PreTrainedModel, since this class inherits from torch.nn.Module and we would like not to.

We could hack it around by adding a name_or_path attribute to the ONNX Runtime models, but I think there is high chance that the rest of the logic will fail anyway since it is designed for transformers.

[ierezell]

Hello,

Sorry to bump (after 2 years...) but is there any news on this topic ? I would also like to use and serve ONNX models from mlflow and this forbid me to do so.