def train_model(self, train_df, output_dir=None, show_running_loss=True, args=None):
"""
Trains the model using 'train_df'
Args:
train_df: Pandas Dataframe containing at least two columns. If the Dataframe has a header, it should contain a 'text' and a 'labels' column. If no header is present,
the Dataframe should contain at least two columns, with the first column containing the text, and the second column containing the label. The model will be trained on this Dataframe.
output_dir: The directory where model files will be saved. If not given, self.args['output_dir'] will be used.
show_running_loss (optional): Set to False to prevent running loss from being printed to console. Defaults to True.
args (optional): Optional changes to the args dict of the model. Any changes made will persist for the model.
Returns:
None
"""
if args:
self.args.update(args)
if not output_dir:
output_dir = self.args['output_dir']
if os.path.exists(output_dir) and os.listdir(output_dir) and not self.args["overwrite_output_dir"]:
raise ValueError("Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(output_dir))
self._move_model_to_device()
if 'text' in train_df.columns and 'labels' in train_df.columns:
train_examples = [InputExample(i, text, None, label) for i, (text, label) in enumerate(zip(train_df['text'], train_df['labels']))]
else:
train_examples = [InputExample(i, text, None, label) for i, (text, label) in enumerate(zip(train_df.iloc[:, 0], train_df.iloc[:, 1]))]
train_dataset = self.load_and_cache_examples(train_examples)
global_step, tr_loss = self.train(train_dataset, output_dir, show_running_loss=show_running_loss)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
model_to_save = self.model.module if hasattr(self.model, "module") else self.model
model_to_save.save_pretrained(output_dir)
self.tokenizer.save_pretrained(output_dir)
torch.save(self.args, os.path.join(output_dir, "training_args.bin"))
if not args["silent"]:
print("Training of {} model complete. Saved to {}.".format(self.args["model_type"], output_dir))
def predict(self, to_predict):
"""
Performs predictions on a list of text.
Args:
to_predict: A python list of text (str) to be sent to the model for prediction.
Returns:
preds: A python list of the predictions (0 or 1) for each text.
model_outputs: A python list of the raw model outputs for each text.
"""
tokenizer = self.tokenizer
device = self.device
model = self.model
args = self.args
self._move_model_to_device()
eval_examples = [InputExample(i, text, None, 0) for i, text in enumerate(to_predict)]
eval_dataset = self.load_and_cache_examples(eval_examples, evaluate=True, no_cache=True)
eval_sampler = SequentialSampler(eval_dataset)
eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args["eval_batch_size"])
eval_loss = 0.0
nb_eval_steps = 0
preds = None
out_label_ids = None
for batch in tqdm(eval_dataloader, disable=args["silent"]):
model.eval()
batch = tuple(t.to(device) for t in batch)
with torch.no_grad():
inputs = self._get_inputs_dict(batch)
outputs = model(**inputs)
tmp_eval_loss, logits = outputs[:2]
eval_loss += tmp_eval_loss.mean().item()
nb_eval_steps += 1
if preds is None:
preds = logits.detach().cpu().numpy()
out_label_ids = inputs["labels"].detach().cpu().numpy()
else:
preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
out_label_ids = np.append(out_label_ids, inputs["labels"].detach().cpu().numpy(), axis=0)
eval_loss = eval_loss / nb_eval_steps
model_outputs = preds
preds = np.argmax(preds, axis=1)
return preds, model_outputs
def predict(self, to_predict):
"""
Performs predictions on a list of text.
Args:
to_predict: A python list of text (str) to be sent to the model for prediction.
Returns:
preds: A python list of the predictions (0 or 1) for each text.
model_outputs: A python list of the raw model outputs for each text.
"""
tokenizer = self.tokenizer
device = self.device
model = self.model
args = self.args
self.model.to(self.device)
eval_examples = [InputExample(i, text, None, 0) for i, text in enumerate(to_predict)]
eval_dataset = self.load_and_cache_examples(eval_examples, evaluate=True, no_cache=True)
eval_sampler = SequentialSampler(eval_dataset)
eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args['eval_batch_size'])
eval_loss = 0.0
nb_eval_steps = 0
preds = None
out_label_ids = None
for batch in tqdm(eval_dataloader):
model.eval()
batch = tuple(t.to(device) for t in batch)
with torch.no_grad():
inputs = {'input_ids': batch[0],
'attention_mask': batch[1],
# XLM don't use segment_ids
'token_type_ids': batch[2] if args['model_type'] in ['bert', 'xlnet'] else None,
'labels': batch[3]}
outputs = model(**inputs)
tmp_eval_loss, logits = outputs[:2]
eval_loss += tmp_eval_loss.mean().item()
nb_eval_steps += 1
if preds is None:
preds = logits.detach().cpu().numpy()
out_label_ids = inputs['labels'].detach().cpu().numpy()
else:
preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
out_label_ids = np.append(
out_label_ids, inputs['labels'].detach().cpu().numpy(), axis=0)
eval_loss = eval_loss / nb_eval_steps
model_outputs = preds
preds = np.argmax(preds, axis=1)
return preds, model_outputs
def train_model(self, train_df, output_dir=None):
"""
Trains the model using 'train_df'
Args:
train_df: Pandas Dataframe (no header) of two columns, first column containing the text, and the second column containing the label. The model will be trained on this Dataframe.
output_dir: The directory where model files will be saved. If not given, self.args['output_dir'] will be used.
Returns:
None
"""
if not output_dir:
output_dir = self.args['output_dir']
if os.path.exists(output_dir) and os.listdir(
output_dir) and not args['overwrite_output_dir']:
raise ValueError(
"Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome."
.format(output_dir))
self.model.to(self.device)
train_examples = [
InputExample(i, text, None, label)
for i, (text, label) in enumerate(
zip(train_df.iloc[:, 0], train_df.iloc[:, 1]))
]
train_dataset = self.load_and_cache_examples(train_examples)
global_step, tr_loss = self.train(train_dataset, output_dir)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
model_to_save = self.model.module if hasattr(self.model,
'module') else self.model
model_to_save.save_pretrained(output_dir)
self.tokenizer.save_pretrained(output_dir)
torch.save(self.args, os.path.join(output_dir, 'training_args.bin'))
print(
f'Training of {self.args["model_type"]} model complete. Saved to {output_dir}.'
)
def evaluate(self, eval_df, output_dir, prefix="", **kwargs):
"""
Evaluates the model on eval_df.
Utility function to be used by the eval_model() method. Not intended to be used directly.
"""
tokenizer = self.tokenizer
device = self.device
model = self.model
args = self.args
eval_output_dir = output_dir
results = {}
if 'text' in eval_df.columns and 'labels' in eval_df.columns:
eval_examples = [InputExample(i, text, None, label) for i, (text, label) in enumerate(zip(eval_df['text'], eval_df['labels']))]
else:
eval_examples = [InputExample(i, text, None, label) for i, (text, label) in enumerate(zip(eval_df.iloc[:, 0], eval_df.iloc[:, 1]))]
eval_dataset = self.load_and_cache_examples(eval_examples, evaluate=True)
if not os.path.exists(eval_output_dir):
os.makedirs(eval_output_dir)
eval_sampler = SequentialSampler(eval_dataset)
eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=args["eval_batch_size"])
eval_loss = 0.0
nb_eval_steps = 0
preds = None
out_label_ids = None
model.eval()
for batch in tqdm(eval_dataloader, disable=args["silent"]):
batch = tuple(t.to(device) for t in batch)
with torch.no_grad():
inputs = self._get_inputs_dict(batch)
outputs = model(**inputs)
tmp_eval_loss, logits = outputs[:2]
eval_loss += tmp_eval_loss.mean().item()
nb_eval_steps += 1
if preds is None:
preds = logits.detach().cpu().numpy()
out_label_ids = inputs["labels"].detach().cpu().numpy()
else:
preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
out_label_ids = np.append(
out_label_ids, inputs["labels"].detach().cpu().numpy(), axis=0)
eval_loss = eval_loss / nb_eval_steps
model_outputs = preds
preds = np.argmax(preds, axis=1)
result, wrong = self.compute_metrics(preds, out_label_ids, eval_examples, **kwargs)
results.update(result)
output_eval_file = os.path.join(eval_output_dir, "eval_results.txt")
with open(output_eval_file, "w") as writer:
for key in sorted(result.keys()):
writer.write("{} = {}\n".format(key, str(result[key])))
return results, model_outputs, wrong
def evaluate(self, eval_df, output_dir, prefix="", **kwargs):
"""
Evaluates the model on eval_df.
Utility function to be used by the eval_model() method. Not intended to be used directly.
"""
tokenizer = self.tokenizer
device = self.device
model = self.model
args = self.args
eval_output_dir = output_dir
results = {}
eval_examples = [
InputExample(i, text, None, label) for i, (
text,
label) in enumerate(zip(eval_df.iloc[:, 0], eval_df.iloc[:,
1]))
]
eval_dataset = self.load_and_cache_examples(eval_examples,
evaluate=True)
if not os.path.exists(eval_output_dir):
os.makedirs(eval_output_dir)
eval_sampler = SequentialSampler(eval_dataset)
eval_dataloader = DataLoader(eval_dataset,
sampler=eval_sampler,
batch_size=args['eval_batch_size'])
eval_loss = 0.0
nb_eval_steps = 0
preds = None
out_label_ids = None
for batch in tqdm(eval_dataloader):
model.eval()
batch = tuple(t.to(device) for t in batch)
with torch.no_grad():
inputs = {
'input_ids': batch[0],
'attention_mask': batch[1],
'labels': batch[3]
}
if args.model_type != 'distilbert':
inputs['token_type_ids'] = batch[2] if args.model_type in [
'bert', 'xlnet'
] else None # XLM, DistilBERT and RoBERTa don't use segment_ids
outputs = model(**inputs)
tmp_eval_loss, logits = outputs[:2]
eval_loss += tmp_eval_loss.mean().item()
nb_eval_steps += 1
if preds is None:
preds = logits.detach().cpu().numpy()
out_label_ids = inputs['labels'].detach().cpu().numpy()
else:
preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)
out_label_ids = np.append(
out_label_ids,
inputs['labels'].detach().cpu().numpy(),
axis=0)
eval_loss = eval_loss / nb_eval_steps
model_outputs = preds
preds = np.argmax(preds, axis=1)
result, wrong = self.compute_metrics(preds, out_label_ids,
eval_examples, **kwargs)
results.update(result)
output_eval_file = os.path.join(eval_output_dir, "eval_results.txt")
with open(output_eval_file, "w") as writer:
for key in sorted(result.keys()):
writer.write("%s = %s\n" % (key, str(result[key])))
return results, model_outputs, wrong