diff --git a/tests/data/__init__.py b/tests/data/__init__.py index 6df7802cd9..f42779830b 100644 --- a/tests/data/__init__.py +++ b/tests/data/__init__.py @@ -20,5 +20,8 @@ ### Constants used for data DATA_DIR = os.path.join(os.path.dirname(__file__)) TWITTER_COMPLAINTS_DATA = os.path.join(DATA_DIR, "twitter_complaints_small.json") +TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT = os.path.join( + DATA_DIR, "twitter_complaints_input_output.json" +) EMPTY_DATA = os.path.join(DATA_DIR, "empty_data.json") MALFORMATTED_DATA = os.path.join(DATA_DIR, "malformatted_data.json") diff --git a/tests/data/twitter_complaints_input_output.json b/tests/data/twitter_complaints_input_output.json new file mode 100644 index 0000000000..2bfba89adb --- /dev/null +++ b/tests/data/twitter_complaints_input_output.json @@ -0,0 +1,50 @@ +{"ID": 0, "Label": 2, "input": "@HMRCcustomers No this is my first job", "output": "no complaint"} +{"ID": 1, "Label": 2, "input": "@KristaMariePark Thank you for your interest! If you decide to cancel, you can call Customer Care at 1-800-NYTIMES.", "output": "no complaint"} +{"ID": 2, "Label": 1, "input": "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", "output": "complaint"} +{"ID": 3, "Label": 1, "input": "@EE On Rosneath Arial having good upload and download speeds but terrible latency 200ms. Why is this.", "output": "complaint"} +{"ID": 4, "Label": 2, "input": "Couples wallpaper, so cute. :) #BrothersAtHome", "output": "no complaint"} +{"ID": 5, "Label": 2, "input": "@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", "output": "no complaint"} +{"ID": 6, "Label": 2, "input": "@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?", "output": "no complaint"} +{"ID": 7, "Label": 1, "input": "@nationalgridus I have no water and the bill is current and paid. Can you do something about this?", "output": "complaint"} +{"ID": 8, "Label": 1, "input": "Never shopping at @MACcosmetics again. Every time I go in there, their employees are super rude/condescending. I'll take my $$ to @Sephora", "output": "complaint"} +{"ID": 9, "Label": 2, "input": "@JenniferTilly Merry Christmas to as well. You get more stunning every year \ufffd\ufffd", "output": "no complaint"} +{"ID": 10, "Label": 2, "input": "@NortonSupport Thanks much.", "output": "no complaint"} +{"ID": 11, "Label": 2, "input": "@VerizonSupport all of a sudden I can't connect to my primary wireless network but guest one works", "output": "no complaint"} +{"ID": 12, "Label": 2, "input": "Aaaahhhhh!!!! My @Razer @PlayOverwatch d.va meka headset came in!!! I didn't even know it had shipped!!! So excited\u2026 https://t.co/4gXy9xED8d", "output": "no complaint"} +{"ID": 13, "Label": 2, "input": "@Lin_Manuel @jmessinaphoto @VAMNit Omg a little squish!!!!! Enjoy and congrats!!!! I miss mine being so young! \ufffd\ufffd\ufffd\ufffd\ufffd\ufffd", "output": "no complaint"} +{"ID": 14, "Label": 2, "input": "@IanJamesPoulter What's your secret to poaching eggs? Mine NEVER look that good.", "output": "no complaint"} +{"ID": 15, "Label": 2, "input": "@AWSSupport When will be able Kinesis Firehose compatible with Elasticsearch 6.0? Thank you!", "output": "no complaint"} +{"ID": 16, "Label": 2, "input": "@NCIS_CBS https://t.co/eeVL9Eu3bE", "output": "no complaint"} +{"ID": 17, "Label": 2, "input": "@msetchell Via the settings? That\u2019s how I do it on master T\u2019s", "output": "no complaint"} +{"ID": 18, "Label": 2, "input": "Today at work there was a low flying duck heading toward a crowd of people, and I yelled \"watch out! and I'm very disappointed with myself.", "output": "no complaint"} +{"ID": 19, "Label": 1, "input": "@NortonSupport @NortonOnline What the hell is a dm 5-10 days to get money back bank account now overdrawn thanks guys", "output": "complaint"} +{"ID": 20, "Label": 1, "input": "@united not happy with this delay from Newark to Manchester tonight :( only 30 mins free Wi-fi sucks ...", "output": "complaint"} +{"ID": 21, "Label": 1, "input": "@ZARA_Care I've been waiting on a reply to my tweets and DMs for days now?", "output": "complaint"} +{"ID": 22, "Label": 2, "input": "New Listing! Large 2 Family Home for Sale in #Passaic Park, #NJ #realestate #homesforsale Great Location!\u2026 https://t.co/IV4OrLXkMk", "output": "no complaint"} +{"ID": 23, "Label": 1, "input": "@SouthwestAir I love you but when sending me flight changes please don't use military time #ignoranceisbliss", "output": "complaint"} +{"ID": 24, "Label": 2, "input": "@JetBlue Completely understand but would prefer being on time to filling out forms....", "output": "no complaint"} +{"ID": 25, "Label": 2, "input": "@nvidiacc I own two gtx 460 in sli. I want to try windows 8 dev preview. Which driver should I use. Can I use the windows 7 one.", "output": "no complaint"} +{"ID": 26, "Label": 2, "input": "Just posted a photo https://t.co/RShFwCjPHu", "output": "no complaint"} +{"ID": 27, "Label": 2, "input": "Love crescent rolls? Try adding pesto @PerdueChicken to them and you\u2019re going to love it! #Promotion #PerdueCrew -\u2026 https://t.co/KBHOfqCukH", "output": "no complaint"} +{"ID": 28, "Label": 1, "input": "@TopmanAskUs please just give me my money back.", "output": "complaint"} +{"ID": 29, "Label": 2, "input": "I just gave 5 stars to Tracee at @neimanmarcus for the great service I received!", "output": "no complaint"} +{"ID": 30, "Label": 2, "input": "@FitbitSupport when are you launching new clock faces for Indian market", "output": "no complaint"} +{"ID": 31, "Label": 1, "input": "@HPSupport my printer will not allow me to choose color instead it only prints monochrome #hppsdr #ijkhelp", "output": "complaint"} +{"ID": 32, "Label": 1, "input": "@DIRECTV can I get a monthly charge double refund when it sprinkles outside and we lose reception? #IamEmbarrasedForYou", "output": "complaint"} +{"ID": 33, "Label": 1, "input": "@AlfaRomeoCares Hi thanks for replying, could be my internet but link doesn't seem to be working", "output": "complaint"} +{"ID": 34, "Label": 2, "input": "Looks tasty! Going to share with everyone I know #FebrezeONE #sponsored https://t.co/4AQI53npei", "output": "no complaint"} +{"ID": 35, "Label": 2, "input": "@OnePlus_IN can OnePlus 5T do front camera portrait?", "output": "no complaint"} +{"ID": 36, "Label": 1, "input": "@sho_help @showtime your arrive is terrible streaming is stop and start every couple mins. Get it together it's xmas", "output": "complaint"} +{"ID": 37, "Label": 2, "input": "@KandraKPTV I just witnessed a huge building fire in Santa Monica California", "output": "no complaint"} +{"ID": 38, "Label": 2, "input": "@fernrocks most definitely the latter for me", "output": "no complaint"} +{"ID": 39, "Label": 1, "input": "@greateranglia Could I ask why the Area in front of BIC Station was not gritted withh all the snow.", "output": "complaint"} +{"ID": 40, "Label": 2, "input": "I'm earning points with #CricketRewards https://t.co/GfpGhqqnhE", "output": "no complaint"} +{"ID": 41, "Label": 2, "input": "@Schrapnel @comcast RIP me", "output": "no complaint"} +{"ID": 42, "Label": 2, "input": "The wait is finally over, just joined @SquareUK, hope to get started real soon!", "output": "no complaint"} +{"ID": 43, "Label": 2, "input": "@WholeFoods what's the best way to give feedback on a particular store to the regional/national office?", "output": "no complaint"} +{"ID": 44, "Label": 2, "input": "@DanielNewman I honestly would believe anything. People are...too much sometimes.", "output": "no complaint"} +{"ID": 45, "Label": 2, "input": "@asblough Yep! It should send you a notification with your driver\u2019s name and what time they\u2019ll be showing up!", "output": "no complaint"} +{"ID": 46, "Label": 2, "input": "@Wavy2Timez for real", "output": "no complaint"} +{"ID": 47, "Label": 1, "input": "@KenyaPower_Care no power in south b area... is it scheduled.", "output": "complaint"} +{"ID": 48, "Label": 1, "input": "Honda won't do anything about water leaking in brand new car. Frustrated! @HondaCustSvc @AmericanHonda", "output": "complaint"} +{"ID": 49, "Label": 1, "input": "@CBSNews @Dodge @ChryslerCares My driver side air bag has been recalled and replaced, but what about the passenger side?", "output": "complaint"} \ No newline at end of file diff --git a/tests/test_sft_trainer.py b/tests/test_sft_trainer.py index bbe91f8908..bbf88819fc 100644 --- a/tests/test_sft_trainer.py +++ b/tests/test_sft_trainer.py @@ -29,7 +29,12 @@ # First Party from scripts.run_inference import TunedCausalLM -from tests.data import EMPTY_DATA, MALFORMATTED_DATA, TWITTER_COMPLAINTS_DATA +from tests.data import ( + EMPTY_DATA, + MALFORMATTED_DATA, + TWITTER_COMPLAINTS_DATA, + TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT, +) # Local from tuning import sft_trainer @@ -436,3 +441,58 @@ def test_bad_torch_dtype(): with pytest.raises(ValueError): sft_trainer.train(model_args, DATA_ARGS, train_args, PEFT_PT_ARGS) + + +### Tests for pretokenized data +def test_pretokenized_dataset(): + """Ensure that we can provide a pretokenized dataset with input/output format.""" + with tempfile.TemporaryDirectory() as tempdir: + train_args = copy.deepcopy(TRAIN_ARGS) + train_args.output_dir = tempdir + data_args = copy.deepcopy(DATA_ARGS) + data_args.dataset_text_field = None + data_args.response_template = None + data_args.training_data_path = TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT + sft_trainer.train(MODEL_ARGS, data_args, train_args, PEFT_PT_ARGS) + _validate_training(tempdir) + + +@pytest.mark.parametrize( + "dataset_text_field,response_template", + [ + ("foo", None), + (None, "bar"), + ], +) +def test_pretokenized_dataset_bad_args(dataset_text_field, response_template): + """Ensure that we can't provide only dataset text field / response template for pretok data.""" + with tempfile.TemporaryDirectory() as tempdir: + train_args = copy.deepcopy(TRAIN_ARGS) + train_args.output_dir = tempdir + + data_args = copy.deepcopy(DATA_ARGS) + data_args.dataset_text_field = dataset_text_field + data_args.response_template = response_template + data_args.training_data_path = TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT + # We should raise an error since we should not have a dataset text + # field or a response template if we're pretokenized data + with pytest.raises(ValueError): + sft_trainer.train(MODEL_ARGS, data_args, train_args, PEFT_PT_ARGS) + + +def test_pretokenized_dataset_wrong_format(): + """Ensure that we fail to generate data if the data is in the wrong format.""" + with tempfile.TemporaryDirectory() as tempdir: + train_args = copy.deepcopy(TRAIN_ARGS) + train_args.output_dir = tempdir + + data_args = copy.deepcopy(DATA_ARGS) + data_args.dataset_text_field = None + data_args.response_template = None + data_args.training_data_path = TWITTER_COMPLAINTS_DATA + + # It would be best to handle this in a way that is more understandable; we might + # need to add directly validation prior to the dataset generation since datasets + # is essentially swallowing a KeyError here. + with pytest.raises(DatasetGenerationError): + sft_trainer.train(MODEL_ARGS, data_args, train_args, PEFT_PT_ARGS) diff --git a/tests/utils/test_preprocessing_utils.py b/tests/utils/test_preprocessing_utils.py new file mode 100644 index 0000000000..bbc02b4d0a --- /dev/null +++ b/tests/utils/test_preprocessing_utils.py @@ -0,0 +1,179 @@ +# Third Party +from datasets import Dataset +from datasets.exceptions import DatasetGenerationError +from transformers import AutoTokenizer, DataCollatorForSeq2Seq +from trl import DataCollatorForCompletionOnlyLM +import pytest + +# First Party +from tests.data import ( + MALFORMATTED_DATA, + TWITTER_COMPLAINTS_DATA, + TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT, +) + +# Local +from tuning.utils.preprocessing_utils import ( + combine_sequence, + get_data_trainer_kwargs, + get_preprocessed_dataset, + load_hf_dataset_from_jsonl_file, +) + + +@pytest.mark.parametrize( + "input_element,output_element,expected_res", + [ + ("foo ", "bar", "foo bar"), + ("foo\n", "bar", "foo\nbar"), + ("foo\t", "bar", "foo\tbar"), + ("foo", "bar", "foo bar"), + ], +) +def test_combine_sequence(input_element, output_element, expected_res): + """Ensure that input / output elements are combined with correct whitespace handling.""" + comb_seq = combine_sequence(input_element, output_element) + assert isinstance(comb_seq, str) + assert comb_seq == expected_res + + +# Tests for loading the dataset from disk +def test_load_hf_dataset_from_jsonl_file(): + input_field_name = "Tweet text" + output_field_name = "text_label" + data = load_hf_dataset_from_jsonl_file( + TWITTER_COMPLAINTS_DATA, + input_field_name=input_field_name, + output_field_name=output_field_name, + ) + # Our dataset should contain dicts that contain the input / output field name types + next_data = next(iter(data)) + assert input_field_name in next_data + assert output_field_name in next_data + + +def test_load_hf_dataset_from_jsonl_file_wrong_keys(): + """Ensure that we explode if the keys are not in the jsonl file.""" + with pytest.raises(DatasetGenerationError): + load_hf_dataset_from_jsonl_file( + TWITTER_COMPLAINTS_DATA, input_field_name="foo", output_field_name="bar" + ) + + +def test_load_hf_dataset_from_malformatted_data(): + """Ensure that we explode if the data is not properly formatted.""" + # NOTE: The actual keys don't matter here + with pytest.raises(DatasetGenerationError): + load_hf_dataset_from_jsonl_file( + MALFORMATTED_DATA, input_field_name="foo", output_field_name="bar" + ) + + +def test_load_hf_dataset_from_jsonl_file_duplicate_keys(): + """Ensure we cannot have the same key for input / output.""" + with pytest.raises(ValueError): + load_hf_dataset_from_jsonl_file( + TWITTER_COMPLAINTS_DATA, + input_field_name="Tweet text", + output_field_name="Tweet text", + ) + + +# Tests for custom masking / preprocessing logic +@pytest.mark.parametrize("max_sequence_length", [1, 10, 100, 1000]) +def test_get_preprocessed_dataset(max_sequence_length): + tokenizer = AutoTokenizer.from_pretrained("Maykeye/TinyLLama-v0") + preprocessed_data = get_preprocessed_dataset( + data_path=TWITTER_COMPLAINTS_DATA, + tokenizer=tokenizer, + max_sequence_length=max_sequence_length, + input_field_name="Tweet text", + output_field_name="text_label", + ) + for tok_res in preprocessed_data: + # Since the padding is left to the collator, there should be no 0s in the attention mask yet + assert sum(tok_res["attention_mask"]) == len(tok_res["attention_mask"]) + # If the source text isn't empty, we start with masked inputs + assert tok_res["labels"][0] == -100 + # All keys in the produced record must be the same length + key_lengths = {len(tok_res[k]) for k in tok_res.keys()} + assert len(key_lengths) == 1 + # And also that length should be less than or equal to the max length depending on if we + # are going up to / over the max size and truncating - padding is handled separately + assert key_lengths.pop() <= max_sequence_length + + +# Tests for fetching train args +@pytest.mark.parametrize( + "use_validation_data, collator_type, packing", + [ + (True, None, True), + (False, None, True), + (True, DataCollatorForCompletionOnlyLM, False), + (False, DataCollatorForCompletionOnlyLM, False), + ], +) +def test_get_trainer_kwargs_with_response_template_and_text_field( + use_validation_data, collator_type, packing +): + training_data_path = TWITTER_COMPLAINTS_DATA + validation_data_path = training_data_path if use_validation_data else None + # Expected columns in the raw loaded dataset for the twitter data + column_names = set(["Tweet text", "ID", "Label", "text_label", "output"]) + trainer_kwargs = get_data_trainer_kwargs( + training_data_path=training_data_path, + validation_data_path=validation_data_path, + packing=packing, + response_template="\n### Label:", + max_sequence_length=100, + tokenizer=AutoTokenizer.from_pretrained("Maykeye/TinyLLama-v0"), + dataset_text_field="output", + ) + assert len(trainer_kwargs) == 3 + # If we are packing, we should not have a data collator + if collator_type is None: + assert trainer_kwargs["data_collator"] is None + else: + assert isinstance(trainer_kwargs["data_collator"], collator_type) + + # We should only have a validation dataset if one is present + if validation_data_path is None: + assert trainer_kwargs["eval_dataset"] is None + else: + assert isinstance(trainer_kwargs["eval_dataset"], Dataset) + assert set(trainer_kwargs["eval_dataset"].column_names) == column_names + + assert isinstance(trainer_kwargs["train_dataset"], Dataset) + assert set(trainer_kwargs["train_dataset"].column_names) == column_names + + +@pytest.mark.parametrize("use_validation_data", [True, False]) +def test_get_trainer_kwargs_with_custom_masking(use_validation_data): + training_data_path = TWITTER_COMPLAINTS_DATA_INPUT_OUTPUT + validation_data_path = training_data_path if use_validation_data else None + # Expected columns in the raw loaded dataset for the twitter data + column_names = set(["input_ids", "attention_mask", "labels"]) + trainer_kwargs = get_data_trainer_kwargs( + training_data_path=training_data_path, + validation_data_path=validation_data_path, + packing=False, + response_template=None, + max_sequence_length=100, + tokenizer=AutoTokenizer.from_pretrained("Maykeye/TinyLLama-v0"), + dataset_text_field=None, + ) + assert len(trainer_kwargs) == 4 + # If we are packing, we should not have a data collator + assert isinstance(trainer_kwargs["data_collator"], DataCollatorForSeq2Seq) + + # We should only have a validation dataset if one is present + if validation_data_path is None: + assert trainer_kwargs["eval_dataset"] is None + else: + assert isinstance(trainer_kwargs["eval_dataset"], Dataset) + assert set(trainer_kwargs["eval_dataset"].column_names) == column_names + + assert isinstance(trainer_kwargs["train_dataset"], Dataset) + assert set(trainer_kwargs["train_dataset"].column_names) == column_names + # Needed to sidestep TRL validation + assert trainer_kwargs["formatting_func"] is not None diff --git a/tuning/sft_trainer.py b/tuning/sft_trainer.py index b307505c07..c6037c096c 100644 --- a/tuning/sft_trainer.py +++ b/tuning/sft_trainer.py @@ -29,8 +29,7 @@ TrainerCallback, ) from transformers.utils import logging -from trl import DataCollatorForCompletionOnlyLM, SFTTrainer -import datasets +from trl import SFTTrainer import fire import transformers @@ -46,6 +45,7 @@ from tuning.trainercontroller import TrainerControllerCallback from tuning.utils.config_utils import get_hf_peft_config from tuning.utils.data_type_utils import get_torch_dtype +from tuning.utils.preprocessing_utils import get_data_trainer_kwargs def train( @@ -165,14 +165,6 @@ def train( } ) - # TODO: near term - how response template ids are parsed out needs to be cleaned. - # The [2:] here applies if response template has \n prefix, it is needed to strip \n, - # otherwise template is not found. We will create issue to clean this out after we discuss - # data formats and collators we will support. - response_template_ids = tokenizer.encode( - data_args.response_template, add_special_tokens=False - )[2:] - max_seq_length = min(train_args.max_seq_length, tokenizer.model_max_length) logger.info("Max sequence length is %s", max_seq_length) if train_args.max_seq_length > tokenizer.model_max_length: @@ -207,59 +199,31 @@ def train( model=model, ) - # Configure the collator and validate args related to packing prior to formatting the dataset - if train_args.packing: - logger.info("Packing is set to True") - data_collator = None - packing = True - else: - logger.info("Packing is set to False") - if data_args.response_template is None: - # TODO: Fix this, currently unreachable due to crashing in batch encoding tokenization - # We should do this validation up front, then do the encoding, then handle the collator - raise ValueError("Response template is None, needs to be set for training") - if data_args.dataset_text_field is None: - raise ValueError("Dataset_text_field is None, needs to be set for training") - data_collator = DataCollatorForCompletionOnlyLM( - response_template_ids, - tokenizer=tokenizer, - ignore_index=configs.IGNORE_INDEX, - ) - packing = False - - # load the data by parsing JSON - data_files = {"train": data_args.training_data_path} - if data_args.validation_data_path: - data_files["validation"] = data_args.validation_data_path - - format_dataset = lambda example: { # pylint: disable=unnecessary-lambda-assignment - f"{data_args.dataset_text_field}": example[f"{data_args.dataset_text_field}"] - + tokenizer.eos_token - } - - json_dataset = datasets.load_dataset("json", data_files=data_files) - formatted_train_dataset = json_dataset["train"].map(format_dataset) - logger.info("Training dataset length is %s", len(formatted_train_dataset)) - - formatted_validation_dataset = None - if data_args.validation_data_path: - formatted_validation_dataset = json_dataset["validation"].map(format_dataset) - logger.info( - "Validation dataset length is %s", len(formatted_validation_dataset) - ) + # Get the trainer args focused on data, specifically those sensitive to + # packing, response template, and dataset text field. Currently this consists of + # - training dataset + # - validation dataset + # - data collator + trainer_data_args = get_data_trainer_kwargs( + training_data_path=data_args.training_data_path, + validation_data_path=data_args.validation_data_path, + packing=train_args.packing, + response_template=data_args.response_template, + max_sequence_length=max_seq_length, + tokenizer=tokenizer, + dataset_text_field=data_args.dataset_text_field, + ) trainer = SFTTrainer( model=model, tokenizer=tokenizer, - train_dataset=formatted_train_dataset, - eval_dataset=formatted_validation_dataset, - packing=packing, - data_collator=data_collator, + packing=train_args.packing, dataset_text_field=data_args.dataset_text_field, args=train_args, max_seq_length=max_seq_length, callbacks=trainer_callbacks, peft_config=peft_config, + **trainer_data_args ) # We track additional metrics and experiment metadata after trainer object creation @@ -322,7 +286,6 @@ def main(**kwargs): # pylint: disable=unused-argument additional, _, ) = parser.parse_args_into_dataclasses(return_remaining_strings=True) - logger = logging.get_logger("__main__") peft_method = additional.peft_method diff --git a/tuning/utils/preprocessing_utils.py b/tuning/utils/preprocessing_utils.py new file mode 100644 index 0000000000..b23b86d6a6 --- /dev/null +++ b/tuning/utils/preprocessing_utils.py @@ -0,0 +1,329 @@ +# Copyright The FMS HF Tuning Authors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# Standard +from typing import Any, Callable, Dict, Optional +import json + +# Third Party +from datasets import Dataset +from transformers import AutoTokenizer, DataCollatorForSeq2Seq +from trl import DataCollatorForCompletionOnlyLM +import datasets + +# Local +from tuning.config import configs + + +def get_data_trainer_kwargs( + training_data_path: str, + validation_data_path: str, + packing: bool, + response_template: Optional[str], + max_sequence_length: int, + tokenizer: AutoTokenizer, + dataset_text_field: Optional[str], +) -> Dict[str, Any]: + """Get trainer args related to data / processing. At the moment, this consists of: + - the training dataset + - the evaluation dataset + - the data collator + - Maybe a formatting a function [only for a special case for validation] + The result can be kwarg expanded into the trainer initialization. + + Args: + training_data_path: str + Path to the training data. + validation_data_path: str + Path to the validation data. + packing: bool + Whether or not we should apply packing or not. + response_template: Optional[str] + Response template to be used for formatting by TRL. + max_sequence_length: int + Max sequence length to be used for sequence tokenization. + tokenizer: AutoTokenizer + Loaded tokenizer object to be used by the collator. + dataset_text_field: Optional[str] + Dataset text field fto be used for formatting by TRL. + + Returns: + Dict[str, Any] + Data related kwargs to be used by the SFT Trainer. + """ + data_collator = get_data_collator( + packing, dataset_text_field, response_template, max_sequence_length, tokenizer + ) + eval_dataset = None + data_kwargs = {} + if isinstance(data_collator, DataCollatorForSeq2Seq): + # HACK: This function is never called, but is needed to sidestep TRL's internal validation. + data_kwargs["formatting_func"] = lambda x: x + train_dataset = get_preprocessed_dataset( + training_data_path, + tokenizer, + max_sequence_length, + input_field_name="input", + output_field_name="output", + ) + if validation_data_path: + eval_dataset = get_preprocessed_dataset( + validation_data_path, + tokenizer, + max_sequence_length, + input_field_name="input", + output_field_name="output", + ) + else: + # Collator is a DataCollatorForCompletionOnlyLM or None; + # Load it as JSON and apply our normal preprocessing logic + train_dataset = get_formatted_dataset( + training_data_path, dataset_text_field, tokenizer + ) + if validation_data_path: + eval_dataset = get_formatted_dataset( + validation_data_path, dataset_text_field, tokenizer + ) + + data_kwargs["data_collator"] = data_collator + data_kwargs["train_dataset"] = train_dataset + data_kwargs["eval_dataset"] = eval_dataset + return data_kwargs + + +def get_data_collator( + packing: bool, + dataset_text_field: Optional[str], + response_template: Optional[str], + max_sequence_length: int, + tokenizer: AutoTokenizer, +) -> Callable: + """Create and return the the appropriate collator type based on the configuration for packing, + response_template, and dataset_text_field. + + Args: + packing: bool + Whether or not we should apply packing or not. + dataset_text_field: Optional[str] + Dataset text field fto be used for formatting by TRL. + response_template: Optional[str] + Response template to be used for formatting by TRL. + max_sequence_length: int + Max sequence length to be used for sequence tokenization. + tokenizer: AutoTokenizer + Loaded tokenizer object to be used by the collator. + + Returns: + Callable + Callable collator to be leveraged by the trainer. + """ + if not packing: + if dataset_text_field is None and response_template is None: + # Use the seq2seq data collator; note that this automatically pads labels with -100 + return DataCollatorForSeq2Seq( + tokenizer=tokenizer, padding=True, max_length=max_sequence_length + ) + if dataset_text_field is None and response_template is not None: + raise ValueError( + "Packing is disabled, but no dataset_text_field is provided" + ) + if response_template is None and dataset_text_field is not None: + raise ValueError( + "Packing is disabled, but no response template is provided" + ) + # TODO: near term - how response template ids are parsed out needs to be cleaned. + # The [2:] here applies if response template has \n prefix, it is needed to strip \n, + # otherwise template is not found. We will create issue to clean this out after we discuss + # data formats and collators we will support. + response_template_ids = tokenizer.encode( + response_template, add_special_tokens=False + )[2:] + return DataCollatorForCompletionOnlyLM( + response_template=response_template_ids, + tokenizer=tokenizer, + ignore_index=configs.IGNORE_INDEX, + ) + + +def get_formatted_dataset( + data_path: str, dataset_text_field: str, tokenizer: AutoTokenizer +) -> Dataset: + """Applies formatting to the loaded dataset instance; does NOT pretokenize data. + + Args: + data_path: str + Path to the file to be loaded. + dataset_text_field: str + Dataset text field fto be used for formatting by TRL. + tokenizer: AutoTokenizer + Loaded tokenizer object to be used by the collator. + + Returns: + Dataset + HF Dataset with formatted [str] data. + """ + format_dataset = lambda example: { # pylint: disable=unnecessary-lambda-assignment + f"{dataset_text_field}": example[f"{dataset_text_field}"] + tokenizer.eos_token + } + json_dataset = datasets.load_dataset("json", data_files=data_path) + return json_dataset.map(format_dataset)[ + "train" + ] # HACK - for now, we just do both datasets separately; train is the default split + + +def get_preprocessed_dataset( + data_path: str, + tokenizer: AutoTokenizer, + max_sequence_length: int, + input_field_name: str, + output_field_name: str, +) -> Dataset: + """Loads the dataset and applies the tokenizer + custom masking logic. + + Args: + data_path: str + Path to the file to be loaded. + tokenizer: AutoTokenizer + Loaded tokenizer object to be used by the collator. + max_sequence_length: int + Max sequence length to be used for sequence tokenization. + input_field_name: str + Name of the input field in the data. + output_field_name: str + Name of the output field in the data. + + Returns: + Dataset + HF Dataset with the pretokenized data. + """ + dataset = load_hf_dataset_from_jsonl_file( + data_path, input_field_name, output_field_name + ) + return dataset.map( + preprocess_and_tokenize, + fn_kwargs={ + "tokenizer": tokenizer, + "max_seq_length": max_sequence_length, + "input_field_name": input_field_name, + "output_field_name": output_field_name, + }, + remove_columns=[input_field_name, output_field_name], + ) + + +### Utils for loading the data from disk in supported formats [currently only jsonl] +def load_hf_dataset_from_jsonl_file( + data_path: str, input_field_name: str, output_field_name: str +) -> Dataset: + """Loads the huggingface datase as a generator. + + Args: + data_path: str + Path to the file to be loaded. + input_field_name: str + Name of the input field in the data. + output_field_name: str + Name of the output field in the data. + + Returns: + Dataset + HF Dataset with the data to be tokenized. + """ + if input_field_name == output_field_name: + raise ValueError("Input field name and output field name should not match!") + + def get_jsonl_object(): + with open(data_path, "r", encoding="utf-8") as jsonl_file: + data_stream = [json.loads(line) for line in jsonl_file] + for data in data_stream: + yield { + input_field_name: data[input_field_name], + output_field_name: data[output_field_name], + } + + return Dataset.from_generator(get_jsonl_object) + + +### Utils for custom masking / manipulating input / output strs, etc +def combine_sequence(input_element: str, output_element: str): + """Combines / concatenates input & output element. + + Args: + input_element: str + Input component of the combined sequence. + output_element: str + Output component of the combined sequence. + + Returns: + str + Sequence combined with whitespace. + """ + if not input_element.endswith((" ", "\n", "\t")) and not output_element.startswith( + (" ", "\n", "\t") + ): + return input_element + " " + output_element + return input_element + output_element + + +def preprocess_and_tokenize( + element: Dict[str, str], + tokenizer: AutoTokenizer, + max_seq_length: int, + input_field_name: str, + output_field_name: str, +) -> Dict[str, Any]: + """Loads the dataset and applies the tokenizer + custom masking logic. + NOTE: Truncation is done in this step, but padding is not, and generally + handled by the collator. + + Args: + element: Dict[str, str] + A single element of the raw Dataset of strings, whose data we would like to apply + custom masking + tokenization logic to. + tokenizer: AutoTokenizer + Loaded tokenizer object to be used by the collator. + max_sequence_length: int + Max sequence length to be used for sequence tokenization. + input_field_name: str + Name of the input field in the data. + output_field_name: str + Name of the output field in the data. + + Returns: + Dict[str, Any] + Dictionary containing the input IDs/labels/attention mask for this record. + """ + combined_seq = combine_sequence( + element[input_field_name], element[output_field_name] + ) + + tokenized_comb_seqs = tokenizer( + combined_seq, max_length=max_seq_length, truncation=True, padding=False + ) + tokenized_input = tokenizer( + element[input_field_name], + max_length=max_seq_length, + truncation=True, + padding=False, + ) + + # mask the prompt part for avoiding loss + masked_labels = [-100] * len( + tokenized_input.input_ids + ) + tokenized_comb_seqs.input_ids[len(tokenized_input.input_ids) :] + + return { + "input_ids": tokenized_comb_seqs.input_ids, + "labels": masked_labels, + "attention_mask": tokenized_comb_seqs.attention_mask, + }