diff --git a/model/supervised_finetuning/custom_datasets/__init__.py b/model/supervised_finetuning/custom_datasets/__init__.py index fcab8a56..907e1a9b 100644 --- a/model/supervised_finetuning/custom_datasets/__init__.py +++ b/model/supervised_finetuning/custom_datasets/__init__.py @@ -2,6 +2,12 @@ from datasets import load_dataset from sklearn.model_selection import train_test_split from torch.utils.data import Dataset, Subset +QA_SPECIAL_TOKENS = { + 'Question': '', + 'Answer': '' +} + + class SquadV2Dataset(Dataset): def __init__(self, cache_dir, split): diff --git a/model/supervised_finetuning/custom_datasets/dialogue_collator.py b/model/supervised_finetuning/custom_datasets/dialogue_collator.py index 17fe1082..f9e1bb5e 100644 --- a/model/supervised_finetuning/custom_datasets/dialogue_collator.py +++ b/model/supervised_finetuning/custom_datasets/dialogue_collator.py @@ -6,6 +6,8 @@ import torch from torch.nn import functional as F from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase +from . import QA_SPECIAL_TOKENS + @dataclass class DialogueDataCollator: @@ -19,22 +21,21 @@ class DialogueDataCollator: pad_to_multiple_of: Optional[int] = None def __call__(self, features): - # TODO add special tokens for question and answer here - # additional_special_tokens = ['', ''] - prompt_tokens = ["Question: ", "Answer: "] - flatten_messages = [] label_masks = [] for messages in features: assert len(messages) % 2 == 0, "Number of messages must be even" messages = [ - (prompt_tokens[0] if i % 2 == 0 else "") + x + ((" " + prompt_tokens[1]) if i % 2 == 0 else "") + (QA_SPECIAL_TOKENS["Question"] if i % 2 == 0 else "") + + x + + (QA_SPECIAL_TOKENS["Answer"] if i % 2 == 0 else "") for i, x in enumerate(messages) ] - # Add a way for the model to terminate generation, reinitialize prompter - messages.append(prompt_tokens[0]) + # Add a way for the model to terminate generation + # When we predict the start of a new expected question, we want to be able to stop generation + messages.append(QA_SPECIAL_TOKENS["Question"]) flatten_messages.append( self.tokenizer( @@ -47,8 +48,10 @@ class DialogueDataCollator: message_change_indices = np.cumsum([len(x) for x in messages[:-1]]) # for each token an integer indicating the index of the message it belongs to. Just to create the label mask. - # TEXT: Question: Hello, how are you? Answer: I am fine. Question: What is your name? Answer: My name is John. - # MESSAGE_INDICES: 0 0 0 0 0 0 1 1 1 2 2 2 2 2 2 3 3 3 3 + # Label mask is true when predicting a token that is part of the answer, false otherwise. + # TEXT: Question: Hello, how are you? Answer: I am fine. Question: What is your name? Answer: My name is John. Question: + # MESSAGE_INDICES: 0 0 0 0 0 0 1 1 1 2 2 2 2 2 2 3 3 3 3 -2 + # LABEL_MASK: 0 0 0 0 0 1 1 1 1 0 0 0 0 0 1 1 1 1 1 0 # If no result in next, we are predicting the last termination token(s) message_indices = list( diff --git a/model/supervised_finetuning/trainer.py b/model/supervised_finetuning/trainer.py index b44890df..dc7b5934 100644 --- a/model/supervised_finetuning/trainer.py +++ b/model/supervised_finetuning/trainer.py @@ -67,6 +67,8 @@ class SFTTrainer(Trainer): optimizers, preprocess_logits_for_metrics, ) + + # By default CrossEntropyLoss ignores padding_index -100, but just in case use our own loss_fct self.loss_fct = get_loss(args.loss_function) def fetch_scheduler(self): @@ -112,7 +114,7 @@ class SFTTrainer(Trainer): with torch.no_grad(): loss, logits, labels, labels_mask = self._compute_loss(model, inputs) - labels[~labels_mask] = -1 + labels[~labels_mask] = -100 # padding_index loss = loss.mean().detach() @@ -159,8 +161,8 @@ def argument_parsing(notebook=False, notebook_args=None): if __name__ == "__main__": training_conf = argument_parsing() - model = get_model(training_conf) tokenizer = get_tokenizer(training_conf) + model = get_model(training_conf, tokenizer) train, evals, collate_fn = get_dataset(training_conf, tokenizer) diff --git a/model/supervised_finetuning/utils.py b/model/supervised_finetuning/utils.py index 4a451bed..6aa5d365 100644 --- a/model/supervised_finetuning/utils.py +++ b/model/supervised_finetuning/utils.py @@ -7,6 +7,7 @@ from losses import CrossEntropyLoss from sklearn.model_selection import train_test_split from torch.utils.data import ConcatDataset, Subset from transformers import AutoModelForCausalLM, AutoTokenizer +from custom_datasets import QA_SPECIAL_TOKENS SUPPORTED_MODELS = ["galactica"] @@ -17,10 +18,19 @@ def get_tokenizer(conf): if "galactica" in conf.model_name: tokenizer.add_special_tokens({"pad_token": "", "eos_token": ""}) + additional_special_tokens = ( + [] + if not "additional_special_tokens" in tokenizer.special_tokens_map + else tokenizer.special_tokens_map["additional_special_tokens"] + ) + additional_special_tokens = list(set(additional_special_tokens + list(QA_SPECIAL_TOKENS.values()))) + + tokenizer.add_special_tokens({"additional_special_tokens": additional_special_tokens}) + return tokenizer -def get_model(conf): +def get_model(conf, tokenizer): if not any([x in conf.model_name for x in SUPPORTED_MODELS]): raise ValueError( f"Model {conf.model_name} not supported. Supported models: {SUPPORTED_MODELS}. " @@ -29,6 +39,11 @@ def get_model(conf): model = AutoModelForCausalLM.from_pretrained(conf.model_name, cache_dir=conf.cache_dir) + if len(tokenizer) != model.get_input_embeddings().num_embeddings: + assert not conf.freeze_layer, "Cannot change the number of embeddings if the model is frozen." + + model.resize_token_embeddings(len(tokenizer)) + if conf.freeze_layer: model = freeze_top_n_layers(model, conf.freeze_layer)