Merge remote-tracking branch 'upstream/main'

Author: dchourasia

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)
 

pyproject.toml

@@ -29,7 +29,7 @@ classifiers=[
 dependencies = [
 "numpy>=1.26.4,<2.0",
 "accelerate>=0.20.3,!=0.34,<1.1",
-"transformers>=4.46,<4.48.2",
+"transformers>=4.49,<5.0",
 "torch>=2.2.0,<2.5",
 "sentencepiece>=0.1.99,<0.3",
 "tokenizers>=0.13.3,<1.0",

tests/artifacts/predefined_data_configs/__init__.py

@@ -49,3 +49,9 @@
 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"
+)
+DATA_CONFIG_SKIP_LARGE_TEXT_HANDLER = os.path.join(
+    PREDEFINED_DATA_CONFIGS, "skip_large_text_data_handler_template.yaml"
+)

tests/artifacts/predefined_data_configs/skip_large_text_data_handler_template.yaml

@@ -0,0 +1,25 @@
+dataprocessor:
+    type: default
+datasets:
+  - name: pre_tokenized
+    data_paths:
+      - "FILE_PATH"
+    data_handlers:
+      - name: tokenize
+        arguments:
+          remove_columns: all
+          batched: true
+          fn_kwargs:
+            dataset_text_field: "output"
+      - name: duplicate_columns
+        arguments:
+          remove_columns: all
+          batched: true
+          fn_kwargs:
+            old_column: "input_ids"
+            new_column: "labels"
+      - name: skip_large_text
+        arguments:
+          fn_kwargs:
+            column_name: "input_ids"
+            max_length: 50

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

@@ -16,7 +16,7 @@
 # https://spdx.dev/learn/handling-license-info/
 
 # Third Party
-from datasets import IterableDatasetDict
+from datasets import Dataset, IterableDatasetDict
 from transformers import AutoTokenizer
 import datasets
 import pytest
@@ -35,7 +35,10 @@
     apply_custom_jinja_template,
     combine_sequence,
     duplicate_columns,
+    skip_large_text,
+    tokenize,
 )
+from tuning.data.setup_dataprocessor import is_pretokenized_dataset
 
 
 def test_apply_custom_formatting_template():
@@ -250,3 +253,62 @@ 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
+
+
+@pytest.mark.parametrize(
+    "column_name, max_length",
+    [
+        (None, None),
+        ("input_ids", None),
+        (1024, 1024),
+        ("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"
+    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)
+
+
+def test_skip_large_text_handler():
+    "Ensure that skip large text handler skips dataset as intended"
+
+    def test_dataset_generator():
+        for i in range(0, 100):
+            yield {"input": list(range(0, i + 1))}
+
+    d = Dataset.from_generator(test_dataset_generator)
+    fn_kwargs = {}
+    fn_kwargs["column_name"] = "input"
+    fn_kwargs["max_length"] = 61
+
+    filtered = d.filter(skip_large_text, fn_kwargs=fn_kwargs)
+    assert len(filtered) == 60

tests/test_sft_trainer.py

@@ -41,7 +41,9 @@
     DATA_CONFIG_MULTIPLE_DATASETS_SAMPLING_YAML,
     DATA_CONFIG_MULTITURN_DATA_YAML,
     DATA_CONFIG_RENAME_RETAIN_COLUMNS,
+    DATA_CONFIG_SKIP_LARGE_TEXT_HANDLER,
     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,
 )
@@ -78,7 +80,11 @@
     DataPreProcessorConfig,
     DataSetConfig,
 )
-from tuning.data.data_handlers import add_tokenizer_eos_token
+from tuning.data.data_handlers import (
+    DataHandler,
+    DataHandlerType,
+    add_tokenizer_eos_token,
+)
 
 MODEL_ARGS = configs.ModelArguments(
     model_name_or_path=MODEL_NAME, use_flash_attn=False, torch_dtype="float32"
@@ -321,14 +327,6 @@ def _get_training_logs_by_epoch(dir_path: str, epoch: int = None):
     return data_list
 
 
-def test_run_train_requires_output_dir():
-    """Check fails when output dir not provided."""
-    updated_output_dir_train_args = copy.deepcopy(TRAIN_ARGS)
-    updated_output_dir_train_args.output_dir = None
-    with pytest.raises(TypeError):
-        sft_trainer.train(MODEL_ARGS, DATA_ARGS, updated_output_dir_train_args, None)
-
-
 def test_run_train_fails_training_data_path_not_exist():
     """Check fails when data path not found."""
     updated_data_path_args = copy.deepcopy(DATA_ARGS)
@@ -996,6 +994,97 @@ 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
+
+
+def test_run_training_with_skip_large_text_handler():
+    """Ensure that we can train succesfully after using skip large text handler."""
+    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_SKIP_LARGE_TEXT_HANDLER
+        datapath = TWITTER_COMPLAINTS_TOKENIZED_JSON
+
+        # 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],
@@ -1656,7 +1745,8 @@ def test_run_with_bad_additional_data_handlers(additional_handlers):
         train_args.output_dir = tempdir
 
         with pytest.raises(
-            ValueError, match="Handlers should be of type Dict, str to callable"
+            ValueError,
+            match="Handler should be of type tuning.data_handler.DataHandler, and name of str",
         ):
             sft_trainer.train(
                 MODEL_ARGS,
@@ -1725,6 +1815,12 @@ def test_handler(element, tokenizer, **kwargs):
             DATA_ARGS,
             train_args,
             PEFT_PT_ARGS,
-            additional_data_handlers={TEST_HANDLER: test_handler},
+            additional_data_handlers={
+                TEST_HANDLER: DataHandler(
+                    op=test_handler,
+                    handler_type=DataHandlerType.MAP,
+                    allows_batching=False,
+                )
+            },
         )
         _validate_training(tempdir)