(test): add deterministic coverage for collator utilities

tests/datasets/test_collator.py

@@ -1,49 +1,93 @@
+import numpy as np
 import torch
-from transformers import PreTrainedTokenizerFast, AutoTokenizer
 
-from roll.datasets.collator import DataCollatorWithPaddingForPaddedKeys
+from roll.datasets.collator import DataCollatorWithPaddingForPaddedKeys, collate_fn_to_dict_list
 
 
-def test_data_collator_with_padding_for_padded_keys():
-    tokenizer = AutoTokenizer.from_pretrained("/Users/pan/Downloads/huggingface/gpt2-imdb", padding_side="left")
+class DummyTokenizer:
+    """Minimal tokenizer stub that behaves like a HF tokenizer for padding."""
 
-    tokenizer.pad_token_id = tokenizer.eos_token_id
+    def __init__(self, pad_token_id: int = 0, padding_side: str = "right"):
+        self.pad_token_id = pad_token_id
+        self.padding_side = padding_side
+        self.model_input_names = ["input_ids", "attention_mask", "labels"]
 
-    max_length = 32
-    data_collator = DataCollatorWithPaddingForPaddedKeys(
-        tokenizer=tokenizer, padding="max_length", max_length=max_length
+    def pad(
+        self,
+        encoded_inputs,
+        padding=True,
+        max_length=None,
+        pad_to_multiple_of=None,
+        return_tensors=None,
+    ):
+        assert padding in [True, "max_length"]
+        target_length = max_length or max(len(feature["input_ids"]) for feature in encoded_inputs)
+        pad_values = {"input_ids": self.pad_token_id, "attention_mask": 0, "labels": -100}
+        padded = {key: [] for key in encoded_inputs[0].keys()}
+        for feature in encoded_inputs:
+            for key, value in feature.items():
+                value_list = list(value)
+                pad_value = pad_values.get(key, 0)
+                padded[key].append(value_list + [pad_value] * (target_length - len(value_list)))
+        if return_tensors == "pt":
+            for key in padded:
+                padded[key] = torch.tensor(padded[key], dtype=torch.long)
+        return padded
+
+
+def test_collate_fn_to_dict_list_merges_tensor_and_python_data():
+    data_list = [
+        {"input_ids": torch.tensor([[1, 2]]), "meta": {"id": "a"}},
+        {"input_ids": torch.tensor([[3, 4]]), "meta": {"id": "b"}},
+    ]
+
+    output = collate_fn_to_dict_list(data_list)
+
+    assert torch.equal(output["input_ids"], torch.tensor([[1, 2], [3, 4]]))
+    assert isinstance(output["meta"], np.ndarray)
+    assert output["meta"].shape == (2,)
+    assert output["meta"][0]["id"] == "a"
+    assert output["meta"][1]["id"] == "b"
+
+
+def test_data_collator_with_padding_for_padded_keys_handles_unpadded_fields():
+    tokenizer = DummyTokenizer(pad_token_id=9)
+    collator = DataCollatorWithPaddingForPaddedKeys(
+        tokenizer=tokenizer,
+        padding="max_length",
+        max_length=6,
     )
 
     features = [
         {
-            "input_ids": tokenizer.encode("Hello, how are you?", return_tensors="pt").squeeze(0),
-            "labels": torch.tensor(1),
+            "input_ids": [1, 2, 3],
+            "attention_mask": [1, 1, 1],
+            "labels": [10, 11, 12],
             "auxiliary": {"type": 1},
         },
         {
-            "input_ids": tokenizer.encode("I'm fine, thank you!", return_tensors="pt").squeeze(0),
-            "labels": torch.tensor(0),
+            "input_ids": [4, 5],
+            "attention_mask": [1, 1],
+            "labels": [13, 14],
             "auxiliary": {"type": 2},
         },
-        {
-            "input_ids": tokenizer.encode("What about you?", return_tensors="pt").squeeze(0),
-            "labels": torch.tensor(1),
-            "auxiliary": {"type": 3},
-        },
     ]
-    for feature in features:
-        feature["attention_mask"] = [1] * len(feature["input_ids"])
-
-    batch = data_collator(features)
-
-    print("Padded input_ids:")
-    print(batch["input_ids"])
-    print("Padded attention_mask:")
-    print(batch["attention_mask"])
-    print("Labels:")
-    print(batch["labels"])
-
-    assert (
-        batch["input_ids"].shape[1] == max_length
-    ), f"Expected max_length {max_length}, got {batch['input_ids'].shape[1]}"
-    print(f"All inputs padded to length {max_length} correctly.")
+
+    batch = collator(features)
+
+    assert batch["input_ids"].shape == (2, 6)
+    assert torch.equal(batch["input_ids"][0, 3:], torch.tensor([9, 9, 9]))
+    assert torch.equal(batch["attention_mask"][1], torch.tensor([1, 1, 0, 0, 0, 0]))
+    assert torch.equal(batch["labels"][1], torch.tensor([13, 14, -100, -100, -100, -100]))
+
+    expected_position_ids = torch.tensor(
+        [
+            [0, 1, 2, 2, 2, 2],
+            [0, 1, 1, 1, 1, 1],
+        ]
+    )
+    assert torch.equal(batch["position_ids"], expected_position_ids)
+
+    assert isinstance(batch["auxiliary"], np.ndarray)
+    assert batch["auxiliary"][0]["type"] == 1
+    assert batch["auxiliary"][1]["type"] == 2