add loss masking col when combining cols

Author: nitsanluke

fast_llm/data/preparator/gpt_memmap/config.py

@@ -77,7 +77,6 @@ class GPTHuggingfaceDatasetConfig(Config):
         hint=FieldHint.optional,
     )
 
-
 @config_class
 class DatasetPreparatorDistributedConfig(Config):
     # TODO: Unify with fast_llm.engine.distributed.config.DistributedConfig
@@ -110,6 +109,22 @@ def _validate(self) -> None:
         super()._validate()
         Assert.in_range(self.rank, 0, self.world_size)
 
+@config_class
+class LossMaskSpansConfig(Config):
+    masking_column: str = Field(
+        default=None, 
+        desc="Field containing character spans to mask for loss computation",
+        hint=FieldHint.optional,
+    )
+    loss_masking_spans: str = Field(
+        default="fast_llm_loss_masking_spans",
+        desc="Field containing character spans to mask for loss computation",
+        hint=FieldHint.optional,
+    )   
+    def _validate(self) -> None:
+        assert isinstance(self.loss_masking_spans, str), "loss_masking_spans col name must be a string."
+        super()._validate()
+
 @config_class
 class FieldCombinePreparatorConfig(Config):
     col_names: typing.List[str] = Field(
@@ -127,13 +142,13 @@ class FieldCombinePreparatorConfig(Config):
         desc="Name of the new field to create.",
         hint=FieldHint.optional,
     )
-    
+    set_masking_span: LossMaskSpansConfig = Field(
+        default=None,
+        desc="Compute loss_masking_spans for the newly combined field.",
+        hint=FieldHint.optional,
+    )
     def _validate(self) -> None:
-        # Assert.gt(len(self.fields), 0)
-        # assert isinstance(self.fields, list), "Fields must be a list."
-        # assert all(isinstance(field, str) for field in self.fields), "All fields must be strings."
         assert isinstance(self.delimiter, str), "Delimiter must be a string."
-        # assert isinstance(self.new_field_name, str), "New field name must be a string."
         super()._validate()
 
 @config_class()
@@ -201,11 +216,13 @@ def _validate(self) -> 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
+            logger.info(
+            f"Setting dataset.field to {self.combine_fields.new_field_name}",
+            )
+            self.dataset.field = self.combine_fields.new_field_name
+            if self.combine_fields.set_masking_span is not None:
+                logger.info(f"Setting dataset.loss_masking_spans to {self.combine_fields.set_masking_span.loss_masking_spans}")
+                self.dataset.loss_masking_spans = self.combine_fields.set_masking_span.loss_masking_spans
         super()._validate()
 
     @classmethod

fast_llm/data/preparator/gpt_memmap/prepare.py

@@ -211,7 +211,9 @@ 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))
+            Assert.eq(len(set(self._config.combine_fields.col_names).intersection(dataset.column_names)), len(self._config.combine_fields.col_names),\
+                msg=f"Some columns to combine are not in the dataset. {set(self._config.combine_fields.col_names).difference(dataset.column_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: {
@@ -224,6 +226,20 @@ def run(self) -> None:
             )
             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
+            
+            if self._config.combine_fields.set_masking_span is not None:
+                Assert.incl(self._config.combine_fields.set_masking_span.masking_column, dataset.column_names)
+ 
+                dataset = dataset.map(
+                    lambda example: {
+                        self._config.dataset.loss_masking_spans: [
+                            (0, len(str(example[self._config.combine_fields.set_masking_span.masking_column])) - 1)
+                        ]# spans are inclusive
+                    },
+                    batched=False,
+                    desc="Setting loss masking spans",
+                )
+                logger.info(f"Sample after setting loss masking spans:\n{dataset[0]}")
 
         dataset = dataset.shard(
             num_shards=self._config.distributed.world_size,