import torch
import argparse
import os
import numpy as np
from pathlib import Path
import logging
from tqdm import tqdm
from huggingface_hub import login
from transformer_deid.data import DeidDataset
from transformer_deid.tokenization import merge_sequences, encodings_to_label_list
from transformers import AutoModelForTokenClassification, AutoTokenizer
from transformer_deid.load_data import load_data, create_deid_dataset, save_labels
from transformer_deid.train import which_transformer_arch
logger = logging.getLogger(__name__)
def annotate(model, test_dataset: DeidDataset, device):
"""Annotates dataset with PHI labels.
Args:
model: trained HuggingFace model
test_dataset: DeidDataset, see data.py
device: either cuda or cpu
Returns:
new_labels: list of lists of [entity type, start, length] objects for each document
"""
model = model.to(device)
model.eval()
logger.info(
f'Running predictions over {len(test_dataset.encodings.encodings)} split sequences from {len(test_dataset.ids)} documents.'
)
result = [
get_logits(encoding, model, device)
for encoding in tqdm(test_dataset.encodings.encodings,
total=len(test_dataset.encodings.encodings))
]
predicted_label = np.argmax(result, axis=2)
id2label = model.config.id2label
predicted_label = [[id2label[token] for token in doc]
for doc in predicted_label]
labels = [encodings_to_label_list(doc, test_dataset.encodings[i], id2label=id2label) for i, doc in enumerate(predicted_label)]
new_labels = merge_sequences(labels, test_dataset.ids)
return new_labels
def get_logits(encodings, model, device):
""" Return predicted labels from the encodings of a *single* text example. """
result = model(input_ids=torch.tensor([encodings.ids]).to(device),
attention_mask=torch.tensor([encodings.attention_mask
]).to(device))
logits = result['logits'].cpu().detach().numpy()
return logits[0]
def parse_args():
parser = argparse.ArgumentParser(
description='Train a transformer-based PHI deidentification model.')
parser.add_argument('-i',
'--test_path',
type=str,
help='string to diretory containing txt directory for the test set.')
parser.add_argument('-m', '--model_path', type=str, help='model directory')
parser.add_argument('-o',
'--output_path',
help='output path in which to optionally save the annotations',
default=None)
args = parser.parse_args()
return args
def main(args):
train_path = args.test_path
out_path = args.output_path
modelDir = args.model_path
baseArchitecture = os.path.basename(modelDir).split('-')[0].lower()
_, tokenizer, _ = which_transformer_arch(baseArchitecture)
data_dict = load_data(Path(train_path))
test_dataset = create_deid_dataset(data_dict, tokenizer)
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
logger.info(f'Running using {device}.')
model = AutoModelForTokenClassification.from_pretrained(modelDir)
annotations = annotate(model, test_dataset, device)
if out_path is not None:
logger.info(f'Saving annotations to {out_path}.')
save_labels(annotations, data_dict['guid'], out_path)
else:
return annotations, test_dataset.ids
if __name__ == '__main__':
args = parse_args()
main(args)