merge update

Author: nitsanluke

fast_llm/data/preparator/gpt_memmap/config.py

@@ -30,7 +30,7 @@ class SourceSchemaConfig(Config):
 
 @config_class()
 class PromptCompletionConfig(SourceSchemaConfig):
-    type: typing.ClassVar[str] = "prompt_completion"
+    type: typing.ClassVar[str] = "prompt_completion" #TODO: Register PromptCompletionConfig for this type for dynamic loading PR #245
     prompt_column: str = Field(
         default="prompt",
         desc="Field of the dataset to use.",
@@ -49,7 +49,7 @@ class PromptCompletionConfig(SourceSchemaConfig):
 
 @config_class()
 class TextColumnConfig(SourceSchemaConfig):
-    type: typing.ClassVar[str] = "text_column"
+    type: typing.ClassVar[str] = "text_column" #TODO: Register TestColumnConfig for this type for dynamic loading PR #245
     input_column: str = Field(
         default="text",
         desc="Field of the dataset to use.",
@@ -88,7 +88,7 @@ class GPTHuggingfaceDatasetConfig(Config):
         hint=FieldHint.optional,
     )
     source_schema: SourceSchemaConfig = Field(
-        default_factory=TextColumnConfig,
+        #TODO: Default should be from subclass TextColumnConfig (waiting for PR #245)
         desc="Configuration for the data source.",
         hint=FieldHint.optional,
     )

fast_llm/data/preparator/gpt_memmap/prepare.py

@@ -37,13 +37,13 @@ class GPTMemmapDatasetPreparator[ConfigType: GPTMemmapDatasetPreparatorConfig](D
 
     _tokenizer: Tokenizer
     _data_type: DataType
-    _data_column: str
+    _text_column: str
     _loss_masking_spans_column: str | None
 
     def _tokenize_batch(self, batch: dict[str, list[typing.Any]]) -> dict[str, list[typing.Any]]:
         input_ids = [
             np.array(self._tokenizer.tokenize(text), dtype=self._data_type.numpy)
-            for text in batch[self._data_column]
+            for text in batch[self._text_column]
         ]
         num_tokens = [len(x) for x in input_ids]
         return {
@@ -63,7 +63,7 @@ def _tokenize_batch_with_spans(self, batch: dict[str, list[typing.Any]]) -> dict
                     for input_ids, token_spans in [
                         self._tokenizer.tokenize_with_spans(text, char_spans)
                         for text, char_spans in zip(
-                            batch[self._data_column], batch[self._loss_masking_spans_column]
+                            batch[self._text_column], batch[self._loss_masking_spans_column]
                         )
                     ]
                 ]
@@ -254,8 +254,8 @@ def run(self) -> None:
             num_shards=self._config.distributed.world_size,
             index=self._config.distributed.rank,
         )
-        if self._data_column not in dataset.column_names:
-            raise ValueError(f"Dataset does not have field '{self._data_column}'.")
+        if self._text_column not in dataset.column_names:
+            raise ValueError(f"Dataset does not have field '{self._text_column}'.")
         if self._loss_masking_spans_column is not None:
             if self._loss_masking_spans_column not in dataset.column_names:
                 raise ValueError(f"Dataset does not have spans field '{self._loss_masking_spans_column}'.")