feat(data loading): GPT2LLMCollateFn can now determine the sub sequence lengths required for inter document attention masking.

Author: BlueCrescent

src/modalities/config/config.py

@@ -443,6 +443,25 @@ class BatchSamplerConfig(BaseModel):
 class GPT2LLMCollateFnConfig(BaseModel):
     sample_key: str
     target_key: str
+    sub_seq_lengths_key: str | None = None
+    eod_token_id: int | None = None
+    padding_token_id: int | None = None
+
+    @model_validator(mode="before")
+    def check_sub_seq_lengths_and_eod_token(cls, values):
+        sub_seq_lengths_key = values.get("sub_seq_lengths_key")
+        eod_token_id = values.get("eod_token_id")
+        if (sub_seq_lengths_key is None) != (eod_token_id is None):
+            raise ValueError("Either both or neither of sub_seq_lengths_key and eod_token_id must be provided.")
+        return values
+
+    @model_validator(mode="before")
+    def check_padding_token_and_sub_seq_lengths(cls, values):
+        padding_token_id = values.get("padding_token_id")
+        sub_seq_lengths_key = values.get("sub_seq_lengths_key")
+        if padding_token_id is not None and sub_seq_lengths_key is None:
+            raise ValueError("If padding_token_id is provided, sub_seq_lengths_key must also be provided.")
+        return values
 
 
 class LLMDataLoaderConfig(BaseModel):

src/modalities/models/gpt2/collator.py

@@ -7,16 +7,31 @@
 class GPT2LLMCollateFn(CollateFnIF):
     """GPT2LLMCollateFn class to define a collate function for GPT2 language model."""
 
-    def __init__(self, sample_key: str, target_key: str):
+    def __init__(
+        self,
+        sample_key: str,
+        target_key: str,
+        sub_seq_lengths_key: str | None = None,
+        eod_token_id: int | None = None,
+        padding_token_id: int | None = None,
+    ):
         """
         Initializes the Collator object.
+        If the eod document token ID and the sub_seq_lengths_key are provided,
+        a list[list[int]] representing the sub-sequence lengths will be created.
 
         Args:
             sample_key (str): The key for accessing the sample data.
             target_key (str): The key for accessing the target data.
+            sub_seq_lengths_key (str | None): The key for accessing the sub-sequence lengths.
+            eod_token_id (int | None): The end-of-document token ID.
+            padding_token_id (int | None): The padding token ID.
         """
         self.sample_key = sample_key
         self.target_key = target_key
+        self.sub_seq_lengths_key = sub_seq_lengths_key
+        self.eod_token_id = eod_token_id
+        self.padding_token_id = padding_token_id
 
     def __call__(self, batch: list[dict[str, torch.Tensor]]) -> DatasetBatch:
         """
@@ -33,4 +48,43 @@ def __call__(self, batch: list[dict[str, torch.Tensor]]) -> DatasetBatch:
         sample_tensor = torch.stack([torch.tensor(d[self.sample_key]) for d in batch])
         samples = {self.sample_key: sample_tensor[:, :-1]}
         targets = {self.target_key: sample_tensor[:, 1:]}
+        if self.sub_seq_lengths_key is not None:
+            # Determine sub sequence lengths by finding the eod tokens in each sequence in the batch.
+            sub_seq_lengths = self._compute_sub_sequence_lengths_for_each_sequence(sample_tensor)
+            samples[self.sub_seq_lengths_key] = sub_seq_lengths
         return DatasetBatch(targets=targets, samples=samples)
+
+    def _compute_sub_sequence_lengths_for_each_sequence(self, sample_tensor: torch.Tensor) -> list[list[int]]:
+        sub_seq_lengths = []
+        for seq in sample_tensor:
+            eod_positions = (seq == self.eod_token_id).nonzero(as_tuple=True)[0]
+            if len(eod_positions) == 0:
+                assert (
+                    self.padding_token_id is None or seq[0] != self.padding_token_id
+                ), "Sequence starts with padding token"
+                sub_seq_lengths.append([len(seq)])
+            else:
+                subseq_lengths = self._compute_subsequence_length(seq, eod_positions)
+                sub_seq_lengths.append(subseq_lengths)
+        return sub_seq_lengths
+
+    def _compute_subsequence_length(self, seq: torch.Tensor, eod_positions: torch.Tensor) -> list[int]:
+        # If the last sequence is cut, i.e. does not end on an eod token,
+        # it should also be included unless the padding token is set and
+        # the last sequence is just padding.
+        last_eod_pos = eod_positions[-1].item()
+        if self._has_cutoff_final_sequence(seq, last_eod_pos):
+            eod_positions = torch.cat([eod_positions, torch.tensor([len(seq) - 1])])
+        # Compute length of each subsequence and add to lengths list.
+        subseq_lengths = []
+        prev_pos = 0
+        for pos in eod_positions:
+            subseq_lengths.append(pos.item() - prev_pos + 1)
+            prev_pos = pos.item() + 1
+        return subseq_lengths
+
+    def _has_cutoff_final_sequence(self, seq: torch.Tensor, last_eod_pos: int) -> bool:
+        # Assumption: If the first token of the last sequence is padding, so is the rest.
+        return last_eod_pos < len(seq) - 1 and (
+            self.padding_token_id is None or seq[last_eod_pos + 1] != self.padding_token_id
+        )

tests/models/test_gpt2_collator.py

@@ -0,0 +1,98 @@
+import pytest
+import torch
+
+from modalities.models.gpt2.collator import GPT2LLMCollateFn
+
+
+def test_gpt2_collate_shifts_samples_and_targets():
+    collator = GPT2LLMCollateFn(sample_key="input_ids", target_key="labels")
+    batch = [
+        {"input_ids": torch.tensor([1, 2, 3, 4])},
+        {"input_ids": torch.tensor([5, 6, 7, 8])},
+    ]
+
+    result = collator(batch)
+
+    assert result.samples["input_ids"].tolist() == [[1, 2, 3], [5, 6, 7]]
+    assert result.targets["labels"].tolist() == [[2, 3, 4], [6, 7, 8]]
+
+
+def test_gpt2_collate_sub_seq_lengths_without_eod():
+    collator = GPT2LLMCollateFn(
+        sample_key="input_ids",
+        target_key="labels",
+        sub_seq_lengths_key="sub_seq_lengths",
+        eod_token_id=99,
+    )
+    batch = [
+        {"input_ids": torch.tensor([10, 11, 12, 13, 14])},
+        {"input_ids": torch.tensor([20, 21, 22, 23, 24])},
+    ]
+
+    result = collator(batch)
+
+    assert result.samples["sub_seq_lengths"] == [[5], [5]]
+
+
+def test_gpt2_collate_sub_seq_lengths_with_eod():
+    collator = GPT2LLMCollateFn(
+        sample_key="input_ids",
+        target_key="labels",
+        sub_seq_lengths_key="sub_seq_lengths",
+        eod_token_id=99,
+    )
+    batch = [
+        {"input_ids": torch.tensor([1, 99, 2, 3, 99])},
+        {"input_ids": torch.tensor([7, 8, 9, 99, 10])},
+    ]
+
+    result = collator(batch)
+
+    assert result.samples["sub_seq_lengths"] == [[2, 3], [4, 1]]
+
+
+def test_gpt2_collate_sub_seq_lengths_with_eod_and_padding():
+    collator = GPT2LLMCollateFn(
+        sample_key="input_ids",
+        target_key="labels",
+        sub_seq_lengths_key="sub_seq_lengths",
+        eod_token_id=99,
+        padding_token_id=0,
+    )
+    batch = [
+        {"input_ids": torch.tensor([1, 99, 2, 3, 4, 5])},
+        {"input_ids": torch.tensor([7, 8, 99, 0, 0, 0])},
+    ]
+
+    result = collator(batch)
+
+    assert result.samples["sub_seq_lengths"] == [[2, 4], [3]]
+
+
+def test_gpt2_collate_sub_seq_lengths_adds_tail_when_not_padding():
+    collator = GPT2LLMCollateFn(
+        sample_key="input_ids",
+        target_key="labels",
+        sub_seq_lengths_key="sub_seq_lengths",
+        eod_token_id=5,
+        padding_token_id=0,
+    )
+    batch = [{"input_ids": torch.tensor([1, 5, 9, 8])}]
+
+    result = collator(batch)
+
+    assert result.samples["sub_seq_lengths"] == [[2, 2]]
+
+
+def test_gpt2_collate_raises_when_sequence_starts_with_padding_and_no_eod():
+    collator = GPT2LLMCollateFn(
+        sample_key="input_ids",
+        target_key="labels",
+        sub_seq_lengths_key="sub_seq_lengths",
+        eod_token_id=99,
+        padding_token_id=0,
+    )
+    batch = [{"input_ids": torch.tensor([0, 1, 2, 3])}]
+
+    with pytest.raises(AssertionError, match="Sequence starts with padding token"):
+        collator(batch)