Update transformers bounds and tests

sdks/python/apache_beam/ml/inference/huggingface_inference_test.py

@@ -47,7 +47,6 @@
   import tensorflow as tf
   import torch
   from transformers import AutoModel
-  from transformers import TFAutoModel
 
   from apache_beam.ml.inference.tensorflow_inference_test import FakeTFTensorModel
   from apache_beam.ml.inference.tensorflow_inference_test import _compare_tensor_prediction_result
@@ -60,7 +59,7 @@
 
 def fake_inference_fn_tensor(
     batch: Sequence[Union['tf.Tensor', 'torch.Tensor']],
-    model: Union['AutoModel', 'TFAutoModel'],
+    model: Any,
     device,
     inference_args: dict[str, Any],
     model_id: Optional[str] = None) -> Iterable[PredictionResult]:
@@ -86,7 +85,7 @@ def test_predict_tensor(self):
     fake_model = FakeTFTensorModel()
     inference_runner = HuggingFaceModelHandlerTensor(
         model_uri='unused',
-        model_class=TFAutoModel,
+        model_class=AutoModel,
         inference_fn=fake_inference_fn_tensor)
     batched_examples = [tf.constant([1]), tf.constant([10]), tf.constant([100])]
     expected_predictions = [
@@ -103,7 +102,7 @@ def test_predict_tensor_with_inference_args(self):
     fake_model = FakeTFTensorModel()
     inference_runner = HuggingFaceModelHandlerTensor(
         model_uri='unused',
-        model_class=TFAutoModel,
+        model_class=AutoModel,
         inference_fn=fake_inference_fn_tensor,
         inference_args={"add": True})
     batched_examples = [tf.constant([1]), tf.constant([10]), tf.constant([100])]
@@ -125,7 +124,7 @@ def test_framework_detection_torch(self):
     fake_model = FakeTorchModel()
     inference_runner = HuggingFaceModelHandlerTensor(
         model_uri='unused',
-        model_class=TFAutoModel,
+        model_class=AutoModel,
         inference_fn=fake_inference_fn_tensor)
     batched_examples = [torch.tensor(1), torch.tensor(10), torch.tensor(100)]
     inference_runner.run_inference(batched_examples, fake_model)
@@ -135,7 +134,7 @@ def test_framework_detection_tensorflow(self):
     fake_model = FakeTFTensorModel()
     inference_runner = HuggingFaceModelHandlerTensor(
         model_uri='unused',
-        model_class=TFAutoModel,
+        model_class=AutoModel,
         inference_fn=fake_inference_fn_tensor)
     batched_examples = [tf.constant(1), tf.constant(10), tf.constant(100)]
     inference_runner.run_inference(batched_examples, fake_model)

sdks/python/setup.py

@@ -665,8 +665,12 @@ def get_portability_package_data():
               'absl-py>=0.12.0'
           ],
           'transformers': [
-              'transformers>=4.28.0,<4.56.0',
+              # Keep in line with the py311-transformers-* environments in
+              # tox.ini.
+              'transformers>=4.41.0,!=4.56.*,<5.21.0',
               'tensorflow>=2.12.0',
+              # torch>=2.12 Linux wheels require triton>=3.7,
+              # which segfaults/aborts when imported after tensorflow.
               'torch>=2.3.0,<2.12.0'
           ],
           'ml_cpu': [

sdks/python/test-suites/tox/py311/build.gradle

@@ -122,11 +122,14 @@ test.dependsOn "testPy311tensorflow-212"
 postCommitPyDep.dependsOn "testPy311tensorflow-212"
 
 // Create a test task for each minor version of transformers.
-// Dropped transformers-428: sentence-transformers==2.2.2 + torch<1.14 is
-// ResolutionImpossible on Python 3.11 (see #30799).
-toxTask "testPy311transformers-447", "py311-transformers-447", "${posargs}"
-test.dependsOn "testPy311transformers-447"
-postCommitPyDep.dependsOn "testPy311transformers-447"
+// Keep in sync with [testenv:py311-transformers-...] in tox.ini.
+toxTask "testPy311transformers-441", "py311-transformers-441", "${posargs}"
+test.dependsOn "testPy311transformers-441"
+postCommitPyDep.dependsOn "testPy311transformers-441"
+
+toxTask "testPy311transformers-457", "py311-transformers-457", "${posargs}"
+test.dependsOn "testPy311transformers-457"
+postCommitPyDep.dependsOn "testPy311transformers-457"
 
 toxTask "testPy311transformers-latest", "py311-transformers-latest", "${posargs}"
 test.dependsOn "testPy311transformers-latest"

sdks/python/tox.ini

@@ -485,17 +485,23 @@ commands =
   # Allow exit code 5 (no tests run) so that we can run this command safely on arbitrary subdirectories.
   /bin/sh -c 'pytest -o junit_suite_name={envname} --junitxml=pytest_{envname}.xml -n 6 -m uses_xgboost {posargs}; ret=$?; [ $ret = 5 ] && exit 0 || exit $ret'
 
-[testenv:py311-transformers-{447,latest}]
+[testenv:py311-transformers-{441,457,latest}]
 deps =
   # Environment dependencies are defined in the `setenv` section and installed in the `commands` section.
 pip_pre = False
 extras = test
 setenv =
-    COMMON_DEPS = tensorflow==2.12.0 protobuf==4.25.5 pip==25.0.1
-    # transformers-428 dropped: sentence-transformers==2.2.2 + torch<1.14 is
-    # ResolutionImpossible on Python 3.11.
-    447: DEPS = 'transformers>=4.47.0,<4.48.0' 'torch>=1.9.0,<1.14.0'
-    latest: DEPS = 'transformers>=4.55.0' 'torch>=2.0.0' 'accelerate>=1.6.0'
+    # Keep in line with the transformers extra in setup.py.
+    # tensorflow-hub is needed by the TensorFlow handler tests; it imports
+    # pkg_resources, which was removed in setuptools 82.
+    COMMON_DEPS = 'tensorflow>=2.16,<2.22' tf-keras tensorflow-hub 'setuptools<82' sentence-transformers pip==25.0.1
+    # Lower bound of the supported range.
+    441: DEPS = 'transformers>=4.41.0,<4.42.0' 'torch>=2.3.0,<2.4.0'
+    # Last transformers release line with TensorFlow support.
+    # torch<2.12: torch>=2.12 pulls in triton>=3.7, which crashes when imported
+    # after tensorflow.
+    457: DEPS = 'transformers>=4.57.0,<4.58.0' 'torch>=2.3.0,<2.12.0'
+    latest: DEPS = 'transformers>=5.20.0,<5.21.0' 'torch>=2.11.0,<2.12.0' 'accelerate>=1.6.0'
 commands =
   /bin/sh -c "pip install .[{extras}] {env:DEPS} {env:COMMON_DEPS}"
   # Log transformers and its dependencies version for debugging