import json
import logging
from typing import List, Any
import copy
import torch
from torch.utils.data import Dataset
from transformers import AutoTokenizer, AutoModelForTokenClassification, Trainer
from util.process_data import Sample, Entity, EntityType, EntityTypeSet, SampleList, Token, Relation
from util.configuration import InferenceConfiguration
valid_relations = { # head : [tail, ...]
"StatedKeyFigure": ["StatedKeyFigure", "Condition", "StatedExpression", "DeclarativeExpression"],
"DeclarativeKeyFigure": ["DeclarativeKeyFigure", "Condition", "StatedExpression", "DeclarativeExpression"],
"StatedExpression": ["Unit", "Factor", "Range", "Condition"],
"DeclarativeExpression": ["DeclarativeExpression", "Unit", "Factor", "Range", "Condition"],
"Condition": ["Condition", "StatedExpression", "DeclarativeExpression"],
"Range": ["Range"]
}
class TokenClassificationDataset(Dataset):
""" Pytorch Dataset """
def __init__(self, encodings, labels):
self.encodings = encodings
self.labels = labels
def __getitem__(self, idx):
item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
item['labels'] = torch.tensor(self.labels[idx])
return item
def __len__(self):
return len(self.labels)
class TransformersInference():
def __init__(self, config: InferenceConfiguration):
super().__init__()
self.__logger = logging.getLogger(self.__class__.__name__)
self.__logger.info(f"Load Configuration: {config.dict()}")
with open(f"classification.json", mode='r', encoding="utf-8") as f:
self.__entity_type_set = EntityTypeSet.parse_obj(json.load(f))
self.__entity_type_label_to_id_mapping = {x.label: x.idx for x in self.__entity_type_set.all_types()}
self.__entity_type_id_to_label_mapping = {x.idx: x.label for x in self.__entity_type_set.all_types()}
self.__logger.info("Load Model: " + config.model_path_keyfigure)
self.__tokenizer = AutoTokenizer.from_pretrained(config.transformer_model,
padding="max_length", max_length=512, truncation=True)
self.__model = AutoModelForTokenClassification.from_pretrained(config.model_path_keyfigure, num_labels=(
len(self.__entity_type_set)))
self.__trainer = Trainer(model=self.__model)
self.__merge_entities = config.merge_entities
self.__split_len = config.split_len
self.__extract_relations = config.extract_relations
# add special tokens
entity_groups = self.__entity_type_set.groups
num_entity_groups = len(entity_groups)
lst_special_tokens = ["[REL]", "[SUB]", "[/SUB]", "[OBJ]", "[/OBJ]"]
for grp_idx, grp in enumerate(entity_groups):
lst_special_tokens.append(f"[GRP-{grp_idx:02d}]")
lst_special_tokens.extend([f"[ENT-{ent:02d}]" for ent in grp if ent != self.__entity_type_set.id_of_non_entity])
lst_special_tokens.extend([f"[/ENT-{ent:02d}]" for ent in grp if ent != self.__entity_type_set.id_of_non_entity])
lst_special_tokens = sorted(list(set(lst_special_tokens)))
special_tokens_dict = {'additional_special_tokens': lst_special_tokens }
num_added_toks = self.__tokenizer.add_special_tokens(special_tokens_dict)
self.__logger.info(f"Added {num_added_toks} new special tokens. All special tokens: '{self.__tokenizer.all_special_tokens}'")
self.__logger.info("Initialization completed.")
def run_inference(self, sample_list: SampleList):
group_predictions = []
group_entity_ids = []
self.__logger.info("Predict Entities ...")
for grp_idx, grp in enumerate(self.__entity_type_set.groups):
token_lists = [[x.text for x in sample.tokens] for sample in sample_list.samples]
predictions = self.__get_predictions(token_lists, f"[GRP-{grp_idx:02d}]")
group_entity_ids_ = []
for sample, prediction_per_tokens in zip(sample_list.samples, predictions):
group_entity_ids_.append(self.generate_response_entities(sample, prediction_per_tokens, grp_idx))
group_predictions.append(predictions)
group_entity_ids.append(group_entity_ids_)
if self.__extract_relations:
self.__logger.info("Predict Relations ...")
self.__do_extract_relations(sample_list, group_predictions, group_entity_ids)
def __do_extract_relations(self, sample_list, group_predictions, group_entity_ids):
id_of_non_entity = self.__entity_type_set.id_of_non_entity
for sample_idx, sample in enumerate(sample_list.samples):
masked_tokens = []
masked_tokens_align = []
# create SUB-Mask for every entity that can be a head
head_entities = [entity_ for entity_ in sample.entities if entity_.ent_type.label in list(valid_relations.keys())]
for entity_ in head_entities:
ent_masked_tokens = []
ent_masked_tokens_align = []
last_preds = [id_of_non_entity for group in group_predictions]
last_ent_ids = [-1 for group in group_entity_ids]
for token_idx, token in enumerate(sample.tokens):
for group, ent_ids, last_pred, last_ent_id in zip(group_predictions, group_entity_ids, last_preds, last_ent_ids):
pred = group[sample_idx][token_idx]
ent_id = ent_ids[sample_idx][token_idx]
if last_pred != pred and last_pred != id_of_non_entity:
mask = "[/SUB]" if last_ent_id == entity_.id else "[/OBJ]"
ent_masked_tokens.extend([f"[/ENT-{last_pred:02d}]", mask])
ent_masked_tokens_align.extend([str(last_ent_id), str(last_ent_id)])
for group, ent_ids, last_pred, last_ent_id in zip(group_predictions, group_entity_ids, last_preds, last_ent_ids):
pred = group[sample_idx][token_idx]
ent_id = ent_ids[sample_idx][token_idx]
if last_pred != pred and pred != id_of_non_entity:
mask = "[SUB]" if ent_id == entity_.id else "[OBJ]"
ent_masked_tokens.extend([mask, f"[ENT-{pred:02d}]"])
ent_masked_tokens_align.extend([str(ent_id), str(ent_id)])
ent_masked_tokens.append(token.text)
ent_masked_tokens_align.append(token.text)
for idx, group in enumerate(group_predictions):
last_preds[idx] = group[sample_idx][token_idx]
for idx, group in enumerate(group_entity_ids):
last_ent_ids[idx] = group[sample_idx][token_idx]
for group, ent_ids, last_pred, last_ent_id in zip(group_predictions, group_entity_ids, last_preds, last_ent_ids):
pred = group[sample_idx][token_idx]
ent_id = ent_ids[sample_idx][token_idx]
if last_pred != id_of_non_entity:
mask = "[/SUB]" if last_ent_id == entity_.id else "[/OBJ]"
ent_masked_tokens.extend([f"[/ENT-{last_pred:02d}]", mask])
ent_masked_tokens_align.extend([str(last_ent_id), str(last_ent_id)])
masked_tokens.append(ent_masked_tokens)
masked_tokens_align.append(ent_masked_tokens_align)
rel_predictions = self.__get_predictions(masked_tokens, "[REL]")
self.generate_response_relations(sample, head_entities, masked_tokens_align, rel_predictions)
def generate_response_entities(self, sample: Sample, predictions_per_tokens: List[int], grp_idx: int):
entities = []
entity_ids = []
id_of_non_entity = self.__entity_type_set.id_of_non_entity
idx = grp_idx * 1000
for token, prediction in zip(sample.tokens, predictions_per_tokens):
if id_of_non_entity == prediction:
entity_ids.append(-1)
continue
idx += 1
entities.append(self.__build_entity(idx, prediction, token))
entity_ids.append(idx)
if self.__merge_entities:
entities = self.__do_merge_entities(copy.deepcopy(entities))
prev_pred = id_of_non_entity
for idx, pred in enumerate(predictions_per_tokens):
if prev_pred == pred and idx > 0:
entity_ids[idx] = entity_ids[idx-1]
prev_pred = pred
sample.entities += entities
tags = sample.tags if len(sample.tags) > 0 else [self.__entity_type_set.id_of_non_entity] * len(sample.tokens)
for tag_id, tok in enumerate(sample.tokens):
for ent in entities:
if tok.start >= ent.start and tok.start < ent.end:
tags[tag_id] = ent.ent_type.idx
logging.info(tags)
sample.tags = tags
return entity_ids
def generate_response_relations(self, sample: Sample, head_entities: List[Entity], masked_tokens_align: List[List[str]], rel_predictions: List[List[int]]):
relations = []
id_of_non_entity = self.__entity_type_set.id_of_non_entity
idx = 0
for entity_, align_per_ent, prediction_per_ent in zip(head_entities, masked_tokens_align, rel_predictions):
for token, prediction in zip(align_per_ent, prediction_per_ent):
if id_of_non_entity == prediction:
continue
try:
tail = int(token)
except:
continue
if not self.__validate_relation(sample.entities, entity_.id, tail, prediction):
continue
idx += 1
relations.append(self.__build_relation(idx, entity_.id, tail, prediction))
sample.relations = relations
def __validate_relation(self, entities: List[Entity], head: int, tail: int, prediction: int):
if head == tail: return False
head_ents = [ent.ent_type.label for ent in entities if ent.id==head]
tail_ents = [ent.ent_type.label for ent in entities if ent.id==tail]
if len(head_ents) > 0:
head_ent = head_ents[0]
else:
return False
if len(tail_ents) > 0:
tail_ent = tail_ents[0]
else:
return False
return tail_ent in valid_relations[head_ent]
def __build_entity(self, idx: int, prediction: int, token: Token) -> Entity:
return Entity(
id=idx,
text=token.text,
start=token.start,
end=token.end,
ent_type=EntityType(
idx=prediction,
label=self.__entity_type_id_to_label_mapping[prediction]
)
)
def __build_relation(self, idx: int, head: int, tail: int, prediction: int) -> Relation:
return Relation(
id=idx,
head=head,
tail=tail,
rel_type=EntityType(
idx=prediction,
label=self.__entity_type_id_to_label_mapping[prediction]
)
)
def __do_merge_entities(self, input_ents_):
out_ents = list()
current_ent = None
for ent in input_ents_:
if current_ent is None:
current_ent = ent
else:
idx_diff = ent.start - current_ent.end
if ent.ent_type.idx == current_ent.ent_type.idx and idx_diff <= 1:
current_ent.end = ent.end
current_ent.text += (" " if idx_diff == 1 else "") + ent.text
else:
out_ents.append(current_ent)
current_ent = ent
if current_ent is not None:
out_ents.append(current_ent)
return out_ents
def __get_predictions(self, token_lists: List[List[str]], trigger: str) -> List[List[int]]:
""" Get predictions of Transformer Sequence Labeling model """
if self.__split_len > 0:
token_lists_split = self.__do_split_sentences(token_lists, self.__split_len)
predictions = []
for sample_token_lists in token_lists_split:
sample_token_lists_trigger = [[trigger]+sample for sample in sample_token_lists]
val_encodings = self.__tokenizer(sample_token_lists_trigger, is_split_into_words=True, padding='max_length', truncation=True) # return_tensors="pt"
val_labels = []
for i in range(len(sample_token_lists_trigger)):
word_ids = val_encodings.word_ids(batch_index=i)
label_ids = [0 for _ in word_ids]
val_labels.append(label_ids)
val_dataset = TokenClassificationDataset(val_encodings, val_labels)
predictions_raw, _, _ = self.__trainer.predict(val_dataset)
predictions_align = self.__align_predictions(predictions_raw, val_encodings)
confidence = [[max(token) for token in sample] for sample in predictions_align]
predictions_sample = [[token.index(max(token)) for token in sample][1:] for sample in predictions_align]
predictions_part = []
for tok, pred in zip(sample_token_lists_trigger, predictions_sample):
if trigger == "[REL]" and "[SUB]" not in tok:
predictions_part += [self.__entity_type_set.id_of_non_entity] * len(pred)
else:
predictions_part += pred
predictions.append(predictions_part)
# predictions.append([j for i in predictions_sample for j in i]))
else:
token_lists_trigger = [[trigger]+sample for sample in token_lists]
val_encodings = self.__tokenizer(token_lists_trigger, is_split_into_words=True, padding='max_length', truncation=True) # return_tensors="pt"
val_labels = []
for i in range(len(token_lists_trigger)):
word_ids = val_encodings.word_ids(batch_index=i)
label_ids = [0 for _ in word_ids]
val_labels.append(label_ids)
val_dataset = TokenClassificationDataset(val_encodings, val_labels)
predictions_raw, _, _ = self.__trainer.predict(val_dataset)
predictions_align = self.__align_predictions(predictions_raw, val_encodings)
confidence = [[max(token) for token in sample] for sample in predictions_align]
predictions = [[token.index(max(token)) for token in sample][1:] for sample in predictions_align]
return predictions
def __do_split_sentences(self, tokens_: List[List[str]], split_len_ = 200) -> List[List[List[str]]]:
# split token lists into shorter lists
res_tokens = []
for tok_lst in tokens_:
res_tokens_sample = []
length = len(tok_lst)
if length > split_len_:
num_lists = length // split_len_ + (1 if (length % split_len_) > 0 else 0)
new_length = int(length / num_lists) + 1
self.__logger.info(f"Splitting a list of {length} elements into {num_lists} lists of length {new_length}..")
start_idx = 0
for i in range(num_lists):
end_idx = min(start_idx + new_length, length)
if "\n" in tok_lst[start_idx]: tok_lst[start_idx] = "."
if "\n" in tok_lst[end_idx-1]: tok_lst[end_idx-1] = "."
res_tokens_sample.append(tok_lst[start_idx:end_idx])
start_idx = end_idx
res_tokens.append(res_tokens_sample)
else:
res_tokens.append([tok_lst])
return res_tokens
def __align_predictions(self, predictions, tokenized_inputs, sum_all_tokens=False) -> List[List[List[float]]]:
""" Align predicted labels from Transformer Tokenizer """
confidence = []
id_of_non_entity = self.__entity_type_set.id_of_non_entity
for i, tagset in enumerate(predictions):
word_ids = tokenized_inputs.word_ids(batch_index=i)
previous_word_idx = None
token_confidence = []
for k, word_idx in enumerate(word_ids):
try:
tok_conf = [value for value in tagset[k]]
except TypeError:
# use the object itself it if's not iterable
tok_conf = tagset[k]
if word_idx is not None:
# add nonentity tokens if there is a gap in word ids (usually caused by a newline token)
if previous_word_idx is not None:
diff = word_idx - previous_word_idx
for i in range(diff - 1):
tmp = [0 for _ in tok_conf]
tmp[id_of_non_entity] = 1.0
token_confidence.append(tmp)
# add confidence value if this is the first token of the word
if word_idx != previous_word_idx:
token_confidence.append(tok_conf)
else:
# if sum_all_tokens=True the confidence for all tokens of one word will be summarized
if sum_all_tokens:
token_confidence[-1] = [a + b for a, b in zip(token_confidence[-1], tok_conf)]
previous_word_idx = word_idx
confidence.append(token_confidence)
return confidence