Add tokenizer data handler and test case via data config (#487)

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

Author: dushyantbehl

docs/advanced-data-preprocessing.md

@@ -233,6 +233,8 @@ This library currently supports the following [preexisting data handlers](https:
     Uses a tokenizer's chat template to preprocess dataset elements, good for single/multi turn chat templates.
  - `duplicate_columns`:
     Duplicates one column of the dataset to another column.
+ - `tokenize`:
+    Tokenizes one column of the dataset passed as input `dataset_text_field`.
 
 These handlers could be requested by their same name and users can lookup the function args from [here](https://github.com/foundation-model-stack/fms-hf-tuning/blob/main/tuning/data/data_handlers.py)
 

tests/artifacts/predefined_data_configs/__init__.py

@@ -49,3 +49,6 @@
 DATA_CONFIG_RENAME_RETAIN_COLUMNS = os.path.join(
     PREDEFINED_DATA_CONFIGS, "rename_retain_columns.yaml"
 )
+DATA_CONFIG_TOKENIZE_AND_TRAIN_WITH_HANDLER = os.path.join(
+    PREDEFINED_DATA_CONFIGS, "tokenize_using_handler_and_train.yaml"
+)

tests/artifacts/predefined_data_configs/tokenize_using_handler_and_train.yaml

@@ -0,0 +1,22 @@
+dataprocessor:
+    type: default
+datasets:
+  - name: non_tokenized_dataset
+    data_paths:
+      - "FILE_PATH"
+    data_handlers:
+      - name: tokenize
+        arguments:
+          remove_columns: all
+          batched: true
+          fn_kwargs:
+            dataset_text_field: "output"
+            truncation: True
+            max_length: 1024
+      - name: duplicate_columns
+        arguments:
+          remove_columns: all
+          batched: true
+          fn_kwargs:
+            old_column: "input_ids"
+            new_column: "labels"

tests/data/test_data_handlers.py

@@ -35,7 +35,9 @@
     apply_custom_jinja_template,
     combine_sequence,
     duplicate_columns,
+    tokenize,
 )
+from tuning.data.setup_dataprocessor import is_pretokenized_dataset
 
 
 def test_apply_custom_formatting_template():
@@ -250,3 +252,26 @@ def test_duplicate_columns_copies_columns():
     assert new in first_element
     assert old in first_element
     assert first_element[new] == first_element[old]
+
+
+def test_tokenizer_data_handler_tokenizes():
+    "Ensure tokenizer data handler tokenizes the input properly with proper truncation"
+    d = datasets.load_dataset("json", data_files=TWITTER_COMPLAINTS_DATA_JSONL)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    dataset_text_field = "output"
+    truncation = True
+    max_length = 10
+
+    updated_dataaset = d.map(
+        tokenize,
+        fn_kwargs={
+            "tokenizer": tokenizer,
+            "dataset_text_field": dataset_text_field,
+            "truncation": truncation,
+            "max_length": max_length,
+        },
+    )
+
+    assert "input_ids" in updated_dataaset["train"][0]
+    for element in updated_dataaset["train"]:
+        assert len(element["input_ids"]) <= max_length

tests/test_sft_trainer.py

@@ -42,6 +42,7 @@
     DATA_CONFIG_MULTITURN_DATA_YAML,
     DATA_CONFIG_RENAME_RETAIN_COLUMNS,
     DATA_CONFIG_TOKENIZE_AND_APPLY_INPUT_MASKING_YAML,
+    DATA_CONFIG_TOKENIZE_AND_TRAIN_WITH_HANDLER,
     DATA_CONFIG_YAML_STREAMING_INPUT_OUTPUT,
     DATA_CONFIG_YAML_STREAMING_PRETOKENIZED,
 )
@@ -996,6 +997,52 @@ def test_run_training_with_pretokenised_dataset_containing_input_ids():
         assert "### Text: @NortonSupport Thanks much.\n\n### Label:" in output_inference
 
 
+def test_run_training_with_data_tokenized_using_tokenizer_handler():
+    """Ensure that we can train on non tokenized dataset works by tokenizing using
+    tokenizer data handler via data config."""
+    with tempfile.TemporaryDirectory() as tempdir:
+
+        data_args = copy.deepcopy(DATA_ARGS)
+
+        # set training_data_path and response_template to none
+        data_args.response_template = None
+        data_args.training_data_path = None
+
+        dataconfigfile = DATA_CONFIG_TOKENIZE_AND_TRAIN_WITH_HANDLER
+        datapath = TWITTER_COMPLAINTS_DATA_JSONL
+
+        # add data_paths in data_config file
+        with tempfile.NamedTemporaryFile(
+            "w", delete=False, suffix=".yaml"
+        ) as temp_yaml_file:
+            with open(dataconfigfile, "r", encoding="utf-8") as f:
+                data = yaml.safe_load(f)
+                datasets = data["datasets"]
+                for _, d in enumerate(datasets):
+                    d["data_paths"] = [datapath]
+                yaml.dump(data, temp_yaml_file)
+                data_args.data_config_path = temp_yaml_file.name
+
+        train_args = copy.deepcopy(TRAIN_ARGS)
+        train_args.output_dir = tempdir
+
+        sft_trainer.train(MODEL_ARGS, data_args, train_args)
+
+        # validate full ft configs
+        _validate_training(tempdir)
+        checkpoint_path = _get_checkpoint_path(tempdir)
+
+        # Load the model
+        loaded_model = TunedCausalLM.load(checkpoint_path, MODEL_NAME)
+
+        # Run inference on the text
+        output_inference = loaded_model.run(
+            "### Text: @NortonSupport Thanks much.\n\n### Label:", max_new_tokens=50
+        )
+        assert len(output_inference) > 0
+        assert "### Text: @NortonSupport Thanks much.\n\n### Label:" in output_inference
+
+
 @pytest.mark.parametrize(
     "dataset_path",
     [CHAT_DATA_SINGLE_TURN, CHAT_DATA_MULTI_TURN],

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, List, Union
 import copy
 import re
 
@@ -257,6 +257,37 @@ def apply_tokenizer_chat_template(
     }
 
 
+def tokenize(
+    element: Union[Dict[str, str], Dict[str, List]],
+    tokenizer: AutoTokenizer,
+    dataset_text_field: str,
+    truncation: Union[bool, str] = None,
+    max_length: int = None,
+    **kwargs,
+):
+    """Function (data handler) to tokenize dataset columns.
+       Expects to be run as a HF Map API function.
+    Args:
+        element: the HF Dataset element.
+        tokenizer: Tokenizer to be used.
+        dataset_text_field: the dataset field to tokenize
+        truncation: Truncation strategy to use, refer the link
+                    (https://huggingface.co/docs/transformers/en/pad_truncation)
+        max_length: Max length to truncate the samples to.
+        kwargs: Any additional kwargs that need to be passed to the tokenizer can be passed as
+                kwargs['tokenizer_kwargs']
+    Returns:
+        tokenized dataset elemenent field "dataset_text_field"
+    """
+    tokenizer_kwargs = kwargs.get("tokenizer_kwargs", {})
+    return tokenizer(
+        element[dataset_text_field],
+        truncation=truncation,
+        max_length=max_length,
+        **tokenizer_kwargs,
+    )
+
+
 def duplicate_columns(
     element: Dict[str, str],
     old_column: str,
@@ -298,4 +329,5 @@ def duplicate_columns(
     "apply_custom_jinja_template": apply_custom_jinja_template,
     "apply_tokenizer_chat_template": apply_tokenizer_chat_template,
     "duplicate_columns": duplicate_columns,
+    "tokenize": tokenize,
 }