Update tuning/data/data_processors.py

Co-authored-by: Will Johnson <[email protected]>
Signed-off-by: Dushyant Behl <[email protected]>

Author: dushyantbehl

tuning/data/data_handlers.py

@@ -15,7 +15,7 @@
 # Definition of some predefined data preprocessing functions that we need.
 
 # Standard
-from typing import Dict, List
+from typing import Dict
 
 # Third Party
 from transformers import AutoTokenizer
@@ -56,12 +56,12 @@ def tokenize_and_apply_input_masking(
 def apply_dataset_formatting(
     element: Dict[str, str], tokenizer: AutoTokenizer, dataset_text_field: str, **kwargs
 ):
-    if isinstance(element[dataset_text_field], list): # batched = True
-      return {
-          f"{dataset_text_field}": [  
-              text + tokenizer.eos_token for text in element[f"{dataset_text_field}"]
-          ]
-      }
+    if isinstance(element[dataset_text_field], list):  # batched = True
+        return {
+            f"{dataset_text_field}": [
+                text + tokenizer.eos_token for text in element[f"{dataset_text_field}"]
+            ]
+        }
     return {
         f"{dataset_text_field}": element[f"{dataset_text_field}"] + tokenizer.eos_token
     }
@@ -77,13 +77,11 @@ def apply_custom_data_formatting_template(
     template += tokenizer.eos_token
 
     # TODO: Eventually move the code here.
-    return custom_data_formatter(
-        element=element, formatted_dataset_field=dataset_text_field, template=template
-    )
+    return custom_data_formatter(element, template, dataset_text_field)
 
 
 AVAILABLE_DATA_HANDLERS = {
     "tokenize_and_apply_instruction_masking": tokenize_and_apply_input_masking,
     "apply_dataset_formatting": apply_dataset_formatting,
-    "apply_custom_data_formatting_template": apply_dataset_formatting,
+    "apply_custom_data_formatting_template": apply_custom_data_formatting_template,
 }

tuning/data/data_processors.py

@@ -19,7 +19,7 @@
 import os
 
 # Third Party
-from datasets import Dataset, IterableDataset, DatasetDict
+from datasets import Dataset, DatasetDict, IterableDataset
 from datasets.exceptions import DatasetNotFoundError
 from transformers import AutoTokenizer
 import datasets
@@ -126,7 +126,7 @@ def _process_dataset_configs(
             if d.sampling:
                 logging.warning("Sampling multiple datasets is not supported yet")
 
-            if d.data_handlers: # Execute the datahandlers
+            if d.data_handlers:  # Execute the datahandlers
                 for data_handler in d.data_handlers:
                     handler_name: str = data_handler.name
                     handler: callable = self.registered_handlers[handler_name]
@@ -157,8 +157,7 @@ def _process_dataset_configs(
 
                     kwargs["fn_kwargs"] = dict(kwargs["fn_kwargs"], **extra_kwargs)
 
-                    # logging.info
-                    # assert ("Applying Handler : {data_handler} Args : {kwargs}") == ""
+                    logging.info("Applying Handler : {data_handler} Args : {kwargs}")
 
                     raw_datasets = raw_datasets.map(handler, **kwargs)
 
@@ -205,12 +204,12 @@ def get_dataprocessor(
 ) -> DataPreProcessor:
     loader = dataloaderconfig.type
     if loader == "default":
-        procesor = HFBasedDataPreProcessor(
+        processor = HFBasedDataPreProcessor(
             dataloaderconfig=dataloaderconfig,
             tokenizer=tokenizer,
             accelerator=accelerator,
         )
     else:
-        procesor = None
-    autoregister_available_handlers(procesor)
-    return procesor
+        processor = None
+    autoregister_available_handlers(processor)
+    return processor

tuning/data/setup_dataprocessor.py

@@ -17,19 +17,17 @@
 import logging
 
 # Third Party
-from datasets import Dataset, DatasetDict
 from transformers import AutoTokenizer
 
 # Local
 from tuning.config.configs import DataArguments
 from tuning.data.data_config import (
-    DataConfig,
     DataHandlerConfig,
     DataLoaderConfig,
     DataSetConfig,
     load_and_validate_data_config,
 )
-from tuning.data.data_processors import DataPreProcessor, get_dataprocessor
+from tuning.data.data_processors import get_dataprocessor
 from tuning.utils.preprocessing_utils import (
     JSON_INPUT_KEY,
     JSON_OUTPUT_KEY,
@@ -45,8 +43,6 @@ def process_dataargs(
     if data_args.validation_data_path:
         validation_dataset = True
 
-    dataset_text_field = data_args.dataset_text_field
-
     # Create a data processor with default loader config
     default_loader_config = DataLoaderConfig()
     data_processor = get_dataprocessor(
@@ -72,6 +68,8 @@ def process_dataargs(
     fn_kwargs = {}
     handlers = None
 
+    dataset_text_field = data_args.dataset_text_field
+
     # Use case specific handlers
     if is_train_data_pretokenized:
         # dataset_text_field is irrelevant to pretokenized datasets
@@ -95,7 +93,6 @@ def process_dataargs(
             )
             handlers = [handler]
     else:
-
         # TODO: These should be called DEFAULT in the name as they are hardcoded.
         fn_kwargs["input_field_name"] = JSON_INPUT_KEY
         fn_kwargs["output_field_name"] = JSON_OUTPUT_KEY
@@ -132,11 +129,10 @@ def process_dataargs(
     if validation_dataset:
         eval_dataset = data_processor.process_dataset_configs([eval_dataset_config])
         logging.info("Validation dataset length is %s", len(eval_dataset))
-    # dataset_text_field is irrelevant to pretokenized datasets
+
     return train_dataset, eval_dataset, dataset_text_field
 
 
-# TODO: This is very basic the handling of validation will be done by adding splitter.
 # For now assume 2 differnet arguments for training and validation dataset config files.
 # This is very limited but is done to keep first implementation minimal
 def process_dataconfig_file(dataconfigfile: str, tokenizer: AutoTokenizer):