import torch
import random
import os
import numpy as np
import logging
import argparse
from pathlib import Path
from transformers import DistilBertForTokenClassification, BertForTokenClassification, RobertaForTokenClassification
from transformers import BertTokenizerFast, DistilBertTokenizerFast, RobertaTokenizerFast
from transformers import Trainer, TrainingArguments
from transformer_deid.load_data import create_deid_dataset, get_labels, load_data
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__)
def seed_everything(seed: int):
random.seed(seed)
os.environ['PYTHONHASHSEED'] = str(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = True
def which_transformer_arch(baseArchitecture):
""" Gets architecture-specific parameters for each supported base architecture. """
if baseArchitecture == 'bert':
load_model = BertForTokenClassification.from_pretrained
tokenizer = BertTokenizerFast.from_pretrained('bert-base-cased')
baseArchitecture = 'bert-base-cased'
elif baseArchitecture == 'roberta':
load_model = RobertaForTokenClassification.from_pretrained
tokenizer = RobertaTokenizerFast.from_pretrained('roberta-base')
baseArchitecture = 'roberta-base'
elif baseArchitecture == 'distilbert':
load_model = DistilBertForTokenClassification.from_pretrained
tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-cased')
baseArchitecture = 'distilbert-base-cased'
else:
raise NotImplementedError(f'{baseArchitecture} not a supported model.')
return load_model, tokenizer, baseArchitecture
def train(train_data_dict, architecture, epochs, out_dir):
""" Trains a transformer-based deidentification model over {epochs} using {architecture}.
Args:
- train_data_dict: dict with 'txt', 'ann', and 'guid' keys with the training data
- architecture: transformer model to use
e.g., bert, roberta, or distilbert
- epochs: int, number of epochs to use
- out_dir: directory to which to save the model; only passed because it's mandatory for TrainingArguments()
Returns: trained HuggingFace Trainer object
"""
seed_everything(42)
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
logger.info(f'Running using {device}.')
load_model, tokenizer, baseArchitecture = which_transformer_arch(
architecture)
unique_labels = get_labels(train_data_dict['ann'])
label2id = {tag: id for id, tag in enumerate(unique_labels)}
id2label = {id: tag for tag, id in label2id.items()}
train_dataset = create_deid_dataset(train_data_dict, tokenizer, label2id)
model = load_model(baseArchitecture,
num_labels=len(unique_labels),
label2id=label2id,
id2label=id2label).to(device)
train_batch_size = 8
training_args = TrainingArguments(
output_dir=out_dir, # this is a mandatory argument
num_train_epochs=epochs,
per_device_train_batch_size=train_batch_size,
per_device_eval_batch_size=8,
warmup_steps=500,
weight_decay=0.01,
logging_dir='./logs',
logging_steps=10,
save_strategy='no', # probably make none? or make a parameter?
)
trainer = Trainer(model=model,
args=training_args,
train_dataset=train_dataset)
logger.info("***** Running training *****")
logger.info(" Num examples = %d", len(train_dataset))
logger.info(" Num Epochs = %d", training_args.num_train_epochs)
trainer.train()
return trainer
def parse_args():
parser = argparse.ArgumentParser(
description='Train a transformer-based PHI deidentification model.')
parser.add_argument('-i',
'--train_path',
type=str,
help=
'string to diretory containing txt and ann directories for the training set.')
parser.add_argument('-e',
'--epochs',
type=int,
help='number of epochs to train over',
default=5)
parser.add_argument('-m',
'--model_architecture',
type=str,
choices=['bert', 'distilbert', 'roberta'],
help='name of model architecture, either bert, roberta, or distilbert',
default='bert')
parser.add_argument('-o',
'--output_path',
help='output path in which to save the model')
args = parser.parse_args()
return args
def main(args):
# arguments
train_path = args.train_path
out_path = args.output_path
model = args.model_architecture
epochs = args.epochs
# read in data from training directory
data_dict = load_data(Path(train_path))
trainer = train(data_dict, model, epochs, out_path)
save_location = f'{out_path}/{model}_model_{epochs}'
trainer.save_model(save_location)
if __name__ == '__main__':
args = parse_args()
main(args)