feat: Support pretokenized

tests/data/__init__.py

@@ -24,6 +24,9 @@
     DATA_DIR, "twitter_complaints_input_output.json"
 )
 TWITTER_COMPLAINTS_JSON_FORMAT = os.path.join(DATA_DIR, "twitter_complaints_json.json")
+TWITTER_COMPLAINTS_TOKENIZED = os.path.join(
+    DATA_DIR, "twitter_complaints_tokenized_with_maykeye_tinyllama_v0.json"
+)
 EMPTY_DATA = os.path.join(DATA_DIR, "empty_data.json")
 MALFORMATTED_DATA = os.path.join(DATA_DIR, "malformatted_data.json")
 MODEL_NAME = "Maykeye/TinyLLama-v0"

tests/data/twitter_complaints_tokenized_with_maykeye_tinyllama_v0.json

@@ -0,0 +1,10 @@
+{"Tweet text":"@HMRCcustomers No this is my first job","ID":0,"Label":2,"text_label":"no complaint","output":"### Text: @HMRCcustomers No this is my first job\n\n### Label: no complaint","input_ids":[1,16121,9211,31871,1662,31866,31856,7416,17632,369,1398,433,322,629,712,1784,13,13,8458,31922,21597,31871,697,9566],"labels":[1,16121,9211,31871,1662,31866,31856,7416,17632,369,1398,433,322,629,712,1784,13,13,8458,31922,21597,31871,697,9566]}
+{"Tweet text":"@KristaMariePark Thank you for your interest! If you decide to cancel, you can call Customer Care at 1-800-NYTIMES.","ID":1,"Label":2,"text_label":"no complaint","output":"### Text: @KristaMariePark Thank you for your interest! If you decide to cancel, you can call Customer Care at 1-800-NYTIMES.\n\n### Label: no complaint","input_ids":[1,16121,9211,31871,1662,31892,1260,31825,11273,503,31857,632,5284,365,329,553,1280,31905,960,365,6194,289,11025,31844,365,473,987,12207,4218,389,31822,31853,31854,31886,31852,31852,31854,11300,31847,3873,1507,31843,13,13,8458,31922,21597,31871,697,9566],"labels":[1,16121,9211,31871,1662,31892,1260,31825,11273,503,31857,632,5284,365,329,553,1280,31905,960,365,6194,289,11025,31844,365,473,987,12207,4218,389,31822,31853,31854,31886,31852,31852,31854,11300,31847,3873,1507,31843,13,13,8458,31922,21597,31871,697,9566]}
+{"Tweet text":"If I can't get my 3rd pair of @beatsbydre powerbeats to work today I'm doneski man. This is a slap in my balls. Your next @Bose @BoseService","ID":2,"Label":1,"text_label":"complaint","output":"### Text: If I can't get my 3rd pair of @beatsbydre powerbeats to work today I'm doneski man. This is a slap in my balls. Your next @Bose @BoseService\n\n### Label: complaint","input_ids":[1,16121,9211,31871,960,312,473,31876,31824,685,629,31822,31878,4449,5861,287,1662,1299,1574,1590,31833,263,1360,1299,1574,289,623,31822,31824,16346,312,31876,31836,994,277,3560,567,31843,672,322,260,29458,288,629,14881,31843,2628,1423,1662,31858,601,1662,31858,601,8378,13,13,8458,31922,21597,31871,9566],"labels":[1,16121,9211,31871,960,312,473,31876,31824,685,629,31822,31878,4449,5861,287,1662,1299,1574,1590,31833,263,1360,1299,1574,289,623,31822,31824,16346,312,31876,31836,994,277,3560,567,31843,672,322,260,29458,288,629,14881,31843,2628,1423,1662,31858,601,1662,31858,601,8378,13,13,8458,31922,21597,31871,9566]}
+{"Tweet text":"@EE On Rosneath Arial having good upload and download speeds but terrible latency 200ms. Why is this.","ID":3,"Label":1,"text_label":"complaint","output":"### Text: @EE On Rosneath Arial having good upload and download speeds but terrible latency 200ms. Why is this.\n\n### Label: complaint","input_ids":[1,16121,9211,31871,1662,7766,1078,8123,17561,308,3456,1833,975,10849,291,4372,15379,504,10011,2368,1512,31822,31855,31852,31852,1243,31843,3007,322,433,31843,13,13,8458,31922,21597,31871,9566],"labels":[1,16121,9211,31871,1662,7766,1078,8123,17561,308,3456,1833,975,10849,291,4372,15379,504,10011,2368,1512,31822,31855,31852,31852,1243,31843,3007,322,433,31843,13,13,8458,31922,21597,31871,9566]}
+{"Tweet text":"Couples wallpaper, so cute. :) #BrothersAtHome","ID":4,"Label":2,"text_label":"no complaint","output":"### Text: Couples wallpaper, so cute. :) #BrothersAtHome\n\n### Label: no complaint","input_ids":[1,16121,9211,31871,12371,2208,26657,31844,560,14138,31843,21994,1257,24870,496,31829,8198,19057,13,13,8458,31922,21597,31871,697,9566],"labels":[1,16121,9211,31871,12371,2208,26657,31844,560,14138,31843,21994,1257,24870,496,31829,8198,19057,13,13,8458,31922,21597,31871,697,9566]}
+{"Tweet text":"@mckelldogs This might just be me, but-- eyedrops? Artificial tears are so useful when you're sleep-deprived and sp\u2026 https:\/\/t.co\/WRtNsokblG","ID":5,"Label":2,"text_label":"no complaint","output":"### Text: @mckelldogs This might just be me, but-- eyedrops? Artificial tears are so useful when you're sleep-deprived and sp\u2026 https:\/\/t.co\/WRtNsokblG\n\n### Label: no complaint","input_ids":[1,16121,9211,31871,1662,31836,651,307,395,13094,672,1467,701,333,515,31844,504,1097,2266,282,305,781,31902,21626,31822,31824,5540,397,560,5253,662,365,31876,263,4985,31854,8903,16801,291,612,31925,2011,1129,31824,31843,1358,31873,19919,31824,31865,31829,469,2131,31874,13,13,8458,31922,21597,31871,697,9566],"labels":[1,16121,9211,31871,1662,31836,651,307,395,13094,672,1467,701,333,515,31844,504,1097,2266,282,305,781,31902,21626,31822,31824,5540,397,560,5253,662,365,31876,263,4985,31854,8903,16801,291,612,31925,2011,1129,31824,31843,1358,31873,19919,31824,31865,31829,469,2131,31874,13,13,8458,31922,21597,31871,697,9566]}
+{"Tweet text":"@Yelp can we get the exact calculations for a business rating (for example if its 4 stars but actually 4.2) or do we use a 3rd party site?","ID":6,"Label":2,"text_label":"no complaint","output":"### Text: @Yelp can we get the exact calculations for a business rating (for example if its 4 stars but actually 4.2) or do we use a 3rd party site?\n\n### Label: no complaint","input_ids":[1,16121,9211,31871,1662,31900,307,31837,473,382,685,266,3195,17532,329,260,1173,9363,352,1671,1881,646,619,31822,31882,5556,504,2091,31822,31882,31843,31855,31861,405,499,382,863,260,31822,31878,4449,2540,2042,31902,13,13,8458,31922,21597,31871,697,9566],"labels":[1,16121,9211,31871,1662,31900,307,31837,473,382,685,266,3195,17532,329,260,1173,9363,352,1671,1881,646,619,31822,31882,5556,504,2091,31822,31882,31843,31855,31861,405,499,382,863,260,31822,31878,4449,2540,2042,31902,13,13,8458,31922,21597,31871,697,9566]}
+{"Tweet text":"@nationalgridus I have no water and the bill is current and paid. Can you do something about this?","ID":7,"Label":1,"text_label":"complaint","output":"### Text: @nationalgridus I have no water and the bill is current and paid. Can you do something about this?\n\n### Label: complaint","input_ids":[1,16121,9211,31871,1662,14390,16373,337,312,435,697,1579,291,266,3925,322,1434,291,3877,31843,1456,365,499,1419,562,433,31902,13,13,8458,31922,21597,31871,9566],"labels":[1,16121,9211,31871,1662,14390,16373,337,312,435,697,1579,291,266,3925,322,1434,291,3877,31843,1456,365,499,1419,562,433,31902,13,13,8458,31922,21597,31871,9566]}
+{"Tweet text":"Never shopping at @MACcosmetics again. Every time I go in there, their employees are super rude\/condescending. I'll take my $$ to @Sephora","ID":8,"Label":1,"text_label":"complaint","output":"### Text: Never shopping at @MACcosmetics again. Every time I go in there, their employees are super rude\/condescending. I'll take my $$ to @Sephora\n\n### Label: complaint","input_ids":[1,16121,9211,31871,7265,7550,389,1662,31856,2226,11596,27771,898,31843,3259,647,312,498,288,635,31844,518,3822,397,2168,28910,31873,13627,4107,1708,31843,312,31876,608,1090,629,10279,289,1662,29966,31831,5605,13,13,8458,31922,21597,31871,9566],"labels":[1,16121,9211,31871,7265,7550,389,1662,31856,2226,11596,27771,898,31843,3259,647,312,498,288,635,31844,518,3822,397,2168,28910,31873,13627,4107,1708,31843,312,31876,608,1090,629,10279,289,1662,29966,31831,5605,13,13,8458,31922,21597,31871,9566]}
+{"Tweet text":"@JenniferTilly Merry Christmas to as well. You get more stunning every year \ufffd\ufffd","ID":9,"Label":2,"text_label":"no complaint","output":"### Text: @JenniferTilly Merry Christmas to as well. You get more stunning every year \ufffd\ufffd\n\n### Label: no complaint","input_ids":[1,16121,9211,31871,1662,31884,1450,7064,31847,6538,30894,4472,289,362,828,31843,864,685,541,9932,843,584,18694,31986,13,13,8458,31922,21597,31871,697,9566],"labels":[1,16121,9211,31871,1662,31884,1450,7064,31847,6538,30894,4472,289,362,828,31843,864,685,541,9932,843,584,18694,31986,13,13,8458,31922,21597,31871,697,9566]}

tests/test_sft_trainer.py

@@ -37,6 +37,7 @@
     TWITTER_COMPLAINTS_DATA,
     TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT,
     TWITTER_COMPLAINTS_JSON_FORMAT,
+    TWITTER_COMPLAINTS_TOKENIZED,
 )
 
 # Local
@@ -416,6 +417,39 @@ def test_run_causallm_ft_and_inference():
         _test_run_inference(tempdir=tempdir)
 
 
+def test_run_causallm_ft_pretokenized():
+    """Check if we can bootstrap and finetune causallm models using pretokenized data"""
+    with tempfile.TemporaryDirectory() as tempdir:
+        data_formatting_args = copy.deepcopy(DATA_ARGS)
+
+        # below args not needed for pretokenized data
+        data_formatting_args.data_formatter_template = None
+        data_formatting_args.dataset_text_field = None
+        data_formatting_args.response_template = None
+
+        # update the training data path to tokenized data
+        data_formatting_args.training_data_path = TWITTER_COMPLAINTS_TOKENIZED
+
+        train_args = copy.deepcopy(TRAIN_ARGS)
+        train_args.output_dir = tempdir
+
+        sft_trainer.train(MODEL_ARGS, data_formatting_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
+
+
 ############################# Helper functions #############################
 def _test_run_causallm_ft(training_args, model_args, data_args, tempdir):
     train_args = copy.deepcopy(training_args)

tests/utils/test_preprocessing_utils.py

@@ -11,6 +11,7 @@
     MODEL_NAME,
     TWITTER_COMPLAINTS_DATA,
     TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT,
+    TWITTER_COMPLAINTS_TOKENIZED,
 )
 
 # Local
@@ -21,6 +22,7 @@
     get_data_collator,
     get_formatted_dataset_with_single_sequence,
     get_preprocessed_dataset,
+    is_pretokenized_dataset,
     load_hf_dataset_from_jsonl_file,
     validate_data_args,
 )
@@ -175,6 +177,29 @@ def test_get_data_collator(
     assert isinstance(collator, expected_collator)
 
 
+@pytest.mark.parametrize(
+    "data, result",
+    [
+        (TWITTER_COMPLAINTS_DATA, False),
+        (
+            Dataset.from_list(
+                [
+                    {
+                        "input_ids": [9437, 29, 210],
+                        "attention_mask": [1, 1, 1],
+                        "labels": [1, 20, 30],
+                    }
+                ]
+            ),
+            True,
+        ),
+    ],
+)
+def test_is_pretokenized_dat(data, result):
+    """Ensure that the correct collator type is fetched based on the data args"""
+    assert is_pretokenized_dataset(data=data) == result
+
+
 # Tests for validating data args
 # Invalid args return ValueError
 @pytest.mark.parametrize(
@@ -211,6 +236,54 @@ def test_get_data_collator(
             ),
             False,
         ),
+        # pretokenized dataset with dataset_text_field
+        (
+            configs.DataArguments(
+                training_data_path=TWITTER_COMPLAINTS_TOKENIZED,
+                dataset_text_field="output",
+            ),
+            False,
+        ),
+        # pretokenized dataset with data formatter
+        (
+            configs.DataArguments(
+                training_data_path=TWITTER_COMPLAINTS_TOKENIZED,
+                data_formatter_template="### Input: {{input}} \n\n### Response: {{output}}",
+            ),
+            False,
+        ),
+        # pretokenized dataset with response template
+        (
+            configs.DataArguments(
+                training_data_path=TWITTER_COMPLAINTS_TOKENIZED,
+                response_template="\n### Label:",
+            ),
+            False,
+        ),
+        # pretokenized training dataset with validation data not pretokenized
+        (
+            configs.DataArguments(
+                training_data_path=TWITTER_COMPLAINTS_TOKENIZED,
+                validation_data_path=TWITTER_COMPLAINTS_DATA,
+            ),
+            False,
+        ),
+        # pretokenized data with dataset_text_field and response template
+        (
+            configs.DataArguments(
+                training_data_path=TWITTER_COMPLAINTS_TOKENIZED,
+                response_template="\n### Label:",
+                dataset_text_field="output",
+            ),
+            False,
+        ),
+        # pretokenized data with packing to True
+        (
+            configs.DataArguments(
+                training_data_path=TWITTER_COMPLAINTS_TOKENIZED,
+            ),
+            True,
+        ),
     ],
 )
 def test_validate_args(data_args, packing):
@@ -219,6 +292,31 @@ def test_validate_args(data_args, packing):
         validate_data_args(data_args, packing)
 
 
+@pytest.mark.parametrize(
+    "data_args, packing",
+    [
+        # pretokenized train dataset and no validation dataset passed
+        (
+            configs.DataArguments(
+                training_data_path=TWITTER_COMPLAINTS_TOKENIZED,
+            ),
+            False,
+        ),
+        # pretokenized train and validation datasets
+        (
+            configs.DataArguments(
+                training_data_path=TWITTER_COMPLAINTS_TOKENIZED,
+                validation_data_path=TWITTER_COMPLAINTS_TOKENIZED,
+            ),
+            False,
+        ),
+    ],
+)
+def test_validate_args_pretokenized(data_args, packing):
+    """Ensure that supported data args do not error out when passing pretokenized datasets"""
+    validate_data_args(data_args, packing)
+
+
 @pytest.mark.parametrize(
     "data_path, dataset_text_field, data_formatter_template",
     [
@@ -286,3 +384,33 @@ def test_format_dataset(data_args):
     else:
         assert dataset_text_field in train_set.column_names
         assert dataset_text_field in eval_set.column_names
+
+
+@pytest.mark.parametrize(
+    "data_args",
+    [
+        # pretokenized train and validation datasets
+        (
+            configs.DataArguments(
+                training_data_path=TWITTER_COMPLAINTS_TOKENIZED,
+                validation_data_path=TWITTER_COMPLAINTS_TOKENIZED,
+            )
+        ),
+        # pretokenized train datasets
+        (
+            configs.DataArguments(
+                training_data_path=TWITTER_COMPLAINTS_TOKENIZED,
+            )
+        ),
+    ],
+)
+def test_format_dataset_pretokenized(data_args):
+    """Ensure that pretokenized datasets are loaded and returned as is"""
+    train_set, eval_set, _ = format_dataset(data_args, None, max_seq_length=1024)
+    assert isinstance(train_set, Dataset)
+    if eval_set:
+        assert isinstance(eval_set, Dataset)
+
+    assert set(["input_ids", "labels"]).issubset(set(train_set.column_names))
+    if eval_set:
+        assert set(["input_ids", "labels"]).issubset(set(eval_set.column_names))