fix collate fn

src/zeroband/data.py

@@ -54,8 +54,8 @@ def _collate_fn_causal_mask(
         elif len(input_ids) > max_seq_length:
             input_ids = input_ids[:max_seq_length]
 
-        batched["input_ids"].append(input_ids[1:])
-        batched["labels"].append(input_ids[:-1])
+        batched["input_ids"].append(input_ids[:-1])
+        batched["labels"].append(input_ids[1:])
 
     return {"input_ids": torch.stack(batched["input_ids"], dim=0), "labels": torch.stack(batched["labels"], dim=0)}
 

tests/test_data.py

@@ -0,0 +1,19 @@
+import torch
+from zeroband.data import collate_causal_mask
+
+
+def test_collate_fn():
+    tensors = [[0, 1, 2, 3, 4], [0, 0, 3, 4, 1, 7]]
+
+    batch = [{"input_ids": torch.Tensor(tensor)} for tensor in tensors]
+
+    collate_fn = collate_causal_mask(max_seq_length=4)
+    collated = collate_fn(batch)
+
+    assert collated is not None
+
+    assert collated["input_ids"][0].tolist() == [0, 1, 2, 3]
+    assert collated["labels"][0].tolist() == [1, 2, 3, 4]
+
+    assert collated["input_ids"][1].tolist() == [0, 0, 3, 4]
+    assert collated["labels"][1].tolist() == [0, 3, 4, 1]