Expose additional data handlers as an argument to the train function.

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

Author: dushyantbehl

tuning/data/data_processors.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 # Standard
-from typing import Dict, List, Union
+from typing import Callable, Dict, List, Union
 import logging
 import os
 
@@ -35,7 +35,7 @@ class DataPreProcessor:
     tokenizer = None
     data_config: DataConfig = None
     processor_config: DataPreProcessorConfig = None
-    registered_handlers: Dict[str, callable] = None
+    registered_handlers: Dict[str, Callable] = None
 
     def __init__(
         self, processor_config: DataPreProcessorConfig, tokenizer: AutoTokenizer
@@ -46,9 +46,20 @@ def __init__(
         # Initialize other objects
         self.registered_handlers = {}
 
-    def register_data_handler(self, name: str, func: callable):
+    def register_data_handler(self, name: str, func: Callable):
+        assert isinstance(name, str), "Handler name should be of str type"
+        assert callable(func), "Handler should be a callable routine"
         self.registered_handlers[name] = func
 
+    def register_data_handlers(self, handlers: Dict[str, Callable]):
+        if handlers is None:
+            return
+        assert isinstance(
+            handlers, Dict
+        ), "Handlers should be of type Dict[str:Callable]"
+        for k, v in handlers.items():
+            self.register_data_handler(name=k, func=v)
+
     def load_dataset(
         self,
         datasetconfig: DataSetConfig,
@@ -238,19 +249,12 @@ def process_dataset_configs(
         return train_dataset
 
 
-def autoregister_available_handlers(processor: DataPreProcessor):
-    if processor is None:
-        return
-    for name, func in AVAILABLE_DATA_HANDLERS.items():
-        processor.register_data_handler(name=name, func=func)
-
-
 def get_datapreprocessor(
     processor_config: DataPreProcessorConfig, tokenizer: AutoTokenizer
 ) -> DataPreProcessor:
     processor = DataPreProcessor(
         processor_config=processor_config,
         tokenizer=tokenizer,
     )
-    autoregister_available_handlers(processor)
+    processor.register_data_handlers(AVAILABLE_DATA_HANDLERS)
     return processor

tuning/data/setup_dataprocessor.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 # Standard
-from typing import Union
+from typing import Callable, Dict, Union
 import logging
 
 # Third Party
@@ -55,11 +55,16 @@ def is_pretokenized_dataset(data: Union[str, Dataset, IterableDataset]):
 
 # TODO: For now assume only training dataset is passed via data config file.
 # This is very limited but is done to keep first implementation minimal
-def _process_dataconfig_file(data_args: DataArguments, tokenizer: AutoTokenizer):
+def _process_dataconfig_file(
+    data_args: DataArguments,
+    tokenizer: AutoTokenizer,
+    additional_data_handlers: Dict[str, Callable] = None,
+):
     data_config = load_and_validate_data_config(data_args.data_config_path)
     processor = get_datapreprocessor(
         processor_config=data_config.dataprocessor, tokenizer=tokenizer
     )
+    processor.register_data_handlers(additional_data_handlers)
     train_dataset = processor.process_dataset_configs(data_config.datasets)
 
     return (train_dataset, None, data_args.dataset_text_field)
@@ -179,14 +184,15 @@ def _process_raw_data_args(
     tokenizer: AutoTokenizer,
     packing: bool,
     max_seq_length: int,
+    additional_data_handlers: Dict[str, Callable] = None,
 ):
 
     # Create a data processor with default processor config
     default_processor_config = DataPreProcessorConfig()
     data_processor = get_datapreprocessor(
         processor_config=default_processor_config, tokenizer=tokenizer
     )
-
+    data_processor.register_data_handlers(additional_data_handlers)
     assert isinstance(
         data_args.training_data_path, str
     ), "Training data path has to be set and str"
@@ -259,7 +265,10 @@ def _process_raw_data_args(
 # If no data config file is specified, process the remaining data arguments
 # to determine the use case based on their presence, as explained in _process_raw_data_args.
 def process_dataargs(
-    data_args: DataArguments, tokenizer: AutoTokenizer, train_args: TrainingArguments
+    data_args: DataArguments,
+    tokenizer: AutoTokenizer,
+    train_args: TrainingArguments,
+    additional_data_handlers: Dict[str, Callable] = None,
 ):
     """
     Args:
@@ -268,11 +277,17 @@ def process_dataargs(
         train_args: TrainingArguments
             Training arguments passed to the library
             Used for packing and max_seq_length
+        additional_data_handlers: A Dict of [str, callable] data handlers
+            which need to be registered with the data preprocessor
     Returns:
         Tuple(Dataset, Dataset, str, DataCollator, int, Dict)
-            tuple containing train_dataset, eval_dataset, dataset_text_field,
-                data_collator, max_seq_length and dataset_kwargs
-
+            tuple containing
+            train_dataset (Dataset/IterableDataset),
+            eval_dataset (Dataset/IterableDataset),
+            dataset_text_field (str),
+            data_collator (DataCollator)
+            max_seq_length(int) and
+            dataset_kwargs (Dict)
     """
 
     max_seq_length = min(train_args.max_seq_length, tokenizer.model_max_length)
@@ -290,26 +305,32 @@ def process_dataargs(
 
     if data_args.data_config_path:
         train_dataset, eval_dataset, dataset_text_field = _process_dataconfig_file(
-            data_args, tokenizer
+            data_args, tokenizer, additional_data_handlers
         )
     else:
         train_dataset, eval_dataset, dataset_text_field = _process_raw_data_args(
-            data_args, tokenizer, train_args.packing, max_seq_length
+            data_args,
+            tokenizer,
+            train_args.packing,
+            max_seq_length,
+            additional_data_handlers,
         )
 
+    # Note: This check should not be removed.
+    #       Its important to recompute this post handling to
+    #       check if we already tokenized the dataset or not.
+    is_tokenized_dataset = is_pretokenized_dataset(train_dataset or eval_dataset)
+
     data_collator = get_data_collator(
         train_args.packing,
         data_args.response_template,
         tokenizer,
-        # Note: This check should not be removed.
-        #       Its important to recompute this post handling to
-        #       check if we already tokenized the dataset or not.
-        is_pretokenized_dataset(train_dataset),
+        is_tokenized_dataset,
         max_seq_length,
     )
 
     dataset_kwargs = {}
-    if is_pretokenized_dataset(train_dataset or eval_dataset):
+    if is_tokenized_dataset:
         dataset_kwargs["skip_prepare_dataset"] = True
 
     return (

tuning/sft_trainer.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 # Standard
-from typing import Dict, List, Optional, Union
+from typing import Callable, Dict, List, Optional, Union
 import dataclasses
 import json
 import logging
@@ -85,6 +85,7 @@ def train(
     attention_and_distributed_packing_config: Optional[
         AttentionAndDistributedPackingConfig
     ] = None,
+    additional_data_handlers: Optional[Dict[str, Callable]] = None,
 ) -> tuple[SFTTrainer, dict]:
     """Call the SFTTrainer
 
@@ -113,7 +114,8 @@ def train(
             Should be used in combination with quantized_lora_config. Also currently 
             fused_lora and fast_kernels must used together (may change in future). \
         attention_and_distributed_packing_config: Used for padding-free attention and multipack.
-
+        additional_data_handlers: Dict [str:Callable] of any extra data handlers \
+                                   to be registered with the data preprocessor
     Returns:
         Tuple: Instance of SFTTrainer , some metadata in a dict
             Metadata contains information on number of added tokens while tuning.
@@ -297,7 +299,7 @@ def train(
         data_collator,
         max_seq_length,
         dataset_kwargs,
-    ) = process_dataargs(data_args, tokenizer, train_args)
+    ) = process_dataargs(data_args, tokenizer, train_args, additional_data_handlers)
     additional_metrics["data_preprocessing_time"] = (
         time.time() - data_preprocessing_time
     )