Allow loss masking for defined spans of characters

fast_llm/data/dataset/gpt/memmap.py

@@ -48,6 +48,25 @@ def _init(self, name: str, prefix: pathlib.Path | str):
             offset=offset + self._document_sizes.nbytes,
         )
 
+        self._num_spans = np.frombuffer(
+            self._index_bin_buffer,
+            dtype=np.int32,
+            count=self._num_documents,
+            offset=offset + self._document_sizes.nbytes + self._pointers.nbytes,
+        )
+        spans = []
+        offset = offset + self._document_sizes.nbytes + self._pointers.nbytes + self._num_spans.nbytes
+        for n_spans in self._num_spans:
+            span = np.frombuffer(
+                self._index_bin_buffer,
+                dtype=np.int32,
+                count=n_spans * 2,
+                offset=offset,
+            ).reshape(-1, 2)
+            spans.append(span)
+            offset += span.nbytes
+        self._spans = spans
+
         self._bin_buffer_mmap = np.memmap(self._prefix.with_suffix(".bin"), mode="r", order="C")
         self._bin_buffer = memoryview(self._bin_buffer_mmap)
 
@@ -64,11 +83,14 @@ def __del__(self):
         del self._index_bin_buffer_mmap
 
     def get(self, idx, offset=0, length=None):
-        return np.frombuffer(
-            self._bin_buffer,
-            dtype=self._dtype,
-            count=self._document_sizes[idx] - offset if length is None else length,
-            offset=self._pointers[idx] + offset * np.dtype(self._dtype).itemsize,
+        return (
+            np.frombuffer(
+                self._bin_buffer,
+                dtype=self._dtype,
+                count=self._document_sizes[idx] - offset if length is None else length,
+                offset=self._pointers[idx] + offset * np.dtype(self._dtype).itemsize,
+            ),
+            self._spans[idx],
         )
 
     @property
@@ -92,20 +114,23 @@ def get_document_sizes(self) -> "np.ndarray":
         return self._document_sizes
 
     @classmethod
-    def write_dataset(cls, prefix: pathlib.Path | str, documents: typing.Iterable[np.ndarray]):
+    def write_dataset(cls, prefix: pathlib.Path | str, documents: typing.Iterable[tuple[np.ndarray, np.ndarray]]):
         # Initialize metadata
         dtype = None
         num_documents = 0
         lengths = []
         pointers = []
         offset = 0
+        # number of spans for each document
+        num_spans = []
+        spans = []
 
         prefix = pathlib.Path(prefix)
         prefix.parent.mkdir(parents=True, exist_ok=True)
 
         # Write the binary data file (.bin) lazily
         with prefix.with_suffix(".bin").open("wb") as bin_stream:
-            for document in documents:
+            for document, mask_spans in documents:
                 # Infer dtype from the first document
                 if dtype is None:
                     dtype = document.dtype
@@ -121,12 +146,16 @@ def write_dataset(cls, prefix: pathlib.Path | str, documents: typing.Iterable[np
                 doc_length = len(document)
                 lengths.append(doc_length)
                 pointers.append(offset)
+                num_spans.append(len(mask_spans))
+                spans.append(mask_spans)
                 offset += doc_length * np.dtype(dtype).itemsize
                 num_documents += 1
 
         # Finalize metadata arrays
         lengths = np.array(lengths, dtype=np.int32)
         pointers = np.array(pointers, dtype=np.int64)
+        num_spans = np.array(num_spans, dtype=np.int32)
+        spans = np.vstack(spans, dtype=np.int32)
 
         # Write the index file (.idx)
         with prefix.with_suffix(".idx").open("wb") as idx_stream:
@@ -142,5 +171,9 @@ def write_dataset(cls, prefix: pathlib.Path | str, documents: typing.Iterable[np
             idx_stream.write(lengths.tobytes(order="C"))
             # Sequence (document) begin offsets in the bin file
             idx_stream.write(pointers.tobytes(order="C"))
+            # Number of spans per document
+            idx_stream.write(num_spans.tobytes(order="C"))
+            # Span indices for each document
+            idx_stream.write(spans.tobytes(order="C"))
             # Document indices, unused but needed for compatibility with Megatron-LM
             idx_stream.write(np.arange(num_documents + 1, dtype=np.int64).tobytes(order="C"))

fast_llm/data/preparator/gpt_memmap/config.py

@@ -44,6 +44,9 @@ class GPTHuggingfaceDatasetConfig(Config):
         desc="Field of the dataset to use.",
         hint=FieldHint.optional,
     )
+    spans_field: None | str = Field(
+        default=None, desc="Field containing character spans to mask for loss computation", hint=FieldHint.optional
+    )
     data_type: DataType | None = Field(
         default=None,
         desc="Data type of the dataset field."

fast_llm/data/preparator/gpt_memmap/prepare.py

@@ -22,14 +22,37 @@ class GPTMemmapDatasetPreparator(DatasetPreparator):
     _tokenizer: Tokenizer
     _data_type: DataType
 
+    def _tokenize_with_spans(self, sample):
+        """
+        Perform span-aware tokenization and return the tokenized input_ids along with token spans.
+        """
+        char_spans = sample.get(self._config.dataset.spans_field, [])
+        text = sample[self._config.dataset.field]
+        input_ids = []
+        token_spans = []
+        char_pos = 0
+        for start, end in char_spans:
+            if char_pos < start:
+                curr_text = text[char_pos:start]
+                tokenized_text = self._tokenizer.tokenize(curr_text)
+                input_ids.extend(tokenized_text)
+            curr_text = text[start : end + 1]
+            tokenized_text = self._tokenizer.tokenize(curr_text)
+            input_ids.extend(tokenized_text)
+            token_spans.append((len(token_spans), len(token_spans) + len(tokenized_text) - 1))
+            char_pos = end + 1
+        if char_pos < len(text):
+            curr_text = text[char_pos:]
+            tokenized_text = self._tokenizer.tokenize(curr_text)
+            input_ids.extend(tokenized_text)
+        return np.array(input_ids, dtype=self._data_type.numpy), np.array(token_spans, dtype=np.int32)
+
     def _tokenize_batch(self, batch):
-        input_ids = [
-            np.array(self._tokenizer.tokenize(text), dtype=self._data_type.numpy)
-            for text in batch[self._config.dataset.field]
-        ]
+        input_ids, token_spans = zip(*[self._tokenize_with_spans(sample) for sample in batch])
         num_tokens = [len(x) for x in input_ids]
         return {
             "input_ids": input_ids,
+            "token_spans": token_spans,
             "num_tokens": num_tokens,
         }
 
@@ -40,7 +63,9 @@ def _save_shard(self, args) -> dict:
 
         def _document_generator():
             for item in tqdm.tqdm(shard_dataset, desc=f"Saving shard {shard_idx}", unit="docs"):
-                yield np.array(item["input_ids"], dtype=self._data_type.numpy)
+                yield np.array(item["input_ids"], dtype=self._data_type.numpy), np.array(
+                    item["token_spans"], dtype=np.int32
+                )
 
         GPTMemmapDataset.write_dataset(prefix=shard_output_path, documents=_document_generator())
 
@@ -126,6 +151,11 @@ def run(self):
         )
         if self._config.dataset.field not in dataset.column_names:
             raise ValueError(f"Dataset does not have field '{self._config.dataset.field}'.")
+        if (
+            self._config.dataset.spans_field is not None
+            and self._config.dataset.spans_field not in dataset.column_names
+        ):
+            raise ValueError(f"Dataset does not have spans field '{self._config.dataset.spans_field}'.")
 
         # Tokenize the dataset in parallel
         tokenized_dataset = dataset.map(

tests/test_memmap_dataset.py

@@ -10,12 +10,26 @@
 
 @pytest.mark.parametrize("dtype", MEMMAP_DTYPES.values())
 def test_gpt_memmap_dataset(dtype):
-    documents = [np.random.randint(1000, size=np.random.randint(1, 100)).astype(dtype) for _ in range(100)]
+    documents = list(
+        zip(
+            [np.random.randint(1000, size=np.random.randint(1, 100)).astype(dtype) for _ in range(100)],
+            np.array([[]] * 100, dtype=np.int32),
+        )
+    )
     with tempfile.TemporaryDirectory() as temp_dir:
         prefix = pathlib.Path(temp_dir)
         GPTMemmapDataset.write_dataset(prefix=prefix, documents=documents)
         dataset = GPTMemmapDataset(name="foo", prefix=prefix)
-        for i, document in enumerate(documents):
+        for i, (document, spans) in enumerate(documents):
+            memmap_document, memmap_spans = dataset.get(i)
             assert np.array_equal(
-                dataset.get(i), document, equal_nan=True
+                memmap_document, document, equal_nan=True
             ), f"Mismatch for document {i}: {document} != {dataset.get(i)}."
+            if len(spans) > 0:
+                assert np.array_equal(
+                    memmap_spans, spans, equal_nan=True
+                ), f"Mismatch for non-empty spans {i}: {spans} != {dataset.get(i)}."
+            else:
+                assert np.array_equal(
+                    memmap_spans.flatten(), spans.flatten(), equal_nan=True
+                ), f"Mismatch for empty spans {i}: {spans} != {dataset.get(i)}."