#!/usr/bin/env python # coding=utf-8 # Copyright 2021 The HuggingFace Team. All rights reserved. # # 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. """ Fine-tuning the library models for sequence to sequence for P-Tuning v2 """ # You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments. # CUDA_VISIBLE_DEVICES=-1 python finetune-p-tuning-v2.py # accelerate launch --cpu --num_machines=1 --num_processes=4 --num_cpu_threads_per_process=1 finetune-p-tuning-v2.py # accelerate launch --num_machines=1 --num_processes=1 --gpu_ids=1 finetune-p-tuning-v2.py # import logging import os import sys import json import numpy as np from datasets import load_dataset import jieba from rouge_chinese import Rouge from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction import torch # import transformers from transformers import ( AutoConfig, AutoModel, AutoTokenizer, DataCollatorForSeq2Seq, # HfArgumentParser, # Seq2SeqTrainingArguments, # set_seed, ) # from typing import Any, Dict, List, Optional, Tuple, Union # import torch # from torch import nn # from torch.utils.data import Dataset # from transformers.deepspeed import is_deepspeed_zero3_enabled # from trainer import PrefixTrainer # from transformers.trainer_utils import PredictionOutput # from transformers.utils import logging # import os # from typing import Optional from transformers import Trainer # import torch # from transformers.modeling_utils import PreTrainedModel, unwrap_model # from transformers.utils import logging # from trainer_seq2seq import Seq2SeqTrainer # from arguments import ModelArguments, DataTrainingArguments # logger = logging.getLogger(__name__) # logger.setLevel(logging.INFO) def main(): # print(torch.backends.mps.is_available()) # print(torch.backends.mps.is_built()) device = torch.device("cpu") if torch.cuda.is_available(): device = torch.device("cuda") elif torch.backends.mps.is_available(): device = torch.device("mps") else: device = torch.device("cpu") print("device:", device) # parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments)) # if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): # # If we pass only one argument to the script and it's the path to a json file, # # let's parse it to get our arguments. # model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1])) # else: # model_args, data_args, training_args = parser.parse_args_into_dataclasses() # Setup logging # logging.basicConfig( # format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", # datefmt="%m/%d/%Y %H:%M:%S", # handlers=[logging.StreamHandler(sys.stdout)], # ) # if training_args.should_log: # # The default of training_args.log_level is passive, so we set log level at info here to have that default. # transformers.utils.logging.set_verbosity_info() # log_level = training_args.get_process_log_level() # logger.setLevel(log_level) # datasets.utils.logging.set_verbosity(log_level) # transformers.utils.logging.set_verbosity(log_level) # transformers.utils.logging.enable_default_handler() # transformers.utils.logging.enable_explicit_format() # Log on each process the small summary: # logger.warning( # f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}" # + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}" # ) # logger.info(f"Training/evaluation parameters {training_args}") # Set seed before initializing model. # set_seed(training_args.seed) # Load dataset data_files_path = "/tmp/hub/dataset/shibing624/AdvertiseGen" print('data_files_path:', data_files_path) data_files = {} data_files["train"] = "train.json" data_files["validation"] = "dev.json" # data_files["test"] = "test.json" print('data_files:', data_files) raw_datasets = load_dataset(data_files_path, data_files=data_files) print("raw_datasets:", raw_datasets) # Load pretrained model and tokenizer model_name_or_path = '/tmp/hub/chatglm2-6b' print('model_name_or_path', model_name_or_path) config = AutoConfig.from_pretrained(model_name_or_path, trust_remote_code=True) print('load autoconfig done') # soft prompt 长度 PRE_SEQ_LEN=128 config.pre_seq_len = PRE_SEQ_LEN config.prefix_projection = False tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, trust_remote_code=True) print('load AutoTokenizer done') # if model_args.ptuning_checkpoint is not None: # # Evaluation # # Loading extra state dict of prefix encoder # model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True) # prefix_state_dict = torch.load(os.path.join(model_args.ptuning_checkpoint, "pytorch_model.bin")) # new_prefix_state_dict = {} # for k, v in prefix_state_dict.items(): # if k.startswith("transformer.prefix_encoder."): # new_prefix_state_dict[k[len("transformer.prefix_encoder."):]] = v # model.transformer.prefix_encoder.load_state_dict(new_prefix_state_dict) # else: # model = AutoModel.from_pretrained(model_args.model_name_or_path, config=config, trust_remote_code=True) model = AutoModel.from_pretrained(model_name_or_path, config=config, trust_remote_code=True) #print('load AutoModel done') # model = model.quantize(4) # if model_args.quantization_bit is not None: # print(f"Quantized to {model_args.quantization_bit} bit") # model = model.quantize(model_args.quantization_bit) # if model_args.pre_seq_len is not None: # # P-tuning v2 # model = model.half() # model.transformer.prefix_encoder.float() # else: # # Finetune # model = model.float() # P-tuning v2, do not work for accelerate model = model.half() model.transformer.prefix_encoder.float() # finetune, work for accelerate # model = model.float() print('model half done') prefix = "" # Preprocessing the datasets. # We need to tokenize inputs and targets. # if training_args.do_train: # column_names = raw_datasets["train"].column_names # elif training_args.do_eval: # column_names = raw_datasets["validation"].column_names # elif training_args.do_predict: # column_names = raw_datasets["test"].column_names # else: # logger.info("There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`.") # return column_names = raw_datasets["train"].column_names # Get the column names for input/target. prompt_column = 'content' response_column = 'summary' history_column = None # Temporarily set max_target_length for training. max_source_length = 64 max_target_length = 128 ignore_pad_token_for_loss = True def preprocess_function_eval(examples): inputs, targets = [], [] for i in range(len(examples[prompt_column])): if examples[prompt_column][i] and examples[response_column][i]: query = examples[prompt_column][i] history = examples[history_column][i] if history_column is not None else None prompt = tokenizer.build_prompt(query, history) inputs.append(prompt) targets.append(examples[response_column][i]) inputs = [prefix + inp for inp in inputs] model_inputs = tokenizer(inputs, max_length=max_source_length, truncation=True, padding=True) labels = tokenizer(text_target=targets, max_length=max_target_length, truncation=True) if ignore_pad_token_for_loss: labels["input_ids"] = [[(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]] model_inputs["labels"] = labels["input_ids"] return model_inputs def preprocess_function_train(examples): max_seq_length = max_source_length + max_target_length + 1 model_inputs = { "input_ids": [], "labels": [] } for i in range(len(examples[prompt_column])): if examples[prompt_column][i] and examples[response_column][i]: query, answer = examples[prompt_column][i], examples[response_column][i] history = examples[history_column][i] if history_column is not None else None prompt = tokenizer.build_prompt(query, history) prompt = prefix + prompt a_ids = tokenizer.encode(text=prompt, add_special_tokens=True, truncation=True, max_length=max_source_length) b_ids = tokenizer.encode(text=answer, add_special_tokens=False, truncation=True, max_length=max_target_length) context_length = len(a_ids) input_ids = a_ids + b_ids + [tokenizer.eos_token_id] labels = [tokenizer.pad_token_id] * context_length + b_ids + [tokenizer.eos_token_id] pad_len = max_seq_length - len(input_ids) input_ids = input_ids + [tokenizer.pad_token_id] * pad_len labels = labels + [tokenizer.pad_token_id] * pad_len if ignore_pad_token_for_loss: labels = [(l if l != tokenizer.pad_token_id else -100) for l in labels] model_inputs["input_ids"].append(input_ids) model_inputs["labels"].append(labels) return model_inputs def print_dataset_example(example): print('******************* print_dataset_example ******************************') print("input_ids:", example["input_ids"]) print("inputs:", tokenizer.decode(example["input_ids"])) print("label_ids:", example["labels"]) print("labels:", tokenizer.decode(example["labels"])) max_train_samples = 5 do_train = True if do_train: train_dataset = raw_datasets["train"] max_train_samples = min(len(train_dataset), max_train_samples) train_dataset = train_dataset.select(range(max_train_samples)) train_dataset = train_dataset.map( preprocess_function_train, batched=True, num_proc=5, remove_columns=column_names, load_from_cache_file=False, desc="Running tokenizer on train dataset", ) # print_dataset_example(train_dataset[0]) max_eval_samples = 5 do_eval = True if do_eval: eval_dataset = raw_datasets["validation"] max_eval_samples = min(len(eval_dataset), max_eval_samples) eval_dataset = eval_dataset.select(range(max_eval_samples)) eval_dataset = eval_dataset.map( preprocess_function_eval, batched=True, num_proc=5, remove_columns=column_names, load_from_cache_file=False, desc="Running tokenizer on validation dataset", ) # print_dataset_example(eval_dataset[0]) # if training_args.do_predict: # max_target_length = data_args.val_max_target_length # if "test" not in raw_datasets: # raise ValueError("--do_predict requires a test dataset") # predict_dataset = raw_datasets["test"] # if data_args.max_predict_samples is not None: # max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples) # predict_dataset = predict_dataset.select(range(max_predict_samples)) # with training_args.main_process_first(desc="prediction dataset map pre-processing"): # predict_dataset = predict_dataset.map( # preprocess_function_eval, # batched=True, # num_proc=data_args.preprocessing_num_workers, # remove_columns=column_names, # load_from_cache_file=not data_args.overwrite_cache, # desc="Running tokenizer on prediction dataset", # ) # print_dataset_example(predict_dataset[0]) # Data collator label_pad_token_id = -100 if ignore_pad_token_for_loss else tokenizer.pad_token_id data_collator = DataCollatorForSeq2Seq( tokenizer, model=model, label_pad_token_id=label_pad_token_id, pad_to_multiple_of=None, padding=False ) print("data_collator done") # Metric def compute_metrics(eval_preds): preds, labels = eval_preds if isinstance(preds, tuple): preds = preds[0] decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True) if ignore_pad_token_for_loss: # Replace -100 in the labels as we can't decode them. labels = np.where(labels != -100, labels, tokenizer.pad_token_id) decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True) score_dict = { "rouge-1": [], "rouge-2": [], "rouge-l": [], "bleu-4": [] } for pred, label in zip(decoded_preds, decoded_labels): hypothesis = list(jieba.cut(pred)) reference = list(jieba.cut(label)) rouge = Rouge() scores = rouge.get_scores(' '.join(hypothesis) , ' '.join(reference)) result = scores[0] for k, v in result.items(): score_dict[k].append(round(v["f"] * 100, 4)) bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3) score_dict["bleu-4"].append(round(bleu_score * 100, 4)) for k, v in score_dict.items(): score_dict[k] = float(np.mean(v)) return score_dict # Override the decoding parameters of Seq2SeqTrainer # training_args.generation_max_length = ( # training_args.generation_max_length # if training_args.generation_max_length is not None # else data_args.val_max_target_length # ) # training_args.generation_num_beams = ( # data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams # ) # Initialize our Trainer # trainer = Seq2SeqTrainer( # model=model, # # args=training_args, # train_dataset=train_dataset, # eval_dataset=eval_dataset, # tokenizer=tokenizer, # data_collator=data_collator, # compute_metrics=compute_metrics, # save_changed=PRE_SEQ_LEN is not None # ) trainer = Trainer( model, train_dataset=train_dataset, eval_dataset=eval_dataset, tokenizer=tokenizer, data_collator=data_collator, compute_metrics=compute_metrics, ) print('build trainer done') # Training if do_train: # checkpoint = False # if training_args.resume_from_checkpoint is not None: # checkpoint = training_args.resume_from_checkpoint # elif last_checkpoint is not None: # checkpoint = last_checkpoint model.gradient_checkpointing_enable() model.enable_input_require_grads() print("begin trainning") # train_result = trainer.train(resume_from_checkpoint=checkpoint) train_result = trainer.train() # trainer.save_model() # Saves the tokenizer too for easy upload print("done trainning") metrics = train_result.metrics max_train_samples = len(train_dataset) metrics["train_samples"] = min(max_train_samples, len(train_dataset)) trainer.log_metrics("train", metrics) trainer.save_metrics("train", metrics) trainer.save_state() print("save state") # trainer.save_model("tmp_trainer/ptuning") print("save model") # # Evaluation # results = {} # max_seq_length = data_args.max_source_length + data_args.max_target_length + 1 # if training_args.do_eval: # logger.info("*** Evaluate ***") # metrics = trainer.evaluate(metric_key_prefix="eval", do_sample=True, top_p=0.7, max_length=max_seq_length, temperature=0.95) # max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset) # metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset)) # trainer.log_metrics("eval", metrics) # trainer.save_metrics("eval", metrics) # if training_args.do_predict: # logger.info("*** Predict ***") # predict_results = trainer.predict(predict_dataset, metric_key_prefix="predict", max_length=max_seq_length, do_sample=True, top_p=0.7, temperature=0.95) # metrics = predict_results.metrics # max_predict_samples = ( # data_args.max_predict_samples if data_args.max_predict_samples is not None else len(predict_dataset) # ) # metrics["predict_samples"] = min(max_predict_samples, len(predict_dataset)) # trainer.log_metrics("predict", metrics) # trainer.save_metrics("predict", metrics) # if trainer.is_world_process_zero(): # if training_args.predict_with_generate: # predictions = tokenizer.batch_decode( # predict_results.predictions, skip_special_tokens=True, clean_up_tokenization_spaces=True # ) # predictions = [pred.strip() for pred in predictions] # labels = tokenizer.batch_decode( # predict_results.label_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True # ) # labels = [label.strip() for label in labels] # output_prediction_file = os.path.join(training_args.output_dir, "generated_predictions.txt") # with open(output_prediction_file, "w", encoding="utf-8") as writer: # for p, l in zip(predictions, labels): # res = json.dumps({"labels": l, "predict": p}, ensure_ascii=False) # writer.write(f"{res}\n") # return results # WEIGHTS_NAME = "pytorch_model.bin" # TRAINING_ARGS_NAME = "training_args.bin" # class PrefixTrainer(Trainer): # def __init__(self, *args, save_changed=False, **kwargs): # self.save_changed = save_changed # super().__init__(*args, **kwargs) # def _save(self, output_dir: Optional[str] = None, state_dict=None): # # If we are executing this function, we are the process zero, so we don't check for that. # output_dir = output_dir if output_dir is not None else self.args.output_dir # os.makedirs(output_dir, exist_ok=True) # logger.info(f"Saving model checkpoint to {output_dir}") # # Save a trained model and configuration using `save_pretrained()`. # # They can then be reloaded using `from_pretrained()` # if not isinstance(self.model, PreTrainedModel): # if isinstance(unwrap_model(self.model), PreTrainedModel): # if state_dict is None: # state_dict = self.model.state_dict() # unwrap_model(self.model).save_pretrained(output_dir, state_dict=state_dict) # else: # logger.info("Trainer.model is not a `PreTrainedModel`, only saving its state dict.") # if state_dict is None: # state_dict = self.model.state_dict() # torch.save(state_dict, os.path.join(output_dir, WEIGHTS_NAME)) # else: # if self.save_changed: # print("Saving PrefixEncoder") # state_dict = self.model.state_dict() # filtered_state_dict = {} # for k, v in self.model.named_parameters(): # if v.requires_grad: # filtered_state_dict[k] = state_dict[k] # self.model.save_pretrained(output_dir, state_dict=filtered_state_dict) # else: # print("Saving the whole model") # self.model.save_pretrained(output_dir, state_dict=state_dict) # if self.tokenizer is not None: # self.tokenizer.save_pretrained(output_dir) # # Good practice: save your training arguments together with the trained model # torch.save(self.args, os.path.join(output_dir, TRAINING_ARGS_NAME)) # class Seq2SeqTrainer(PrefixTrainer): # def evaluate( # self, # eval_dataset: Optional[Dataset] = None, # ignore_keys: Optional[List[str]] = None, # metric_key_prefix: str = "eval", # **gen_kwargs # ) -> Dict[str, float]: # """ # Run evaluation and returns metrics. # The calling script will be responsible for providing a method to compute metrics, as they are task-dependent # (pass it to the init `compute_metrics` argument). # You can also subclass and override this method to inject custom behavior. # Args: # eval_dataset (`Dataset`, *optional*): # Pass a dataset if you wish to override `self.eval_dataset`. If it is an [`~datasets.Dataset`], columns # not accepted by the `model.forward()` method are automatically removed. It must implement the `__len__` # method. # ignore_keys (`List[str]`, *optional*): # A list of keys in the output of your model (if it is a dictionary) that should be ignored when # gathering predictions. # metric_key_prefix (`str`, *optional*, defaults to `"eval"`): # An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named # "eval_bleu" if the prefix is `"eval"` (default) # max_length (`int`, *optional*): # The maximum target length to use when predicting with the generate method. # num_beams (`int`, *optional*): # Number of beams for beam search that will be used when predicting with the generate method. 1 means no # beam search. # gen_kwargs: # Additional `generate` specific kwargs. # Returns: # A dictionary containing the evaluation loss and the potential metrics computed from the predictions. The # dictionary also contains the epoch number which comes from the training state. # """ # gen_kwargs = gen_kwargs.copy() # if gen_kwargs.get("max_length") is None and gen_kwargs.get("max_new_tokens") is None: # gen_kwargs["max_length"] = self.args.generation_max_length # gen_kwargs["num_beams"] = ( # gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.args.generation_num_beams # ) # self._gen_kwargs = gen_kwargs # return super().evaluate(eval_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix) # def predict( # self, # test_dataset: Dataset, # ignore_keys: Optional[List[str]] = None, # metric_key_prefix: str = "test", # **gen_kwargs # ) -> PredictionOutput: # """ # Run prediction and returns predictions and potential metrics. # Depending on the dataset and your use case, your test dataset may contain labels. In that case, this method # will also return metrics, like in `evaluate()`. # Args: # test_dataset (`Dataset`): # Dataset to run the predictions on. If it is a [`~datasets.Dataset`], columns not accepted by the # `model.forward()` method are automatically removed. Has to implement the method `__len__` # ignore_keys (`List[str]`, *optional*): # A list of keys in the output of your model (if it is a dictionary) that should be ignored when # gathering predictions. # metric_key_prefix (`str`, *optional*, defaults to `"eval"`): # An optional prefix to be used as the metrics key prefix. For example the metrics "bleu" will be named # "eval_bleu" if the prefix is `"eval"` (default) # max_length (`int`, *optional*): # The maximum target length to use when predicting with the generate method. # num_beams (`int`, *optional*): # Number of beams for beam search that will be used when predicting with the generate method. 1 means no # beam search. # gen_kwargs: # Additional `generate` specific kwargs. # # If your predictions or labels have different sequence lengths (for instance because you're doing dynamic # padding in a token classification task) the predictions will be padded (on the right) to allow for # concatenation into one array. The padding index is -100. # # Returns: *NamedTuple* A namedtuple with the following keys: # - predictions (`np.ndarray`): The predictions on `test_dataset`. # - label_ids (`np.ndarray`, *optional*): The labels (if the dataset contained some). # - metrics (`Dict[str, float]`, *optional*): The potential dictionary of metrics (if the dataset contained # labels). # """ # gen_kwargs = gen_kwargs.copy() # if gen_kwargs.get("max_length") is None and gen_kwargs.get("max_new_tokens") is None: # gen_kwargs["max_length"] = self.args.generation_max_length # gen_kwargs["num_beams"] = ( # gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.args.generation_num_beams # ) # self._gen_kwargs = gen_kwargs # return super().predict(test_dataset, ignore_keys=ignore_keys, metric_key_prefix=metric_key_prefix) # def prediction_step( # self, # model: nn.Module, # inputs: Dict[str, Union[torch.Tensor, Any]], # prediction_loss_only: bool, # ignore_keys: Optional[List[str]] = None, # ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: # """ # Perform an evaluation step on `model` using `inputs`. # Subclass and override to inject custom behavior. # Args: # model (`nn.Module`): # The model to evaluate. # inputs (`Dict[str, Union[torch.Tensor, Any]]`): # The inputs and targets of the model. # The dictionary will be unpacked before being fed to the model. Most models expect the targets under the # argument `labels`. Check your model's documentation for all accepted arguments. # prediction_loss_only (`bool`): # Whether or not to return the loss only. # Return: # Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]: A tuple with the loss, logits and # labels (each being optional). # """ # if not self.args.predict_with_generate or prediction_loss_only: # return super().prediction_step( # model, inputs, prediction_loss_only=prediction_loss_only, ignore_keys=ignore_keys # ) # has_labels = "labels" in inputs # inputs = self._prepare_inputs(inputs) # # XXX: adapt synced_gpus for fairscale as well # gen_kwargs = self._gen_kwargs.copy() # if gen_kwargs.get("max_length") is None and gen_kwargs.get("max_new_tokens") is None: # gen_kwargs["max_length"] = self.model.config.max_length # gen_kwargs["num_beams"] = ( # gen_kwargs["num_beams"] if gen_kwargs.get("num_beams") is not None else self.model.config.num_beams # ) # default_synced_gpus = True if is_deepspeed_zero3_enabled() else False # gen_kwargs["synced_gpus"] = ( # gen_kwargs["synced_gpus"] if gen_kwargs.get("synced_gpus") is not None else default_synced_gpus # ) # if "attention_mask" in inputs: # gen_kwargs["attention_mask"] = inputs.get("attention_mask", None) # if "position_ids" in inputs: # gen_kwargs["position_ids"] = inputs.get("position_ids", None) # if "global_attention_mask" in inputs: # gen_kwargs["global_attention_mask"] = inputs.get("global_attention_mask", None) # # prepare generation inputs # # some encoder-decoder models can have varying encoder's and thus # # varying model input names # if hasattr(self.model, "encoder") and self.model.encoder.main_input_name != self.model.main_input_name: # generation_inputs = inputs[self.model.encoder.main_input_name] # else: # generation_inputs = inputs[self.model.main_input_name] # gen_kwargs["input_ids"] = generation_inputs # generated_tokens = self.model.generate(**gen_kwargs) # generated_tokens = generated_tokens[:, generation_inputs.size()[-1]:] # # in case the batch is shorter than max length, the output should be padded # if gen_kwargs.get("max_length") is not None and generated_tokens.shape[-1] < gen_kwargs["max_length"]: # generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs["max_length"]) # elif gen_kwargs.get("max_new_tokens") is not None and generated_tokens.shape[-1] < ( # gen_kwargs["max_new_tokens"] + 1 # ): # generated_tokens = self._pad_tensors_to_max_len(generated_tokens, gen_kwargs["max_new_tokens"] + 1) # loss = None # if self.args.prediction_loss_only: # return (loss, None, None) # if has_labels: # labels = inputs["labels"] # if gen_kwargs.get("max_length") is not None and labels.shape[-1] < gen_kwargs["max_length"]: # labels = self._pad_tensors_to_max_len(labels, gen_kwargs["max_length"]) # elif gen_kwargs.get("max_new_tokens") is not None and labels.shape[-1] < ( # gen_kwargs["max_new_tokens"] + 1 # ): # labels = self._pad_tensors_to_max_len(labels, (gen_kwargs["max_new_tokens"] + 1)) # else: # labels = None # return (loss, generated_tokens, labels) # def _pad_tensors_to_max_len(self, tensor, max_length): # if self.tokenizer is not None and hasattr(self.tokenizer, "pad_token_id"): # # If PAD token is not defined at least EOS token has to be defined # pad_token_id = ( # self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id # ) # else: # if self.model.config.pad_token_id is not None: # pad_token_id = self.model.config.pad_token_id # else: # raise ValueError("Pad_token_id must be set in the configuration of the model, in order to pad tensors") # padded_tensor = pad_token_id * torch.ones( # (tensor.shape[0], max_length), dtype=tensor.dtype, device=tensor.device # ) # padded_tensor[:, : tensor.shape[-1]] = tensor # return padded_tensor if __name__ == "__main__": main()