lfqa1 / training /run_seq2seq_no_trainer.py
Achyut Tiwari
Add files via upload
e49e418 unverified
raw
history blame
18.5 kB
import argparse
import logging
import math
import re
import numpy as np
import torch
from accelerate import Accelerator
from accelerate.utils import set_seed
from torch.utils.data import DataLoader
from tqdm.auto import tqdm
from transformers import get_scheduler, AutoTokenizer, AdamW, SchedulerType, AutoModelForSeq2SeqLM, \
DataCollatorWithPadding
from datasets import load_dataset
logger = logging.getLogger(__name__)
def get_parser():
parser = argparse.ArgumentParser(description="Train ELI5 seq2seq answer generation model")
parser.add_argument(
"--dataset_name",
type=str,
default="vblagoje/lfqa",
help="The name of the dataset to use (via the datasets library).",
)
parser.add_argument(
"--per_device_train_batch_size",
type=int,
default=4,
)
parser.add_argument(
"--per_device_eval_batch_size",
type=int,
default=4,
help="Batch size (per device) for the evaluation dataloader.",
)
parser.add_argument(
"--pretrained_model_name",
type=str,
default="facebook/bart-large",
)
parser.add_argument(
"--model_save_name",
type=str,
default="eli5_bart_model",
)
parser.add_argument(
"--learning_rate",
type=float,
default=2e-4,
)
parser.add_argument(
"--weight_decay",
type=float,
default=0.0,
help="Weight decay to use."
)
parser.add_argument(
"--log_freq",
type=int,
default=100,
help="Log train/validation loss every log_freq update steps"
)
parser.add_argument(
"--ignore_pad_token_for_loss",
type=bool,
default=True,
help="Whether to ignore the tokens corresponding to " "padded labels in the loss computation or not.",
)
parser.add_argument(
"--num_train_epochs",
type=int,
default=3,
)
parser.add_argument(
"--max_train_steps",
type=int,
default=None,
help="Total number of training steps to perform. If provided, overrides num_train_epochs.",
)
parser.add_argument(
"--gradient_accumulation_steps",
type=int,
default=16,
help="Number of updates steps to accumulate before performing a backward/update pass.",
)
parser.add_argument(
"--pad_to_max_length",
action="store_true",
help="If passed, pad all samples to `max_length`. Otherwise, dynamic padding is used.",
)
parser.add_argument(
"--overwrite_cache", type=bool, default=None, help="Overwrite the cached training and evaluation sets"
)
parser.add_argument(
"--max_source_length",
type=int,
default=1024,
help="The maximum total input sequence length after "
"tokenization.Sequences longer than this will be truncated, sequences shorter will be padded.",
)
parser.add_argument(
"--max_target_length",
type=int,
default=360,
help="The maximum total sequence length for target text after "
"tokenization. Sequences longer than this will be truncated, sequences shorter will be padded."
)
parser.add_argument(
"--lr_scheduler_type",
type=SchedulerType,
default="linear", # this is linear with warmup
help="The scheduler type to use.",
choices=["linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"],
)
parser.add_argument(
"--num_warmup_steps",
type=int,
default=None,
help="Number of steps for the warmup in the lr scheduler."
)
parser.add_argument(
"--warmup_percentage",
type=float,
default=0.08,
help="Number of steps for the warmup in the lr scheduler."
)
return parser
def cleanup_references(text):
# URL reference where we need to remove both the link text and URL
# ...and this letter is used by most biographers as the cornerstone of Lee's personal
# views on slavery ([1](_URL_2_ & pg=PA173), [2](_URL_1_), [3](_URL_5_)).
# ...and this letter is used by most biographers as the cornerstone of Lee's personal views on slavery.
result = re.sub(r"[\(\s]*\[\d+\]\([^)]+\)[,)]*", "", text, 0, re.MULTILINE)
# URL reference where we need to preserve link text but remove URL
# At the outbreak of the Civil War, [Leyburn left his church](_URL_19_) and joined the South.
# At the outbreak of the Civil War, Leyburn left his church and joined the South.
result = re.sub(r"\[([^]]+)\]\([^)]+\)", "\\1", result, 0, re.MULTILINE)
# lastly remove just dangling _URL_[0-9]_ URL references
result = re.sub(r"_URL_\d_", "", result, 0, re.MULTILINE)
return result
def clean_answer(text):
result = cleanup_references(text)
result = result.replace("\n", " ")
result = re.sub(r"\s\s+", " ", result)
result = re.sub(r"BULLET::::-", "", result)
return result.strip()
def clean_question(text):
result = cleanup_references(text)
result = result.replace("\n", " ")
result = re.sub(r"\s\s+", " ", result)
result = result.replace("[deleted]", "")
return result.lower().strip()
def prepare_support_docs(example):
provenances = example["output"][-1]["provenance"]
context = "<P> " + " <P> ".join([p["text"] for p in provenances])
return {"context": context}
def preprocess_eli5(examples, **fn_kwargs):
document_cache = fn_kwargs["document_cache"]
training = fn_kwargs.get("training", True)
extra_answer_threshold = fn_kwargs.get("extra_answer_threshold", 3)
include_selftext = fn_kwargs.get("include_selftext", False)
exclude_answer_patterns = fn_kwargs.get("exclude_answer_patterns", [])
questions, contexts, answers = [], [], []
for q_id, question, selftext, answer in zip(examples["q_id"], examples["title"], examples["selftext"],
examples["answers"]):
accepted_answer_idx = []
if training:
accepted_answer_idx = [idx for idx, score in enumerate(answer["score"]) if
score > extra_answer_threshold]
if not training or not accepted_answer_idx:
accepted_answer_idx = [0]
document = document_cache[q_id]
for idx in accepted_answer_idx:
skip_answer = any([p.search(answer["text"][idx]) for p in exclude_answer_patterns])
if skip_answer:
continue
if include_selftext:
questions.append(clean_question(f"{question} {selftext}"))
else:
questions.append(clean_question(question))
contexts.append(document.lower().strip())
answers.append(clean_answer(answer["text"][idx]))
return {"question": questions, "context": contexts, "answer": answers}
def eval_qa_s2s_epoch(model, dataloader, accelerator, args):
model.eval()
num_eval_steps = math.ceil(len(dataloader))
progress_bar = tqdm(range(num_eval_steps), disable=not accelerator.is_local_main_process)
total_loss = 0.
with torch.no_grad():
for step, batch in enumerate(dataloader):
outputs = model(**batch)
loss = outputs.loss
total_loss += loss.item()
progress_bar.update(1)
progress_bar.set_postfix(loss=round((total_loss / (step + 1)), 3))
return total_loss / (step + 1)
def train(config):
set_seed(42)
args = config["args"]
eli5 = load_dataset(args.dataset_name)
support_docs = load_dataset("vblagoje/lfqa_support_docs")
# Initialize the accelerator. We will let the accelerator handle device placement for us in this example.
accelerator = Accelerator()
# Make one log on every process with the configuration for debugging.
logging.basicConfig(
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
datefmt="%m/%d/%Y %H:%M:%S",
level=logging.INFO,
)
logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR)
logger.info(accelerator.state)
tokenizer = AutoTokenizer.from_pretrained(args.pretrained_model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(args.pretrained_model_name)
# Optimizer
# Split weights in two groups, one with weight decay and the other not.
no_decay = ["bias", "LayerNorm.weight"]
optimizer_grouped_parameters = [
{
"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
"weight_decay": args.weight_decay,
},
{
"params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
"weight_decay": 0.0,
},
]
optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, weight_decay=args.weight_decay)
processed_datasets = {}
support_docs_prepared = {}
with accelerator.main_process_first():
for split in ["train", "validation"]:
support_docs_prepared[split] = support_docs[split].map(prepare_support_docs,
batched=False,
cache_file_name=f"./support_docs_{split}.arrow",
load_from_cache_file=not args.overwrite_cache,
desc="Preparing support docs",
)
column_names = eli5["train"].column_names
for split in ["train", "validation"]:
d_cache = dict([(e["id"], e["context"]) for e in tqdm(support_docs_prepared[split],
desc=f"Adding support docs to LFQA {split}")])
processed_datasets[split] = eli5[split].map(preprocess_eli5,
batched=True,
remove_columns=column_names,
cache_file_name=f"./processed_datasets_{split}.arrow",
load_from_cache_file=not args.overwrite_cache,
desc="Preparing dataset for tokenization",
fn_kwargs={"document_cache": d_cache,
"training": split == "train",
"exclude_answer_patterns": [re.compile("not sure what you"),
re.compile("\n\n >")]}
)
padding = "max_length" if args.pad_to_max_length else False
# Temporarily set max_target_length for training.
max_target_length = args.max_target_length
label_pad_token_id = -100 if args.ignore_pad_token_for_loss else tokenizer.pad_token_id
def tokenize_dataset(examples):
inputs = ["question: {} context: {}".format(q, c) for q, c in zip(examples["question"], examples["context"])]
targets = examples["answer"]
model_inputs = tokenizer(inputs, max_length=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=True, truncation=True,
return_tensors="np")
model_inputs["decoder_input_ids"] = labels["input_ids"][:, :-1].tolist()
# replace pad_token_id with label_pad_token_id to avoid loss calculation on those tokens
labels["input_ids"] = np.where(labels["input_ids"] == tokenizer.pad_token_id,
label_pad_token_id, labels["input_ids"])
model_inputs["labels"] = labels["input_ids"][:, 1:].tolist()
return model_inputs
tokenized_datasets = {}
with accelerator.main_process_first():
for split, dataset in processed_datasets.items():
tokenized_datasets[split] = dataset.map(
tokenize_dataset,
batched=True,
cache_file_name=f"./tokenized_dataset_{split}.arrow",
remove_columns=dataset.column_names,
load_from_cache_file=not args.overwrite_cache,
desc="Running tokenizer on dataset"
)
train_dataset = tokenized_datasets["train"]
eval_dataset = tokenized_datasets["validation"]
train_dataset.set_format(type='torch')
eval_dataset.set_format(type='torch')
data_collator = DataCollatorWithPadding(tokenizer, "max_length")
# first epoch we don't shuffle
train_dataloader = DataLoader(train_dataset, shuffle=False, batch_size=args.per_device_train_batch_size,
collate_fn=data_collator)
eval_dataloader = DataLoader(eval_dataset, batch_size=args.per_device_eval_batch_size, collate_fn=data_collator)
# train the model
model, optimizer, train_dataloader, eval_dataloader = accelerator.prepare(model, optimizer, train_dataloader,
eval_dataloader)
# Scheduler and math around the number of training steps.
num_update_steps_per_epoch = math.ceil(len(train_dataloader) / args.gradient_accumulation_steps)
if args.max_train_steps is None:
args.max_train_steps = args.num_train_epochs * num_update_steps_per_epoch
else:
args.num_train_epochs = math.ceil(args.max_train_steps / num_update_steps_per_epoch)
num_warmup_steps = args.num_warmup_steps if args.num_warmup_steps else math.ceil(args.max_train_steps *
args.warmup_percentage)
scheduler = get_scheduler(
name=args.lr_scheduler_type,
optimizer=optimizer,
num_warmup_steps=num_warmup_steps,
num_training_steps=args.max_train_steps,
)
# Train!
total_batch_size = args.per_device_train_batch_size * accelerator.num_processes * args.gradient_accumulation_steps
logger.info("***** Running training *****")
logger.info(f" Num examples = {len(train_dataset)}")
logger.info(f" Num eval examples = {len(eval_dataset)}")
logger.info(f" Num Epochs = {args.num_train_epochs}")
logger.info(f" Instantaneous batch size per device = {args.per_device_train_batch_size}")
logger.info(f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
logger.info(f" Gradient Accumulation steps = {args.gradient_accumulation_steps}")
logger.info(f" Total optimization steps = {args.max_train_steps}")
logger.info(f" Warmup steps = {num_warmup_steps}")
logger.info(f" Logging training progress every {args.log_freq} optimization steps")
# Only show the progress bar once on each machine.
progress_bar = tqdm(range(args.max_train_steps), disable=not accelerator.is_local_main_process)
completed_steps = 0
switched_train_dataloader = False
for epoch in range(args.num_train_epochs):
model.train()
if epoch > 0 and not switched_train_dataloader:
train_dataloader = DataLoader(train_dataset, batch_size=args.per_device_train_batch_size,
shuffle=True, collate_fn=data_collator)
train_dataloader = accelerator.prepare(train_dataloader)
switched_train_dataloader = True
for step, batch in enumerate(train_dataloader):
outputs = model(**batch)
loss = torch.mean(outputs.loss)
accelerator.backward(loss)
if ((step + 1) % args.gradient_accumulation_steps == 0) or (step + 1 == len(train_dataloader)):
optimizer.step()
scheduler.step()
optimizer.zero_grad()
progress_bar.update(1)
progress_bar.set_postfix(loss=round(loss.item(), 3))
completed_steps += 1
if completed_steps >= args.max_train_steps:
break
if step % (args.log_freq * args.gradient_accumulation_steps) == 0:
validation_loss = eval_qa_s2s_epoch(model, eval_dataloader, accelerator, args)
model.train()
logger.info(f"Train loss {loss.item()} , validation loss {validation_loss}")
if args.wandb and accelerator.is_local_main_process:
import wandb
wandb.log({"loss": loss.item(),
"lr": scheduler.get_last_lr()[0],
"validation_loss": validation_loss,
"completed_steps": completed_steps})
logger.info("Saving model {}".format(args.model_save_name))
accelerator.wait_for_everyone()
unwrapped_model = accelerator.unwrap_model(model)
accelerator.save(unwrapped_model.state_dict(), "{}_{}.bin".format(args.model_save_name, epoch))
# Calculating the validation loss over epoch
validation_loss = eval_qa_s2s_epoch(model, eval_dataloader, accelerator, args)
logger.info("Epoch: {}".format(epoch))
logger.info("Validation loss: {}".format(validation_loss))
def main():
parser = get_parser()
parser.add_argument(
"--wandb",
action="store_true",
help="If true, use W&B logging",
)
main_args, _ = parser.parse_known_args()
config = {"args": main_args}
if main_args.wandb:
import wandb
wandb.init(project="Bart_ELI5")
train(config=config)
main()