Repository: huggingface/transformers
Pull request: fix: support tensor labels in DataCollatorWithFlattening

## Description

Fixes #42502

This PR adds support for `torch.Tensor` labels in `DataCollatorWithFlattening`. Previously, the collator only worked with list labels and raised a `TypeError` when tensor labels were provided.

## Problem

When users passed `torch.Tensor` objects as labels to `DataCollatorWithFlattening`, it failed with:
```python
TypeError: can only concatenate list (not "Tensor") to list
```

This happened because the collator used list concatenation (`+=`) without checking if inputs were tensors.

## Solution

- Added type checking using `hasattr(obj, "tolist")` 
- Convert tensors to lists before concatenation
- Applied to both `input_ids` and `labels`
- Maintains backward compatibility with list inputs

## Changes Made

- Modified `src/transformers/data/data_collator.py`:
  - Added tensor-to-list conversion in `DataCollatorWithFlattening.__call__()`
- Added tests in `tests/trainer/test_data_collator.py`:
  - `test_flattening_with_tensor_labels`: Verifies tensor labels work
  - `test_flattening_with_list_labels`: Regression test for list labels

## Testing

### Before this PR ❌
```python
features = [
    {"input_ids": torch.tensor([1, 2, 3, 4]), "labels": torch.tensor([10, 11, 12, 13])},
    {"input_ids": torch.tensor([5, 6, 7]), "labels": torch.tensor([14, 15, 16])},
]
collator = DataCollatorWithFlattening(return_tensors="pt")
batch = collator(features)  # TypeError!
```

### After this PR ✅
```python
features = [
    {"input_ids": torch.tensor([1, 2, 3, 4]), "labels": torch.tensor([10, 11, 12, 13])},
    {"input_ids": torch.tensor([5, 6, 7]), "labels": torch.tensor([14, 15, 16])},
]
collator = DataCollatorWithFlattening(return_tensors="pt")
batch = collator(features)  # Works perfectly!
```

## Checklist

- [x] Bug fix (non-breaking change which fixes an issue)
- [x] Tests added for the fix
- [x] All new and existing tests passed
- [x] Code follows the project's style guidelines

Changed files:
filename: src/transformers/data/data_collator.py
patch: @@ -1413,9 +1413,17 @@ def __call__(self, features, return_tensors=None, separator_id=None):
             max_length = 0
         for seq_idx, sample in enumerate(features):
             input_ids = sample["input_ids"]
+            # Convert to list if tensor
+            if hasattr(input_ids, "tolist"):
+                input_ids = input_ids.tolist()
             batch["input_ids"] += input_ids
+
             if is_labels_provided:
-                batch["labels"] += [separator_id] + sample["labels"][1:]
+                labels = sample["labels"]
+                # Convert to list if tensor
+                if hasattr(labels, "tolist"):
+                    labels = labels.tolist()
+                batch["labels"] += [separator_id] + labels[1:]
             else:
                 batch["labels"] += [separator_id] + input_ids[1:]
             if self.return_position_ids:

filename: tests/trainer/test_data_collator.py
patch: @@ -1965,3 +1965,55 @@ def test__whole_word_mask(self):
         ).astype(bool)
 
         np.testing.assert_array_equal(output_mask, expected_mask)
+
+
+class DataCollatorWithFlatteningTest(unittest.TestCase):
+    """Tests for DataCollatorWithFlattening"""
+
+    def test_flattening_with_tensor_labels(self):
+        """Test that DataCollatorWithFlattening supports tensor labels (fixes issue #42599)."""
+        features = [
+            {
+                "input_ids": torch.tensor([1, 2, 3, 4]),
+                "labels": torch.tensor([10, 11, 12, 13]),
+            },
+            {
+                "input_ids": torch.tensor([5, 6, 7]),
+                "labels": torch.tensor([14, 15, 16]),
+            },
+        ]
+        collator = DataCollatorWithFlattening(return_tensors="pt")
+
+        # This should not raise TypeError anymore
+        batch = collator(features)
+
+        # Verify the output
+        self.assertIsInstance(batch, dict)
+        self.assertIn("input_ids", batch)
+        self.assertIn("labels", batch)
+        self.assertIn("position_ids", batch)
+
+        # Check shapes
+        self.assertEqual(batch["input_ids"].shape, (1, 7))  # 4 + 3 tokens
+        self.assertEqual(batch["labels"].shape, (1, 7))
+        self.assertEqual(batch["position_ids"].shape, (1, 7))
+
+    def test_flattening_with_list_labels(self):
+        """Test that DataCollatorWithFlattening still works with list labels."""
+        features = [
+            {
+                "input_ids": torch.tensor([1, 2, 3, 4]),
+                "labels": [10, 11, 12, 13],
+            },
+            {
+                "input_ids": torch.tensor([5, 6, 7]),
+                "labels": [14, 15, 16],
+            },
+        ]
+        collator = DataCollatorWithFlattening(return_tensors="pt")
+        batch = collator(features)
+
+        # Verify it still works with lists
+        self.assertIsInstance(batch, dict)
+        self.assertEqual(batch["input_ids"].shape, (1, 7))
+        self.assertEqual(batch["labels"].shape, (1, 7))