combined cols to create a new col

Author: nitsanluke

fast_llm/data/preparator/gpt_memmap/config.py

@@ -1,6 +1,7 @@
 import os
 import pathlib
 import typing
+import dataclasses
 
 from fast_llm.config import Config, Field, FieldHint, check_field, config_class
 from fast_llm.data.config import TokenizerConfig
@@ -109,6 +110,31 @@ def _validate(self) -> None:
         super()._validate()
         Assert.in_range(self.rank, 0, self.world_size)
 
+@config_class
+class FieldCombinePreparatorConfig(Config):
+    fields: typing.List[str] = Field(
+        default_factory=list,
+        desc="Fields of the dataset to combine.",
+        hint=FieldHint.core,
+    )
+    delimiter: str = Field(
+        default=" ",
+        desc="Delimiter to use when combining fields.",
+        hint=FieldHint.optional,
+    )
+    new_field_name: str = Field(
+        default="fast_llm_combined_field",
+        desc="Name of the new field to create.",
+        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()
 class GPTMemmapDatasetPreparatorConfig(DatasetPreparatorConfig):
@@ -164,6 +190,11 @@ class GPTMemmapDatasetPreparatorConfig(DatasetPreparatorConfig):
         " Does not shuffle samples.",
         hint=FieldHint.optional,
     )
+    combine_fields: FieldCombinePreparatorConfig = Field(
+        default=None,
+        desc="Combine all files into a single file.",
+        hint=FieldHint.optional,
+    )
 
     def _validate(self) -> None:
         assert self.tokenizer.path is not None

fast_llm/data/preparator/gpt_memmap/prepare.py

@@ -208,6 +208,22 @@ def run(self) -> None:
                 torch.distributed.barrier()
 
         assert isinstance(dataset, datasets.Dataset)
+        
+        # Check for combining fields
+        if self._config.combine_fields: 
+            Assert.eq(len(set(self._config.combine_fields.fields).intersection(dataset.column_names)), len(self._config.combine_fields.fields))
+            dataset = dataset.map(
+                lambda example: {
+                    self._config.combine_fields.new_field_name: self._config.combine_fields.delimiter.join(
+                        str(example[column]) for column in self._config.combine_fields.fields
+                    )
+                },
+                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
+        
         dataset = dataset.shard(
             num_shards=self._config.distributed.world_size,
             index=self._config.distributed.rank,