set dataset.field to the new field

Author: nitsanluke

fast_llm/data/preparator/gpt_memmap/config.py

@@ -1,7 +1,7 @@
 import os
 import pathlib
 import typing
-import dataclasses
+import logging
 
 from fast_llm.config import Config, Field, FieldHint, check_field, config_class
 from fast_llm.data.config import TokenizerConfig
@@ -24,7 +24,7 @@
 MEMMAP_DTYPES_INV = {y: x for x, y in MEMMAP_DTYPES.items()}
 MEMMAP_INDEX_HEADER = b"MMIDIDX\x00\x00"
 
-
+logger = logging.getLogger(__name__)
 @config_class
 class GPTHuggingfaceDatasetConfig(Config):
     path: str = Field(
@@ -200,6 +200,12 @@ def _validate(self) -> None:
         assert self.tokenizer.path is not None
         if self.dataset.data_type is not None:
             Assert.incl(DataType.from_numpy(self.dataset.data_type.numpy), MEMMAP_DTYPES_INV)
+        if self.combine_fields is not None:
+            if self.dataset.field != self.combine_fields.new_field_name:
+                 logger.warning(
+                    f"Combine mode activated yet dataset.field != combine_fields.new_field_name ({self.dataset.field} != {self.combine_fields.new_field_name}). Setting dataset.field to {self.combine_fields.new_field_name}",
+                 )
+                 self.dataset.field = self.combine_fields.new_field_name
         super()._validate()
 
     @classmethod

fast_llm/data/preparator/gpt_memmap/prepare.py

@@ -212,6 +212,7 @@ def run(self) -> None:
         # Check for combining fields
         if self._config.combine_fields: 
             Assert.eq(len(set(self._config.combine_fields.col_names).intersection(dataset.column_names)), len(self._config.combine_fields.col_names))
+            logger.info(f"Combining fields {self._config.combine_fields.col_names} into {self._config.combine_fields.new_field_name}")
             dataset = dataset.map(
                 lambda example: {
                     self._config.combine_fields.new_field_name: self._config.combine_fields.delimiter.join(
@@ -221,9 +222,9 @@ def run(self) -> None:
                 batched=False,
                 desc="Combining fields",
             )
-            # Set the new field name in the config for following operations
-            self._config.dataset.field = self._config.combine_fields.new_field_name
-        
+            logger.info(f"Sample after combining fields:\n{dataset[0]}")
+            # Note: self.dataset.field is set to new_field_name for the rest of the operation see config validation
+
         dataset = dataset.shard(
             num_shards=self._config.distributed.world_size,
             index=self._config.distributed.rank,