Repository: mlflow/mlflow

Issue:
[BUG] Cant load logged transformer model that has remote code 
### Issues Policy acknowledgement

- [X] I have read and agree to submit bug reports in accordance with the [issues policy](https://www.github.com/mlflow/mlflow/blob/master/ISSUE_POLICY.md)

### Where did you encounter this bug?

Other

### Willingness to contribute

Yes. I would be willing to contribute a fix for this bug with guidance from the MLflow community.

### MLflow version

- Client: 2.8.1



### System information

- Amazon linux
- python 3.9.10


### Describe the problem

When attempting to load a custom transformers model using MLflow's transformers integration, the process fails with an AttributeError. The specific error is: module transformers has no attribute LSGDistilBertForSequenceClassification. This suggests an issue with how MLflow is handling or recognizing the custom model class.

The error occurs within the load_model function, indicating a possible issue with the way MLflow interfaces with the transformers library, particularly around custom models.
It probably has to do with `_add_code_from_conf_to_system_path` function at mlflow/transformers/__init__.py line 866, which has a default for the CODE_KEY which does not match up (it should the 'auto_map' in this case) but i cant set that.


### Tracking information

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```shell
REPLACE_ME
```


### Code to reproduce issue

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
from transformers import AutoModelForSequenceClassification
import mlflow

# Attempt to load the custom transformers model
model = AutoModelForSequenceClassification.from_pretrained("ccdv/lsg-distilbert-base-uncased-4096",trust_remote_code=True)
with mlflow.start_run():
    mlflow.transformers.log_model(model, "custom_transformer_model",code_paths=["/home/ec2-user/.cache/huggingface/modules/transformers_modules/ccdv/lsg-distilbert-base-uncased-4096/79cac77ade3abac244e8037579f4cba2cf62736c/modeling_lsg_distilbert.py"])

# Load model with MLflow
loaded_model = mlflow.transformers.load_model("path/to/custom_transformer_model")

```


### Stack trace

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
Traceback (most recent call last):
  ...
  File "/path/to/mlflow/transformers/__init__.py", line 928, in _load_model
    model_instance = getattr(transformers, flavor_config[_PIPELINE_MODEL_TYPE_KEY])
  ...
AttributeError: module transformers has no attribute LSGDistilBertForSequenceClassification

```


### Other info / logs

<!-- PLEASE KEEP BACKTICKS AND CHECK PREVIEW -->
```
REPLACE_ME
```


### What component(s) does this bug affect?

- [X] `area/artifacts`: Artifact stores and artifact logging
- [ ] `area/build`: Build and test infrastructure for MLflow
- [ ] `area/deployments`: MLflow Deployments client APIs, server, and third-party Deployments integrations
- [ ] `area/docs`: MLflow documentation pages
- [ ] `area/examples`: Example code
- [X] `area/model-registry`: Model Registry service, APIs, and the fluent client calls for Model Registry
- [X] `area/models`: MLmodel format, model serialization/deserialization, flavors
- [ ] `area/recipes`: Recipes, Recipe APIs, Recipe configs, Recipe Templates
- [ ] `area/projects`: MLproject format, project running backends
- [ ] `area/scoring`: MLflow Model server, model deployment tools, Spark UDFs
- [ ] `area/server-infra`: MLflow Tracking server backend
- [ ] `area/tracking`: Tracking Service, tracking client APIs, autologging

### What interface(s) does this bug affect?

- [ ] `area/uiux`: Front-end, user experience, plotting, JavaScript, JavaScript dev server
- [ ] `area/docker`: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
- [ ] `area/sqlalchemy`: Use of SQLAlchemy in the Tracking Service or Model Registry
- [ ] `area/windows`: Windows support

### What language(s) does this bug affect?

- [ ] `language/r`: R APIs and clients
- [ ] `language/java`: Java APIs and clients
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/azure`: Azure and Azure ML integrations
- [ ] `integrations/sagemaker`: SageMaker integrations
- [ ] `integrations/databricks`: Databricks integrations


Discussion:
<!-- assign-maintainer -->
@mlflow/mlflow-team Please assign a maintainer and start triaging this issue.
I just found a simply temporary solution and it worked for me! 
In your module, where you are loading the mlflow model, put the following code:

```
import transformers 
transformers.LSGDistilBertForSequenceClassification = LSGDistilBertForSequenceClassification

``` 



Resolving patch:
diff --git a/mlflow/transformers/__init__.py b/mlflow/transformers/__init__.py
--- a/mlflow/transformers/__init__.py
+++ b/mlflow/transformers/__init__.py
@@ -1,4 +1,5 @@
 """MLflow module for HuggingFace/transformer support."""
+
 from __future__ import annotations
 
 import ast
@@ -280,7 +281,9 @@ def save_model(
     **kwargs,  # pylint: disable=unused-argument
 ) -> None:
     """
-    Save a trained transformers model to a path on the local file system.
+    Save a trained transformers model to a path on the local file system. Note that
+    saving transformers models with custom code (i.e. models that require
+    ``trust_remote_code=True``) requires ``transformers >= 4.26.0``.
 
     Args:
         transformers_model:
@@ -708,7 +711,9 @@ def log_model(
     **kwargs,
 ):
     """
-    Log a ``transformers`` object as an MLflow artifact for the current run.
+    Log a ``transformers`` object as an MLflow artifact for the current run. Note that
+    logging transformers models with custom code (i.e. models that require
+    ``trust_remote_code=True``) requires ``transformers >= 4.26.0``.
 
     Args:
         transformers_model:
diff --git a/mlflow/transformers/model_io.py b/mlflow/transformers/model_io.py
--- a/mlflow/transformers/model_io.py
+++ b/mlflow/transformers/model_io.py
@@ -123,6 +123,56 @@ def _load_component(flavor_conf, name, local_path=None):
         return cls.from_pretrained(repo, revision=revision)
 
 
+def _load_class_from_transformers_config(model_name_or_path, revision=None):
+    """
+    This method retrieves the Transformers AutoClass from the transformers config.
+    Using the correct AutoClass allows us to leverage Transformers' model loading
+    machinery, which is necessary for supporting models using custom code.
+    """
+    import transformers
+    from transformers import AutoConfig
+
+    config = AutoConfig.from_pretrained(
+        model_name_or_path,
+        revision=revision,
+        # trust_remote_code is set to True in order to
+        # make sure the config gets loaded as the correct
+        # class. if this is not set for custom models, the
+        # base class will be loaded instead of the custom one.
+        trust_remote_code=True,
+    )
+
+    # the model's class name (e.g. "MPTForCausalLM")
+    # is stored in the `architectures` field. it
+    # seems to usually just have one element.
+    class_name = config.architectures[0]
+
+    # if the class is available in transformers natively,
+    # then we don't need to execute any custom code.
+    if hasattr(transformers, class_name):
+        cls = getattr(transformers, class_name)
+        return cls, False
+    else:
+        # else, we need to fetch the correct AutoClass.
+        # this is defined in the `auto_map` field. there
+        # should only be one AutoClass that maps to the
+        # model's class name.
+        auto_classes = [
+            auto_class
+            for auto_class, module in config.auto_map.items()
+            if module.split(".")[-1] == class_name
+        ]
+
+        if len(auto_classes) == 0:
+            raise MlflowException(f"Couldn't find a loader class for {class_name}")
+
+        auto_class = auto_classes[0]
+        cls = getattr(transformers, auto_class)
+
+        # we will need to trust remote code when loading the model
+        return cls, True
+
+
 def _load_model(model_name_or_path, flavor_conf, accelerate_conf, device, revision=None):
     """
     Try to load a model with various loading strategies.
@@ -132,9 +182,17 @@ def _load_model(model_name_or_path, flavor_conf, accelerate_conf, device, revisi
     """
     import transformers
 
-    cls = getattr(transformers, flavor_conf[FlavorKey.MODEL_TYPE])
+    if hasattr(transformers, flavor_conf[FlavorKey.MODEL_TYPE]):
+        cls = getattr(transformers, flavor_conf[FlavorKey.MODEL_TYPE])
+        trust_remote = False
+    else:
+        cls, trust_remote = _load_class_from_transformers_config(
+            model_name_or_path, revision=revision
+        )
 
     load_kwargs = {"revision": revision} if revision else {}
+    if trust_remote:
+        load_kwargs.update({"trust_remote_code": True})
 
     if model := _try_load_model_with_accelerate(
         cls, model_name_or_path, {**accelerate_conf, **load_kwargs}


Test patch:
diff --git a/tests/transformers/conftest.py b/tests/transformers/conftest.py
--- a/tests/transformers/conftest.py
+++ b/tests/transformers/conftest.py
@@ -4,6 +4,7 @@
     load_audio_classification_pipeline,
     load_component_multi_modal,
     load_conversational_pipeline,
+    load_custom_code_pipeline,
     load_feature_extraction_pipeline,
     load_fill_mask_pipeline,
     load_ner_pipeline,
@@ -50,6 +51,11 @@ def component_multi_modal():
     return load_component_multi_modal()
 
 
+@pytest.fixture
+def custom_code_pipeline():
+    return load_custom_code_pipeline()
+
+
 @pytest.fixture
 def small_conversational_model():
     return load_small_conversational_model()
diff --git a/tests/transformers/helper.py b/tests/transformers/helper.py
--- a/tests/transformers/helper.py
+++ b/tests/transformers/helper.py
@@ -287,6 +287,20 @@ def load_peft_pipeline():
         )
 
 
+@prefetch
+@flaky()
+def load_custom_code_pipeline():
+    model = transformers.AutoModel.from_pretrained(
+        "hf-internal-testing/test_dynamic_model_with_util", trust_remote_code=True
+    )
+    tokenizer = transformers.AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+    return transformers.pipeline(
+        task="feature-extraction",
+        model=model,
+        tokenizer=tokenizer,
+    )
+
+
 def prefetch_models():
     """
     Prefetches models used in the test suite to avoid downloading them during the test run.
diff --git a/tests/transformers/test_transformers_model_export.py b/tests/transformers/test_transformers_model_export.py
--- a/tests/transformers/test_transformers_model_export.py
+++ b/tests/transformers/test_transformers_model_export.py
@@ -1369,6 +1369,32 @@ def test_table_question_answering_pipeline(
     assert pd_inference == result
 
 
+@pytest.mark.skipif(
+    Version(transformers.__version__) < Version("4.26"), reason="Feature is not available"
+)
+def test_custom_code_pipeline(custom_code_pipeline, model_path):
+    data = "hello"
+
+    signature = infer_signature(
+        data, mlflow.transformers.generate_signature_output(custom_code_pipeline, data)
+    )
+
+    mlflow.transformers.save_model(
+        custom_code_pipeline,
+        path=model_path,
+        signature=signature,
+    )
+
+    # just test that it doens't blow up when performing inference
+    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
+    pyfunc_pred = pyfunc_loaded.predict(data)
+    assert isinstance(pyfunc_pred[0][0], float)
+
+    transformers_loaded = mlflow.transformers.load_model(model_path)
+    transformers_pred = transformers_loaded(data)
+    assert pyfunc_pred[0][0] == transformers_pred[0][0][0]
+
+
 @pytest.mark.parametrize(
     ("data", "result"),
     [
@@ -2846,8 +2872,8 @@ def test_whisper_model_with_malformed_audio(whisper_pipeline):
 def test_save_model_card_with_non_utf_characters(tmp_path, model_name):
     # non-ascii unicode characters
     test_text = (
-        "Emoji testing! \u2728 \U0001F600 \U0001F609 \U0001F606 "
-        "\U0001F970 \U0001F60E \U0001F917 \U0001F9D0"
+        "Emoji testing! \u2728 \U0001f600 \U0001f609 \U0001f606 "
+        "\U0001f970 \U0001f60e \U0001f917 \U0001f9d0"
     )
 
     card_data: ModelCard = huggingface_hub.ModelCard.load(model_name)