| \n", + " | text | \n", + "
|---|---|
| 0 | \n", + "Question: To whom did the Virgin Mary allegedl... | \n", + "
| 1 | \n", + "Question: What is in front of the Notre Dame M... | \n", + "
| 2 | \n", + "Question: The Basilica of the Sacred heart at ... | \n", + "
| 3 | \n", + "Question: What is the Grotto at Notre Dame? An... | \n", + "
| 4 | \n", + "Question: What sits on top of the Main Buildin... | \n", + "
| ... | \n", + "... | \n", + "
| 87594 | \n", + "Question: In what US state did Kathmandu first... | \n", + "
| 87595 | \n", + "Question: What was Yangon previously known as?... | \n", + "
| 87596 | \n", + "Question: With what Belorussian city does Kath... | \n", + "
| 87597 | \n", + "Question: In what year did Kathmandu create it... | \n", + "
| 87598 | \n", + "Question: What is KMC an initialism of? Answer... | \n", + "
87599 rows × 1 columns
\n", + "" + ], + "text/plain": [ + " text\n", + "0 Question: To whom did the Virgin Mary allegedl...\n", + "1 Question: What is in front of the Notre Dame M...\n", + "2 Question: The Basilica of the Sacred heart at ...\n", + "3 Question: What is the Grotto at Notre Dame? An...\n", + "4 Question: What sits on top of the Main Buildin...\n", + "... ...\n", + "87594 Question: In what US state did Kathmandu first...\n", + "87595 Question: What was Yangon previously known as?...\n", + "87596 Question: With what Belorussian city does Kath...\n", + "87597 Question: In what year did Kathmandu create it...\n", + "87598 Question: What is KMC an initialism of? Answer...\n", + "\n", + "[87599 rows x 1 columns]" + ] + }, + "execution_count": 21, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "train_df = pd.DataFrame(train_texts, columns=[\"text\"])\n", + "eval_df = pd.DataFrame({\"text\": eval_texts, \"answer\": eval_answers})\n", + "\n", + "train_df" + ] + }, + { + "cell_type": "code", + "execution_count": 22, + "metadata": {}, + "outputs": [], + "source": [ + "train_df.to_json(\"../data/squad-train.jsonl\", orient=\"records\", lines=True)\n", + "eval_df.to_json(\"../data/squad-eval.jsonl\", orient=\"records\", lines=True)" + ] + }, + { + "cell_type": "code", + "execution_count": 28, + "metadata": {}, + "outputs": [], + "source": [ + "train_df[\"rag_query\"] = train_data[\"question\"]\n", + "eval_df[\"rag_query\"] = eval_data[\"question\"]" + ] + }, + { + "cell_type": "code", + "execution_count": 29, + "metadata": {}, + "outputs": [], + "source": [ + "train_df.to_json(\"../data/squad-train.jsonl\", orient=\"records\", lines=True)\n", + "eval_df.to_json(\"../data/squad-eval.jsonl\", orient=\"records\", lines=True)\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "st", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.13" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/examples/llms/train.py b/examples/llms/train.py new file mode 100644 index 00000000..c8eb2bf2 --- /dev/null +++ b/examples/llms/train.py @@ -0,0 +1,98 @@ +import logging +import pandas as pd + +from simpletransformers.language_modeling import LanguageModelingModel, LanguageModelingArgs, GenerationArgs +from simpletransformers.retrieval import ( + RetrievalModel, + RetrievalArgs, +) + + +logging.basicConfig(level=logging.INFO) +transformers_logger = logging.getLogger("transformers") +transformers_logger.setLevel(logging.WARNING) +trainsformers_modules_logger = logging.getLogger("transformers_modules") +trainsformers_modules_logger.setLevel(logging.ERROR) + + +rag = True + +if rag: + from rag_setup import model as rag_model + + +train_file = "../data/squad-train.jsonl" + + +model_args = LanguageModelingArgs() +model_args.reprocess_input_data = True +model_args.overwrite_output_dir = True +model_args.num_train_epochs = 1 +model_args.save_eval_checkpoints = False +model_args.save_model_every_epoch = False +model_args.train_batch_size = 2 +model_args.eval_batch_size = 4 +model_args.gradient_accumulation_steps = 1 +model_args.manual_seed = 4 +model_args.fp16 = True +model_args.dataset_type = "simple" +model_args.logging_steps = 100 +model_args.evaluate_during_training = False +model_args.mlm = False +model_args.use_multiprocessing = False +model_args.use_hf_datasets = True +model_args.peft = True +model_args.qlora = False +model_args.nf4 = True +model_args.loftq_bits = 4 +model_args.lora_config = {"r": 8} +model_args.data_format = "jsonl" +model_args.trust_remote_code = True +model_args.save_steps = 1000 +model_args.optimizer = "Adam8bit" +model_args.chunk_text = False +model_args.max_seq_length = 500 + +if not rag: + model_args.wandb_project = "llama-adapter-tuning-squad" + model_args.wandb_kwargs = {"name": "squad-llama-2-7b-vanilla"} + model_args.output_dir = "squad-llama-2-7b" + +if rag: + model_args.rag = rag + model_args.wandb_project = "llama-adapter-tuning-squad" + model_args.wandb_kwargs = {"name": "squad-llama-2-7b-rag"} + model_args.output_dir = "squad-llama-2-7b-rag" + + +model = LanguageModelingModel( + "causal", + # "outputs", + # "stabilityai/stablelm-zephyr-3b", + "meta-llama/Llama-2-7b-hf", + args=model_args, + retrieval_model=rag_model if rag else None, +) + +model.train_model( + train_file, +) + +# generation_args = GenerationArgs() +# generation_args.max_length = None +# generation_args.max_new_tokens = 100 + +# test_df = pd.read_json("data/test.jsonl", lines=True) + +# to_predict = test_df["input_text"].tolist() + +# responses, _ = model.predict( +# to_predict, +# generation_args=generation_args, +# ) + +# print(responses[:5]) + +# test_df["generated_text"] = responses + +# test_df.to_json("data/test_output-finetuned.jsonl", orient="records", lines=True) \ No newline at end of file diff --git a/examples/retrieval/download_msmarco.py b/examples/retrieval/download_msmarco.py new file mode 100644 index 00000000..7976fcda --- /dev/null +++ b/examples/retrieval/download_msmarco.py @@ -0,0 +1,22 @@ +import os +from datasets import load_dataset + + +os.makedirs("data/msmarco", exist_ok=True) + +print("=== Downloading MSMARCO ===") +print("Downloading MSMARCO training triples...") +dataset = load_dataset("thilina/negative-sampling")["train"] + +print("Dataset loaded. Sample:") +print(dataset[0]) + +qrels = load_dataset("BeIR/msmarco-qrels")["validation"] + +print("Saving dataset to disk...") +# Save the dataset to disk +dataset.to_csv("data/msmarco/msmarco-train.tsv", sep="\t", index=False) +qrels.to_csv("data/msmarco/devs.tsv", sep="\t", index=False) + +print("Done.") +print("=== MSMARCO download complete ===") diff --git a/examples/retrieval/train_dpr_base.py b/examples/retrieval/train_dpr_base.py new file mode 100644 index 00000000..c6562704 --- /dev/null +++ b/examples/retrieval/train_dpr_base.py @@ -0,0 +1,64 @@ +import logging + +import pandas as pd +from simpletransformers.retrieval import RetrievalModel, RetrievalArgs + +# Configuring logging +logging.basicConfig(level=logging.INFO) +transformers_logger = logging.getLogger("transformers") +transformers_logger.setLevel(logging.WARNING) + +# Specifying the path to the training data +train_data_path = "data/msmarco/msmarco-train.tsv" + +# Loading the training data +if train_data_path.endswith(".tsv"): + train_data = pd.read_csv(train_data_path, sep="\t") +else: + train_data = train_data_path + +# Configuring the model arguments +model_args = RetrievalArgs() +model_args.reprocess_input_data = True +model_args.overwrite_output_dir = True +model_args.use_cached_eval_features = False +model_args.include_title = False if "msmarco" in train_data_path else True +model_args.max_seq_length = 256 +model_args.num_train_epochs = 40 +model_args.train_batch_size = 16 +model_args.use_hf_datasets = True +model_args.learning_rate = 1e-6 +model_args.warmup_steps = 5000 +model_args.save_steps = 300000 +model_args.evaluate_during_training = True +model_args.evaluate_during_training_steps = False +model_args.save_model_every_epoch = False +model_args.wandb_project = "Retrieval training example" +model_args.hard_negatives = False +model_args.n_gpu = 1 +model_args.data_format = "beir" +model_args.output_dir = f"trained_models/pretrained/DPR-base-msmarco" +model_args.wandb_kwargs = {"name": f"DPR-base-msmarco"} + +# Defining the model type and names +model_type = "custom" +model_name = None +context_name = "bert-base-multilingual-cased" +question_name = "bert-base-multilingual-cased" + +# Main execution +if __name__ == "__main__": + # Creating the model + model = RetrievalModel( + model_type, + model_name, + context_name, + question_name, + args=model_args, + ) + + # Training the model + model.train_model( + train_data, + eval_set="dev", + ) diff --git a/setup.py b/setup.py index 7d924cc3..657771f1 100755 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name="simpletransformers", - version="0.65.1", + version="0.70.0", author="Thilina Rajapakse", author_email="chaturangarajapakshe@gmail.com", description="An easy-to-use wrapper library for the Transformers library.", diff --git a/simpletransformers/classification/classification_model.py b/simpletransformers/classification/classification_model.py index 66b01238..6ea168b4 100755 --- a/simpletransformers/classification/classification_model.py +++ b/simpletransformers/classification/classification_model.py @@ -22,12 +22,14 @@ from scipy.special import softmax from sklearn.metrics import ( confusion_matrix, + f1_score, label_ranking_average_precision_score, matthews_corrcoef, mean_squared_error, roc_curve, auc, average_precision_score, + accuracy_score, ) from torch.utils.tensorboard import SummaryWriter from torch.nn import CrossEntropyLoss @@ -1155,8 +1157,7 @@ def train( epoch_number += 1 output_dir_current = os.path.join( output_dir, - "checkpoint-{}-epoch-{}".format(global_step, epoch_number) -, + "checkpoint-{}-epoch-{}".format(global_step, epoch_number), ) if args.save_model_every_epoch or args.evaluate_during_training: @@ -1924,6 +1925,8 @@ def compute_metrics( return {**extra_metrics}, wrong mcc = matthews_corrcoef(labels, preds) + accuracy = accuracy_score(labels, preds) + f1 = f1_score(labels, preds) if self.model.num_labels == 2: tn, fp, fn, tp = confusion_matrix(labels, preds, labels=[0, 1]).ravel() if self.args.sliding_window: @@ -1943,6 +1946,8 @@ def compute_metrics( { **{ "mcc": mcc, + "accuracy": accuracy, + "f1_score": f1, "tp": tp, "tn": tn, "fp": fp, diff --git a/simpletransformers/classification/multi_modal_classification_model.py b/simpletransformers/classification/multi_modal_classification_model.py index b59acca7..9ec83919 100644 --- a/simpletransformers/classification/multi_modal_classification_model.py +++ b/simpletransformers/classification/multi_modal_classification_model.py @@ -814,8 +814,7 @@ def train( epoch_number += 1 output_dir_current = os.path.join( output_dir, - "checkpoint-{}-epoch-{}".format(global_step, epoch_number) -, + "checkpoint-{}-epoch-{}".format(global_step, epoch_number), ) if args.save_model_every_epoch or args.evaluate_during_training: diff --git a/simpletransformers/config/model_args.py b/simpletransformers/config/model_args.py index 2ab9da85..7e55cc57 100644 --- a/simpletransformers/config/model_args.py +++ b/simpletransformers/config/model_args.py @@ -93,6 +93,7 @@ class ModelArgs: tokenizer_type: str = None train_batch_size: int = 8 train_custom_parameters_only: bool = False + trust_remote_code: bool = False use_cached_eval_features: bool = False use_early_stopping: bool = False use_hf_datasets: bool = False @@ -246,6 +247,36 @@ class T5Args(ModelArgs): use_multiprocessed_decoding: bool = True +@dataclass +class GenerationArgs: + """ + Args for language generation. + """ + + max_length: int = 20 + max_new_tokens: int = None + min_length: int = 0 + min_new_tokens: int = None + early_stopping: bool = False + max_time: float = None + + do_sample: bool = False + num_beams: int = 1 + num_beam_groups: int = 1 + penalty_alpha: float = None + use_cache: bool = True + + temperature: float = 1.0 + top_k: int = 50 + top_p: float = 1.0 + repetition_penalty: float = 1.0 + diversity_penalty: float = 0.0 + + def get_dict(self): + d = asdict(self) + return {k: v for k, v in d.items() if v is not None} + + @dataclass class LanguageModelingArgs(ModelArgs): """ @@ -254,6 +285,7 @@ class LanguageModelingArgs(ModelArgs): model_class: str = "LanguageModelingModel" block_size: int = -1 + chunk_text: bool = True config_name: str = None dataset_class: Dataset = None dataset_type: str = "None" @@ -276,6 +308,14 @@ class LanguageModelingArgs(ModelArgs): special_tokens_list: list = field(default_factory=list) strip_accents: bool = True local_rank: int = -1 + loftq_bits: int = 4 + loftq_config: dict = field(default_factory=dict) + lora_config: dict = field(default_factory=dict) + peft: bool = False + qlora: bool = False + rag: bool = False + rag_replace_method: str = "prepend" + nf4: bool = False use_autoencoder: bool = False stream_hf_datasets: bool = False diff --git a/simpletransformers/conv_ai/conv_ai_model.py b/simpletransformers/conv_ai/conv_ai_model.py index b8162f9f..f7a33da7 100644 --- a/simpletransformers/conv_ai/conv_ai_model.py +++ b/simpletransformers/conv_ai/conv_ai_model.py @@ -723,8 +723,7 @@ def train( epoch_number += 1 output_dir_current = os.path.join( output_dir, - "checkpoint-{}-epoch-{}".format(global_step, epoch_number) -, + "checkpoint-{}-epoch-{}".format(global_step, epoch_number), ) if args.save_model_every_epoch or args.evaluate_during_training: diff --git a/simpletransformers/language_modeling/__init__.py b/simpletransformers/language_modeling/__init__.py index bde1ccf1..264f6e5e 100755 --- a/simpletransformers/language_modeling/__init__.py +++ b/simpletransformers/language_modeling/__init__.py @@ -1,4 +1,4 @@ -from simpletransformers.config.model_args import LanguageModelingArgs +from simpletransformers.config.model_args import LanguageModelingArgs, GenerationArgs from simpletransformers.language_modeling.language_modeling_model import ( LanguageModelingModel, ) diff --git a/simpletransformers/language_modeling/language_modeling_model.py b/simpletransformers/language_modeling/language_modeling_model.py index c9cdb307..31dea2f4 100755 --- a/simpletransformers/language_modeling/language_modeling_model.py +++ b/simpletransformers/language_modeling/language_modeling_model.py @@ -49,6 +49,7 @@ from simpletransformers.custom_models.models import RobertaWithAutoEncoderForMaskedLM from transformers import DummyObject, requires_backends +from torch.utils.data import DataLoader, TensorDataset class NystromformerTokenizer(metaclass=DummyObject): @@ -63,6 +64,7 @@ def __init__(self, *args, **kwargs): AutoConfig, AutoModelWithLMHead, AutoTokenizer, + AutoModelForCausalLM, BertConfig, BertForMaskedLM, BertTokenizer, @@ -102,6 +104,7 @@ def __init__(self, *args, **kwargs): XLMRobertaConfig, XLMRobertaForMaskedLM, XLMRobertaTokenizer, + GenerationConfig, ) from transformers.data.datasets.language_modeling import ( LineByLineTextDataset, @@ -109,7 +112,7 @@ def __init__(self, *args, **kwargs): ) from simpletransformers.config.global_args import global_args -from simpletransformers.config.model_args import LanguageModelingArgs +from simpletransformers.config.model_args import LanguageModelingArgs, GenerationArgs from simpletransformers.config.utils import sweep_config_to_sweep_values from simpletransformers.custom_models.models import ElectraForLanguageModelingModel from simpletransformers.language_modeling.language_modeling_utils import ( @@ -132,6 +135,7 @@ def __init__(self, *args, **kwargs): "bert": (BertConfig, BertForMaskedLM, BertTokenizer), "bigbird": (BigBirdConfig, BigBirdForMaskedLM, BigBirdTokenizer), "camembert": (CamembertConfig, CamembertForMaskedLM, CamembertTokenizer), + "causal": (AutoConfig, AutoModelForCausalLM, AutoTokenizer), "distilbert": (DistilBertConfig, DistilBertForMaskedLM, DistilBertTokenizer), "electra": (ElectraConfig, ElectraForLanguageModelingModel, ElectraTokenizer), "gpt2": (GPT2Config, GPT2LMHeadModel, GPT2Tokenizer), @@ -159,7 +163,9 @@ def __init__( train_files=None, args=None, use_cuda=True, - autoencoder_model=None, + retrieval_model=None, + adapter_name=None, + # autoencoder_model=None, cuda_device=-1, **kwargs, ): @@ -173,6 +179,8 @@ def __init__( discriminator_name (optional): A pretrained model name or path to a directory containing an ELECTRA discriminator model. args (optional): Default args will be used if this parameter is not provided. If provided, it should be a dict containing the args that should be changed in the default args. train_files (optional): List of files to be used when training the tokenizer. + rag_corpus (optional): A collection of documents to be used for Retrieval-Augmented Generation. This may + retrieval_model (optional): A pretrained model name or path to a directory containing a retrieval model. This should be preloaded with a knowledge index. use_cuda (optional): Use GPU if available. Setting to False will force model to use CPU only. cuda_device (optional): Specific GPU that should be used. Will use the first available GPU by default. **kwargs (optional): For providing proxies, force_download, resume_download, cache_dir and other options specific to the 'from_pretrained' implementation where this will be supplied. @@ -258,11 +266,17 @@ def __init__( if self.args.config_name: self.config = config_class.from_pretrained( - self.args.config_name, cache_dir=self.args.cache_dir + self.args.config_name, + cache_dir=self.args.cache_dir, + trust_remote_code=self.args.trust_remote_code, + **kwargs, ) elif self.args.model_name and self.args.model_name != "electra": self.config = config_class.from_pretrained( - model_name, cache_dir=self.args.cache_dir, **kwargs + model_name, + cache_dir=self.args.cache_dir, + trust_remote_code=self.args.trust_remote_code, + **kwargs, ) else: self.config = config_class(**self.args.config, **kwargs) @@ -359,12 +373,65 @@ def __init__( ) ) else: - self.model = model_class.from_pretrained( - model_name, - config=self.config, - cache_dir=self.args.cache_dir, - **kwargs, - ) + if self.args.nf4: + from transformers import BitsAndBytesConfig + + + nf4_config = BitsAndBytesConfig( + load_in_4bit=True, + bnb_4bit_quant_type="nf4", + bnb_4bit_use_double_quant=True, + bnb_4bit_compute_dtype=torch.bfloat16, + ) + self.model = model_class.from_pretrained( + model_name, + quantization_config=nf4_config, + trust_remote_code=self.args.trust_remote_code, + ) + else: + self.model = model_class.from_pretrained( + model_name, + config=self.config, + cache_dir=self.args.cache_dir, + **kwargs, + ) + + if self.args.peft: + from peft import LoraConfig, get_peft_model, LoftQConfig + from peft.peft_model import PeftModel + + if self.args.qlora: + if self.args.nf4: + raise ValueError( + "PEFT and QLORA cannot be used together with NF4" + ) + loftq_config = LoftQConfig( + loftq_bits=self.args.loftq_bits, **self.args.loftq_config + ) + self.lora_config = LoraConfig( + init_lora_weights="loftq", + target_modules="all-linear", + loftq_config=loftq_config, + **self.args.lora_config, + ) + self.args.fp16 = False + else: + self.lora_config = LoraConfig( + use_rslora=True, target_modules="all-linear" + ) + self.model.gradient_checkpointing_enable() + self.model.enable_input_require_grads() + if adapter_name is not None: + self.model = PeftModel.from_pretrained( + self.model, + model_id=adapter_name, + adapter_name=adapter_name, + ) + self.adapter_name = adapter_name + else: + self.model = get_peft_model(self.model, self.lora_config) + self.model.print_trainable_parameters() + else: logger.info(" Training language model from scratch") if self.args.model_type == "electra": @@ -432,6 +499,15 @@ def __init__( ) self.args.wandb_project = None + if self.args.rag: + if retrieval_model: + self.retrieval_model = retrieval_model + else: + raise ValueError( + "RAG is enabled but no retrieval model is specified." + " Pass a retrieval model when instantiating the LanguageModelingModel to use RAG." + ) + def train_model( self, train_file, @@ -608,6 +684,15 @@ def collate(examples: List[torch.Tensor]): relative_step=args.adafactor_relative_step, warmup_init=args.adafactor_warmup_init, ) + elif args.optimizer == "Adam8bit": + from bitsandbytes.optim import Adam8bit + + optimizer = Adam8bit( + optimizer_grouped_parameters, + lr=args.learning_rate, + eps=args.adam_epsilon, + betas=args.adam_betas, + ) else: raise ValueError( @@ -769,7 +854,6 @@ def collate(examples: List[torch.Tensor]): # TODO: Move this to _get_inputs_dict and keep the attention masks if self.args.use_hf_datasets: if self.args.stream_hf_datasets: - # BUG: TOKENIZATION IS BUGGED FOR HF DATASETS batch["input_ids"] = torch.stack(batch["input_ids"]) batch["attention_mask"] = torch.stack(batch["attention_mask"]) if self.args.model_type in [ @@ -806,16 +890,22 @@ def collate(examples: List[torch.Tensor]): ) token_type_ids = ( batch["token_type_ids"].to(self.device) - if self.args.use_hf_datasets + if self.args.use_hf_datasets and "token_type_ids" in batch else None ) labels = labels.to(self.device) - inputs_dict = { - "input_ids": inputs, - "attention_mask": attention_mask, - "token_type_ids": token_type_ids, - } + if token_type_ids is None: + inputs_dict = { + "input_ids": inputs, + "attention_mask": attention_mask, + } + else: + inputs_dict = { + "input_ids": inputs, + "attention_mask": attention_mask, + "token_type_ids": token_type_ids, + } if args.fp16: with amp.autocast(): @@ -1100,8 +1190,7 @@ def collate(examples: List[torch.Tensor]): epoch_number += 1 output_dir_current = os.path.join( output_dir, - "checkpoint-{}-epoch-{}".format(global_step, epoch_number) -, + "checkpoint-{}-epoch-{}".format(global_step, epoch_number), ) if args.save_model_every_epoch or args.evaluate_during_training: @@ -1249,7 +1338,13 @@ def collate(examples: List[torch.Tensor]): ) def eval_model( - self, eval_file, output_dir=None, verbose=True, silent=False, **kwargs + self, + eval_file, + output_dir=None, + evaluate_generated_text=False, + verbose=True, + silent=False, + **kwargs, ): """ Evaluates the model on eval_df. Saves results to args.output_dir @@ -1261,26 +1356,30 @@ def eval_model( self._move_model_to_device() - eval_dataset = self.load_and_cache_examples( - eval_file, evaluate=True, verbose=verbose, silent=silent - ) - os.makedirs(output_dir, exist_ok=True) + if evaluate_generated_text: + raise NotImplementedError( + "evaluate_generated_text is not yet implemented for this model type." + ) + else: + eval_dataset = self.load_and_cache_examples( + eval_file, evaluate=True, verbose=verbose, silent=silent + ) + os.makedirs(output_dir, exist_ok=True) - result = self.evaluate( - eval_dataset, output_dir, verbose=verbose, silent=silent, **kwargs - ) - self.results.update(result) + result = self.evaluate( + eval_dataset, output_dir, verbose=verbose, silent=silent, **kwargs + ) + self.results.update(result) - if verbose: - logger.info(self.results) + if verbose: + logger.info(self.results) - return result + return result def evaluate( self, eval_dataset, output_dir, - multi_label=False, prefix="", verbose=True, silent=False, @@ -1331,18 +1430,39 @@ def collate(examples: List[torch.Tensor]): eval_dataloader, disable=args.silent or silent, desc="Running Evaluation" ): if self.args.use_hf_datasets: - batch = batch["input_ids"] + input_ids = batch["input_ids"] inputs, labels = ( - mask_tokens(batch, tokenizer, args) if args.mlm else (batch, batch) + mask_tokens(batch, tokenizer, args) + if args.mlm + else (input_ids, input_ids) ) inputs = inputs.to(self.device) labels = labels.to(self.device) + + if "token_type_ids" in batch: + inputs_dict = { + "input_ids": inputs, + "attention_mask": batch["attention_mask"].to(self.device) + if self.args.use_hf_datasets + else None, + "token_type_ids": batch["token_type_ids"].to(self.device) + if self.args.use_hf_datasets + else None, + } + else: + inputs_dict = { + "input_ids": inputs, + "attention_mask": batch["attention_mask"].to(self.device) + if self.args.use_hf_datasets + else None, + } + with torch.no_grad(): outputs = ( - model(inputs, labels=labels) + model(**inputs_dict, labels=labels) if args.mlm - else model(inputs, labels=labels) + else model(**inputs_dict, labels=labels) ) if args.model_type == "electra": g_loss = outputs[0] @@ -1368,6 +1488,127 @@ def collate(examples: List[torch.Tensor]): return results + def predict( + self, + to_predict, + generation_args=None, + rag_queries=None, + knowledge_dataset=None, + **kwargs, + ): + """ + Performs text completions on a list of text. To be used with language models. + + Args: + + to_predict: A list of text to make predictions on. + generation_args: An instance of the `GenerationArgs` class containing the generation arguments for the model. + rag_queries (optional): A list of text to be used as queries for the RAG model. Only applicable if rag is enabled. + knowledge_dataset (optional): A list of text to be used as knowledge for the RAG model. Only applicable if the model is a RAG model. + **kwargs: Additional arguments to be passed to the models `generate()` method during inference. + + Returns: + preds: A list of the predicted sequences. + """ + self._move_model_to_device() + + if not generation_args: + generation_args = GenerationArgs() + + if self.args.peft and self.adapter_name: + logger.info( + "Merging adapter with model for faster inference. Contunuing training from this point may result in unexpected behavior." + ) + self.model = self.model.merge_and_unload() + + self.tokenizer.padding_side = "left" + + if self.args.rag: + if not rag_queries: + rag_queries = to_predict + raise Warning( + "No `rag_queries` provided. Using `to_predict` as `rag_queries`." + ) + + context_docs = self.retrieval_model.predict( + rag_queries, + passages_only=True, + prediction_passages=knowledge_dataset, + ) + + to_predict = [ + f"Context: {' '.join(context_doc)} {text}" + for context_doc, text in zip(context_docs, to_predict) + ] + # TODO: + # - Simplest option is to just prepend context: context_docs to to_predict + # - Advanced option is to have