Return combined decorator in transformers 5.2.0

Author: oliverholworthy

nemo_automodel/shared/import_utils.py

@@ -497,6 +497,8 @@ def get_check_model_inputs_decorator():
 
     In transformers >= 4.57.3, check_model_inputs became a function that returns a decorator.
     In older versions, it was directly a decorator.
+    In transformers >= 5.2.0, check_model_inputs was removed and split into
+    ``merge_with_config_defaults`` and ``capture_outputs``.
 
     Returns:
         Decorator function to validate model inputs.
@@ -511,5 +513,19 @@ def get_check_model_inputs_decorator():
             # Old API: check_model_inputs is directly a decorator
             return check_model_inputs
     except ImportError:
-        # If transformers is not available, return a no-op decorator
-        return null_decorator
+        pass
+
+    # transformers >= 5.2.0: check_model_inputs was split into two decorators
+    try:
+        from transformers.utils.generic import merge_with_config_defaults
+        from transformers.utils.output_capturing import capture_outputs
+
+        def _combined_decorator(func):
+            return merge_with_config_defaults(capture_outputs(func))
+
+        return _combined_decorator
+    except ImportError:
+        pass
+
+    # No transformers decorator available — return a no-op decorator
+    return null_decorator

tests/unit_tests/shared/test_import_utils.py

@@ -202,6 +202,25 @@ def test_get_check_model_inputs_decorator():
     assert callable(decorator)
 
 
+def test_get_check_model_inputs_decorator_fallback_with_kwargs():
+    """
+    The ``null_decorator`` fallback (returned by ``get_check_model_inputs_decorator``
+    when transformers decorators are unavailable) must work as a plain ``@decorator``
+    on functions called with keyword arguments.
+
+    This guards against a bug where a ``@contextmanager``-wrapped fallback
+    produced a ``ContextDecorator`` whose ``__call__`` collided with model
+    forward kwargs like ``input_ids``.
+    """
+    decorator = si.null_decorator
+
+    @decorator
+    def dummy(x=None):
+        return x
+
+    assert dummy(x=42) == 42
+
+
 def test_null_decorator_as_direct_decorator():
     """
     ``null_decorator`` must be a valid no-op decorator in ``@decorator`` form.