Fix transformers tests (mlflow#14177)

Signed-off-by: serena-ruan <[email protected]>
Signed-off-by: Serena Ruan <[email protected]>

mlflow/ml-package-versions.yml

@@ -561,7 +561,6 @@ transformers:
           "torch",
           "torchvision",
           "tensorflow",
-          "accelerate",
           "setfit",
           "optuna",
           "accelerate<1",
@@ -576,11 +575,13 @@ transformers:
       # Installing `transformers==4.34.1` with `datasets` results in installing `huggingface_hub==0.17.3`.
       # `accelerate>=0.32.0` is incompatible with `huggingface_hub==0.17.3`.
       "== 4.34.1": ["accelerate<0.32.0", "sentence-transformers<3.1.0"]
-      "< 4.38": ["tensorflow<2.14"] # Older versions of transformers are incompatible with tensorflow >= 2.14
+      "< 4.38": ["tensorflow<2.14", "pydantic<2"] # Older versions of transformers are incompatible with tensorflow >= 2.14
       ">= 4.38": ["tf-keras"] # Transformers doesn't support Keras 3.0. tf-keras needs to be installed.
       # hf_hub>=0.24.0 is broken with setfit<=1.0.3. the incompatibility was fixed in setfit>=1.1, but
       # that version of setfit requires tranformers>=4.41
-      "< 4.41": ["huggingface_hub<0.24.0"]
+      # datasets > 2.4.0 is incompatible with huggingface_hub<0.24.0
+      # evaluate >= 0.4.3 is incompatible with datasets<=2.4.0
+      "< 4.41": ["huggingface_hub<0.24.0", "datasets<=2.4.0", "evaluate<0.4.3"]
     run: |
       pytest tests/transformers/test_transformers_autolog.py
 

mlflow/models/utils.py

@@ -27,13 +27,13 @@
 from mlflow.tracking.artifact_utils import _download_artifact_from_uri
 from mlflow.types import DataType, ParamSchema, ParamSpec, Schema, TensorSpec
 from mlflow.types.schema import AnyType, Array, Map, Object, Property
-from mlflow.types.type_hints import PYDANTIC_V1_OR_OLDER
 from mlflow.types.utils import (
     TensorsNotSupportedException,
     _infer_param_schema,
     _is_none_or_nan,
     clean_tensor_type,
 )
+from mlflow.utils import IS_PYDANTIC_V2_OR_NEWER
 from mlflow.utils.annotations import experimental
 from mlflow.utils.file_utils import create_tmp_dir, get_local_path_or_none
 from mlflow.utils.proto_json_utils import (
@@ -304,7 +304,9 @@ def __init__(self, input_example: ModelInputExample):
         model_input = deepcopy(self._inference_data)
 
         if isinstance(model_input, pydantic.BaseModel):
-            model_input = model_input.dict() if PYDANTIC_V1_OR_OLDER else model_input.model_dump()
+            model_input = (
+                model_input.model_dump() if IS_PYDANTIC_V2_OR_NEWER else model_input.dict()
+            )
 
         is_unified_llm_input = False
         if isinstance(model_input, dict):