diff --git a/convlab2/base_models/t5/create_data.py b/convlab2/base_models/t5/create_data.py
new file mode 100644
index 0000000000000000000000000000000000000000..01b353f20890657236eef8c9faf8f4fa46054fb3
--- /dev/null
+++ b/convlab2/base_models/t5/create_data.py
@@ -0,0 +1,36 @@
+import os
+import json
+from tqdm import tqdm
+from convlab2.util import load_dataset, load_nlu_data, load_dst_data, load_policy_data, load_nlg_data, load_e2e_data, load_rg_data
+
+def create_rg_data(data_by_split, data_dir):
+ data_splits = data_by_split.keys()
+ file_name = os.path.join(data_dir, f"source_prefix.txt")
+ with open(file_name, "w") as f:
+ f.write("generate a system response according to the context: ")
+ for data_split in data_splits:
+ data = []
+ for sample in tqdm(data_by_split[data_split], desc=f'{data_split} sample', leave=False):
+ context = ' '.join([f"{turn['speaker']}: {turn['utterance']}" for turn in sample['context']])
+ response = f"{sample['speaker']}: {sample['utterance']}"
+ data.append(json.dumps({'context': context, 'response': response})+'\n')
+
+ file_name = os.path.join(data_dir, f"{data_split}.json")
+ with open(file_name, "w") as f:
+ f.writelines(data)
+
+if __name__ == '__main__':
+ from argparse import ArgumentParser
+ parser = ArgumentParser(description="create data for seq2seq training")
+ parser.add_argument('--tasks', metavar='task_name', nargs='*', choices=['rg'], help='names of tasks')
+ parser.add_argument('--datasets', metavar='dataset_name', nargs='*', help='names of unified datasets')
+ parser.add_argument('--save_dir', metavar='save_directory', type=str, default='data', help='directory to save the data, default: data/$task_name/$dataset_name')
+ args = parser.parse_args()
+ print(args)
+ for dataset_name in tqdm(args.datasets, desc='datasets'):
+ dataset = load_dataset(dataset_name)
+ for task_name in tqdm(args.tasks, desc='tasks', leave=False):
+ data_by_split = eval(f"load_{task_name}_data")(dataset)
+ data_dir = os.path.join(args.save_dir, task_name, dataset_name)
+ os.makedirs(data_dir, exist_ok=True)
+ eval(f"create_{task_name}_data")(data_by_split, data_dir)
\ No newline at end of file
diff --git a/convlab2/base_models/t5/run_rg.sh b/convlab2/base_models/t5/run_rg.sh
new file mode 100644
index 0000000000000000000000000000000000000000..60c88cfb385a1b3babc67c9f80f1f7fb663e0456
--- /dev/null
+++ b/convlab2/base_models/t5/run_rg.sh
@@ -0,0 +1,50 @@
+n_gpus=8
+task_name="rg"
+dataset_name="multiwoz21"
+data_dir="data/${task_name}/${dataset_name}"
+output_dir="output/${task_name}/${dataset_name}"
+cache_dir="cache"
+logging_dir="${output_dir}/runs"
+train_file="${data_dir}/train.json"
+validation_file="${data_dir}/validation.json"
+test_file="${data_dir}/test.json"
+source_prefix="${data_dir}/source_prefix.txt"
+source_column="context"
+target_column="response"
+model_name_or_path="t5-small"
+per_device_train_batch_size=32
+per_device_eval_batch_size=128
+gradient_accumulation_steps=1
+lr=1e-3
+num_train_epochs=5
+
+python -m torch.distributed.launch \
+ --nproc_per_node ${n_gpus} run_seq2seq.py \
+ --task_name ${task_name} \
+ --train_file ${train_file} \
+ --validation_file ${validation_file} \
+ --test_file ${test_file} \
+ --source_column ${source_column} \
+ --target_column ${target_column} \
+ --source_prefix ${source_prefix} \
+ --model_name_or_path ${model_name_or_path} \
+ --do_train \
+ --do_eval \
+ --do_predict \
+ --save_strategy epoch \
+ --evaluation_strategy epoch \
+ --load_best_model_at_end \
+ --predict_with_generate \
+ --cache_dir ${cache_dir} \
+ --output_dir ${output_dir} \
+ --logging_dir ${logging_dir} \
+ --overwrite_output_dir \
+ --preprocessing_num_workers 4 \
+ --per_device_train_batch_size ${per_device_train_batch_size} \
+ --per_device_eval_batch_size ${per_device_eval_batch_size} \
+ --gradient_accumulation_steps ${gradient_accumulation_steps} \
+ --learning_rate ${lr} \
+ --num_train_epochs ${num_train_epochs} \
+ --debug underflow_overflow \
+ --adafactor \
+ --gradient_checkpointing
diff --git a/convlab2/base_models/t5/run_seq2seq.py b/convlab2/base_models/t5/run_seq2seq.py
new file mode 100644
index 0000000000000000000000000000000000000000..30e2b1d5e840292b030a626eacc8ac3628b89f5f
--- /dev/null
+++ b/convlab2/base_models/t5/run_seq2seq.py
@@ -0,0 +1,639 @@
+#!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.
+Modified from https://github.com/huggingface/transformers/blob/master/examples/pytorch/summarization/run_summarization.py
+"""
+# You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments.
+
+import logging
+import os
+import sys
+import json
+from dataclasses import dataclass, field
+from typing import Optional
+
+import datasets
+import nltk # Here to have a nice missing dependency error message early on
+import numpy as np
+from datasets import load_dataset, load_metric
+
+import transformers
+from filelock import FileLock
+from transformers import (
+ AutoConfig,
+ AutoModelForSeq2SeqLM,
+ AutoTokenizer,
+ DataCollatorForSeq2Seq,
+ HfArgumentParser,
+ Seq2SeqTrainer,
+ Seq2SeqTrainingArguments,
+ set_seed,
+)
+from transformers.file_utils import is_offline_mode
+from transformers.trainer_utils import EvalPrediction, get_last_checkpoint
+from transformers.utils import check_min_version
+from transformers.utils.versions import require_version
+
+
+# Will error if the minimal version of Transformers is not installed. Remove at your own risks.
+# check_min_version("4.16.0.dev0")
+
+# require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/summarization/requirements.txt")
+
+logger = logging.getLogger(__name__)
+os.environ["WANDB_DISABLED"] = "true"
+
+
+@dataclass
+class ModelArguments:
+ """
+ Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+ """
+
+ model_name_or_path: str = field(
+ metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+ )
+ config_name: Optional[str] = field(
+ default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+ )
+ tokenizer_name: Optional[str] = field(
+ default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+ )
+ cache_dir: Optional[str] = field(
+ default=None,
+ metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
+ )
+ use_fast_tokenizer: bool = field(
+ default=True,
+ metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+ )
+ model_revision: str = field(
+ default="main",
+ metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+ )
+ use_auth_token: bool = field(
+ default=False,
+ metadata={
+ "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script "
+ "with private models)."
+ },
+ )
+ resize_position_embeddings: Optional[bool] = field(
+ default=None,
+ metadata={
+ "help": "Whether to automatically resize the position embeddings if `max_source_length` exceeds "
+ "the model's position embeddings."
+ },
+ )
+
+
+@dataclass
+class DataTrainingArguments:
+ """
+ Arguments pertaining to what data we are going to input our model for training and eval.
+ """
+
+
+ task_name: Optional[str] = field(
+ default=None, metadata={"help": "The name of the task, e.g., response generation."}
+ )
+ dataset_name: Optional[str] = field(
+ default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+ )
+ dataset_config_name: Optional[str] = field(
+ default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+ )
+ source_column: Optional[str] = field(
+ default=None,
+ metadata={"help": "The name of the column in the datasets containing the source texts."},
+ )
+ target_column: Optional[str] = field(
+ default=None,
+ metadata={"help": "The name of the column in the datasets containing the target texts."},
+ )
+ train_file: Optional[str] = field(
+ default=None, metadata={"help": "The input training data file (a jsonlines or csv file)."}
+ )
+ validation_file: Optional[str] = field(
+ default=None,
+ metadata={
+ "help": "An optional input evaluation data file to evaluate the metrics on (a jsonlines or csv file)."
+ },
+ )
+ test_file: Optional[str] = field(
+ default=None,
+ metadata={
+ "help": "An optional input test data file to evaluate the metrics on (a jsonlines or csv file)."
+ },
+ )
+ overwrite_cache: bool = field(
+ default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+ )
+ preprocessing_num_workers: Optional[int] = field(
+ default=None,
+ metadata={"help": "The number of processes to use for the preprocessing."},
+ )
+ max_source_length: Optional[int] = field(
+ default=1024,
+ metadata={
+ "help": "The maximum total input sequence length after tokenization. Sequences longer "
+ "than this will be truncated, sequences shorter will be padded."
+ },
+ )
+ max_target_length: Optional[int] = field(
+ default=128,
+ metadata={
+ "help": "The maximum total sequence length for target text after tokenization. Sequences longer "
+ "than this will be truncated, sequences shorter will be padded."
+ },
+ )
+ val_max_target_length: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": "The maximum total sequence length for validation target text after tokenization. Sequences longer "
+ "than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`."
+ "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
+ "during ``evaluate`` and ``predict``."
+ },
+ )
+ pad_to_max_length: bool = field(
+ default=False,
+ metadata={
+ "help": "Whether to pad all samples to model maximum sentence length. "
+ "If False, will pad the samples dynamically when batching to the maximum length in the batch. More "
+ "efficient on GPU but very bad for TPU."
+ },
+ )
+ max_train_samples: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
+ "value if set."
+ },
+ )
+ max_eval_samples: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": "For debugging purposes or quicker training, truncate the number of evaluation examples to this "
+ "value if set."
+ },
+ )
+ max_predict_samples: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": "For debugging purposes or quicker training, truncate the number of prediction examples to this "
+ "value if set."
+ },
+ )
+ num_beams: Optional[int] = field(
+ default=None,
+ metadata={
+ "help": "Number of beams to use for evaluation. This argument will be passed to ``model.generate``, "
+ "which is used during ``evaluate`` and ``predict``."
+ },
+ )
+ ignore_pad_token_for_loss: bool = field(
+ default=True,
+ metadata={
+ "help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."
+ },
+ )
+ source_prefix_filepath: Optional[str] = field(
+ default=None, metadata={"help": "A file whose first line is the prefix to add before every source text (useful for T5 models)."}
+ )
+
+ def __post_init__(self):
+ if (
+ self.dataset_name is None
+ and self.train_file is None
+ and self.validation_file is None
+ and self.test_file is None
+ ):
+ raise ValueError("Need either a dataset name or a training/validation file/test_file.")
+ else:
+ if self.train_file is not None:
+ extension = self.train_file.split(".")[-1]
+ assert extension in ["csv", "json"], "`train_file` should be a csv or a json file."
+ if self.validation_file is not None:
+ extension = self.validation_file.split(".")[-1]
+ assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file."
+ if self.test_file is not None:
+ extension = self.test_file.split(".")[-1]
+ assert extension in ["csv", "json"], "`test_file` should be a csv or a json file."
+ if self.val_max_target_length is None:
+ self.val_max_target_length = self.max_target_length
+
+
+
+
+def main():
+ # See all possible arguments in src/transformers/training_args.py
+ # or by passing the --help flag to this script.
+ # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+ 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)],
+ )
+ 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}")
+
+ if data_args.source_prefix_filepath is None and model_args.model_name_or_path in [
+ "t5-small",
+ "t5-base",
+ "t5-large",
+ "t5-3b",
+ "t5-11b",
+ ]:
+ logger.warning(
+ "You're running a t5 model but didn't provide a source prefix, which is the expected, e.g. with "
+ "`--source_prefix_filepath 'path_to_prefix_file' ` whose first line is the source prefix"
+ )
+
+ # Detecting last checkpoint.
+ last_checkpoint = None
+ if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+ last_checkpoint = get_last_checkpoint(training_args.output_dir)
+ if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+ raise ValueError(
+ f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+ "Use --overwrite_output_dir to overcome."
+ )
+ elif last_checkpoint is not None and training_args.resume_from_checkpoint is None:
+ logger.info(
+ f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+ "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+ )
+
+ # Set seed before initializing model.
+ set_seed(training_args.seed)
+
+ # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below)
+ # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+ # (the dataset will be downloaded automatically from the datasets Hub).
+ #
+ # For CSV/JSON files this script will use the first column for the source texts and the second column for the
+ # target texts (unless you specify column names for this with the `source_column` and `target_column` arguments).
+ #
+ # In distributed training, the load_dataset function guarantee that only one local process can concurrently
+ # download the dataset.
+ if data_args.dataset_name is not None:
+ # Downloading and loading a dataset from the hub.
+ raw_datasets = load_dataset(
+ data_args.dataset_name, data_args.dataset_config_name, cache_dir=model_args.cache_dir
+ )
+ else:
+ data_files = {}
+ if data_args.train_file is not None:
+ data_files["train"] = data_args.train_file
+ extension = data_args.train_file.split(".")[-1]
+ if data_args.validation_file is not None:
+ data_files["validation"] = data_args.validation_file
+ extension = data_args.validation_file.split(".")[-1]
+ if data_args.test_file is not None:
+ data_files["test"] = data_args.test_file
+ extension = data_args.test_file.split(".")[-1]
+ raw_datasets = load_dataset(extension, data_files=data_files, cache_dir=model_args.cache_dir)
+ # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
+ # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+ # Load pretrained model and tokenizer
+ #
+ # Distributed training:
+ # The .from_pretrained methods guarantee that only one local process can concurrently
+ # download model & vocab.
+ config = AutoConfig.from_pretrained(
+ model_args.config_name if model_args.config_name else model_args.model_name_or_path,
+ cache_dir=model_args.cache_dir,
+ revision=model_args.model_revision,
+ use_auth_token=True if model_args.use_auth_token else None,
+ )
+ tokenizer = AutoTokenizer.from_pretrained(
+ model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
+ cache_dir=model_args.cache_dir,
+ use_fast=model_args.use_fast_tokenizer,
+ revision=model_args.model_revision,
+ use_auth_token=True if model_args.use_auth_token else None,
+ )
+ model = AutoModelForSeq2SeqLM.from_pretrained(
+ model_args.model_name_or_path,
+ from_tf=bool(".ckpt" in model_args.model_name_or_path),
+ config=config,
+ cache_dir=model_args.cache_dir,
+ revision=model_args.model_revision,
+ use_auth_token=True if model_args.use_auth_token else None,
+ )
+
+ model.resize_token_embeddings(len(tokenizer))
+
+ if training_args.gradient_checkpointing:
+ # use_cache=True is incompatible with gradient checkpointing.
+ config.use_cache = False
+
+ if model.config.decoder_start_token_id is None:
+ raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")
+
+ if (
+ hasattr(model.config, "max_position_embeddings")
+ and model.config.max_position_embeddings < data_args.max_source_length
+ ):
+ if model_args.resize_position_embeddings is None:
+ logger.warning(
+ f"Increasing the model's number of position embedding vectors from {model.config.max_position_embeddings} "
+ f"to {data_args.max_source_length}."
+ )
+ model.resize_position_embeddings(data_args.max_source_length)
+ elif model_args.resize_position_embeddings:
+ model.resize_position_embeddings(data_args.max_source_length)
+ else:
+ raise ValueError(
+ f"`--max_source_length` is set to {data_args.max_source_length}, but the model only has {model.config.max_position_embeddings}"
+ f" position encodings. Consider either reducing `--max_source_length` to {model.config.max_position_embeddings} or to automatically "
+ "resize the model's position encodings by passing `--resize_position_embeddings`."
+ )
+
+ if data_args.source_prefix_filepath is not None:
+ prefix = open(data_args.source_prefix_filepath, 'r', encoding='utf-8').readline().strip()
+ else:
+ 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
+
+ # Get the column names for input/target.
+ if data_args.source_column is None:
+ source_column = column_names[0]
+ else:
+ source_column = data_args.source_column
+ if source_column not in column_names:
+ raise ValueError(
+ f"--source_column' value '{data_args.source_column}' needs to be one of: {', '.join(column_names)}"
+ )
+ if data_args.target_column is None:
+ target_column = column_names[1]
+ else:
+ target_column = data_args.target_column
+ if target_column not in column_names:
+ raise ValueError(
+ f"--target_column' value '{data_args.target_column}' needs to be one of: {', '.join(column_names)}"
+ )
+
+ # Temporarily set max_target_length for training.
+ max_target_length = data_args.max_target_length
+ padding = "max_length" if data_args.pad_to_max_length else False
+
+ if training_args.label_smoothing_factor > 0 and not hasattr(model, "prepare_decoder_input_ids_from_labels"):
+ logger.warning(
+ "label_smoothing is enabled but the `prepare_decoder_input_ids_from_labels` method is not defined for"
+ f"`{model.__class__.__name__}`. This will lead to loss being calculated twice and will take up more memory"
+ )
+
+ def preprocess_function(examples):
+
+ # remove pairs where at least one record is None
+ inputs, targets = [], []
+ for i in range(len(examples[source_column])):
+ if examples[source_column][i] is not None and examples[target_column][i] is not None:
+ inputs.append(examples[source_column][i])
+ targets.append(examples[target_column][i])
+
+ inputs = [prefix + inp for inp in inputs]
+ model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, padding=padding, truncation=True)
+
+ # Setup the tokenizer for targets
+ with tokenizer.as_target_tokenizer():
+ labels = tokenizer(targets, max_length=max_target_length, padding=padding, truncation=True)
+
+ # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
+ # padding in the loss.
+ if padding == "max_length" and data_args.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
+
+ if training_args.do_train:
+ if "train" not in raw_datasets:
+ raise ValueError("--do_train requires a train dataset")
+ train_dataset = raw_datasets["train"]
+ if data_args.max_train_samples is not None:
+ train_dataset = train_dataset.select(range(data_args.max_train_samples))
+ with training_args.main_process_first(desc="train dataset map pre-processing"):
+ train_dataset = train_dataset.map(
+ preprocess_function,
+ 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 train dataset",
+ )
+
+ if training_args.do_eval:
+ max_target_length = data_args.val_max_target_length
+ if "validation" not in raw_datasets:
+ raise ValueError("--do_eval requires a validation dataset")
+ eval_dataset = raw_datasets["validation"]
+ if data_args.max_eval_samples is not None:
+ eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
+ with training_args.main_process_first(desc="validation dataset map pre-processing"):
+ eval_dataset = eval_dataset.map(
+ preprocess_function,
+ 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 validation dataset",
+ )
+
+ 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:
+ predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))
+ with training_args.main_process_first(desc="prediction dataset map pre-processing"):
+ predict_dataset = predict_dataset.map(
+ preprocess_function,
+ 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",
+ )
+
+ # Data collator
+ label_pad_token_id = -100 if data_args.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=8 if training_args.fp16 else None,
+ )
+
+ # TODO: compute custom metric at evaluation.
+ # Must take a EvalPrediction and return a dictionary string to metric values.
+ def compute_metrics(p: EvalPrediction):
+ preds, labels = p.predictions, p.label_ids
+ if isinstance(preds, tuple):
+ preds = preds[0]
+ decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
+ if data_args.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)
+
+ # TODO: compute metric using decoded_preds & decoded_labels
+ result = {}
+
+ prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
+ result["gen_len"] = np.mean(prediction_lens)
+ result = {k: round(v, 4) for k, v in result.items()}
+ return result
+
+ # Initialize our Trainer
+ trainer = Seq2SeqTrainer(
+ model=model,
+ args=training_args,
+ train_dataset=train_dataset if training_args.do_train else None,
+ eval_dataset=eval_dataset if training_args.do_eval else None,
+ tokenizer=tokenizer,
+ data_collator=data_collator,
+ compute_metrics=compute_metrics if training_args.predict_with_generate else None,
+ )
+
+ # Training
+ if training_args.do_train:
+ checkpoint = None
+ if training_args.resume_from_checkpoint is not None:
+ checkpoint = training_args.resume_from_checkpoint
+ elif last_checkpoint is not None:
+ checkpoint = last_checkpoint
+ train_result = trainer.train(resume_from_checkpoint=checkpoint)
+ trainer.save_model() # Saves the tokenizer too for easy upload
+
+ metrics = train_result.metrics
+ max_train_samples = (
+ data_args.max_train_samples if data_args.max_train_samples is not None else 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()
+
+ # Evaluation
+ results = {}
+ max_length = (
+ training_args.generation_max_length
+ if training_args.generation_max_length is not None
+ else data_args.val_max_target_length
+ )
+ num_beams = data_args.num_beams if data_args.num_beams is not None else training_args.generation_num_beams
+ if training_args.do_eval:
+ logger.info("*** Evaluate ***")
+ metrics = trainer.evaluate(max_length=max_length, num_beams=num_beams, metric_key_prefix="eval")
+ 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_length, num_beams=num_beams
+ )
+ 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]
+ output_prediction_file = os.path.join(training_args.output_dir, "generated_predictions.json")
+ with open(output_prediction_file, "w") as writer:
+ for sample, pred in zip(raw_datasets["test"], predictions):
+ sample["predictions"] = pred
+ writer.write(json.dumps(sample)+'\n')
+
+ kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": data_args.task_name}
+ if data_args.dataset_name is not None:
+ kwargs["dataset_tags"] = data_args.dataset_name
+ if data_args.dataset_config_name is not None:
+ kwargs["dataset_args"] = data_args.dataset_config_name
+ kwargs["dataset"] = f"{data_args.dataset_name} {data_args.dataset_config_name}"
+ else:
+ kwargs["dataset"] = data_args.dataset_name
+
+ if training_args.push_to_hub:
+ trainer.push_to_hub(**kwargs)
+ else:
+ trainer.create_model_card(**kwargs)
+
+ return results
+
+
+def _mp_fn(index):
+ # For xla_spawn (TPUs)
+ main()
+
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file