adam/cartman
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

retrain

Browse source
This commit is contained in:
Adam 2023-02-12 15:03:47 -05:00
parent 667d98c017
commit 7a0e2b1411
3 changed files with 8931 additions and 8874 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
train/clean.py Normal file
View file

@ -0,0 +1,18 @@
import pandas as pd
df = pd.read_csv('./data/All-seasons.csv')
cleanlines = pd.Series(
[cell
.replace('\n', '')
.replace('(', '')
.replace(')', '')
.replace(' ', ' ')
.strip()
for cell in df.Line
]
)
train = pd.DataFrame(df.Character)
train['line'] = cleanlines
train.columns = ['name', 'line']
train.to_csv('./data/train.csv', index=False)

17530
train/data/train.csv
View file

File diff suppressed because it is too large Load diff

251
train/train.py
View file

@ -1,5 +1,20 @@
# all the imports # all the imports
from transformers.models.auto.configuration_auto import AutoConfig
from transformers.modeling_utils import PreTrainedModel
from transformers.tokenization_utils import PreTrainedTokenizer
from transformers.models.auto.tokenization_auto import AutoTokenizer
from transformers.models.auto.modeling_auto import (
AutoModelForCausalLM,
MODEL_WITH_LM_HEAD_MAPPING,
)
from transformers.utils import WEIGHTS_NAME
from transformers.optimization import (
AdamW,
get_linear_schedule_with_warmup,
)
import torch
import glob import glob
import logging import logging
import os import os
@ -15,33 +30,22 @@ import pandas as pd
from sklearn.model_selection import train_test_split from sklearn.model_selection import train_test_split
from torch.nn.utils.rnn import pad_sequence from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler from torch.utils.data import (
from torch.utils.data.distributed import DistributedSampler DataLoader,
from tqdm.notebook import tqdm, trange Dataset,
RandomSampler,
from pathlib import Path SequentialSampler,
from transformers import (
MODEL_WITH_LM_HEAD_MAPPING,
WEIGHTS_NAME,
AdamW,
AutoConfig,
PreTrainedModel,
PreTrainedTokenizer,
get_linear_schedule_with_warmup,
) )
from torch.utils.data.distributed import DistributedSampler
from tqdm import tqdm, trange
from torch.utils.tensorboard.writer import SummaryWriter
try:
from torch.utils.tensorboard import SummaryWriter
except ImportError:
from tensorboardX import SummaryWriter
# -------------------------------------------------------------------------- # --------------------------------------------------------------------------
data = pd.read_csv('data/train.csv') data = pd.read_csv('data/train.csv')
CHARACTER_NAME = 'TARGET' CHARACTER_NAME = 'Cartman'
contexted = [] contexted = []
# context window of size 7 # context window of size 7
@ -51,7 +55,7 @@ for i in data[data.name == CHARACTER_NAME].index:
if i < n: if i < n:
continue continue
row = [] row = []
prev = i - 1 - n # we additionally substract 1, so row will contain current response and 7 previous responses prev = i - 1 - n # we additionally substract 1, so row will contain current response and 7 previous responses
for j in range(i, prev, -1): for j in range(i, prev, -1):
row.append(data.line[j]) row.append(data.line[j])
contexted.append(row) contexted.append(row)
@ -64,16 +68,21 @@ df = pd.DataFrame.from_records(contexted, columns=columns)
trn_df, val_df = train_test_split(df, test_size=0.1) trn_df, val_df = train_test_split(df, test_size=0.1)
# create dataset suitable for our model # create dataset suitable for our model
def construct_conv(row, tokenizer, eos = True):
flatten = lambda l: [item for sublist in l for item in sublist]
conv = list(reversed([tokenizer.encode(x) + [tokenizer.eos_token_id] for x in row])) def construct_conv(row, tokenizer, eos=True):
def flatten(l): return [item for sublist in l for item in sublist]
conv = list(
reversed([tokenizer.encode(x) + [tokenizer.eos_token_id] for x in row]))
conv = flatten(conv) conv = flatten(conv)
return conv return conv
class ConversationDataset(Dataset): class ConversationDataset(Dataset):
def __init__(self, tokenizer: PreTrainedTokenizer, args, df, block_size=512): def __init__(self, tokenizer: PreTrainedTokenizer, args, df, block_size=512):
block_size = block_size - (tokenizer.model_max_length - tokenizer.max_len_single_sentence) block_size = block_size - \
(tokenizer.model_max_length - tokenizer.max_len_single_sentence)
directory = args.cache_dir directory = args.cache_dir
cached_features_file = os.path.join( cached_features_file = os.path.join(
@ -81,7 +90,8 @@ class ConversationDataset(Dataset):
) )
if os.path.exists(cached_features_file) and not args.overwrite_cache: if os.path.exists(cached_features_file) and not args.overwrite_cache:
logger.info("Loading features from cached file %s", cached_features_file) logger.info("Loading features from cached file %s",
cached_features_file)
with open(cached_features_file, "rb") as handle: with open(cached_features_file, "rb") as handle:
self.examples = pickle.load(handle) self.examples = pickle.load(handle)
else: else:
@ -92,9 +102,11 @@ class ConversationDataset(Dataset):
conv = construct_conv(row, tokenizer) conv = construct_conv(row, tokenizer)
self.examples.append(conv) self.examples.append(conv)
logger.info("Saving features into cached file %s", cached_features_file) logger.info("Saving features into cached file %s",
cached_features_file)
with open(cached_features_file, "wb") as handle: with open(cached_features_file, "wb") as handle:
pickle.dump(self.examples, handle, protocol=pickle.HIGHEST_PROTOCOL) pickle.dump(self.examples, handle,
protocol=pickle.HIGHEST_PROTOCOL)
def __len__(self): def __len__(self):
return len(self.examples) return len(self.examples)
@ -102,7 +114,8 @@ class ConversationDataset(Dataset):
def __getitem__(self, item): def __getitem__(self, item):
return torch.tensor(self.examples[item], dtype=torch.long) return torch.tensor(self.examples[item], dtype=torch.long)
# Cacheing and storing of data/checkpoints # Caching and storing of data/checkpoints
def load_and_cache_examples(args, tokenizer, df_trn, df_val, evaluate=False): def load_and_cache_examples(args, tokenizer, df_trn, df_val, evaluate=False):
return ConversationDataset(tokenizer, args, df_val if evaluate else df_trn) return ConversationDataset(tokenizer, args, df_val if evaluate else df_trn)
@ -119,15 +132,18 @@ def set_seed(args):
def _sorted_checkpoints(args, checkpoint_prefix="checkpoint", use_mtime=False) -> List[str]: def _sorted_checkpoints(args, checkpoint_prefix="checkpoint", use_mtime=False) -> List[str]:
ordering_and_checkpoint_path = [] ordering_and_checkpoint_path = []
glob_checkpoints = glob.glob(os.path.join(args.output_dir, "{}-*".format(checkpoint_prefix))) glob_checkpoints = glob.glob(os.path.join(
args.output_dir, "{}-*".format(checkpoint_prefix)))
for path in glob_checkpoints: for path in glob_checkpoints:
if use_mtime: if use_mtime:
ordering_and_checkpoint_path.append((os.path.getmtime(path), path)) ordering_and_checkpoint_path.append((os.path.getmtime(path), path))
else: else:
regex_match = re.match(".*{}-([0-9]+)".format(checkpoint_prefix), path) regex_match = re.match(
".*{}-([0-9]+)".format(checkpoint_prefix), path)
if regex_match and regex_match.groups(): if regex_match and regex_match.groups():
ordering_and_checkpoint_path.append((int(regex_match.groups()[0]), path)) ordering_and_checkpoint_path.append(
(int(regex_match.groups()[0]), path))
checkpoints_sorted = sorted(ordering_and_checkpoint_path) checkpoints_sorted = sorted(ordering_and_checkpoint_path)
checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted] checkpoints_sorted = [checkpoint[1] for checkpoint in checkpoints_sorted]
@ -141,18 +157,19 @@ def _rotate_checkpoints(args, checkpoint_prefix="checkpoint", use_mtime=False) -
return return
# Check if we should delete older checkpoint(s) # Check if we should delete older checkpoint(s)
checkpoints_sorted = _sorted_checkpoints(args, checkpoint_prefix, use_mtime) checkpoints_sorted = _sorted_checkpoints(
args, checkpoint_prefix, use_mtime)
if len(checkpoints_sorted) <= args.save_total_limit: if len(checkpoints_sorted) <= args.save_total_limit:
return return
number_of_checkpoints_to_delete = max(0, len(checkpoints_sorted) - args.save_total_limit) 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] checkpoints_to_be_deleted = checkpoints_sorted[:number_of_checkpoints_to_delete]
for checkpoint in checkpoints_to_be_deleted: for checkpoint in checkpoints_to_be_deleted:
logger.info("Deleting older checkpoint [{}] due to args.save_total_limit".format(checkpoint)) logger.info(
"Deleting older checkpoint [{}] due to args.save_total_limit".format(checkpoint))
shutil.rmtree(checkpoint) shutil.rmtree(checkpoint)
from transformers import AutoModelWithLMHead, AutoModelForCausalLM, AutoTokenizer
import torch
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large") tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-large")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large") model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-large")
@ -170,13 +187,15 @@ MODEL_CONFIG_CLASSES = list(MODEL_WITH_LM_HEAD_MAPPING.keys())
MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES) MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
# Args to allow for easy conversion of python script to notebook # Args to allow for easy conversion of python script to notebook
class Args(): class Args():
def __init__(self): def __init__(self):
self.output_dir = 'models/output-medium' self.output_dir = 'cartman/models/output-medium-3ep'
self.model_type = 'gpt2' self.model_type = 'gpt2'
self.model_name_or_path = 'microsoft/DialoGPT-large' self.model_name_or_path = 'microsoft/DialoGPT-medium'
self.config_name = 'microsoft/DialoGPT-large' self.config_name = 'microsoft/DialoGPT-medium'
self.tokenizer_name = 'microsoft/DialoGPT-large' self.tokenizer_name = 'microsoft/DialoGPT-medium'
self.cache_dir = 'cached' self.cache_dir = 'cached'
self.block_size = 512 self.block_size = 512
self.do_train = True self.do_train = True
@ -189,7 +208,7 @@ class Args():
self.weight_decay = 0.0 self.weight_decay = 0.0
self.adam_epsilon = 1e-8 self.adam_epsilon = 1e-8
self.max_grad_norm = 1.0 self.max_grad_norm = 1.0
self.num_train_epochs = 4 self.num_train_epochs = 3
self.max_steps = -1 self.max_steps = -1
self.warmup_steps = 0 self.warmup_steps = 0
self.logging_steps = 1000 self.logging_steps = 1000
@ -205,8 +224,10 @@ class Args():
self.fp16 = False self.fp16 = False
self.fp16_opt_level = 'O1' self.fp16_opt_level = 'O1'
args = Args() args = Args()
def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedTokenizer) -> Tuple[int, float]: def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedTokenizer) -> Tuple[int, float]:
""" Train the model """ """ Train the model """
if args.local_rank in [-1, 0]: if args.local_rank in [-1, 0]:
@ -219,22 +240,25 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
return pad_sequence(examples, batch_first=True) return pad_sequence(examples, batch_first=True)
return pad_sequence(examples, batch_first=True, padding_value=tokenizer.pad_token_id) return pad_sequence(examples, batch_first=True, padding_value=tokenizer.pad_token_id)
train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset) train_sampler = RandomSampler(
train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
train_dataloader = DataLoader( train_dataloader = DataLoader(
train_dataset, sampler=train_sampler, batch_size=args.train_batch_size, collate_fn=collate, drop_last = True train_dataset, sampler=train_sampler, batch_size=args.train_batch_size, collate_fn=collate, drop_last=True
) )
if args.max_steps > 0: if args.max_steps > 0:
t_total = args.max_steps t_total = args.max_steps
args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1 args.num_train_epochs = args.max_steps // (
len(train_dataloader) // args.gradient_accumulation_steps) + 1
else: else:
t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs t_total = len(
train_dataloader) // args.gradient_accumulation_steps * args.num_train_epochs
model = model.module if hasattr(model, "module") else model # Take care of distributed/parallel training # Take care of distributed/parallel training
model = model.module if hasattr(model, "module") else model
model.resize_token_embeddings(len(tokenizer)) model.resize_token_embeddings(len(tokenizer))
# add_special_tokens_(model, tokenizer) # add_special_tokens_(model, tokenizer)
# Prepare optimizer and schedule (linear warmup and decay) # Prepare optimizer and schedule (linear warmup and decay)
no_decay = ["bias", "LayerNorm.weight"] no_decay = ["bias", "LayerNorm.weight"]
optimizer_grouped_parameters = [ optimizer_grouped_parameters = [
@ -242,9 +266,11 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
"params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)], "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, "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}, {"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) optimizer = AdamW(optimizer_grouped_parameters,
lr=args.learning_rate, eps=args.adam_epsilon)
scheduler = get_linear_schedule_with_warmup( scheduler = get_linear_schedule_with_warmup(
optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total
) )
@ -256,15 +282,19 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
and os.path.isfile(os.path.join(args.model_name_or_path, "scheduler.pt")) and os.path.isfile(os.path.join(args.model_name_or_path, "scheduler.pt"))
): ):
# Load in optimizer and scheduler states # Load in optimizer and scheduler states
optimizer.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "optimizer.pt"))) optimizer.load_state_dict(torch.load(
scheduler.load_state_dict(torch.load(os.path.join(args.model_name_or_path, "scheduler.pt"))) os.path.join(args.model_name_or_path, "optimizer.pt")))
scheduler.load_state_dict(torch.load(
os.path.join(args.model_name_or_path, "scheduler.pt")))
if args.fp16: if args.fp16:
try: try:
from apex import amp from apex import amp
except ImportError: except ImportError:
raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use fp16 training.") raise ImportError(
model, optimizer = amp.initialize(model, optimizer, opt_level=args.fp16_opt_level) "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)
# multi-gpu training (should be after apex fp16 initialization) # multi-gpu training (should be after apex fp16 initialization)
if args.n_gpu > 1: if args.n_gpu > 1:
@ -273,21 +303,24 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
# Distributed training (should be after apex fp16 initialization) # Distributed training (should be after apex fp16 initialization)
if args.local_rank != -1: if args.local_rank != -1:
model = torch.nn.parallel.DistributedDataParallel( model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True model, device_ids=[
args.local_rank], output_device=args.local_rank, find_unused_parameters=True
) )
# Train! # Train!
logger.info("***** Running training *****") logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_dataset)) logger.info(" Num examples = %d", len(train_dataset))
logger.info(" Num Epochs = %d", args.num_train_epochs) 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(" Instantaneous batch size per GPU = %d",
args.per_gpu_train_batch_size)
logger.info( logger.info(
" Total train batch size (w. parallel, distributed & accumulation) = %d", " Total train batch size (w. parallel, distributed & accumulation) = %d",
args.train_batch_size args.train_batch_size
* args.gradient_accumulation_steps * args.gradient_accumulation_steps
* (torch.distributed.get_world_size() if args.local_rank != -1 else 1), * (torch.distributed.get_world_size() if args.local_rank != -1 else 1),
) )
logger.info(" Gradient Accumulation steps = %d", args.gradient_accumulation_steps) logger.info(" Gradient Accumulation steps = %d",
args.gradient_accumulation_steps)
logger.info(" Total optimization steps = %d", t_total) logger.info(" Total optimization steps = %d", t_total)
global_step = 0 global_step = 0
@ -297,15 +330,21 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
if args.model_name_or_path and os.path.exists(args.model_name_or_path): if args.model_name_or_path and os.path.exists(args.model_name_or_path):
try: try:
# set global_step to gobal_step of last saved checkpoint from model path # set global_step to gobal_step of last saved checkpoint from model path
checkpoint_suffix = args.model_name_or_path.split("-")[-1].split("/")[0] checkpoint_suffix = args.model_name_or_path.split(
"-")[-1].split("/")[0]
global_step = int(checkpoint_suffix) global_step = int(checkpoint_suffix)
epochs_trained = global_step // (len(train_dataloader) // args.gradient_accumulation_steps) epochs_trained = global_step // (len(train_dataloader) //
steps_trained_in_current_epoch = global_step % (len(train_dataloader) // args.gradient_accumulation_steps) args.gradient_accumulation_steps)
steps_trained_in_current_epoch = global_step % (
len(train_dataloader) // args.gradient_accumulation_steps)
logger.info(" Continuing training from checkpoint, will skip to saved global_step") logger.info(
" Continuing training from checkpoint, will skip to saved global_step")
logger.info(" Continuing training from epoch %d", epochs_trained) logger.info(" Continuing training from epoch %d", epochs_trained)
logger.info(" Continuing training from global step %d", global_step) logger.info(
logger.info(" Will skip the first %d steps in the first epoch", steps_trained_in_current_epoch) " Continuing training from global step %d", global_step)
logger.info(" Will skip the first %d steps in the first epoch",
steps_trained_in_current_epoch)
except ValueError: except ValueError:
logger.info(" Starting fine-tuning.") logger.info(" Starting fine-tuning.")
@ -317,7 +356,8 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
) )
set_seed(args) # Added here for reproducibility set_seed(args) # Added here for reproducibility
for _ in train_iterator: for _ in train_iterator:
epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0]) epoch_iterator = tqdm(train_dataloader, desc="Iteration",
disable=args.local_rank not in [-1, 0])
for step, batch in enumerate(epoch_iterator): for step, batch in enumerate(epoch_iterator):
# Skip past any already trained steps if resuming training # Skip past any already trained steps if resuming training
@ -326,12 +366,14 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
continue continue
inputs, labels = (batch, batch) inputs, labels = (batch, batch)
if inputs.shape[1] > 1024: continue if inputs.shape[1] > 1024:
continue
inputs = inputs.to(args.device) inputs = inputs.to(args.device)
labels = labels.to(args.device) labels = labels.to(args.device)
model.train() model.train()
outputs = model(inputs, labels=labels) outputs = model(inputs, labels=labels)
loss = outputs[0] # model outputs are always tuple in transformers (see doc) # model outputs are always tuple in transformers (see doc)
loss = outputs[0]
if args.n_gpu > 1: if args.n_gpu > 1:
loss = loss.mean() # mean() to average on multi-gpu parallel training loss = loss.mean() # mean() to average on multi-gpu parallel training
@ -347,9 +389,11 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
tr_loss += loss.item() tr_loss += loss.item()
if (step + 1) % args.gradient_accumulation_steps == 0: if (step + 1) % args.gradient_accumulation_steps == 0:
if args.fp16: if args.fp16:
torch.nn.utils.clip_grad_norm_(amp.master_params(optimizer), args.max_grad_norm) torch.nn.utils.clip_grad_norm_(
amp.master_params(optimizer), args.max_grad_norm)
else: else:
torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm) torch.nn.utils.clip_grad_norm_(
model.parameters(), args.max_grad_norm)
optimizer.step() optimizer.step()
scheduler.step() # Update learning rate schedule scheduler.step() # Update learning rate schedule
model.zero_grad() model.zero_grad()
@ -362,15 +406,19 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
): # Only evaluate when single GPU otherwise metrics may not average well ): # Only evaluate when single GPU otherwise metrics may not average well
results = evaluate(args, model, tokenizer) results = evaluate(args, model, tokenizer)
for key, value in results.items(): for key, value in results.items():
tb_writer.add_scalar("eval_{}".format(key), value, global_step) tb_writer.add_scalar(
tb_writer.add_scalar("lr", scheduler.get_lr()[0], global_step) "eval_{}".format(key), value, global_step)
tb_writer.add_scalar("loss", (tr_loss - logging_loss) / args.logging_steps, 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)
logging_loss = tr_loss logging_loss = tr_loss
if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0: if args.local_rank in [-1, 0] and args.save_steps > 0 and global_step % args.save_steps == 0:
checkpoint_prefix = "checkpoint" checkpoint_prefix = "checkpoint"
# Save model checkpoint # Save model checkpoint
output_dir = os.path.join(args.output_dir, "{}-{}".format(checkpoint_prefix, global_step)) output_dir = os.path.join(
args.output_dir, "{}-{}".format(checkpoint_prefix, global_step))
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
model_to_save = ( model_to_save = (
model.module if hasattr(model, "module") else model model.module if hasattr(model, "module") else model
@ -378,14 +426,18 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
model_to_save.save_pretrained(output_dir) model_to_save.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir) tokenizer.save_pretrained(output_dir)
torch.save(args, os.path.join(output_dir, "training_args.bin")) torch.save(args, os.path.join(
output_dir, "training_args.bin"))
logger.info("Saving model checkpoint to %s", output_dir) logger.info("Saving model checkpoint to %s", output_dir)
_rotate_checkpoints(args, checkpoint_prefix) _rotate_checkpoints(args, checkpoint_prefix)
torch.save(optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt")) torch.save(optimizer.state_dict(), os.path.join(
torch.save(scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt")) output_dir, "optimizer.pt"))
logger.info("Saving optimizer and scheduler states to %s", output_dir) torch.save(scheduler.state_dict(), os.path.join(
output_dir, "scheduler.pt"))
logger.info(
"Saving optimizer and scheduler states to %s", output_dir)
if args.max_steps > 0 and global_step > args.max_steps: if args.max_steps > 0 and global_step > args.max_steps:
epoch_iterator.close() epoch_iterator.close()
@ -401,11 +453,13 @@ def train(args, train_dataset, model: PreTrainedModel, tokenizer: PreTrainedToke
# Evaluation of some model # Evaluation of some model
def evaluate(args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, df_trn, df_val, prefix="") -> Dict: def evaluate(args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, df_trn, df_val, prefix="") -> Dict:
# Loop to handle MNLI double evaluation (matched, mis-matched) # Loop to handle MNLI double evaluation (matched, mis-matched)
eval_output_dir = args.output_dir eval_output_dir = args.output_dir
eval_dataset = load_and_cache_examples(args, tokenizer, df_trn, df_val, evaluate=True) eval_dataset = load_and_cache_examples(
args, tokenizer, df_trn, df_val, evaluate=True)
os.makedirs(eval_output_dir, exist_ok=True) os.makedirs(eval_output_dir, exist_ok=True)
args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu) args.eval_batch_size = args.per_gpu_eval_batch_size * max(1, args.n_gpu)
# Note that DistributedSampler samples randomly # Note that DistributedSampler samples randomly
@ -417,7 +471,7 @@ def evaluate(args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, df_tr
eval_sampler = SequentialSampler(eval_dataset) eval_sampler = SequentialSampler(eval_dataset)
eval_dataloader = DataLoader( eval_dataloader = DataLoader(
eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size, collate_fn=collate, drop_last = True eval_dataset, sampler=eval_sampler, batch_size=args.eval_batch_size, collate_fn=collate, drop_last=True
) )
# multi-gpu evaluate # multi-gpu evaluate
@ -448,7 +502,8 @@ def evaluate(args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, df_tr
result = {"perplexity": perplexity} result = {"perplexity": perplexity}
output_eval_file = os.path.join(eval_output_dir, prefix, "eval_results.txt") output_eval_file = os.path.join(
eval_output_dir, prefix, "eval_results.txt")
with open(output_eval_file, "w") as writer: with open(output_eval_file, "w") as writer:
logger.info("***** Eval results {} *****".format(prefix)) logger.info("***** Eval results {} *****".format(prefix))
for key in sorted(result.keys()): for key in sorted(result.keys()):
@ -459,13 +514,15 @@ def evaluate(args, model: PreTrainedModel, tokenizer: PreTrainedTokenizer, df_tr
# Main runner # Main runner
def main(df_trn, df_val): def main(df_trn, df_val):
args = Args() args = Args()
if args.should_continue: if args.should_continue:
sorted_checkpoints = _sorted_checkpoints(args) sorted_checkpoints = _sorted_checkpoints(args)
if len(sorted_checkpoints) == 0: if len(sorted_checkpoints) == 0:
raise ValueError("Used --should_continue but no checkpoint was found in --output_dir.") raise ValueError(
"Used --should_continue but no checkpoint was found in --output_dir.")
else: else:
args.model_name_or_path = sorted_checkpoints[-1] args.model_name_or_path = sorted_checkpoints[-1]
@ -505,8 +562,10 @@ def main(df_trn, df_val):
# Set seed # Set seed
set_seed(args) set_seed(args)
config = AutoConfig.from_pretrained(args.config_name, cache_dir=args.cache_dir) config = AutoConfig.from_pretrained(
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name, cache_dir=args.cache_dir) args.config_name, cache_dir=args.cache_dir)
tokenizer = AutoTokenizer.from_pretrained(
args.tokenizer_name, cache_dir=args.cache_dir)
model = AutoModelForCausalLM.from_pretrained( model = AutoModelForCausalLM.from_pretrained(
args.model_name_or_path, args.model_name_or_path,
from_tf=False, from_tf=False,
@ -519,10 +578,12 @@ def main(df_trn, df_val):
# Training # Training
if args.do_train: if args.do_train:
train_dataset = load_and_cache_examples(args, tokenizer, df_trn, df_val, evaluate=False) train_dataset = load_and_cache_examples(
args, tokenizer, df_trn, df_val, evaluate=False)
global_step, tr_loss = train(args, train_dataset, model, tokenizer) global_step, tr_loss = train(args, train_dataset, model, tokenizer)
logger.info(" global_step = %s, average loss = %s", global_step, tr_loss) 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() # Saving best-practices: if you use save_pretrained for the model and tokenizer, you can reload them using from_pretrained()
if args.do_train: if args.do_train:
@ -546,26 +607,4 @@ def main(df_trn, df_val):
tokenizer = AutoTokenizer.from_pretrained(args.output_dir) tokenizer = AutoTokenizer.from_pretrained(args.output_dir)
model.to(args.device) 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 = AutoModelForCausalLM.from_pretrained(checkpoint)
model.to(args.device)
result = evaluate(args, model, tokenizer, df_trn, df_val, prefix=prefix)
result = dict((k + "_{}".format(global_step), v) for k, v in result.items())
results.update(result)
return results
main(trn_df, val_df) main(trn_df, val_df)