fix(train): pad training batches with the model's pad id

model2vec/train/base.py

@@ -402,7 +402,7 @@ def _prepare_dataset(self, X: list[str], y: torch.Tensor, max_length: int | None
             encoded = self.tokenizer.encode_batch_fast(batch, add_special_tokens=False)
             tokenized.extend([encoding.ids[:max_length] for encoding in encoded])
 
-        return TextDataset(tokenized, y)
+        return TextDataset(tokenized, y, pad_id=self.pad_id)
 
     def _labels_to_tensor(self, labels: Any) -> torch.Tensor:
         """Turn the labels into a tensor."""

model2vec/train/dataset.py

@@ -4,17 +4,19 @@
 
 
 class TextDataset(Dataset):
-    def __init__(self, tokenized_texts: list[list[int]], targets: torch.Tensor) -> None:
+    def __init__(self, tokenized_texts: list[list[int]], targets: torch.Tensor, pad_id: int = 0) -> None:
         """A dataset of texts.
 
         :param tokenized_texts: The tokenized texts. Each text is a list of token ids.
         :param targets: The targets.
+        :param pad_id: The id used to pad batches. Must match the `pad_id` of the model being trained.
         :raises ValueError: If the number of targets does not match the number of texts.
         """
         if len(targets) != len(tokenized_texts):
             raise ValueError("Number of targets does not match number of texts.")
         self.tokenized_texts = tokenized_texts
         self.targets = targets
+        self.pad_id = pad_id
 
     def __len__(self) -> int:
         """Return the length of the dataset."""
@@ -24,13 +26,12 @@ def __getitem__(self, index: int) -> tuple[list[int], torch.Tensor]:
         """Gets an item."""
         return self.tokenized_texts[index], self.targets[index]
 
-    @staticmethod
-    def collate_fn(batch: list[tuple[list[list[int]], int]]) -> tuple[torch.Tensor, torch.Tensor]:
+    def collate_fn(self, batch: list[tuple[list[list[int]], int]]) -> tuple[torch.Tensor, torch.Tensor]:
         """Collate function."""
         texts, targets = zip(*batch)
 
         tensors: list[torch.Tensor] = [torch.LongTensor(x) for x in texts]
-        padded = pad_sequence(tensors, batch_first=True, padding_value=0)
+        padded = pad_sequence(tensors, batch_first=True, padding_value=self.pad_id)
 
         return padded, torch.stack(targets)
 

tests/test_trainable.py

@@ -147,6 +147,19 @@ def test_textdataset_init_incorrect() -> None:
         TextDataset([[0]], torch.arange(2))
 
 
+def test_training_batch_padding_is_masked(mock_vectors: np.ndarray, mock_tokenizer: Tokenizer) -> None:
+    """Training batches should pad with the model's pad id, so padding stays masked and out of the mean."""
+    s = StaticModelForClassification(vectors=torch.from_numpy(mock_vectors).float(), tokenizer=mock_tokenizer, pad_id=1)
+    texts = ["word2", "word2 word3"]
+
+    dataset = s._prepare_dataset(texts, torch.arange(2), max_length=None)
+    batch, _ = next(iter(dataset.to_dataloader(shuffle=False, batch_size=2)))
+
+    assert torch.equal(batch, s.tokenize(texts))
+    with torch.no_grad():
+        assert torch.allclose(s._encode(batch)[0], s._encode(s.tokenize(texts[:1]))[0])
+
+
 def test_predict(mock_trained_pipeline: StaticModelForClassification) -> None:
     """Test the predict function."""
     result = mock_trained_pipeline.predict(["dog cat", "dog"]).tolist()