Repository: huggingface/transformers
Pull request: Patch `__setitem__` on `ModelOutput` even if the parameter was previously `None`

# What does this PR do?

* Patch `__setitem__` on `ModelOutput` even if the parameter was previously `None`

Fixes #44079, follow-up from #44050.

Essentially, it brings behaviour to the expected as described in #44079:

> If I 1) initialize a `ModelClass` subclass with some parameter set to None, and then 2) set that parameter to some non-None value, then I would expect that 3a) the parameter is added to `.keys()` and 3b) I can access it with `outputs[parameter]`.

This was caused by this: https://github.com/huggingface/transformers/blob/16a3bea3b88e0530f78d4d7a2fcc0f6387ac72b9/src/transformers/utils/generic.py#L429-L433

Here, `__setitem__` is only set if the name is already in `self.keys()`, but `self.keys()` excludes keys whose value have been set to `None`. This PR instead checks whether the name is a field name.

## Reproduction
When running the script from #44079, I now get:
```python
import torch
from transformers.modeling_outputs import BaseModelOutputWithPooling

# An example BaseModelOutputWithPooling with last_hidden_state and pooler_output:
outputs = BaseModelOutputWithPooling(
    last_hidden_state=torch.randn(1, 2, 4),
    pooler_output=torch.randn(1, 4),
    hidden_states=None,
    attentions=None,
)
print(outputs)
print(outputs.keys())
"""
BaseModelOutputWithPooling(last_hidden_state=tensor([[[ 0.1468,  0.8285,  1.9449,  0.3687],
         [ 1.1413, -1.6430, -2.5313, -0.9286]]]), pooler_output=tensor([[-0.2389,  1.8526,  0.9567,  0.7564]]), hidden_states=None, attentions=None)
dict_keys(['last_hidden_state', 'pooler_output'])
"""
# ✅

# We can override the pooler output with a new tensor:
outputs.pooler_output = torch.arange(4)
print(outputs)
print(outputs.keys())
"""
BaseModelOutputWithPooling(last_hidden_state=tensor([[[ 0.1468,  0.8285,  1.9449,  0.3687],
         [ 1.1413, -1.6430, -2.5313, -0.9286]]]), pooler_output=tensor([0, 1, 2, 3]), hidden_states=None, attentions=None)
dict_keys(['last_hidden_state', 'pooler_output'])
"""
# ✅

# An example BaseModelOutputWithPooling without pooler_output:
no_pooler_outputs = BaseModelOutputWithPooling(
    last_hidden_state=torch.randn(1, 2, 4),
    pooler_output=None,
    hidden_states=None,
    attentions=None,
)
print(no_pooler_outputs)
print(no_pooler_outputs.keys())
"""
BaseModelOutputWithPooling(last_hidden_state=tensor([[[ 1.2644, -0.6101, -0.8053, -0.2578],
         [ 1.5456,  1.2600, -1.4123,  1.4299]]]), pooler_output=None, hidden_states=None, attentions=None)
dict_keys(['last_hidden_state'])
"""
# ✅

# Now let's try and override the pooler output on the no_pooler_outputs:
no_pooler_outputs.pooler_output = torch.arange(4)
print(no_pooler_outputs)
print(no_pooler_outputs.keys())
"""
BaseModelOutputWithPooling(last_hidden_state=tensor([[[-0.0020,  0.2205,  0.9009,  0.3189], 
         [-0.9656,  1.6390, -0.9418,  0.7022]]]), pooler_output=tensor([0, 1, 2, 3]), hidden_states=None, attentions=None)
dict_keys(['last_hidden_state', 'pooler_output'])
"""
# ✅
print(no_pooler_outputs.pooler_output)
# tensor([0, 1, 2, 3]) ✅
print(no_pooler_outputs["pooler_output"])
# tensor([0, 1, 2, 3]) ✅
```

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@zucchini-nlp

- Tom Aarsen

Changed files:
filename: src/transformers/utils/generic.py
patch: @@ -383,8 +383,9 @@ 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
+                # reset first field to None and remove it from the internal dictionary
                 setattr(self, class_fields[0].name, None)
+                super().__delitem__(class_fields[0].name)
                 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:
@@ -427,7 +428,8 @@ def __getitem__(self, k):
             return self.to_tuple()[k]
 
     def __setattr__(self, name, value):
-        if name in self.keys() and value is not None:
+        field_names = {field.name for field in fields(self)}
+        if name in field_names and value is not None:
             # Don't call self.__setitem__ to avoid recursion errors
             super().__setitem__(name, value)
         super().__setattr__(name, value)

filename: tests/utils/test_generic.py
patch: @@ -142,15 +142,27 @@ def test_to_py_obj_torch(self):
     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)
+        self.assertNotEqual(out["logits"], None)
+        self.assertEqual(out.loss, None)
+        self.assertEqual(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)
+        self.assertNotEqual(out["logits"], None)
+        self.assertEqual(out.loss, None)
+        self.assertEqual(len(out.to_tuple()), 1)
+
+        # testing with updating a previously-None key after init with attribute assignment
+        out = CausalLMOutputWithPast(logits=torch.ones(2, 3, 4))
+        out.loss = torch.tensor(0.5)
+        self.assertEqual(out.loss, torch.tensor(0.5))
+        self.assertEqual(len(out.to_tuple()), 2)
+
+        # testing with updating a previously-None key after init with dictionary assignment
+        out = CausalLMOutputWithPast(logits=torch.ones(2, 3, 4))
+        out["loss"] = torch.tensor(0.5)
+        self.assertEqual(out.loss, torch.tensor(0.5))
+        self.assertEqual(len(out.to_tuple()), 2)
 
 
 class ValidationDecoratorTester(unittest.TestCase):