fix dict like init for ModelOutput (#41002)

* fix dict like init

* style

Author: SunMarc

src/transformers/utils/generic.py

@@ -380,6 +380,8 @@ def __post_init__(self):
             # if we provided an iterator as first field and the iterator is a (key, value) iterator
             # set the associated fields
             if first_field_iterator:
+                # reset first field to None
+                setattr(self, class_fields[0].name, None)
                 for idx, element in enumerate(iterator):
                     if not isinstance(element, (list, tuple)) or len(element) != 2 or not isinstance(element[0], str):
                         if idx == 0:

tests/utils/test_generic.py

@@ -19,7 +19,7 @@
 import pytest
 
 from transformers.configuration_utils import PretrainedConfig
-from transformers.modeling_outputs import BaseModelOutput
+from transformers.modeling_outputs import BaseModelOutput, CausalLMOutputWithPast
 from transformers.testing_utils import require_torch
 from transformers.utils import (
     can_return_tuple,
@@ -139,6 +139,19 @@ def test_to_py_obj_torch(self):
 
         self.assertTrue(to_py_obj([t1, t2]) == [x1, x2])
 
+    def test_model_output_subclass(self):
+        # testing with “dict-like init” case
+        out = CausalLMOutputWithPast({"logits": torch.ones(2, 3, 4)})
+        self.assertTrue(out["logits"] is not None)
+        self.assertTrue(out.loss is None)
+        self.assertTrue(len(out.to_tuple()) == 1)
+
+        # testing with dataclass init case
+        out = CausalLMOutputWithPast(logits=torch.ones(2, 3, 4))
+        self.assertTrue(out["logits"] is not None)
+        self.assertTrue(out.loss is None)
+        self.assertTrue(len(out.to_tuple()) == 1)
+
 
 class ValidationDecoratorTester(unittest.TestCase):
     def test_cases_no_warning(self):