import logging
import os
import time
import json
import torch
import argparse
import numpy as np
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, NewType, NamedTuple, Union, Tuple
from tqdm import tqdm
from torch import nn
from torch.utils.data.dataset import Dataset
from torch.utils.data.dataloader import DataLoader
from torch.utils.data.sampler import SequentialSampler
from transformers import (
AutoConfig,
AutoTokenizer,
set_seed,
PreTrainedTokenizer,
BertTokenizerFast
)
from ops import (
json_to_sent,
input_form,
get_prob,
detokenize,
preprocess,
Profile,
)
from models import RoBERTaMultiNER2, BERTMultiNER2
logger = logging.getLogger(__name__)
InputDataClass = NewType("InputDataClass", Any)
DataCollator = NewType("DataCollator", Callable[[List[InputDataClass]], Dict[str, torch.Tensor]])
@dataclass
class InputExample:
"""
A single training/test example for token classification.
Args:
guid: Unique id for the example.
words: list. The words of the sequence.
labels: (Optional) list. The labels for each word of the sequence. This should be
specified for train and dev examples, but not for test examples.
"""
guid: str
words: List[str]
labels: Optional[List[str]]
entity_labels: Optional[List[int]]
@dataclass
class InputFeatures:
"""
A single set of features of data.
Property names are the same names as the corresponding inputs to a model.
"""
input_ids: List[int]
attention_mask: List[int]
token_type_ids: Optional[List[int]] = None
label_ids: Optional[List[int]] = None
entity_type_ids: Optional[List[int]] = None
class DataProcessor(object):
"""Base class for data converters for sequence classification data sets."""
def get_train_examples(self, data_dir):
"""Gets a collection of `InputExample`s for the train set."""
raise NotImplementedError()
def get_dev_examples(self, data_dir):
"""Gets a collection of `InputExample`s for the dev set."""
raise NotImplementedError()
def get_labels(self):
"""Gets the list of labels for this data set."""
raise NotImplementedError()
@classmethod
def _read_data(cls, data, pmids):
"""Reads a BIO data."""
lines = []
words = []
labels = []
entity_labels = []
for pmid in pmids:
for sent in data[pmid]['words']:
words = sent[:]
labels = ['O'] * len(words)
entity_labels = [str(0)] * len(words)
if len(words) >= 30:
while len(words) >= 30:
tmplabel = labels[:30]
l = ' '.join([label for label
in labels[:len(tmplabel)]
if len(label) > 0])
w = ' '.join([word for word
in words[:len(tmplabel)]
if len(word) > 0])
e = ' '.join([el for el
in entity_labels[:len(tmplabel)]
if len(el) > 0])
lines.append([l, w, e])
words = words[len(tmplabel):]
labels = labels[len(tmplabel):]
entity_labels = entity_labels[len(tmplabel):]
if len(words) == 0:
continue
l = ' '.join([label for label in labels if len(label) > 0])
w = ' '.join([word for word in words if len(word) > 0])
e = ' '.join([el for el in entity_labels if len(entity_labels) > 0])
lines.append([l, w, e])
words = []
labels = []
entity_labels = []
continue
return lines
class NerDataset(Dataset):
"""
This will be superseded by a framework-agnostic approach soon.
"""
features: List[InputFeatures]
pad_token_label_id: int = nn.CrossEntropyLoss().ignore_index
def __init__(
self,
predict_examples,
labels: List[str],
tokenizer: PreTrainedTokenizer,
config,
params,
base_name
):
logger.info(f"Creating features from dataset file")
self.labels = labels
self.predict_examples = predict_examples
self.tokenizer = tokenizer
self.config = config
self.params = params
self.features = convert_examples_to_features(
self.predict_examples,
self.labels,
self.params.max_seq_length,
self.tokenizer,
cls_token_at_end=bool(self.config.model_type in ["xlnet"]),
cls_token=self.tokenizer.cls_token,
cls_token_segment_id=2 if self.config.model_type in ["xlnet"] else 0,
sep_token=self.tokenizer.sep_token,
sep_token_extra=False,
pad_on_left=bool(self.tokenizer.padding_side=="left"),
pad_token=self.tokenizer.pad_token_id,
pad_token_segment_id=self.tokenizer.pad_token_type_id,
pad_token_label_id=self.pad_token_label_id,
base_name=base_name,
)
def __len__(self):
return len(self.features)
def __getitem__(self, i) -> InputFeatures:
return self.features[i]
class PredictionOutput(NamedTuple):
predictions: np.ndarray
label_ids: Optional[np.ndarray]
def default_data_collator(features: List[InputDataClass]) -> Dict[str, torch.Tensor]:
"""
Very simple data collator that:
- simply collates batches of dict-like objects
- Performs special handling for potential keys named:
- `label`: handles a single value (int or float) per object
- `label_ids`: handles a list of values per object
- does not do any additional preprocessing
i.e., Property names of the input object will be used as corresponding inputs to the model.
See glue and ner for example of how it's useful.
"""
# In this function we'll make the assumption that all `features` in the batch
# have the same attributes.
# So we will look at the first element as a proxy for what attributes exist
# on the whole batch.
if not isinstance(features[0], dict):
features = [vars(f) for f in features]
first = features[0]
batch = {}
# Special handling for labels.
# Ensure that tensor is created with the correct type
# (it should be automatically the case, but let's make sure of it.)
if "label" in first and first["label"] is not None:
dtype = torch.long if type(first["label"]) is int else torch.float
batch["labels"] = torch.tensor([f["label"] for f in features], dtype=dtype)
elif "label_ids" in first and first["label_ids"] is not None:
if isinstance(first["label_ids"], torch.Tensor):
batch["labels"] = torch.stack([f["label_ids"] for f in features])
else:
dtype = torch.long if type(first["label_ids"][0]) is int else torch.float
batch["labels"] = torch.tensor([f["label_ids"] for f in features], dtype=dtype)
# Handling of all other possible keys.
# Again, we will use the first element to figure out which key/values are not None for this model.
for k, v in first.items():
if k not in ("label", "label_ids") and v is not None and not isinstance(v, str):
if isinstance(v, torch.Tensor):
batch[k] = torch.stack([f[k] for f in features])
else:
batch[k] = torch.tensor([f[k] for f in features], dtype=torch.long)
return batch
def convert_examples_to_features(
examples: List[InputExample],
label_list: List[str],
max_seq_length: int,
tokenizer: PreTrainedTokenizer,
cls_token_at_end=False,
cls_token="[CLS]",
cls_token_segment_id=1,
sep_token="[SEP]",
sep_token_extra=False,
pad_on_left=False,
pad_token=0,
pad_token_segment_id=0,
pad_token_label_id=-100,
sequence_a_segment_id=0,
mask_padding_with_zero=True,
base_name="",
) -> List[InputFeatures]:
""" Loads a data file into a list of `InputFeatures`
`cls_token_at_end` define the location of the CLS token:
- False (Default, BERT/XLM pattern): [CLS] + A + [SEP] + B + [SEP]
- True (XLNet/GPT pattern): A + [SEP] + B + [SEP] + [CLS]
`cls_token_segment_id` define the segment id associated to the CLS token (0 for BERT, 2 for XLNet)
"""
# TODO clean up all this to leverage built-in features of tokenizers
label_map = {label: i for i, label in enumerate(label_list)}
features = []
for (ex_index, example) in tqdm(enumerate(examples)):
if ex_index % 10_000 == 0:
logger.info("Writing example %d of %d", ex_index, len(examples))
tokens, label_ids, = [], []
det_tokens = []
for word_idx, (word, label) in enumerate(zip(example.words.split(), example.labels.split())):
word_tokens = tokenizer.tokenize(word)
# bert-base-multilingual-cased sometimes output "nothing ([]) when calling tokenize with just a space.
if len(word_tokens) > 0:
tokens.extend(word_tokens)
# Use the real label id for the first token of the word, and padding ids for the remaining tokens
label_ids.extend([label_map[label]] + [pad_token_label_id] * (len(word_tokens) - 1))
if len(word_tokens) == 1:
det_tokens.extend(word_tokens)
elif len(word_tokens) > 1:
for det_idx, det_word in enumerate(word_tokens):
if det_idx > 0:
det_word = '##' + det_word
det_tokens.append(det_word)
else:
det_tokens.append(det_word)
# calculate temperature with length : temp = 1 - 0.02 * length
# temperature = [1 - sharpening * i if i > 1 else i for _, i in enumerate(entity_length)]
# Account for [CLS] and [SEP] with "- 2" and with "- 3" for RoBERTa.
special_tokens_count = tokenizer.num_special_tokens_to_add()
## truncating tokens with max_seq_length
# if len(tokens) > max_seq_length - special_tokens_count:
# tokens = tokens[: (max_seq_length - special_tokens_count)]
# label_ids = label_ids[: (max_seq_length - special_tokens_count)]
# det_tokens = det_tokens[: (max_seq_length - special_tokens_count)]
# for sliding window tokens - update 23.11.13
for i in range(0, (len(tokens) // max_seq_length) + 1):
if i == 0:
window_tokens = tokens[i*max_seq_length:(i+1)*max_seq_length-special_tokens_count]
window_label_ids = label_ids[i*max_seq_length:(i+1)*max_seq_length-special_tokens_count]
window_det_tokens = det_tokens[i*max_seq_length:(i+1)*max_seq_length-special_tokens_count]
elif i >= 1:
window_tokens = tokens[i*max_seq_length-special_tokens_count:(i+1)*max_seq_length-special_tokens_count]
window_label_ids = label_ids[i*max_seq_length-special_tokens_count:(i+1)*max_seq_length-special_tokens_count]
window_det_tokens = det_tokens[i*max_seq_length-special_tokens_count:(i+1)*max_seq_length-special_tokens_count]
# The convention in BERT is:
# (a) For sequence pairs:
# tokens: [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
# type_ids: 0 0 0 0 0 0 0 0 1 1 1 1 1 1
# (b) For single sequences:
# tokens: [CLS] the dog is hairy . [SEP]
# type_ids: 0 0 0 0 0 0 0
#
# Where "type_ids" are used to indicate whether this is the first
# sequence or the second sequence. The embedding vectors for `type=0` and
# `type=1` were learned during pre-training and are added to the wordpiece
# embedding vector (and position vector). This is not *strictly* necessary
# since the [SEP] token unambiguously separates the sequences, but it makes
# it easier for the model to learn the concept of sequences.
#
# For classification tasks, the first vector (corresponding to [CLS]) is
# used as as the "sentence vector". Note that this only makes sense because
# the entire model is fine-tuned.
window_tokens += [sep_token]
window_label_ids += [pad_token_label_id]
window_det_tokens += [sep_token]
if sep_token_extra:
# roberta uses an extra separator b/w pairs of sentences
window_tokens += [sep_token]
window_label_ids += [pad_token_label_id]
window_det_tokens += [sep_token]
# make entity type label index for multiner
entity_type_ids = [int(example.entity_labels[0])] * len(window_tokens)
segment_ids = [sequence_a_segment_id] * len(window_tokens)
if cls_token_at_end:
window_tokens += [cls_token]
window_label_ids += [pad_token_label_id]
segment_ids += [cls_token_segment_id]
entity_type_ids += [int(example.entity_labels[0])]
window_det_tokens += [cls_token]
else:
window_tokens = [cls_token] + window_tokens
window_label_ids = [pad_token_label_id] + window_label_ids
segment_ids = [cls_token_segment_id] + segment_ids
entity_type_ids = [int(example.entity_labels[0])] + entity_type_ids
window_det_tokens = [cls_token] + window_det_tokens
input_ids = tokenizer.convert_tokens_to_ids(window_tokens)
# The mask has 1 for real tokens and 0 for padding tokens. Only real
# tokens are attended to.
input_mask = [1 if mask_padding_with_zero else 0] * len(input_ids)
# Zero-pad up to the sequence length.
padding_length = max_seq_length - len(input_ids)
if pad_on_left:
input_ids = ([pad_token] * padding_length) + input_ids
input_mask = ([0 if mask_padding_with_zero else 1] * padding_length) + input_mask
segment_ids = ([pad_token_segment_id] * padding_length) + segment_ids
window_label_ids = ([pad_token_label_id] * padding_length) + window_label_ids
entity_type_ids = ([int(example.entity_labels[0])] * padding_length) + entity_type_ids
window_tokens = (["**NULL**"] * padding_length) + window_tokens
window_det_tokens = (["**NULL**"] * padding_length) + window_det_tokens
else:
input_ids += [pad_token] * padding_length
input_mask += [0 if mask_padding_with_zero else 1] * padding_length
segment_ids += [pad_token_segment_id] * padding_length
window_label_ids += [pad_token_label_id] * padding_length
entity_type_ids += [int(example.entity_labels[0])] * padding_length
window_tokens += ["**NULL**"] * padding_length
window_det_tokens += ["**NULL**"] * padding_length
assert len(input_ids) == max_seq_length
assert len(input_mask) == max_seq_length
assert len(segment_ids) == max_seq_length
assert len(window_label_ids) == max_seq_length
assert len(entity_type_ids) == max_seq_length
assert len(window_tokens) == max_seq_length
if ex_index < 1:
logger.info("*** Example ***")
logger.info("guid: %s", example.guid)
logger.info("tokens: %s", " ".join([str(x) for x in window_tokens]))
logger.info("input_ids: %s", " ".join([str(x) for x in input_ids]))
logger.info("input_mask: %s", " ".join([str(x) for x in input_mask]))
logger.info("segment_ids: %s", " ".join([str(x) for x in segment_ids]))
logger.info("label_ids: %s", " ".join([str(x) for x in window_label_ids]))
logger.info("entity_type_ids: %s", " ".join([str(x) for x in entity_type_ids]))
if "token_type_ids" not in tokenizer.model_input_names:
segment_ids = None
features.append(
InputFeatures(
input_ids=input_ids, attention_mask=input_mask, token_type_ids=segment_ids, \
label_ids=window_label_ids, entity_type_ids=entity_type_ids, \
)
)
write_tokens(window_tokens, window_det_tokens, 'test', base_name)
return features
def write_tokens(tokens, det_tokens, mode, base_name):
if mode == "test":
tmp_path = os.path.join('multi_ner', 'tmp')
if not os.path.exists(tmp_path):
os.makedirs(tmp_path)
path = os.path.join("multi_ner", "tmp",
"token_{}_{}.txt".format(mode, base_name))
with open(path, 'a') as wf:
for token in tokens:
if token != "**NULL**":
wf.write(token + '\n')
det_path = os.path.join("multi_ner", "tmp",
"det_token_{}_{}.txt".format(mode, base_name))
with open(det_path, 'a') as wf:
for token in det_tokens:
if token != "**NULL**":
wf.write(token + '\n')
class NerProcessor(DataProcessor):
def get_test_examples(self, data_dir):
data = list()
pmids = list()
with open(data_dir, 'r') as in_:
for line in in_:
line = line.strip()
tmp = json.loads(line)
tmp['title'] = preprocess(tmp['title'])
tmp['abstract'] = preprocess(tmp['abstract'])
data.append(tmp)
pmids.append(tmp["pmid"])
json_file = input_form(json_to_sent(data))
return \
self._create_example(self._read_data(json_file, pmids), "test"), \
json_file, data
def get_test_dict_list(self, dict_list):
pmids = list()
for d in dict_list:
pmids.append(d["pmid"])
json_file = input_form(json_to_sent(dict_list))
return \
self._create_example(self._read_data(json_file, pmids), "test"), \
json_file
def get_labels(self):
return ["B", "I", "O"]
def _create_example(self, lines, set_type):
examples = []
for (i,line) in enumerate(lines):
guid = "%s-%s" % (set_type, i)
text = line[1]
label = line[0]
entity_labels = line[2]
examples.append(InputExample(guid=guid, words=text, labels=label, entity_labels=entity_labels))
return examples
class BioMedNER:
def __init__(self, params):
# See all possible arguments in src/transformers/training_args.py
# or by passing the --help flag to this script.
# We now keep distinct sets of args, for a cleaner separation of concerns.
init_start_t = time.time()
# Set ner processor
self.processor = NerProcessor()
# Setup parsing
self.params = params
self.prediction_loss_only = False
# Set seed
set_seed(self.params.seed)
# Prepare Labels
self.labels = self.processor.get_labels()
self.id2label: Dict[int, str] = {i: label for i, label in enumerate(self.labels)}
self.label2id = {label:i for i, label in enumerate(self.labels)}
self.num_labels = len(self.labels)
self.config = AutoConfig.from_pretrained(
self.params.model_name_or_path,
num_labels=self.num_labels,
id2label=self.id2label,
label2id=self.label2id,
)
self.tokenizer = BertTokenizerFast.from_pretrained(
self.params.model_name_or_path,
)
self.model = BERTMultiNER2.from_pretrained(
self.params.model_name_or_path,
num_labels=self.num_labels,
config=self.config,
)
if not self.params.no_cuda:
self.model = self.model.cuda()
self.entity_types = ['disease', 'drug', 'gene', 'species', 'cell_line', 'DNA', 'RNA', 'cell_type']
# 'biological_structure', 'diagnostic_procedure', 'duration', 'date', 'therapeutic_procedure',
# 'sign_symptom', 'lab_value']
self.estimator_dict = {}
for etype in self.entity_types:
self.estimator_dict[etype] = {}
self.estimator_dict[etype]['prediction'] = []
self.estimator_dict[etype]['log_probs'] = []
self.counter = 0
self.pad_token_label_id:int = nn.CrossEntropyLoss().ignore_index
init_end_t = time.time()
print('BioMedNER init_t {:.3f} sec.'.format(init_end_t - init_start_t))
@Profile(__name__)
def recognize(self, input_dl, base_name, indent=None):
if type(input_dl) is str:
predict_examples, self.json_dict, self.data_list = \
self.processor.get_test_examples(input_dl)
elif type(input_dl) is list:
predict_examples, self.json_dict = \
self.processor.get_test_dict_list(input_dl)
self.data_list = input_dl
else:
raise ValueError('Wrong type')
token_path = os.path.join("multi_ner", "tmp",
"token_test_{}.txt".format(base_name))
det_token_path = os.path.join("multi_ner", "tmp",
"det_token_test_{}.txt".format(base_name))
if os.path.exists(token_path):
os.remove(token_path)
if os.path.exists(det_token_path):
os.remove(det_token_path)
predict_example_list = (NerDataset(predict_examples, self.labels,\
self.tokenizer, self.config, self.params, base_name))
tokens, tot_tokens = list(), list()
"""
Aggregate label results with detokenized tokens
words: <s> Auto phagy main tain s tumour growth ... </s>
label: O O O O O O B O ... O
detok_words: <s> Authophagy maintains tumour growth ... </s>
detok_label: O O O B O ... </s>
"""
with open(det_token_path, 'r') as reader:
for line_idx, line in enumerate(reader):
tok = line.strip()
tot_tokens.append(tok)
if tok == '[CLS]' or tok == '<s>':
tmp_toks = [tok]
elif tok == '[SEP]' or tok == '</s>':
tmp_toks.append(tok)
tokens.append(tmp_toks)
else:
tmp_toks.append(tok)
self.predict_dict, self.prob_dict = dict(), dict()
threads, self.out_tag_dict = list(), dict()
all_type = self._predict(predict_example_list)
# disease, drug, gene, spec, cell_line, dna, rna, cell_type
for etype_idx, etype in enumerate(self.entity_types):
predictions, label_ids = all_type[etype_idx] # batch, seq, labels
preds_array = self.align_predictions(predictions) # batch, seq
self.out_tag_dict[etype] = (False, None)
self.recognize_etype(etype, tokens, tot_tokens, predictions, preds_array)
for etype in self.entity_types:
if self.out_tag_dict[etype][0]:
if type(input_dl) is str:
print(os.path.split(input_dl)[1],
'Found an error:', self.out_tag_dict[etype][1])
else:
print('Found an error:', self.out_tag_dict[etype][1])
if os.path.exists(token_path):
os.remove(token_path)
return None
# get probability of all mentions
data_list = get_prob(self.data_list, self.json_dict, self.predict_dict,
self.prob_dict, entity_types=self.entity_types)
if type(input_dl) is str:
output_path = os.path.join('result/', os.path.splitext(
os.path.basename(input_dl))[0] + '_NER_{}.json'.format(base_name))
print('pred', output_path)
with open(output_path, 'w') as resultf:
for paper in data_list:
paper['ner_model'] = "MULTI-TASK NER v.20210707"
resultf.write(
json.dumps(paper, sort_keys=True, indent=indent) + '\n'
)
# delete temp files
if os.path.exists(token_path):
os.remove(token_path)
if os.path.exists(det_token_path):
os.remove(det_token_path)
return data_list
@Profile(__name__)
def recognize_etype(self, etype, tokens, tot_tokens, predictions, preds_array):
result = []
for one_batch in range(predictions.shape[0]):
result.append({'prediction':preds_array[one_batch],
'log_probs':predictions[one_batch]})
predicts = list()
logits = list()
for pidx, prediction in enumerate(result):
slen = len(tokens[pidx])
for p in prediction['prediction'][:slen]:
predicts.append(self.id2label[p])
for l in prediction['log_probs'][:slen]:
logits.append(l)
de_toks, de_labels, de_logits = detokenize(tot_tokens, predicts, logits)
self.predict_dict[etype] = dict()
self.prob_dict[etype] = dict()
piv = 0
for data in self.data_list:
pmid = data['pmid']
self.predict_dict[etype][pmid] = list()
self.prob_dict[etype][pmid] = list()
sent_lens = list()
for sent in self.json_dict[pmid]['words']:
sent_lens.append(len(sent))
sent_idx = 0
de_i = 0
overlen = False
while True:
if overlen:
try:
self.predict_dict[etype][pmid][-1].extend(
de_labels[piv + de_i])
except Exception as e:
self.out_tag_dict[etype] = (True, e)
break
self.prob_dict[etype][pmid][-1].extend(de_logits[piv + de_i])
de_i += 1
if len(self.predict_dict[etype][pmid][-1]) == len(
self.json_dict[pmid]['words'][
len(self.predict_dict[etype][pmid]) - 1]):
sent_idx += 1
overlen = False
else:
self.predict_dict[etype][pmid].append(de_labels[piv + de_i])
self.prob_dict[etype][pmid].append(de_logits[piv + de_i])
de_i += 1
if len(self.predict_dict[etype][pmid][-1]) == len(
self.json_dict[pmid]['words'][
len(self.predict_dict[etype][pmid]) - 1]):
sent_idx += 1
overlen = False
else:
overlen = True
if sent_idx == len(self.json_dict[pmid]['words']):
piv += de_i
break
if self.out_tag_dict[etype][0]:
break
def _predict(self, test_dataset:Dataset):
sampler = SequentialSampler(test_dataset)
data_loader = DataLoader(
test_dataset,
sampler=sampler,
batch_size=32, # you can adjust evaluation batch size, we prefer using 32
collate_fn=default_data_collator,
drop_last=False,
)
return self._prediction_loop(data_loader, description="Prediction")
def _prediction_loop(
self, dataloader: DataLoader, description: str, prediction_loss_only: Optional[bool] = None
) -> PredictionOutput:
"""
Prediction/evaluation loop, shared by `evaluate()` and `predict()`.
Works both with or without labels.
"""
prediction_loss_only = prediction_loss_only if prediction_loss_only is not None else self.prediction_loss_only
model = self.model
eval_losses: List[float] = []
dise_preds: torch.Tensor = None
chem_preds: torch.Tensor = None
gene_preds: torch.Tensor = None
spec_preds: torch.Tensor = None
cl_preds: torch.Tensor = None
dna_preds: torch.Tensor = None
rna_preds: torch.Tensor = None
ct_preds: torch.Tensor = None
# biological_preds: torch.Tensor = None
# diagnostic_preds: torch.Tensor = None
# duration_preds: torch.Tensor = None
# date_preds: torch.Tensor = None
# therapeutic_preds: torch.Tensor = None
# sign_symptom_preds: torch.Tensor = None
# lab_value_preds: torch.Tensor = None
label_ids: torch.Tensor = None
model.eval()
for inputs in tqdm(dataloader, desc=description):
has_labels = any(inputs.get(k) is not None for k in ["labels", "lm_labels", "masked_lm_labels"])
for k, v in inputs.items():
if isinstance(v, torch.Tensor):
inputs[k] = v.to(self.model.device)
with torch.no_grad():
outputs = model(**inputs)
if has_labels:
step_eval_loss, logits = outputs[:2]
eval_losses += [step_eval_loss.mean().item()]
else:
logits = outputs[0]
if not prediction_loss_only:
(dise_logits, chem_logits, gene_logits, spec_logits, cl_logits, dna_logits, rna_logits, ct_logits) = logits
# biological_logits, diagnostic_logits, duration_logits, date_logits, therapeutic_logits,
# sign_symptom_logits, lab_value_logits) = logits
if dise_preds is None \
and chem_preds is None \
and gene_preds is None \
and spec_preds is None \
and cl_preds is None \
and dna_preds is None \
and rna_preds is None \
and ct_preds is None :
# and biological_preds is None \
# and diagnostic_preds is None \
# and duration_preds is None \
# and date_preds is None \
# and therapeutic_preds is None \
# and sign_symptom_preds is None \
# and lab_value_preds is None:
dise_preds = dise_logits.detach()
chem_preds = chem_logits.detach()
gene_preds = gene_logits.detach()
spec_preds = spec_logits.detach()
cl_preds = cl_logits.detach()
dna_preds = dna_logits.detach()
rna_preds = rna_logits.detach()
ct_preds = ct_logits.detach()
# biological_preds = biological_logits.detach()
# diagnostic_preds = diagnostic_logits.detach()
# duration_preds = duration_logits.detach()
# date_preds = date_logits.detach()
# therapeutic_preds = therapeutic_logits.detach()
# sign_symptom_preds = sign_symptom_logits.detach()
# lab_value_preds = lab_value_logits.detach()
else:
dise_preds = torch.cat((dise_preds, dise_logits.detach()), dim=0)
chem_preds = torch.cat((chem_preds, chem_logits.detach()), dim=0)
gene_preds = torch.cat((gene_preds, gene_logits.detach()), dim=0)
spec_preds = torch.cat((spec_preds, spec_logits.detach()), dim=0)
cl_preds = torch.cat((cl_preds, cl_logits.detach()), dim=0)
dna_preds = torch.cat((dna_preds, dna_logits.detach()), dim=0)
rna_preds = torch.cat((rna_preds, rna_logits.detach()), dim=0)
ct_preds = torch.cat((ct_preds, ct_logits.detach()), dim=0)
# biological_preds = torch.cat((biological_preds, biological_logits.detach()), dim=0)
# diagnostic_preds = torch.cat((diagnostic_preds, diagnostic_logits.detach()), dim=0)
# duration_preds = torch.cat((duration_preds, duration_logits.detach()), dim=0)
# date_preds = torch.cat((date_preds, date_logits.detach()), dim=0)
# therapeutic_preds = torch.cat((therapeutic_preds, therapeutic_logits.detach()), dim=0)
# sign_symptom_preds = torch.cat((sign_symptom_preds, sign_symptom_logits.detach()), dim=0)
# lab_value_preds = torch.cat((lab_value_preds, lab_value_logits.detach()), dim=0)
if inputs.get("labels") is not None:
if label_ids is None:
label_ids = inputs["labels"].detach()
else:
label_ids = torch.cat((label_ids, inputs["labels"].detach()), dim=0)
# Finally, turn the aggregated tensors into numpy arrays.
if dise_preds is not None \
and chem_preds is not None \
and gene_preds is not None \
and spec_preds is not None \
and cl_preds is not None \
and dna_preds is not None \
and rna_preds is not None \
and ct_preds is not None :
# and biological_preds is not None \
# and diagnostic_preds is not None \
# and duration_preds is not None \
# and date_preds is not None \
# and therapeutic_preds is not None \
# and sign_symptom_preds is not None \
# and lab_value_preds is not None:
dise_preds = dise_preds.cpu().numpy()
chem_preds = chem_preds.cpu().numpy()
gene_preds = gene_preds.cpu().numpy()
spec_preds = spec_preds.cpu().numpy()
cl_preds = cl_preds.cpu().numpy()
dna_preds = dna_preds.cpu().numpy()
rna_preds = rna_preds.cpu().numpy()
ct_preds = ct_preds.cpu().numpy()
# biological_preds = biological_preds.cpu().numpy()
# diagnostic_preds = diagnostic_preds.cpu().numpy()
# duration_preds = duration_preds.cpu().numpy()
# date_preds = date_preds.cpu().numpy()
# therapeutic_preds = therapeutic_preds.cpu().numpy()
# sign_symptom_preds = sign_symptom_preds.cpu().numpy()
# lab_value_preds = lab_value_preds.cpu().numpy()
if label_ids is not None:
label_ids = label_ids.cpu().numpy()
return_output = (PredictionOutput(predictions=dise_preds, label_ids=label_ids), \
PredictionOutput(predictions=chem_preds, label_ids=label_ids), \
PredictionOutput(predictions=gene_preds, label_ids=label_ids), \
PredictionOutput(predictions=spec_preds, label_ids=label_ids), \
PredictionOutput(predictions=cl_preds, label_ids=label_ids), \
PredictionOutput(predictions=dna_preds, label_ids=label_ids), \
PredictionOutput(predictions=rna_preds, label_ids=label_ids), \
PredictionOutput(predictions=ct_preds, label_ids=label_ids))
# PredictionOutput(predictions=biological_preds, label_ids=label_ids),
# PredictionOutput(predictions=diagnostic_preds, label_ids=label_ids),
# PredictionOutput(predictions=duration_preds, label_ids=label_ids),
# PredictionOutput(predictions=date_preds, label_ids=label_ids),
# PredictionOutput(predictions=therapeutic_preds, label_ids=label_ids),
# PredictionOutput(predictions=sign_symptom_preds, label_ids=label_ids),
# PredictionOutput(predictions=lab_value_preds, label_ids=label_ids))
return return_output
def align_predictions(self, predictions: np.ndarray) -> List[int]:
preds = np.argmax(predictions, axis=2)
batch_size, seq_len = preds.shape
preds_list = [[] for _ in range(batch_size)]
for i in range(batch_size):
for j in range(seq_len):
preds_list[i].append(preds[i][j])
return np.array(preds_list)
def main():
os.environ["CUDA_VISIBLE_DEVICES"]="6"
argparser = argparse.ArgumentParser()
argparser.add_argument('--model_name_or_path', default='dmis-lab/bern2-ner')
argparser.add_argument('--max_seq_length', type=int, help='The maximum total input sequence length after tokenization. Sequences longer than this will be truncated, sequences shorter will be padded.',
default=128)
argparser.add_argument('--seed', type=int, help='random seed for initialization',
default=1)
args = argparser.parse_args()
biomedner = BioMedNER(args)
if __name__ == "__main__":
main()
Datasets
Models
