You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
735 lines
37 KiB
735 lines
37 KiB
from __future__ import absolute_import, division, print_function
|
|
|
|
import argparse
|
|
import glob
|
|
import logging
|
|
import os
|
|
import pickle
|
|
import random
|
|
import re
|
|
import shutil
|
|
import time
|
|
import json
|
|
|
|
import numpy as np
|
|
import torch
|
|
from torch.utils.data import DataLoader, Dataset, SequentialSampler, RandomSampler
|
|
from torch.utils.data.distributed import DistributedSampler
|
|
|
|
try:
|
|
from torch.utils.tensorboard import SummaryWriter
|
|
except:
|
|
from tensorboardX import SummaryWriter
|
|
|
|
from tqdm import tqdm, trange
|
|
|
|
from transformers import (WEIGHTS_NAME, AdamW, get_linear_schedule_with_warmup,GPT2Config, GPT2DoubleHeadsModel, GPT2Tokenizer)
|
|
|
|
|
|
logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s - %(message)s',
|
|
datefmt = '%m/%d/%Y %H:%M:%S',
|
|
level = logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
MODEL_CLASSES = {
|
|
'gpt2': (GPT2Config, GPT2DoubleHeadsModel, GPT2Tokenizer)
|
|
}
|
|
s = 0
|
|
s_time = time.time()
|
|
def log_every_n_interval(interval, msg):
|
|
global s, s_time
|
|
s += 1
|
|
if s % interval == 0:
|
|
if s_time is None:
|
|
s_time = time.time()
|
|
iter_p_s = 0
|
|
else:
|
|
e_time = time.time()
|
|
elapse = e_time - s_time
|
|
iter_p_s = interval / elapse
|
|
s_time = e_time
|
|
logger.info(f'MSG: {msg}; ITER: {iter_p_s:.2f}/s')
|
|
|
|
class JsonDataset(Dataset):
|
|
def __init__(self, tokenizer, args, file_path='train', max_seq=80, max_turn=1, seperator=' & '):
|
|
assert os.path.isfile(file_path)
|
|
directory, filename = os.path.split(file_path)
|
|
cached_features_file = os.path.join(directory, args.output_dir + '_cached_lm' + '_seqlen_' + str(max_seq) + '_' + filename)
|
|
|
|
if os.path.exists(cached_features_file) and not args.overwrite_cache:
|
|
logger.info("Loading features from cached file %s", cached_features_file)
|
|
with open(cached_features_file, 'rb') as handle:
|
|
self.examples = pickle.load(handle)
|
|
else:
|
|
logger.info(f"Creating features from dataset file at {directory}")
|
|
|
|
self.examples = []
|
|
self.labels = []
|
|
self.token_ids = []
|
|
self.attention_masks = []
|
|
|
|
self.mc_token_ids = []
|
|
self.mc_labels = []
|
|
|
|
system_token_id = tokenizer.convert_tokens_to_ids(['system'])
|
|
user_token_id = tokenizer.convert_tokens_to_ids(['user'])
|
|
induction_token_id = tokenizer.convert_tokens_to_ids(['=>'])
|
|
|
|
examples = json.load(open(file_path))
|
|
|
|
response_pool = []
|
|
belief_pool = []
|
|
idxs = list(range(len(examples)))
|
|
for i in examples:
|
|
response = i['reply'] + ' '+tokenizer.eos_token
|
|
response_id = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(response))
|
|
response_pool.append(response_id)
|
|
|
|
belief_id = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(i['belief']))
|
|
belief_pool.append(belief_id)
|
|
|
|
for example in examples:
|
|
history = example['history']
|
|
context = history[-max_turn:]
|
|
context_ids = []
|
|
token_ids_for_context = []
|
|
for cxt in context:
|
|
ids = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(cxt))
|
|
context_ids += ids
|
|
if 'user :' in cxt:
|
|
token_ids_for_context += user_token_id * len(ids)
|
|
else:
|
|
token_ids_for_context += system_token_id * len(ids)
|
|
|
|
history = ' '.join(history[-max_turn:])
|
|
kb = ' '#example['kb']
|
|
belief = example['belief']
|
|
|
|
belief_id = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(belief))
|
|
response = example['reply'] + ' '+tokenizer.eos_token
|
|
response_id = tokenizer.convert_tokens_to_ids(tokenizer.tokenize(response))
|
|
|
|
token_id = token_ids_for_context + system_token_id + system_token_id * len(belief_id) + system_token_id * len(response_id)
|
|
source = context_ids + induction_token_id + belief_id + response_id
|
|
if args.with_LM:
|
|
target = source
|
|
else:
|
|
target = [-1] * len(context_ids) + [-1] * len(induction_token_id) + belief_id + response_id
|
|
|
|
|
|
if len(source) < max_seq:
|
|
attention_mask = [0] * max_seq
|
|
attention_mask[:len(source)] = [1] * len(source)
|
|
self.mc_token_ids.append(len(source) - 1)
|
|
source += [0] * (max_seq - len(source))
|
|
target += [-1] * (max_seq - len(target))
|
|
token_id += [0] * (max_seq - len(token_id))
|
|
else:
|
|
attention_mask = [1] * max_seq
|
|
self.mc_token_ids.append(max_seq - 1)
|
|
source = source[-max_seq:]
|
|
target = target[-max_seq:]
|
|
token_id = token_id[-max_seq:]
|
|
|
|
self.mc_labels.append(0)
|
|
|
|
if not len(source) == len(target) == len(token_id) == len(attention_mask):
|
|
import pdb
|
|
pdb.set_trace()
|
|
|
|
self.examples.append(source)
|
|
self.labels.append(target)
|
|
self.token_ids.append(token_id)
|
|
self.attention_masks.append(attention_mask)
|
|
|
|
if args.add_same_belief_response_prediction:
|
|
for _ in range(args.num_candidates):
|
|
|
|
random_idx = random.choice(idxs)
|
|
new_response_id = response_pool[random_idx]
|
|
new_belief_id = belief_pool[random_idx]
|
|
|
|
source = context_ids + induction_token_id + new_belief_id + new_response_id
|
|
token_id = token_ids_for_context + system_token_id + system_token_id * len(new_belief_id) + system_token_id * len(new_response_id)
|
|
if len(source) < max_seq:
|
|
attention_mask = [0] * max_seq
|
|
attention_mask[:len(source)] = [1] * len(source)
|
|
self.mc_token_ids.append(len(source)-1)
|
|
source += [0] * (max_seq - len(source))
|
|
token_id += [0] * (max_seq - len(token_id))
|
|
else:
|
|
attention_mask = [1] * max_seq
|
|
self.mc_token_ids.append(max_seq - 1)
|
|
source = source[-max_seq:]
|
|
target = target[-max_seq:]
|
|
token_id = token_id[-max_seq:]
|
|
|
|
|
|
self.examples.append(source)
|
|
self.labels.append([-1] * len(source))
|
|
self.token_ids.append(token_id)
|
|
self.mc_labels.append(1)
|
|
self.attention_masks.append(attention_mask)
|
|
|
|
if args.add_response_prediction:
|
|
for _ in range(args.num_candidates):
|
|
|
|
random_idx = random.choice(idxs)
|
|
new_response_id = response_pool[random_idx]
|
|
new_belief_id = belief_pool[random_idx]
|
|
|
|
source = context_ids + induction_token_id + belief_id + new_response_id
|
|
token_id = token_ids_for_context + system_token_id + system_token_id * len(belief_id) + system_token_id * len(new_response_id)
|
|
if len(source) < max_seq:
|
|
attention_mask = [0] * max_seq
|
|
attention_mask[:len(source)] = [1] * len(source)
|
|
self.mc_token_ids.append(len(source)-1)
|
|
source += [0] * (max_seq - len(source))
|
|
token_id += [0] * (max_seq - len(token_id))
|
|
else:
|
|
attention_mask = [1] * max_seq
|
|
self.mc_token_ids.append(max_seq - 1)
|
|
source = source[-max_seq:]
|
|
target = target[-max_seq:]
|
|
token_id = token_id[-max_seq:]
|
|
|
|
|
|
self.examples.append(source)
|
|
self.labels.append([-1] * len(source))
|
|
self.token_ids.append(token_id)
|
|
self.mc_labels.append(1)
|
|
self.attention_masks.append(attention_mask)
|
|
|
|
if args.add_belief_prediction:
|
|
for _ in range(args.num_candidates):
|
|
random_idx = random.choice(idxs)
|
|
new_response_id = response_pool[random_idx]
|
|
new_belief_id = belief_pool[random_idx]
|
|
|
|
source = context_ids + induction_token_id + new_belief_id + response_id
|
|
token_id = token_ids_for_context + system_token_id + system_token_id * len(new_belief_id) + system_token_id * len(response_id)
|
|
if len(source) < max_seq:
|
|
attention_mask = [0] * max_seq
|
|
attention_mask[:len(source)] = [1] * len(source)
|
|
self.mc_token_ids.append(len(source)-1)
|
|
source += [0] * (max_seq - len(source))
|
|
token_id += [0] * (max_seq - len(token_id))
|
|
else:
|
|
attention_mask = [1] * max_seq
|
|
self.mc_token_ids.append(max_seq - 1)
|
|
source = source[-max_seq:]
|
|
target = target[-max_seq:]
|
|
token_id = token_id[-max_seq:]
|
|
|
|
|
|
self.examples.append(source)
|
|
self.labels.append([-1] * len(source))
|
|
self.token_ids.append(token_id)
|
|
self.mc_labels.append(1)
|
|
self.attention_masks.append(attention_mask)
|
|
def __len__(self):
|
|
return len(self.examples)
|
|
|
|
def __getitem__(self, item):
|
|
return torch.tensor(self.examples[item]), torch.tensor(self.token_ids[item]), torch.tensor(self.labels[item]), torch.tensor(self.attention_masks[item]), torch.tensor(self.mc_labels[item]), torch.tensor(self.mc_token_ids[item]),
|
|
|
|
def load_and_cache_examples(args, tokenizer, evaluate=False):
|
|
dataset = JsonDataset(tokenizer, args, file_path=args.eval_data_file if evaluate else args.train_data_file, max_seq=args.max_seq, max_turn=args.max_turn)
|
|
return dataset
|
|
|
|
def set_seed(args):
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
if args.n_gpu > 0:
|
|
torch.cuda.manual_seed_all(args.seed)
|
|
|
|
|
|
def _rotate_checkpoints(args, checkpoint_prefix, use_mtime=False):
|
|
if not args.save_total_limit:
|
|
return
|
|
if args.save_total_limit <= 0:
|
|
return
|
|
|
|
# Check if we should delete older checkpoint(s)
|
|
glob_checkpoints = glob.glob(os.path.join(args.output_dir, '{}-*'.format(checkpoint_prefix)))
|
|
if len(glob_checkpoints) <= args.save_total_limit:
|
|
return
|
|
|
|
ordering_and_checkpoint_path = []
|
|
for path in glob_checkpoints:
|
|
if use_mtime:
|
|
ordering_and_checkpoint_path.append((os.path.getmtime(path), path))
|
|
else:
|
|
regex_match = re.match('.*{}-([0-9]+)'.format(checkpoint_prefix), path)
|
|
if regex_match and regex_match.groups():
|
|
ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path))
|
|
|
|
checkpoints_sorted = sorted(ordering_and_checkpoint_path)
|
|
checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
|
|
number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - args.save_total_limit)
|
|
checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
|
|
for checkpoint in checkpoints_to_be_deleted:
|
|
logger.info("Deleting older checkpoint [{}] due to args.save_total_limit".format(checkpoint))
|
|
shutil.rmtree(checkpoint)
|
|
|
|
|
|
def mask_tokens(inputs, tokenizer, args):
|
|
""" Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original. """
|
|
labels = inputs.clone()
|
|
# We sample a few tokens in each sequence for masked-LM training (with probability args.mlm_probability defaults to 0.15 in Bert/RoBERTa)
|
|
probability_matrix = torch.full(labels.shape, args.mlm_probability)
|
|
special_tokens_mask = [tokenizer.get_special_tokens_mask(val, already_has_special_tokens=True) for val in labels.tolist()]
|
|
probability_matrix.masked_fill_(torch.tensor(special_tokens_mask, dtype=torch.bool), value=0.0)
|
|
masked_indices = torch.bernoulli(probability_matrix).bool()
|
|
labels[~masked_indices] = -1 # We only compute loss on masked tokens
|
|
|
|
# 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
|
|
indices_replaced = torch.bernoulli(torch.full(labels.shape, 0.8)).bool() & masked_indices
|
|
inputs[indices_replaced] = tokenizer.convert_tokens_to_ids(tokenizer.mask_token)
|
|
|
|
# 10% of the time, we replace masked input tokens with random word
|
|
indices_random = torch.bernoulli(torch.full(labels.shape, 0.5)).bool() & masked_indices & ~indices_replaced
|
|
random_words = torch.randint(len(tokenizer), labels.shape, dtype=torch.long)
|
|
inputs[indices_random] = random_words[indices_random]
|
|
|
|
# The rest of the time (10% of the time) we keep the masked input tokens unchanged
|
|
return inputs, labels
|
|
|
|
|
|
def train(args, train_dataset, model, tokenizer):
|
|
""" Train the model """
|
|
if args.local_rank in [-1, 0]:
|
|
tb_writer = SummaryWriter()
|
|
|
|
args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
|
|
train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
|
|
train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)
|
|
|
|
if args.max_steps > 0:
|
|
t_total = args.max_steps
|
|
args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
|
|
else:
|
|
t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
|
|
|
|
# Prepare optimizer and schedule (linear warmup and decay)
|
|
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, eps=args.adam_epsilon)
|
|
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total)
|
|
if args.fp16:
|
|
try:
|
|
from apex import amp
|
|
except ImportError:
|
|
raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.")
|
|
model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level)
|
|
model.resize_token_embeddings(len(tokenizer))
|
|
# multi-gpu training (should be after apex fp16 initialization)
|
|
if args.n_gpu > 1:
|
|
model = torch.nn.DataParallel(model)
|
|
|
|
# Distributed training (should be after apex fp16 initialization)
|
|
if args.local_rank != -1:
|
|
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank],
|
|
output_device=args.local_rank,
|
|
find_unused_parameters=True)
|
|
|
|
# Train!
|
|
logger.info("***** Running training *****")
|
|
logger.info(" Num examples = %d", len(train_dataset))
|
|
logger.info(" Num Epochs = %d", args.num_train_epochs)
|
|
logger.info(" Instantaneous batch size per GPU = %d", args.per_gpu_train_batch_size)
|
|
logger.info(" Total train batch size (w. parallel, distributed & accumulation) = %d",
|
|
args.train_batch_size * args.gradient_accumulation_steps * (torch.distributed.get_world_size() if args.local_rank != -1 else 1))
|
|
logger.info(" Gradient Accumulation steps = %d", args.gradient_accumulation_steps)
|
|
logger.info(" Total optimization steps = %d", t_total)
|
|
|
|
global_step = 0
|
|
tr_loss, logging_loss = 0.0, 0.0
|
|
model.zero_grad()
|
|
train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
|
|
set_seed(args) # Added here for reproducibility (even between python 2 and 3)
|
|
for e in train_iterator:
|
|
# epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
|
|
for step, batch in enumerate(train_dataloader):
|
|
# inputs, labels = mask_tokens(batch, tokenizer, args) if args.mlm else (batch, batch)
|
|
# logger.info(f" PROGRESS: {float(global_step)/t_total*100}%")
|
|
log_every_n_interval(500, f" PROGRESS: {int(float(global_step)/t_total*100)}%")
|
|
if step % 500 == 0:
|
|
logger.info(f" PROGRESS: {int(float(global_step)/t_total*100)}%")
|
|
inputs, tokens, labels, masks,mc_labels, mc_token_ids = batch
|
|
|
|
inputs = inputs.to(args.device)
|
|
tokens = tokens.to(args.device)
|
|
labels = labels.to(args.device)
|
|
masks = masks.to(args.device)
|
|
mc_labels = mc_labels.to(args.device)
|
|
mc_token_ids = mc_token_ids.to(args.device)
|
|
|
|
model.train()
|
|
outputs = model(inputs, lm_labels=labels, mc_labels=mc_labels, mc_token_ids=mc_token_ids, token_type_ids=tokens, attention_mask=masks)
|
|
lm_loss = outputs[0] # model outputs are always tuple in transformers (see doc)
|
|
mc_loss = outputs[1]
|
|
|
|
loss = lm_loss + args.mc_loss_efficient * mc_loss
|
|
|
|
if args.n_gpu > 1:
|
|
loss = loss.mean() # mean() to average on multi-gpu parallel training
|
|
if args.gradient_accumulation_steps > 1:
|
|
loss = loss / args.gradient_accumulation_steps
|
|
|
|
if args.fp16:
|
|
with amp.scale_loss(loss, optimizer) as scaled_loss:
|
|
scaled_loss.backward()
|
|
else:
|
|
loss.backward()
|
|
|
|
tr_loss += loss.item()
|
|
if (step + 1) % args.gradient_accumulation_steps == 0:
|
|
if args.fp16:
|
|
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm)
|
|
else:
|
|
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
|
|
optimizer.step()
|
|
scheduler.step() # Update learning rate schedule
|
|
model.zero_grad()
|
|
global_step += 1
|
|
|
|
if args.local_rank in [-1, 0] and args.logging_steps > 0 and global_step % args.logging_steps == 0:
|
|
# Log metrics
|
|
if args.local_rank == -1 and args.evaluate_during_training: # Only evaluate when single GPU otherwise metrics may not average well
|
|
results = evaluate(args, model, tokenizer)
|
|
for key, value in results.items():
|
|
tb_writer.add_scalar('eval_{}'.format(key), value, global_step)
|
|
tb_writer.add_scalar('lr', scheduler.get_lr()[0], global_step)
|
|
tb_writer.add_scalar('loss', (tr_loss - logging_loss)/args.logging_steps, global_step)
|
|
logger.info(f" EVALERR: {(tr_loss - logging_loss)/float(args.logging_steps)}")
|
|
logging_loss = tr_loss
|
|
|
|
if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
|
|
checkpoint_prefix = 'checkpoint'
|
|
# Save model checkpoint
|
|
output_dir = os.path.join(args.output_dir, '{}-{}'.format(checkpoint_prefix, global_step))
|
|
if not os.path.exists(output_dir):
|
|
os.makedirs(output_dir)
|
|
model_to_save = model.module if hasattr(model, 'module') else model # Take care of distributed/parallel training
|
|
model_to_save.save_pretrained(output_dir)
|
|
tokenizer.save_pretrained(output_dir)
|
|
torch.save(args, os.path.join(output_dir, 'training_args.bin'))
|
|
logger.info("Saving model checkpoint to %s", output_dir)
|
|
|
|
_rotate_checkpoints(args, checkpoint_prefix)
|
|
|
|
# if args.max_steps > 0 and global_step > args.max_steps:
|
|
# epoch_iterator.close()
|
|
# break
|
|
if args.max_steps > 0 and global_step > args.max_steps:
|
|
train_iterator.close()
|
|
break
|
|
|
|
if args.local_rank in [-1, 0]:
|
|
tb_writer.close()
|
|
|
|
return global_step, tr_loss / global_step
|
|
|
|
|
|
def evaluate(args, model, tokenizer, prefix=""):
|
|
# Loop to handle MNLI double evaluation (matched, mis-matched)
|
|
eval_output_dir = args.output_dir
|
|
|
|
eval_dataset = load_and_cache_examples(args, tokenizer, evaluate=True)
|
|
|
|
if not os.path.exists(eval_output_dir) and args.local_rank in [-1, 0]:
|
|
os.makedirs(eval_output_dir)
|
|
|
|
args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
|
|
# Note that DistributedSampler samples randomly
|
|
eval_sampler = SequentialSampler(eval_dataset) if args.local_rank == -1 else DistributedSampler(eval_dataset)
|
|
eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size)
|
|
|
|
# multi-gpu evaluate
|
|
if args.n_gpu > 1:
|
|
model = torch.nn.DataParallel(model)
|
|
|
|
# Eval!
|
|
logger.info("***** Running evaluation {} *****".format(prefix))
|
|
logger.info(" Num examples = %d", len(eval_dataset))
|
|
logger.info(" Batch size = %d", args.eval_batch_size)
|
|
eval_loss = 0.0
|
|
nb_eval_steps = 0
|
|
model.eval()
|
|
|
|
for batch in tqdm(eval_dataloader, desc="Evaluating"):
|
|
# inputs, labels = mask_tokens(batch, tokenizer, args) if args.mlm else (batch, batch)
|
|
|
|
inputs, tokens, labels, masks = batch
|
|
# import pdb
|
|
# pdb.set_trace()
|
|
inputs = inputs.to(args.device)
|
|
tokens = tokens.to(args.device)
|
|
labels = labels.to(args.device)
|
|
masks = masks.to(args.device)
|
|
# inputs = inputs.to(args.device)
|
|
# labels = labels.to(args.device)
|
|
|
|
with torch.no_grad():
|
|
outputs = model(inputs, masked_lm_labels=labels, token_type_ids=tokens) if args.mlm else model(inputs, labels=labels)
|
|
lm_loss = outputs[0]
|
|
eval_loss += lm_loss.mean().item()
|
|
nb_eval_steps += 1
|
|
|
|
eval_loss = eval_loss / nb_eval_steps
|
|
perplexity = torch.exp(torch.tensor(eval_loss))
|
|
|
|
result = {
|
|
"perplexity": perplexity
|
|
}
|
|
|
|
output_eval_file = os.path.join(eval_output_dir, prefix, "eval_results.txt")
|
|
with open(output_eval_file, "w") as writer:
|
|
logger.info("***** Eval results {} *****".format(prefix))
|
|
for key in sorted(result.keys()):
|
|
logger.info(" %s = %s", key, str(result[key]))
|
|
writer.write("%s = %s\n" % (key, str(result[key])))
|
|
|
|
return result
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
|
|
## Required parameters
|
|
parser.add_argument("--train_data_file", default=None, type=str, required=True,
|
|
help="The input training data file (a text file).")
|
|
parser.add_argument("--output_dir", default=None, type=str, required=True,
|
|
help="The output directory where the model predictions and checkpoints will be written.")
|
|
|
|
## Other parameters
|
|
parser.add_argument("--eval_data_file", default=None, type=str,
|
|
help="An optional input evaluation data file to evaluate the perplexity on (a text file).")
|
|
|
|
parser.add_argument("--model_type", default="bert", type=str,
|
|
help="The model architecture to be fine-tuned.")
|
|
parser.add_argument("--model_name_or_path", default="bert-base-cased", type=str,
|
|
help="The model checkpoint for weights initialization.")
|
|
|
|
parser.add_argument("--mlm", action='store_true',
|
|
help="Train with masked-language modeling loss instead of language modeling.")
|
|
parser.add_argument("--mlm_probability", type=float, default=0.15,
|
|
help="Ratio of tokens to mask for masked language modeling loss")
|
|
|
|
parser.add_argument("--config_name", default="", type=str,
|
|
help="Optional pretrained config name or path if not the same as model_name_or_path")
|
|
parser.add_argument("--tokenizer_name", default="", type=str,
|
|
help="Optional pretrained tokenizer name or path if not the same as model_name_or_path")
|
|
parser.add_argument("--cache_dir", default="", type=str,
|
|
help="Optional directory to store the pre-trained models downloaded from s3 (instread of the default one)")
|
|
parser.add_argument("--block_size", default=80, type=int,
|
|
help="Optional input sequence length after tokenization."
|
|
"The training dataset will be truncated in block of this size for training."
|
|
"Default to the model max input length for single sentence inputs (take into account special tokens).")
|
|
parser.add_argument("--do_train", action='store_true',
|
|
help="Whether to run training.")
|
|
parser.add_argument("--do_eval", action='store_true',
|
|
help="Whether to run eval on the dev set.")
|
|
parser.add_argument("--evaluate_during_training", action='store_true',
|
|
help="Run evaluation during training at each logging step.")
|
|
parser.add_argument("--do_lower_case", action='store_true',
|
|
help="Set this flag if you are using an uncased model.")
|
|
|
|
parser.add_argument("--per_gpu_train_batch_size", default=1, type=int,
|
|
help="Batch size per GPU/CPU for training.")
|
|
parser.add_argument("--per_gpu_eval_batch_size", default=1, type=int,
|
|
help="Batch size per GPU/CPU for evaluation.")
|
|
parser.add_argument('--gradient_accumulation_steps', type=int, default=1,
|
|
help="Number of updates steps to accumulate before performing a backward/update pass.")
|
|
parser.add_argument("--learning_rate", default=5e-5, type=float,
|
|
help="The initial learning rate for Adam.")
|
|
parser.add_argument("--weight_decay", default=0.0, type=float,
|
|
help="Weight deay if we apply some.")
|
|
parser.add_argument("--adam_epsilon", default=1e-8, type=float,
|
|
help="Epsilon for Adam optimizer.")
|
|
parser.add_argument("--max_grad_norm", default=1.0, type=float,
|
|
help="Max gradient norm.")
|
|
parser.add_argument("--num_train_epochs", default=1.0, type=float,
|
|
help="Total number of training epochs to perform.")
|
|
parser.add_argument("--max_steps", default=-1, type=int,
|
|
help="If > 0: set total number of training steps to perform. Override num_train_epochs.")
|
|
parser.add_argument("--warmup_steps", default=0, type=int,
|
|
help="Linear warmup over warmup_steps.")
|
|
|
|
parser.add_argument('--logging_steps', type=int, default=10,
|
|
help="Log every X updates steps.")
|
|
parser.add_argument('--save_steps', type=int, default=5000,
|
|
help="Save checkpoint every X updates steps.")
|
|
parser.add_argument('--save_total_limit', type=int, default=None,
|
|
help='Limit the total amount of checkpoints, delete the older checkpoints in the output_dir, does not delete by default')
|
|
parser.add_argument("--eval_all_checkpoints", action='store_true',
|
|
help="Evaluate all checkpoints starting with the same prefix as model_name_or_path ending and ending with step number")
|
|
parser.add_argument("--no_cuda", action='store_true',
|
|
help="Avoid using CUDA when available")
|
|
parser.add_argument('--overwrite_output_dir', action='store_true',
|
|
help="Overwrite the content of the output directory")
|
|
parser.add_argument('--overwrite_cache', action='store_true',
|
|
help="Overwrite the cached training and evaluation sets")
|
|
parser.add_argument('--seed', type=int, default=42,
|
|
help="random seed for initialization")
|
|
|
|
parser.add_argument('--fp16', action='store_true',
|
|
help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit")
|
|
parser.add_argument('--fp16_opt_level', type=str, default='O1',
|
|
help="For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']."
|
|
"See details at https://nvidia.github.io/apex/amp.html")
|
|
parser.add_argument("--local_rank", type=int, default=-1,
|
|
help="For distributed training: local_rank")
|
|
parser.add_argument('--server_ip', type=str, default='', help="For distant debugging.")
|
|
parser.add_argument('--server_port', type=str, default='', help="For distant debugging.")
|
|
parser.add_argument('--text_chunk', action='store_true', help="")
|
|
parser.add_argument('--with_LM', type=bool, default=True, help="")
|
|
|
|
parser.add_argument("--max_seq", default=200, type=int, help="")
|
|
parser.add_argument("--max_turn", default=1, type=int, help="")
|
|
parser.add_argument("--mc_loss_efficient", default=1, type=float, help="")
|
|
parser.add_argument("--num_candidates", default=1, type=int, help="")
|
|
parser.add_argument("--add_special_action_tokens", default='', type=str)
|
|
parser.add_argument("--add_same_belief_response_prediction", action='store_true')
|
|
parser.add_argument("--add_response_prediction", action='store_true')
|
|
parser.add_argument("--add_belief_prediction", action='store_true')
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.model_type in ["bert", "roberta", "distilbert"] and not args.mlm:
|
|
raise ValueError("BERT and RoBERTa do not have LM heads but masked LM heads. They must be run using the --mlm "
|
|
"flag (masked language modeling).")
|
|
if args.eval_data_file is None and args.do_eval:
|
|
raise ValueError("Cannot do evaluation without an evaluation data file. Either supply a file to --eval_data_file "
|
|
"or remove the --do_eval argument.")
|
|
|
|
if os.path.exists(args.output_dir) and os.listdir(args.output_dir) and args.do_train and not args.overwrite_output_dir:
|
|
raise ValueError("Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(args.output_dir))
|
|
|
|
# Setup distant debugging if needed
|
|
if args.server_ip and args.server_port:
|
|
# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
|
|
import ptvsd
|
|
print("Waiting for debugger attach")
|
|
ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)
|
|
ptvsd.wait_for_attach()
|
|
|
|
# Setup CUDA, GPU & distributed training
|
|
if args.local_rank == -1 or args.no_cuda:
|
|
device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
|
|
args.n_gpu = torch.cuda.device_count()
|
|
else: # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
|
|
torch.cuda.set_device(args.local_rank)
|
|
device = torch.device("cuda", args.local_rank)
|
|
torch.distributed.init_process_group(backend='nccl')
|
|
args.n_gpu = 1
|
|
args.device = device
|
|
|
|
# Setup logging
|
|
logging.basicConfig(format = '%(asctime)s - %(levelname)s - %(name)s - %(message)s',
|
|
datefmt = '%m/%d/%Y %H:%M:%S',
|
|
level = logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
|
|
logger.warning("Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
|
|
args.local_rank, device, args.n_gpu, bool(args.local_rank != -1), args.fp16)
|
|
|
|
# Set seed
|
|
set_seed(args)
|
|
|
|
# Load pretrained model and tokenizer
|
|
if args.local_rank not in [-1, 0]:
|
|
torch.distributed.barrier() # Barrier to make sure only the first process in distributed training download model & vocab
|
|
|
|
config_class, model_class, tokenizer_class = MODEL_CLASSES[args.model_type]
|
|
config = config_class.from_pretrained(args.config_name if args.config_name else args.model_name_or_path,
|
|
cache_dir=args.cache_dir if args.cache_dir else None)
|
|
config.num_labels = 2
|
|
tokenizer = tokenizer_class.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
|
|
do_lower_case=args.do_lower_case,
|
|
cache_dir=args.cache_dir if args.cache_dir else None)
|
|
if args.block_size <= 0:
|
|
args.block_size = tokenizer.max_len_single_sentence # Our input block size will be the max possible for the model
|
|
args.block_size = min(args.block_size, tokenizer.max_len_single_sentence)
|
|
model = model_class.from_pretrained(args.model_name_or_path,
|
|
from_tf=bool('.ckpt' in args.model_name_or_path),
|
|
config=config,
|
|
cache_dir=args.cache_dir if args.cache_dir else None)
|
|
|
|
if args.add_special_action_tokens:
|
|
special_tokens = []
|
|
for line in open(args.add_special_action_tokens):
|
|
special_tokens.append(line.strip())
|
|
tokenizer.add_tokens(special_tokens)
|
|
model.resize_token_embeddings(len(tokenizer))
|
|
model.to(args.device)
|
|
|
|
if args.local_rank == 0:
|
|
torch.distributed.barrier() # End of barrier to make sure only the first process in distributed training download model & vocab
|
|
|
|
logger.info("Training/evaluation parameters %s", args)
|
|
|
|
# Training
|
|
if args.do_train:
|
|
if args.local_rank not in [-1, 0]:
|
|
torch.distributed.barrier() # Barrier to make sure only the first process in distributed training process the dataset, and the others will use the cache
|
|
|
|
train_dataset = load_and_cache_examples(args, tokenizer, evaluate=False)
|
|
|
|
if args.local_rank == 0:
|
|
torch.distributed.barrier()
|
|
|
|
global_step, tr_loss = train(args, train_dataset, model, tokenizer)
|
|
logger.info(" global_step = %s, average loss = %s", global_step, tr_loss)
|
|
|
|
|
|
# Saving best-practices: if you use save_pretrained for the model and tokenizer, you can reload them using from_pretrained()
|
|
if args.do_train and (args.local_rank == -1 or torch.distributed.get_rank() == 0):
|
|
# Create output directory if needed
|
|
if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
|
|
os.makedirs(args.output_dir)
|
|
|
|
logger.info("Saving model checkpoint to %s", args.output_dir)
|
|
# Save a trained model, configuration and tokenizer using `save_pretrained()`.
|
|
# They can then be reloaded using `from_pretrained()`
|
|
model_to_save = model.module if hasattr(model, 'module') else model # Take care of distributed/parallel training
|
|
model_to_save.save_pretrained(args.output_dir)
|
|
tokenizer.save_pretrained(args.output_dir)
|
|
|
|
# Good practice: save your training arguments together with the trained model
|
|
torch.save(args, os.path.join(args.output_dir, 'training_args.bin'))
|
|
|
|
# Load a trained model and vocabulary that you have fine-tuned
|
|
model = model_class.from_pretrained(args.output_dir)
|
|
tokenizer = tokenizer_class.from_pretrained(args.output_dir, do_lower_case=args.do_lower_case)
|
|
model.to(args.device)
|
|
|
|
|
|
# Evaluation
|
|
results = {}
|
|
if args.do_eval and args.local_rank in [-1, 0]:
|
|
checkpoints = [args.output_dir]
|
|
if args.eval_all_checkpoints:
|
|
checkpoints = list(os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + '/**/' + WEIGHTS_NAME, recursive=True)))
|
|
logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN) # Reduce logging
|
|
logger.info("Evaluate the following checkpoints: %s", checkpoints)
|
|
for checkpoint in checkpoints:
|
|
global_step = checkpoint.split('-')[-1] if len(checkpoints) > 1 else ""
|
|
prefix = checkpoint.split('/')[-1] if checkpoint.find('checkpoint') != -1 else ""
|
|
|
|
model = model_class.from_pretrained(checkpoint)
|
|
model.to(args.device)
|
|
result = evaluate(args, model, tokenizer, prefix=prefix)
|
|
result = dict((k + '_{}'.format(global_step), v) for k, v in result.items())
|
|
results.update(result)
|
|
|
|
return results
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|