simple emebdding calculation

app.py

@@ -1,21 +1,84 @@
 import json
-from typing import Tuple
+import logging
+from functools import partial
+from typing import Optional, Tuple
 
 import gradio as gr
+from pie_modules.document.processing import tokenize_document
+from pie_modules.documents import TokenDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
 from pie_modules.models import *  # noqa: F403
 from pie_modules.taskmodules import *  # noqa: F403
+from pytorch_ie import Pipeline
 from pytorch_ie.annotations import LabeledSpan
 from pytorch_ie.auto import AutoPipeline
 from pytorch_ie.documents import TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
 from pytorch_ie.models import *  # noqa: F403
 from pytorch_ie.taskmodules import *  # noqa: F403
 from rendering_utils import render_displacy, render_pretty_table
+from transformers import AutoModel, AutoTokenizer, PreTrainedModel, PreTrainedTokenizer
+
+logger = logging.getLogger(__name__)
 
 RENDER_WITH_DISPLACY = "displaCy + highlighted arguments"
 RENDER_WITH_PRETTY_TABLE = "Pretty Table"
 
 
-def predict(text: str) -> Tuple[dict, str]:
+def embed_text_annotations(
+    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+    model: PreTrainedModel,
+    tokenizer: PreTrainedTokenizer,
+    text_layer_name: str,
+) -> dict:
+    # to not modify the original document
+    document = document.copy()
+    # tokenize_document does not yet consider predictions, so we need to add them manually
+    document[text_layer_name].extend(document[text_layer_name].predictions.clear())
+    added_annotations = []
+    # TODO: set return_overflowing_tokens=True and max_length=...?
+    tokenizer_kwargs = {}
+    tokenized_documents = tokenize_document(
+        document,
+        tokenizer=tokenizer,
+        result_document_type=TokenDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+        partition_layer="labeled_partitions",
+        added_annotations=added_annotations,
+        **tokenizer_kwargs,
+    )
+    # just tokenize again to get tensors in the correct format for the model
+    model_inputs = tokenizer(document.text, return_tensors="pt", **tokenizer_kwargs)
+    assert len(model_inputs.encodings) == len(tokenized_documents)
+    model_output = model(**model_inputs)
+
+    # get embeddings for all text annotations
+    embeddings = {}
+    for batch_idx in range(len(model_output.last_hidden_state)):
+        text2tok_ann = added_annotations[batch_idx][text_layer_name]
+        tok2text_ann = {v: k for k, v in text2tok_ann.items()}
+        for tok_ann in tokenized_documents[batch_idx].labeled_spans:
+            # skip "empty" annotations
+            if tok_ann.start == tok_ann.end:
+                continue
+            # use the max pooling strategy to get a single embedding for the annotation text
+            embedding = model_output.last_hidden_state[batch_idx, tok_ann.start : tok_ann.end].max(
+                dim=0
+            )[0]
+            text_ann = tok2text_ann[tok_ann]
+
+            if text_ann in embeddings:
+                logger.warning(
+                    f"Overwriting embedding for annotation '{text_ann}' (do you use striding?)"
+                )
+            embeddings[text_ann] = embedding
+
+    return embeddings
+
+
+def predict(
+    text: str,
+    pipeline: Pipeline,
+    embedding_model: Optional[PreTrainedModel] = None,
+    embedding_tokenizer: Optional[PreTrainedTokenizer] = None,
+) -> Tuple[dict, str]:
     document = TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions(text=text)
 
     # add single partition from the whole text (the model only considers text in partitions)
@@ -25,11 +88,27 @@ def predict(text: str) -> Tuple[dict, str]:
     pipeline(document)
 
     document_dict = document.asdict()
+
+    if embedding_model is not None and embedding_tokenizer is not None:
+        adu_embeddings = embed_text_annotations(
+            document=document,
+            model=embedding_model,
+            tokenizer=embedding_tokenizer,
+            text_layer_name="labeled_spans",
+        )
+        # convert keys to str because JSON keys must be strings
+        adu_embeddings_dict = {str(k._id): v.detach().tolist() for k, v in adu_embeddings.items()}
+        document_dict["embeddings"] = adu_embeddings_dict
+
+    # Return as dict and JSON string. The latter is required because the JSON component converts floats
+    # to ints which destroys de-serialization of the document (the scores of the annotations need to be floats)
     return document_dict, json.dumps(document_dict)
 
 
 def render(document_txt: str, render_with: str, render_kwargs_json: str) -> str:
     document_dict = json.loads(document_txt)
+    # remove embeddings from document_dict to make it de-serializable
+    document_dict.pop("embeddings", None)
     document = TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions.fromdict(
         document_dict
     )
@@ -52,16 +131,18 @@ def close_accordion():
     return gr.Accordion(open=False)
 
 
-if __name__ == "__main__":
+def main():
 
     model_name_or_path = "ArneBinder/sam-pointer-bart-base-v0.3"
+    # W&B run: https://wandb.ai/arne/dataset-sciarg-task-ner_re-v0.3-training/runs/prik91di
     revision = "76300f8e534e2fcf695f00cb49bba166739b8d8a"
     # local path
-    # model_name_or_path = "models/dataset-sciarg/task-ner_re/v0.3/2024-03-01_18-25-32"
+    # model_name_or_path = "models/dataset-sciarg/task-ner_re/v0.3/2024-05-28_23-33-46"
     # revision = None
 
     example_text = "Scholarly Argumentation Mining (SAM) has recently gained attention due to its potential to help scholars with the rapid growth of published scientific literature. It comprises two subtasks: argumentative discourse unit recognition (ADUR) and argumentative relation extraction (ARE), both of which are challenging since they require e.g. the integration of domain knowledge, the detection of implicit statements, and the disambiguation of argument structure. While previous work focused on dataset construction and baseline methods for specific document sections, such as abstract or results, full-text scholarly argumentation mining has seen little progress. In this work, we introduce a sequential pipeline model combining ADUR and ARE for full-text SAM, and provide a first analysis of the performance of pretrained language models (PLMs) on both subtasks. We establish a new SotA for ADUR on the Sci-Arg corpus, outperforming the previous best reported result by a large margin (+7% F1). We also present the first results for ARE, and thus for the full AM pipeline, on this benchmark dataset. Our detailed error analysis reveals that non-contiguous ADUs as well as the interpretation of discourse connectors pose major challenges and that data annotation needs to be more consistent."
 
+    print("loading argumentation mining model ...")
     pipeline = AutoPipeline.from_pretrained(
         model_name_or_path,
         device=-1,
@@ -70,6 +151,10 @@ if __name__ == "__main__":
         model_kwargs=dict(revision=revision),
     )
 
+    print("loading SciBERT embedding model ...")
+    embedding_model = AutoModel.from_pretrained("allenai/scibert_scivocab_uncased")
+    embedding_tokenizer = AutoTokenizer.from_pretrained("allenai/scibert_scivocab_uncased")
+
     default_render_kwargs = {
         "entity_options": {
             # we need to convert the keys to uppercase because the spacy rendering function expects them in uppercase
@@ -132,7 +217,15 @@ if __name__ == "__main__":
             fn=render, inputs=[output_txt, render_as, render_kwargs], outputs=rendered_output
         )
         predict_btn.click(open_accordion, inputs=[], outputs=[output_accordion]).then(
-            fn=predict, inputs=text, outputs=[output_json, output_txt], api_name="predict"
+            fn=partial(
+                predict,
+                pipeline=pipeline,
+                embedding_model=embedding_model,
+                embedding_tokenizer=embedding_tokenizer,
+            ),
+            inputs=text,
+            outputs=[output_json, output_txt],
+            api_name="predict",
         ).success(**render_button_kwargs).success(
             close_accordion, inputs=[], outputs=[output_accordion]
         )
@@ -223,3 +316,10 @@ if __name__ == "__main__":
         rendered_output.change(fn=None, js=js, inputs=[], outputs=[])
 
     demo.launch()
+
+
+if __name__ == "__main__":
+    # configure logging
+    logging.basicConfig()
+
+    main()

requirements.txt

@@ -1,5 +1,4 @@
-transformers==4.39.3
 gradio==4.31.4
 prettytable==3.10.0
-pie-modules>=0.11.0,<0.12.0
+pie-modules==0.12.0
 beautifulsoup4==4.12.3