Repository: huggingface/transformers

Issue:
`ModelOutput` keys aren't correctly assigned if key was previously None
Related issue: https://github.com/huggingface/transformers/pull/44050#discussion_r2815826882

### System Info

- `transformers` version: 5.2.0.dev0
- Platform: Windows-10-10.0.26200-SP0
- Python version: 3.11.13
- Huggingface_hub version: 1.3.1
- Safetensors version: 0.6.2
- Accelerate version: 1.11.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.9.0+cu126 (CUDA)
- Using distributed or parallel set-up in script?: no
- Using GPU in script?: no
- GPU type: NVIDIA GeForce RTX 3090

### Who can help?

@zucchini-nlp

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [x] My own task or dataset (give details below)

### Reproduction

If a `ModelOutput` (subclass) has been initialized with e.g. `pooler_output=None` and `pooler_output` is a valid dataclass entry for this subclass, then setting the value for this `pooler_output` with `outputs.pooler_output = tensor` will *not* add `pooler_output` to the `outputs` keys anymore. It's easiest to explain with this code example:

```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'])
"""
# ❌ The dict_keys doesn't have `pooler_output` in its keys, because __setitem__ is only called on keys already in
# self.keys(), otherwise __setattr__ is called instead.
print(no_pooler_outputs.pooler_output)
# tensor([0, 1, 2, 3]) ✅
print(no_pooler_outputs["pooler_output"])
# KeyError: 'pooler_output' ❌ because pooler_output is not in self.keys()
```

This is 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`. We can, or should, instead check whether the name is a field name. I'll open a PR.

### Expected behavior

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]`.

- Tom Aarsen

Discussion hints:

Solution patch:
diff --git a/src/transformers/utils/generic.py b/src/transformers/utils/generic.py
index 23eaa418cfc2..d6eec7433a79 100644
--- a/src/transformers/utils/generic.py
+++ b/src/transformers/utils/generic.py
@@ -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)


Test patch:
diff --git a/tests/utils/test_generic.py b/tests/utils/test_generic.py
index e99621d2d7ec..73d8fec05f23 100644
--- a/tests/utils/test_generic.py
+++ b/tests/utils/test_generic.py
@@ -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):