minor changes

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

Author: dushyantbehl

scripts/offline_data_processing.py

@@ -29,11 +29,17 @@ def save_dataset_shards(
         dataset_name (str): Name of the dataset (used for logging).
     """
     os.makedirs(output_dir, exist_ok=True)
+    logging.info(
+        "Dumping processesd dataaset %s at %s in %d shards",
+        dataset_name,
+        output_dir,
+        num_shards,
+    )
     for shard_idx in range(num_shards):
         shard = dataset.shard(index=shard_idx, num_shards=num_shards)
         shard_path = os.path.join(output_dir, f"ds_{shard_idx:05d}.parquet")
         shard.to_parquet(shard_path)
-    logging.info("Dumped %d shards of %s at %s", num_shards, dataset_name, output_dir)
+    logging.info("Dumped %d shards", num_shards)
 
 
 def process_datasets_offline(
@@ -53,10 +59,13 @@ def process_datasets_offline(
         tuple: A tuple containing the formatted training dataset and validation dataset.
     """
     # Set log level for this function
-    train_args, logger = set_log_level(train_args, "get_processed_dataset")
+    train_args, logger = set_log_level(train_args, "process_datasets_offline")
 
     logger.info(
-        "Starting dataset processing with model_args: %s, data_args: %s, training_args: %s",
+        "Starting offline dataset processing with \n\
+         model_args: %s, \n\
+         data_args: %s, \n\
+         training_args: %s",
         model_args,
         data_args,
         train_args,

tuning/data/data_handlers.py

@@ -609,11 +609,23 @@ def tokenize_and_apply_chat_template_with_masking(
 
 
 AVAILABLE_DATA_HANDLERS = {
-    "tokenize_and_apply_input_masking": DataHandler(
-        op=tokenize_and_apply_input_masking,
+    "remove_columns": DataHandler(
+        # Native function
+        handler_type=DataHandlerType.REMOVE,
+    ),
+    "select_columns": DataHandler(
+        # Native function
+        handler_type=DataHandlerType.SELECT,
+    ),
+    "rename_columns": DataHandler(
+        # Native function
+        handler_type=DataHandlerType.RENAME,
+    ),
+    "tokenize": DataHandler(
+        op=tokenize,
         handler_type=DataHandlerType.MAP,
-        allows_batching=False,
-        desc="Combining and tokenizing instruction and response, masking instructions",
+        allows_batching=True,
+        desc="Tokenizing the dataset",
     ),
     "add_tokenizer_eos_token": DataHandler(
         op=add_tokenizer_eos_token,
@@ -625,51 +637,42 @@ def tokenize_and_apply_chat_template_with_masking(
         op=apply_custom_jinja_template,
         handler_type=DataHandlerType.MAP,
         allows_batching=False,
-        desc="Formatting dataset with given jinja template",
+        desc="Formatting dataset with given formatting template",
+    ),
+    "tokenize_and_apply_input_masking": DataHandler(
+        op=tokenize_and_apply_input_masking,
+        handler_type=DataHandlerType.MAP,
+        allows_batching=False,
+        desc="Combining and tokenizing instruction and response, masking instructions",
     ),
     "apply_tokenizer_chat_template": DataHandler(
         op=apply_tokenizer_chat_template,
         handler_type=DataHandlerType.MAP,
         allows_batching=False,
-        desc="Applying tokenizers chat template to dataset",
+        desc="Applying chat template to dataset",
     ),
     "tokenize_and_apply_chat_template_with_masking": DataHandler(
         op=tokenize_and_apply_chat_template_with_masking,
         handler_type=DataHandlerType.MAP,
         allows_batching=False,
-        desc="Applying chat template to dataset with tokenization",
+        desc="Applying chat template to dataset and tokenizing",
     ),
     "duplicate_columns": DataHandler(
         op=duplicate_columns,
         handler_type=DataHandlerType.MAP,
         allows_batching=True,
         desc="Duplicating columns",
     ),
-    "prepare_multimodal_data_processor": DataHandler(
-        op=prepare_multimodal_data_processor,
-        handler_type=DataHandlerType.MAP,
-        allows_batching=False,
-        desc="Processing text+image data",
-    ),
-    "tokenize": DataHandler(
-        op=tokenize,
-        handler_type=DataHandlerType.MAP,
-        allows_batching=True,
-        desc="Tokenizing the dataset",
-    ),
     "skip_samples_with_large_columns": DataHandler(
         op=skip_samples_with_large_columns,
         handler_type=DataHandlerType.FILTER,
         allows_batching=False,
         desc="Skipping large samples",
     ),
-    "remove_columns": DataHandler(
-        handler_type=DataHandlerType.REMOVE,
-    ),
-    "select_columns": DataHandler(
-        handler_type=DataHandlerType.SELECT,
-    ),
-    "rename_columns": DataHandler(
-        handler_type=DataHandlerType.RENAME,
+    "prepare_multimodal_data_processor": DataHandler(
+        op=prepare_multimodal_data_processor,
+        handler_type=DataHandlerType.MAP,
+        allows_batching=False,
+        desc="Processing multimodal data",
     ),
 }