add demonstrator

README.md

@@ -1,6 +1,6 @@
 ---
 title: NLP Legal Texts
-emoji: 🐠
+emoji: ⚖‍
 colorFrom: red
 colorTo: gray
 sdk: streamlit

app.py

@@ -0,0 +1,209 @@
+import os
+import streamlit as st
+from spacy import displacy
+from PIL import Image
+import json
+import requests
+from pyvis.network import Network
+import streamlit.components.v1 as components
+
+from util.process_data import Entity, EntityType, Relation, Sample, SampleList
+from util.tokenizer import Tokenizer
+from model_inference import TransformersInference
+from util.configuration import InferenceConfiguration
+
+inference_config = InferenceConfiguration()
+tokenizer = Tokenizer(inference_config.spacy_model)
+
+SAMPLE_66 = "EStG § 66 Höhe des Kindergeldes, Zahlungszeitraum (1) Das Kindergeld beträgt monatlich für das erste und zweite Kind jeweils 219 Euro, für das dritte Kind 225 Euro und für das vierte und jedes weitere Kind jeweils 250 Euro."
+SAMPLE_9 = "EStG § 9 Werbungskosten ... Zur Abgeltung dieser Aufwendungen ist für jeden Arbeitstag, an dem der Arbeitnehmer die erste Tätigkeitsstätte aufsucht eine Entfernungspauschale für jeden vollen Kilometer der Entfernung zwischen Wohnung und erster Tätigkeitsstätte von 0,30 Euro anzusetzen, höchstens jedoch 4 500 Euro im Kalenderjahr; ein höherer Betrag als 4 500 Euro ist anzusetzen, soweit der Arbeitnehmer einen eigenen oder ihm zur Nutzung überlassenen Kraftwagen benutzt."
+
+
+############################################################
+## Constants
+############################################################
+max_width_str = f"max-width: 60%;"
+paragraph = None
+style = "<style>mark.entity { display: inline-block }</style>"
+graph_options = ''' 
+var options = {
+  "edges": {
+    "arrows": {
+      "to": {
+        "enabled": true,
+        "scaleFactor": 1.2
+      }
+    }
+  }
+}
+'''
+
+legend_content = {
+    "text": "StatedKeyFigure StatedExpression Unit Range Factor Condition DeclarativeKeyFigure DeclarativeExpression",
+    "ents": [
+        {"start": 0, "end": 15, "label": "K"}, 
+        {"start": 16, "end": 32, "label": "E"},
+        {"start": 33, "end": 37, "label": "U"}, 
+        {"start": 38, "end": 43, "label": "R"}, 
+        {"start": 44, "end": 50, "label": "F"}, 
+        {"start": 51, "end": 60, "label": "C"}, 
+        {"start": 61, "end": 81, "label": "DK"}, 
+        {"start": 82, "end": 103, "label": "DE"}, 
+    ]}
+legend_options = {
+    "ents": ["K","U","E","R","F","C","DK","DE"], 
+    "colors": {'K': '#46d000',"U": "#e861ef", "E": "#538cff", "R": "#ffbe00", "F": "#0fd5dc", "C":"#ff484b", "DK":"#46d000", "DE":"#538cff"}
+    }
+legend_mapping = {"StatedKeyFigure": "K","Unit": "U","StatedExpression": "E","Range": "R","Factor": "F","Condition": "C","DeclarativeKeyFigure": "DK","DeclarativeExpression": "DE"}
+edge_colors = {'hasKeyFigure': '#46d000',"hasUnit": "#e861ef", "hasExpression": "#538cff", "hasRange": "#ffbe00", "hasFactor": "#0fd5dc", "hasCondition":"#ff484b", "join":"#aaa", "Typ":"#aaa", "hasParagraph": "#FF8B15"}
+
+
+############################################################
+## Function definitions
+############################################################
+
+def get_html(html: str, legend=False):
+    """Convert HTML so it can be rendered."""
+    WRAPPER = """<div style="overflow-x: auto; border: 1px solid #e6e9ef; border-radius: 0.25rem; padding: 1.5rem">{}</div>"""
+    if legend: WRAPPER = """<div style="overflow-x: auto; padding: 1rem">{}</div>"""
+    # Newlines seem to mess with the rendering
+    html = html.replace("\n", " ")
+    return WRAPPER.format(html)
+
+
+def get_displacy_ent_obj(paragraph, bedingungen=False, send_request=False):
+    entities = []
+    for entity in paragraph['entities']:
+        label = entity["entity"] if not send_request else entity["ent_type"]["label"]
+        if (bedingungen and label == "Condition") or (not bedingungen and label != "Condition") :
+            entities.append({
+                'start': entity['start'],
+                'end': entity["end"],
+                'label': legend_mapping[label]
+            })
+    return [{'text': paragraph['text'], 'ents': entities}]
+
+
+def request_extractor(text_data): 
+    try:
+        data = SampleList(
+                samples=[
+                    Sample(
+                        idx=0,
+                        text=str(text_data),
+                        entities=[],
+                        relations=[]
+                    )
+                ]
+            )
+        tokenizer.run(data)
+
+        model_inference = TransformersInference(inference_config)
+        model_inference.run_inference(data)
+        return data.dict()["samples"][0]
+    except Exception as e:
+        result = e
+    return {"text":"error","entities":[], "relations":[]}
+
+
+def generate_graph(nodes, edges, send_request=False):
+    net = Network(height="450px", width="100%")#, bgcolor="#222222", font_color="white", select_menu=True, filter_menu=True)
+    for node in nodes:
+        if "id" in node:
+            label = node["entity"] if not send_request else node["ent_type"]["label"]
+            node_color = legend_options["colors"][legend_mapping[label]]
+            node_label = node["text"] if len(node["text"]) < 30 else (node["text"][:27]+" ...")
+            if label in ["Kennzahl", "Kennzahlumschreibung"]:
+                net.add_node(node["id"], label=node_label, title=node["text"], mass=2, shape="ellipse", color=node_color, physics=False)
+            else:
+                net.add_node(node["id"], label=node_label, title=node["text"], mass=1, shape="ellipse", color=node_color) 
+    for edge in edges:
+        label = edge["relation"] if not send_request else edge["rel_type"]["label"]
+        net.add_edge(edge["head"], edge["tail"], width=1, title=label, arrowStrikethrough=False, color=edge_colors[label])
+    # net.force_atlas_2based()    # barnes_hut()  force_atlas_2based()  hrepulsion()   repulsion()
+    net.toggle_physics(True)
+    net.set_edge_smooth("dynamic")  # dynamic, continuous, discrete, diagonalCross, straightCross, horizontal, vertical, curvedCW, curvedCCW, cubicBezier
+    net.set_options(graph_options)
+    html_graph = net.generate_html()
+    return html_graph
+
+############################################################
+## Page configuration
+############################################################
+st.set_page_config(
+    page_title="NLP Gesetzestexte", 
+    menu_items={
+         'Get Help': None,
+         'Report a bug': None,
+         'About': "## Demonstrator NLP"
+        }
+        #  layout="wide")
+    )
+
+st.markdown(
+    f"""
+    <style>
+    .appview-container .main .block-container{{
+        {max_width_str}
+    }}
+    </style>    
+    """,
+    unsafe_allow_html=True,
+    )
+
+# radio button formatting in line
+st.write('<style>div.row-widget.stRadio > div{flex-direction:row;justify-content: left;} </style>', unsafe_allow_html=True)
+
+############################################################
+## Page formating
+############################################################
+col3, col4 = st.columns([2.4,1.6])
+st.write('\n')
+st.write('\n')
+
+
+with col3:
+    st.subheader("Extraction of Key Figures")
+    st.write("Demonstrator Application for Paper 'Semantic Extraction of Key Figures and Their Properties From Tax Legal Texts using Neural Models'")
+with col4:
+    st.caption("Semantic Model")
+    image = Image.open('util/ontology.png')
+    st.image(image, width=350)
+
+
+text_option = st.radio("Select Example", ["Insert your paragraph", "EStG § 66 Kindergeld", "EStG § 9 Werbungskosten"])
+st.write('\n')
+if text_option == "EStG § 66 Kindergeld":
+    text_area_input = st.text_area("Given paragraph", SAMPLE_66, height=200)
+elif text_option == "EStG § 9 Werbungskosten":
+    text_area_input = st.text_area("Given paragraph", SAMPLE_9, height=200)
+else:
+    text_area_input = st.text_area("Given paragraph", "", height=200)
+
+if st.button("Start Extraction") and text_area_input != "":
+    with st.spinner('Executing Extraction ...'):
+        paragraph = request_extractor(text_area_input)
+    if paragraph["text"] == "error":
+        st.error("Error while executing extraction.")
+    else:
+        legend = displacy.render([legend_content], style="ent", options=legend_options, manual=True)
+        st.write(f"{style}{get_html(legend, True)}", unsafe_allow_html=True)
+
+        st.caption("Entities:")
+        extracted_data = get_displacy_ent_obj(paragraph, False, True)
+        html = displacy.render(extracted_data, style="ent", options=legend_options, manual=True)
+        st.write(f"{style}{get_html(html)}", unsafe_allow_html=True)
+        
+        st.write('\n')
+        st.caption("Conditions:")
+        extracted_data = get_displacy_ent_obj(paragraph, True, True)
+        html = displacy.render(extracted_data, style="ent", options=legend_options, manual=True)
+        st.write(f"{style}{get_html(html)}", unsafe_allow_html=True)
+
+        st.write('\n')
+        st.caption("\n\nRelations:")
+        html_graph_req = generate_graph(paragraph["entities"], paragraph["relations"], send_request=True)
+        components.html(html_graph_req, height=500)
+        st.write('\n')
+        with st.expander("Show JSON"):
+            st.json(paragraph)

classification.json

@@ -0,0 +1,83 @@
+{
+    "entity_types": [
+        {
+            "idx": 0,
+            "label": "O"
+        },
+        {
+            "idx": 1,
+            "label": "StatedKeyFigure"
+        },
+        {
+            "idx": 2,
+            "label": "Condition"
+        },
+        {
+            "idx": 3,
+            "label": "StatedExpression"
+        },
+        {
+            "idx": 4,
+            "label": "Unit"
+        },
+        {
+            "idx": 5,
+            "label": "Range"
+        },
+        {
+            "idx": 6,
+            "label": "DeclarativeKeyFigure"
+        },
+        {
+            "idx": 7,
+            "label": "Factor"
+        },
+        {
+            "idx": 8,
+            "label": "DeclarativeExpression"
+        }
+    ],
+    "relation_types": [
+        {
+            "idx": 9,
+            "label": "hasCondition"
+        },
+        {
+            "idx": 10,
+            "label": "hasExpression"
+        },
+        {
+            "idx": 11,
+            "label": "hasUnit"
+        },
+        {
+            "idx": 12,
+            "label": "join"
+        },
+        {
+            "idx": 13,
+            "label": "hasRange"
+        },
+        {
+            "idx": 14,
+            "label": "hasFactor"
+        }
+    ],
+    "id_of_non_entity": 0,
+    "groups": [
+        [
+            0,
+            2
+        ],
+        [
+            0,
+            1,
+            3,
+            4,
+            5,
+            6,
+            7,
+            8
+        ]
+    ]
+}

model_inference.py

@@ -0,0 +1,380 @@
+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

requirements.txt

@@ -0,0 +1,11 @@
+# Python 3.8.11
+numpy~=1.23.5
+PyYAML~=5.4.1
+pydantic==1.8.2
+tqdm~=4.56.2
+scikit-learn~=0.24.2
+spacy==3.2.0
+# MarkupSafe==2.0.1
+torch==1.6.0
+transformers[sentencepiece]==4.26.1
+pyvis==0.3.2

util/configuration.py

@@ -0,0 +1,9 @@
+from pydantic import BaseModel
+
+class InferenceConfiguration(BaseModel):
+    model_path_keyfigure: str = "danielsteinigen/KeyFiTax"
+    spacy_model: str = "de_core_news_sm"
+    transformer_model: str = "xlm-roberta-large"
+    merge_entities: bool = True
+    split_len: int = 200
+    extract_relations: bool = True

util/process_data.py

@@ -0,0 +1,55 @@
+from typing import Optional, List
+
+from pydantic import BaseModel, Extra
+
+class EntityType(BaseModel):
+    idx: int
+    label: str
+
+
+class EntityTypeSet(BaseModel):
+    entity_types: List[EntityType]
+    relation_types: List[EntityType]
+    id_of_non_entity: int
+    groups: List[List[int]]
+
+    def __len__(self):
+        return len(self.entity_types) + len(self.relation_types)
+
+    def all_types(self):
+        return [*self.entity_types, *self.relation_types]
+
+
+class Token(BaseModel):
+    text: str
+    start: int
+    end: int
+
+
+class Entity(BaseModel):
+    id: int
+    text: str
+    start: int
+    end: int
+    ent_type: EntityType 
+    confidence: Optional[float]
+
+
+class Relation(BaseModel):
+    id: int
+    head: int
+    tail: int
+    rel_type: EntityType
+
+
+class Sample(BaseModel):
+    idx: int
+    text: str
+    entities: List[Entity] = []
+    relations: List[Relation] = []
+    tokens: List[Token] = []
+    tags: List[int] = []
+
+
+class SampleList(BaseModel):
+    samples: List[Sample]

util/tokenizer.py

@@ -0,0 +1,24 @@
+from typing import List
+
+import spacy
+
+from util.process_data import Token, Sample, SampleList
+
+class Tokenizer():
+
+    def __init__(self, spacy_model: str):
+        self.__spacy_model = spacy.load(spacy_model)
+
+    def run(self, sample_list: SampleList):
+        self.__tokenize(sample_list.samples, self.__spacy_model)
+
+    def __tokenize(self, samples: List[Sample], spacy_model):
+        doc_pipe = spacy_model.pipe([sample.text.replace('\xa0', ' ') for sample in samples])
+        for sample, doc in zip(samples, doc_pipe):
+            sample.tokens = [Token(
+                text=x.text,
+                start=x.idx,
+                end=x.idx + len(x.text)
+            ) for x in doc]
+            while '\n' in sample.tokens[-1].text or ' ' in sample.tokens[-1].text:
+                sample.tokens = sample.tokens[:-1]