Clean up datahandler processing.

Rename data handlers to match HF API names.
Remove rename/retain features and make them data handlers under the new
framework.

Signed-off-by: Dushyant Behl <dushyantbehl@in.ibm.com>

Author: dushyantbehl

.gitignore

@@ -8,6 +8,7 @@ coverage*.xml
 dist
 htmlcov
 test
+error.log
 
 # IDEs
 .vscode/

tests/artifacts/predefined_data_configs/__init__.py

@@ -64,12 +64,12 @@
 DATA_CONFIG_DUPLICATE_COLUMNS = os.path.join(
     PREDEFINED_DATA_CONFIGS, "duplicate_columns.yaml"
 )
-DATA_CONFIG_RENAME_RETAIN_COLUMNS = os.path.join(
-    PREDEFINED_DATA_CONFIGS, "rename_retain_columns.yaml"
+DATA_CONFIG_RENAME_SELECT_COLUMNS = os.path.join(
+    PREDEFINED_DATA_CONFIGS, "rename_select_columns.yaml"
 )
 DATA_CONFIG_TOKENIZE_AND_TRAIN_WITH_HANDLER = os.path.join(
     PREDEFINED_DATA_CONFIGS, "tokenize_using_handler_and_train.yaml"
 )
-DATA_CONFIG_SKIP_LARGE_TEXT_HANDLER = os.path.join(
-    PREDEFINED_DATA_CONFIGS, "skip_large_text_data_handler_template.yaml"
+DATA_CONFIG_SKIP_LARGE_COLUMNS_HANDLER = os.path.join(
+    PREDEFINED_DATA_CONFIGS, "skip_large_columns_data_handler_template.yaml"
 )

…_data_configs/rename_retain_columns.yaml …_data_configs/rename_select_columns.yamltests/artifacts/predefined_data_configs/rename_retain_columns.yaml renamed to tests/artifacts/predefined_data_configs/rename_select_columns.yaml tests/artifacts/predefined_data_configs/rename_retain_columns.yaml renamed to tests/artifacts/predefined_data_configs/rename_select_columns.yaml

@@ -2,15 +2,19 @@ dataprocessor:
     type: default
 datasets:
   - name: text_dataset_input_output_masking
-    rename_columns:
-      "input"  : "instruction"
-      "output" : "response"
-    retain_columns:
-      - "instruction"
-      - "response"
     data_paths:
       - "FILE_PATH"
     data_handlers:
+      - name: rename_columns
+        arguments:
+          column_mapping:
+            "input"  : "instruction"
+            "output" : "response"
+      - name: select_columns
+        arguments:
+          column_names:
+            - "instruction"
+            - "response"
       - name: tokenize_and_apply_input_masking
         arguments:
           remove_columns: all

…ip_large_text_data_handler_template.yaml …large_columns_data_handler_template.yamltests/artifacts/predefined_data_configs/skip_large_text_data_handler_template.yaml renamed to tests/artifacts/predefined_data_configs/skip_large_columns_data_handler_template.yaml tests/artifacts/predefined_data_configs/skip_large_text_data_handler_template.yaml renamed to tests/artifacts/predefined_data_configs/skip_large_columns_data_handler_template.yaml

@@ -1,7 +1,7 @@
 dataprocessor:
     type: default
 datasets:
-  - name: pre_tokenized
+  - name: non_tokenized
     data_paths:
       - "FILE_PATH"
     data_handlers:
@@ -17,7 +17,7 @@ datasets:
           fn_kwargs:
             old_column: "input_ids"
             new_column: "labels"
-      - name: skip_large_text
+      - name: skip_large_columns
         arguments:
           fn_kwargs:
             column_name: "input_ids"

tests/data/test_data_handlers.py

@@ -35,10 +35,9 @@
     apply_custom_jinja_template,
     combine_sequence,
     duplicate_columns,
-    skip_large_text,
+    skip_large_columns,
     tokenize,
 )
-from tuning.data.setup_dataprocessor import is_pretokenized_dataset
 
 
 def test_apply_custom_formatting_template():
@@ -287,19 +286,19 @@ def test_tokenizer_data_handler_tokenizes():
         ("not_existing", "not_existing"),
     ],
 )
-def test_skip_large_text_handler_throws_error_on_bad_args(column_name, max_length):
-    "Ensure that skip large text handler throws error on bad arguments"
+def test_skip_large_columns_handler_throws_error_on_bad_args(column_name, max_length):
+    "Ensure that skip large columns handler throws error on bad arguments"
     d = datasets.load_dataset("json", data_files=TWITTER_COMPLAINTS_DATA_JSONL)
     fn_kwargs = {}
     fn_kwargs["column_name"] = column_name
     fn_kwargs["max_length"] = max_length
 
     with pytest.raises(ValueError):
-        filtered = d.filter(skip_large_text, fn_kwargs=fn_kwargs)
+        filtered = d.filter(skip_large_columns, fn_kwargs=fn_kwargs)
 
 
-def test_skip_large_text_handler():
-    "Ensure that skip large text handler skips dataset as intended"
+def test_skip_large_columns_handler():
+    "Ensure that skip large columns handler skips dataset as intended"
 
     def test_dataset_generator():
         for i in range(0, 100):
@@ -308,7 +307,7 @@ def test_dataset_generator():
     d = Dataset.from_generator(test_dataset_generator)
     fn_kwargs = {}
     fn_kwargs["column_name"] = "input"
-    fn_kwargs["max_length"] = 61
+    fn_kwargs["max_length"] = 60
 
-    filtered = d.filter(skip_large_text, fn_kwargs=fn_kwargs)
+    filtered = d.filter(skip_large_columns, fn_kwargs=fn_kwargs)
     assert len(filtered) == 60

tests/data/test_data_preprocessing.py

@@ -37,7 +37,7 @@
     DATA_CONFIG_MULTIPLE_DATASETS_SAMPLING_YAML,
     DATA_CONFIG_MULTITURN_DATA_YAML,
     DATA_CONFIG_PRETOKENIZE_JSON_DATA_YAML,
-    DATA_CONFIG_RENAME_RETAIN_COLUMNS,
+    DATA_CONFIG_RENAME_SELECT_COLUMNS,
     DATA_CONFIG_TOKENIZE_AND_APPLY_INPUT_MASKING_YAML,
     DATA_CONFIG_YAML_STREAMING_INPUT_OUTPUT,
     DATA_CONFIG_YAML_STREAMING_PRETOKENIZED,
@@ -70,7 +70,11 @@
 from tuning.config import configs
 from tuning.config.acceleration_configs import AttentionAndDistributedPackingConfig
 from tuning.data.collators import VisionDataCollator
-from tuning.data.data_config import DataPreProcessorConfig, DataSetConfig
+from tuning.data.data_config import (
+    DataHandlerConfig,
+    DataPreProcessorConfig,
+    DataSetConfig,
+)
 from tuning.data.data_preprocessing_utils import get_data_collator
 from tuning.data.data_processors import DataPreProcessor, get_datapreprocessor
 from tuning.data.setup_dataprocessor import (
@@ -1674,33 +1678,33 @@ def test_process_dataset_configs_with_sampling_error(
 
 
 @pytest.mark.parametrize(
-    "datafile, rename, retain, final, datasetconfigname",
+    "datafile, rename, select, final, datasetconfigname",
     [
         (
             TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_JSON,
             {"input": "instruction", "output": "response"},
             None,
             ["ID", "Label", "instruction", "response"],
-            DATA_CONFIG_RENAME_RETAIN_COLUMNS,
+            DATA_CONFIG_RENAME_SELECT_COLUMNS,
         ),
         (
             TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_JSON,
             None,
             ["ID", "input", "output"],
             ["ID", "input", "output"],
-            DATA_CONFIG_RENAME_RETAIN_COLUMNS,
+            DATA_CONFIG_RENAME_SELECT_COLUMNS,
         ),
         (
             TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_JSON,
             {"input": "instruction", "output": "response"},
             ["Label", "instruction", "response"],
             ["Label", "instruction", "response"],
-            DATA_CONFIG_RENAME_RETAIN_COLUMNS,
+            DATA_CONFIG_RENAME_SELECT_COLUMNS,
         ),
     ],
 )
-def test_rename_and_retain_dataset_columns(
-    datafile, rename, retain, final, datasetconfigname
+def test_rename_and_select_dataset_columns(
+    datafile, rename, select, final, datasetconfigname
 ):
     """Test process_dataset_configs for expected output."""
     dataprocessor_config = DataPreProcessorConfig()
@@ -1709,12 +1713,23 @@ def test_rename_and_retain_dataset_columns(
         processor_config=dataprocessor_config,
         tokenizer=tokenizer,
     )
+
+    handlers = []
+    if rename:
+        handlers.append(
+            DataHandlerConfig(
+                name="rename_columns", arguments={"column_mapping": rename}
+            )
+        )
+    if select:
+        handlers.append(
+            DataHandlerConfig(name="select_columns", arguments={"column_names": select})
+        )
+    data_paths = [datafile]
+
     datasetconfig = [
         DataSetConfig(
-            name=datasetconfigname,
-            data_paths=[datafile],
-            rename_columns=rename,
-            retain_columns=retain,
+            name=datasetconfigname, data_paths=data_paths, data_handlers=handlers
         )
     ]
     train_dataset = processor.process_dataset_configs(dataset_configs=datasetconfig)

tests/test_sft_trainer.py

@@ -43,8 +43,8 @@
     DATA_CONFIG_MULTITURN_CHAT_TOKENIZE_AND_MASKING_DATA_HANDLER,
     DATA_CONFIG_MULTITURN_DATA_YAML,
     DATA_CONFIG_MULTITURN_GRANITE_3_1B_DATA_YAML,
-    DATA_CONFIG_RENAME_RETAIN_COLUMNS,
-    DATA_CONFIG_SKIP_LARGE_TEXT_HANDLER,
+    DATA_CONFIG_RENAME_SELECT_COLUMNS,
+    DATA_CONFIG_SKIP_LARGE_COLUMNS_HANDLER,
     DATA_CONFIG_TOKENIZE_AND_APPLY_INPUT_MASKING_YAML,
     DATA_CONFIG_TOKENIZE_AND_TRAIN_WITH_HANDLER,
     DATA_CONFIG_VALID_BASE64_CHAT_TEMPLATE,
@@ -925,7 +925,7 @@ def test_run_causallm_ft_and_inference_streaming(datasetconfigname, datafiles):
         ),
         (
             [TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT_JSON],
-            DATA_CONFIG_RENAME_RETAIN_COLUMNS,
+            DATA_CONFIG_RENAME_SELECT_COLUMNS,
         ),
     ],
 )
@@ -1064,8 +1064,8 @@ def test_run_training_with_data_tokenized_using_tokenizer_handler():
         assert "### Text: @NortonSupport Thanks much.\n\n### Label:" in output_inference
 
 
-def test_run_training_with_skip_large_text_handler():
-    """Ensure that we can train succesfully after using skip large text handler."""
+def test_run_training_with_skip_large_column_handler():
+    """Ensure that we can train succesfully after using skip large column handler."""
     with tempfile.TemporaryDirectory() as tempdir:
 
         data_args = copy.deepcopy(DATA_ARGS)
@@ -1074,8 +1074,8 @@ def test_run_training_with_skip_large_text_handler():
         data_args.response_template = None
         data_args.training_data_path = None
 
-        dataconfigfile = DATA_CONFIG_SKIP_LARGE_TEXT_HANDLER
-        datapath = TWITTER_COMPLAINTS_TOKENIZED_JSON
+        dataconfigfile = DATA_CONFIG_SKIP_LARGE_COLUMNS_HANDLER
+        datapath = TWITTER_COMPLAINTS_DATA_JSONL
 
         # add data_paths in data_config file
         with tempfile.NamedTemporaryFile(

tuning/config/configs.py

@@ -72,7 +72,9 @@ class DataArguments:
     dataset_text_field: str = field(
         default=None,
         metadata={
-            "help": "Training dataset text field containing single sequence. \
+            "help": "[DEPRECATED] "
+                    "Use text_column_name to specify this argument going forward\n"\
+                    "Training dataset text field containing single sequence. \
                     Either the dataset_text_field \
                     or data_formatter_template need to be supplied. \
                     For running vision language model tuning pass the column name for text data."
@@ -85,6 +87,14 @@ class DataArguments:
                     Used as key to point multi-turn data field."
         },
     )
+    text_column_name : str = field(
+        default=None,
+        metadata={
+            "help": "Training dataset text column name containing single sequence. \
+                    Either the text_column_name \
+                    or data_formatter_template need to be supplied."
+        },
+    )
     validation_data_path: str = field(
         default=None,
         metadata={"help": "Path to the validation data in JSON/JSONL format."},
@@ -157,6 +167,9 @@ def unescape(s):
         self.response_template = unescape(self.response_template)
         self.instruction_template = unescape(self.instruction_template)
 
+        # Initialise deprecated field
+        if self.dataset_text_field:
+            self.text_column_name = self.dataset_text_field
 
 @dataclass
 class TrainingArguments(transformers.TrainingArguments):

tuning/data/data_config.py

@@ -37,8 +37,6 @@ class DataSetConfig:
     data_paths: List[str]
     builder: Optional[str] = None  # Referring to Hugging Face dataset builder
     sampling: Optional[float] = None
-    rename_columns: Optional[Dict] = None
-    retain_columns: Optional[List] = None
     data_handlers: Optional[List[DataHandlerConfig]] = None