Upload 7 files

__init__.py

@@ -3,3 +3,9 @@ import sys
 
 # add current folder to the python path
 sys.path.append(os.path.dirname(__file__))
+
+# this is required to dynamically load the PIE models
+from pie_modules.models import *  # noqa: F403
+from pie_modules.taskmodules import *  # noqa: F403
+from pytorch_ie.models import *  # noqa: F403
+from pytorch_ie.taskmodules import *  # noqa: F403

app.py

@@ -1,25 +1,19 @@
 import json
 import logging
 import os.path
-from collections import defaultdict
 from functools import partial
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import gradio as gr
 import pandas as pd
-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 backend import get_annotation_from_document, get_relevant_adus, get_similar_adus, process_text
+from pie_modules.taskmodules import PointerNetworkTaskModuleForEnd2EndRE
 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
-from vector_store import SimpleVectorStore
+from vector_store import SimpleVectorStore, VectorStore
 
 logger = logging.getLogger(__name__)
 
@@ -34,91 +28,6 @@ DEFAULT_MODEL_REVISION = "76300f8e534e2fcf695f00cb49bba166739b8d8a"
 DEFAULT_EMBEDDING_MODEL_NAME = "allenai/scibert_scivocab_uncased"
 
 
-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 = []
-    tokenizer_kwargs = {
-        "max_length": 512,
-        "stride": 64,
-        "truncation": True,
-        "return_overflowing_tokens": True,
-    }
-    tokenized_documents = tokenize_document(
-        document,
-        tokenizer=tokenizer,
-        result_document_type=TokenDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-        partition_layer="labeled_partitions",
-        added_annotations=added_annotations,
-        strict_span_conversion=False,
-        **tokenizer_kwargs,
-    )
-    # just tokenize again to get tensors in the correct format for the model
-    # TODO: fix for A34.txt from sciarg corpus
-    model_inputs = tokenizer(document.text, return_tensors="pt", **tokenizer_kwargs)
-    # this is added when using return_overflowing_tokens=True, but the model does not accept it
-    model_inputs.pop("overflow_to_sample_mapping", None)
-    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 annotate(
-    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-    pipeline: Pipeline,
-    embedding_model: Optional[PreTrainedModel] = None,
-    embedding_tokenizer: Optional[PreTrainedTokenizer] = None,
-) -> None:
-
-    # execute prediction pipeline
-    pipeline(document)
-
-    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.metadata["embeddings"] = adu_embeddings_dict
-    else:
-        gr.Warning(
-            "No embedding model provided. Skipping embedding extraction. You can load an embedding "
-            "model in the 'Model Configuration' section."
-        )
-
-
 def render_annotated_document(
     document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
     render_with: str,
@@ -135,57 +44,6 @@ def render_annotated_document(
     return html
 
 
-def add_to_index(
-    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-    processed_documents: dict,
-    vector_store: SimpleVectorStore,
-) -> None:
-    try:
-        if document.id in processed_documents:
-            gr.Warning(f"Document '{document.id}' already in index. Overwriting.")
-        # save the processed document to the index
-        processed_documents[document.id] = document
-        # save the embeddings to the vector store
-        for adu_id, embedding in document.metadata["embeddings"].items():
-            vector_store.save((document.id, adu_id), embedding)
-        gr.Info(
-            f"Added document {document.id} to index (index contains {len(processed_documents)} "
-            f"documents and {len(vector_store)} embeddings)."
-        )
-    except Exception as e:
-        raise gr.Error(f"Failed to add document {document.id} to index: {e}")
-
-
-def process_text(
-    text: str,
-    doc_id: str,
-    models: Tuple[Pipeline, Optional[PreTrainedModel], Optional[PreTrainedTokenizer]],
-    processed_documents: dict[
-        str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
-    ],
-    vector_store: SimpleVectorStore,
-) -> TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions:
-    try:
-        document = TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions(
-            id=doc_id, text=text, metadata={}
-        )
-        # add single partition from the whole text (the model only considers text in partitions)
-        document.labeled_partitions.append(LabeledSpan(start=0, end=len(text), label="text"))
-        # annotate the document
-        annotate(
-            document=document,
-            pipeline=models[0],
-            embedding_model=models[1],
-            embedding_tokenizer=models[2],
-        )
-        # add the document to the index
-        add_to_index(document, processed_documents, vector_store)
-
-        return document
-    except Exception as e:
-        raise gr.Error(f"Failed to process text: {e}")
-
-
 def wrapped_process_text(
     text: str,
     doc_id: str,
@@ -193,7 +51,7 @@ def wrapped_process_text(
     processed_documents: dict[
         str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
     ],
-    vector_store: SimpleVectorStore,
+    vector_store: VectorStore[Tuple[str, str]],
 ) -> Tuple[dict, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions]:
     document = process_text(
         text=text,
@@ -206,13 +64,13 @@ def wrapped_process_text(
     return document.asdict(), document
 
 
-def process_uploaded_file(
+def process_uploaded_files(
     file_names: List[str],
     models: Tuple[Pipeline, Optional[PreTrainedModel], Optional[PreTrainedTokenizer]],
     processed_documents: dict[
         str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
     ],
-    vector_store: SimpleVectorStore,
+    vector_store: VectorStore[Tuple[str, str]],
 ) -> None:
     try:
         for file_name in file_names:
@@ -229,164 +87,6 @@ def process_uploaded_file(
         raise gr.Error(f"Failed to process uploaded files: {e}")
 
 
-def _get_annotation_from_document(
-    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-    annotation_id: str,
-    annotation_layer: str,
-) -> LabeledSpan:
-    # use predictions
-    annotations = document[annotation_layer].predictions
-    id2annotation = {str(annotation._id): annotation for annotation in annotations}
-    annotation = id2annotation.get(annotation_id)
-    if annotation is None:
-        raise gr.Error(
-            f"annotation '{annotation_id}' not found in document '{document.id}'. Available "
-            f"annotations: {id2annotation}"
-        )
-    return annotation
-
-
-def _get_annotation(
-    doc_id: str,
-    annotation_id: str,
-    annotation_layer: str,
-    processed_documents: dict[
-        str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
-    ],
-) -> LabeledSpan:
-    document = processed_documents.get(doc_id)
-    if document is None:
-        raise gr.Error(
-            f"Document '{doc_id}' not found in index. Available documents: {list(processed_documents)}"
-        )
-    return _get_annotation_from_document(document, annotation_id, annotation_layer)
-
-
-def _get_similar_entries_from_vector_store(
-    ref_annotation_id: str,
-    ref_document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-    vector_store: SimpleVectorStore[Tuple[str, str]],
-    **retrieval_kwargs,
-) -> List[Tuple[Tuple[str, str], float]]:
-    embeddings = ref_document.metadata["embeddings"]
-    ref_embedding = embeddings.get(ref_annotation_id)
-    if ref_embedding is None:
-        raise gr.Error(
-            f"Embedding for annotation '{ref_annotation_id}' not found in metadata of "
-            f"document '{ref_document.id}'. Annotations with embeddings: {list(embeddings)}"
-        )
-
-    try:
-        similar_entries = vector_store.retrieve_similar(
-            ref_id=(ref_document.id, ref_annotation_id), **retrieval_kwargs
-        )
-    except Exception as e:
-        raise gr.Error(f"Failed to retrieve similar ADUs: {e}")
-
-    return similar_entries
-
-
-def get_similar_adus(
-    ref_annotation_id: str,
-    ref_document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-    vector_store: SimpleVectorStore,
-    processed_documents: dict[
-        str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
-    ],
-    min_similarity: float,
-) -> pd.DataFrame:
-    similar_entries = _get_similar_entries_from_vector_store(
-        ref_annotation_id=ref_annotation_id,
-        ref_document=ref_document,
-        vector_store=vector_store,
-        min_similarity=min_similarity,
-    )
-
-    similar_annotations = [
-        _get_annotation(
-            doc_id=doc_id,
-            annotation_id=annotation_id,
-            annotation_layer="labeled_spans",
-            processed_documents=processed_documents,
-        )
-        for (doc_id, annotation_id), _ in similar_entries
-    ]
-    df = pd.DataFrame(
-        [
-            # unpack the tuple (doc_id, annotation_id) to separate columns
-            # and add the similarity score and the text of the annotation
-            (doc_id, annotation_id, score, str(annotation))
-            for ((doc_id, annotation_id), score), annotation in zip(
-                similar_entries, similar_annotations
-            )
-        ],
-        columns=["doc_id", "adu_id", "sim_score", "text"],
-    )
-
-    return df
-
-
-def get_relevant_adus(
-    ref_annotation_id: str,
-    ref_document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-    vector_store: SimpleVectorStore,
-    processed_documents: dict[
-        str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
-    ],
-    min_similarity: float,
-) -> pd.DataFrame:
-    similar_entries = _get_similar_entries_from_vector_store(
-        ref_annotation_id=ref_annotation_id,
-        ref_document=ref_document,
-        vector_store=vector_store,
-        min_similarity=min_similarity,
-    )
-    ref_annotation = _get_annotation(
-        doc_id=ref_document.id,
-        annotation_id=ref_annotation_id,
-        annotation_layer="labeled_spans",
-        processed_documents=processed_documents,
-    )
-    result = []
-    for (doc_id, annotation_id), score in similar_entries:
-        # skip entries from the same document
-        if doc_id == ref_document.id:
-            continue
-        document = processed_documents[doc_id]
-        tail2rels = defaultdict(list)
-        head2rels = defaultdict(list)
-        for rel in document.binary_relations.predictions:
-            # skip non-argumentative relations
-            if rel.label in ["parts_of_same", "semantically_same"]:
-                continue
-            head2rels[rel.head].append(rel)
-            tail2rels[rel.tail].append(rel)
-
-        id2annotation = {
-            str(annotation._id): annotation for annotation in document.labeled_spans.predictions
-        }
-        annotation = id2annotation.get(annotation_id)
-        # note: we do not need to check if the annotation is different from the reference annotation,
-        # because they com from different documents and we already skip entries from the same document
-        for rel in head2rels.get(annotation, []):
-            result.append(
-                {
-                    "doc_id": doc_id,
-                    "reference_adu": str(annotation),
-                    "sim_score": score,
-                    "rel_score": rel.score,
-                    "relation": rel.label,
-                    "text": str(rel.tail),
-                }
-            )
-
-    # define column order
-    df = pd.DataFrame(
-        result, columns=["text", "relation", "doc_id", "reference_adu", "sim_score", "rel_score"]
-    )
-    return df
-
-
 def open_accordion():
     return gr.Accordion(open=True)
 
@@ -463,6 +163,24 @@ def select_processed_document(
     return doc
 
 
+def set_relation_types(
+    models: Tuple[Pipeline, Optional[PreTrainedModel], Optional[PreTrainedTokenizer]],
+    default: Optional[List[str]] = None,
+) -> gr.Dropdown:
+    arg_pipeline = models[0]
+    if isinstance(arg_pipeline.taskmodule, PointerNetworkTaskModuleForEnd2EndRE):
+        relation_types = arg_pipeline.taskmodule.labels_per_layer["binary_relations"]
+    else:
+        raise gr.Error("Unsupported taskmodule for relation types")
+
+    return gr.Dropdown(
+        choices=relation_types,
+        label="Relation Types",
+        value=default,
+        multiselect=True,
+    )
+
+
 def main():
 
     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."
@@ -526,7 +244,7 @@ def main():
                     )
                     embedding_model_name = gr.Textbox(
                         label=f"Embedding Model Name (e.g. {DEFAULT_EMBEDDING_MODEL_NAME})",
-                        value="",
+                        value=DEFAULT_EMBEDDING_MODEL_NAME,
                     )
                     load_models_btn = gr.Button("Load Models")
                     load_models_btn.click(
@@ -583,8 +301,18 @@ def main():
                         step=0.01,
                         value=0.8,
                     )
+                    top_k = gr.Slider(
+                        label="Top K",
+                        minimum=2,
+                        maximum=50,
+                        step=1,
+                        value=20,
+                    )
                     retrieve_similar_adus_btn = gr.Button("Retrieve similar ADUs")
                     similar_adus = gr.DataFrame(headers=["doc_id", "adu_id", "score", "text"])
+                    relation_types = set_relation_types(
+                        models_state.value, default=["supports", "contradicts"]
+                    )
 
                 # retrieve_relevant_adus_btn = gr.Button("Retrieve relevant ADUs")
                 relevant_adus = gr.DataFrame(
@@ -626,7 +354,7 @@ def main():
         )
 
         upload_btn.upload(
-            fn=process_uploaded_file,
+            fn=process_uploaded_files,
             inputs=[upload_btn, models_state, processed_documents_state, vector_store_state],
             outputs=[],
         ).success(
@@ -648,12 +376,14 @@ def main():
                 vector_store_state,
                 processed_documents_state,
                 min_similarity,
+                top_k,
+                relation_types,
             ],
             outputs=[relevant_adus],
         )
 
         reference_adu_id.change(
-            fn=partial(_get_annotation_from_document, annotation_layer="labeled_spans"),
+            fn=partial(get_annotation_from_document, annotation_layer="labeled_spans"),
             inputs=[document_state, reference_adu_id],
             outputs=[reference_adu_text],
         ).success(**retrieve_relevant_adus_event_kwargs)
@@ -666,10 +396,17 @@ def main():
                 vector_store_state,
                 processed_documents_state,
                 min_similarity,
+                top_k,
             ],
             outputs=[similar_adus],
         )
 
+        models_state.change(
+            fn=set_relation_types,
+            inputs=[models_state],
+            outputs=[relation_types],
+        )
+
         # retrieve_relevant_adus_btn.click(
         #     **retrieve_relevant_adus_event_kwargs
         # )

backend.py

@@ -0,0 +1,300 @@
+import logging
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple
+
+import gradio as gr
+import pandas as pd
+from pie_modules.document.processing import tokenize_document
+from pie_modules.documents import TokenDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
+from pytorch_ie import Pipeline
+from pytorch_ie.annotations import LabeledSpan, Span
+from pytorch_ie.documents import TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
+from rendering_utils import labeled_span_to_id
+from transformers import PreTrainedModel, PreTrainedTokenizer
+from vector_store import VectorStore
+
+logger = logging.getLogger(__name__)
+
+
+def embed_text_annotations(
+    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+    model: PreTrainedModel,
+    tokenizer: PreTrainedTokenizer,
+    text_layer_name: str,
+) -> Dict[Span, List[float]]:
+    # 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 = []
+    tokenizer_kwargs = {
+        "max_length": 512,
+        "stride": 64,
+        "truncation": True,
+        "return_overflowing_tokens": True,
+    }
+    tokenized_documents = tokenize_document(
+        document,
+        tokenizer=tokenizer,
+        result_document_type=TokenDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+        partition_layer="labeled_partitions",
+        added_annotations=added_annotations,
+        strict_span_conversion=False,
+        **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)
+    # this is added when using return_overflowing_tokens=True, but the model does not accept it
+    model_inputs.pop("overflow_to_sample_mapping", None)
+    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 annotate(
+    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+    pipeline: Pipeline,
+    embedding_model: Optional[PreTrainedModel] = None,
+    embedding_tokenizer: Optional[PreTrainedTokenizer] = None,
+) -> None:
+
+    # execute prediction pipeline
+    pipeline(document)
+
+    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 = {
+            labeled_span_to_id(k): v.detach().tolist() for k, v in adu_embeddings.items()
+        }
+        document.metadata["embeddings"] = adu_embeddings_dict
+    else:
+        gr.Warning(
+            "No embedding model provided. Skipping embedding extraction. You can load an embedding "
+            "model in the 'Model Configuration' section."
+        )
+
+
+def add_to_index(
+    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+    processed_documents: dict,
+    vector_store: VectorStore[Tuple[str, str]],
+) -> None:
+    try:
+        if document.id in processed_documents:
+            gr.Warning(f"Document '{document.id}' already in index. Overwriting.")
+
+        # save the processed document to the index
+        processed_documents[document.id] = document
+
+        # save the embeddings to the vector store
+        for adu_id, embedding in document.metadata["embeddings"].items():
+            vector_store.save((document.id, adu_id), embedding)
+
+        gr.Info(
+            f"Added document {document.id} to index (index contains {len(processed_documents)} "
+            f"documents and {len(vector_store)} embeddings)."
+        )
+    except Exception as e:
+        raise gr.Error(f"Failed to add document {document.id} to index: {e}")
+
+
+def process_text(
+    text: str,
+    doc_id: str,
+    models: Tuple[Pipeline, Optional[PreTrainedModel], Optional[PreTrainedTokenizer]],
+    processed_documents: dict[
+        str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
+    ],
+    vector_store: VectorStore[Tuple[str, str]],
+) -> TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions:
+    """Create a TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions from the provided
+    text, annotate it, and add it to the index.
+
+    Parameters:
+        text: The text to process.
+        doc_id: The ID of the document.
+        models: A tuple containing the prediction pipeline and the embedding model and tokenizer.
+        processed_documents: The index of processed documents.
+        vector_store: The vector store to save the embeddings.
+
+    Returns:
+        The processed document.
+    """
+
+    try:
+        document = TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions(
+            id=doc_id, text=text, metadata={}
+        )
+        # add single partition from the whole text (the model only considers text in partitions)
+        document.labeled_partitions.append(LabeledSpan(start=0, end=len(text), label="text"))
+        # annotate the document
+        annotate(
+            document=document,
+            pipeline=models[0],
+            embedding_model=models[1],
+            embedding_tokenizer=models[2],
+        )
+        # add the document to the index
+        add_to_index(document, processed_documents, vector_store)
+
+        return document
+    except Exception as e:
+        raise gr.Error(f"Failed to process text: {e}")
+
+
+def get_annotation_from_document(
+    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+    annotation_id: str,
+    annotation_layer: str,
+) -> LabeledSpan:
+    # use predictions
+    annotations = document[annotation_layer].predictions
+    id2annotation = {labeled_span_to_id(annotation): annotation for annotation in annotations}
+    annotation = id2annotation.get(annotation_id)
+    if annotation is None:
+        raise gr.Error(
+            f"annotation '{annotation_id}' not found in document '{document.id}'. Available "
+            f"annotations: {id2annotation}"
+        )
+    return annotation
+
+
+def get_annotation_from_processed_documents(
+    doc_id: str,
+    annotation_id: str,
+    annotation_layer: str,
+    processed_documents: dict[
+        str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
+    ],
+) -> LabeledSpan:
+    document = processed_documents.get(doc_id)
+    if document is None:
+        raise gr.Error(
+            f"Document '{doc_id}' not found in index. Available documents: {list(processed_documents)}"
+        )
+    return get_annotation_from_document(document, annotation_id, annotation_layer)
+
+
+def get_similar_adus(
+    ref_annotation_id: str,
+    ref_document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+    vector_store: VectorStore[Tuple[str, str]],
+    processed_documents: dict[
+        str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
+    ],
+    min_similarity: float,
+    top_k: int,
+) -> pd.DataFrame:
+    similar_entries = vector_store.retrieve_similar(
+        ref_id=(ref_document.id, ref_annotation_id),
+        min_similarity=min_similarity,
+        top_k=top_k,
+    )
+
+    similar_annotations = [
+        get_annotation_from_processed_documents(
+            doc_id=doc_id,
+            annotation_id=annotation_id,
+            annotation_layer="labeled_spans",
+            processed_documents=processed_documents,
+        )
+        for (doc_id, annotation_id), _ in similar_entries
+    ]
+    df = pd.DataFrame(
+        [
+            # unpack the tuple (doc_id, annotation_id) to separate columns
+            # and add the similarity score and the text of the annotation
+            (doc_id, annotation_id, score, str(annotation))
+            for ((doc_id, annotation_id), score), annotation in zip(
+                similar_entries, similar_annotations
+            )
+        ],
+        columns=["doc_id", "adu_id", "sim_score", "text"],
+    )
+
+    return df
+
+
+def get_relevant_adus(
+    ref_annotation_id: str,
+    ref_document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+    vector_store: VectorStore[Tuple[str, str]],
+    processed_documents: dict[
+        str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
+    ],
+    min_similarity: float,
+    top_k: int,
+    relation_types: List[str],
+) -> pd.DataFrame:
+    similar_entries = vector_store.retrieve_similar(
+        ref_id=(ref_document.id, ref_annotation_id),
+        min_similarity=min_similarity,
+        top_k=top_k,
+    )
+    result = []
+    for (doc_id, annotation_id), score in similar_entries:
+        # skip entries from the same document
+        if doc_id == ref_document.id:
+            continue
+        document = processed_documents[doc_id]
+        tail2rels = defaultdict(list)
+        head2rels = defaultdict(list)
+        for rel in document.binary_relations.predictions:
+            # skip non-argumentative relations
+            if rel.label not in relation_types:
+                continue
+            head2rels[rel.head].append(rel)
+            tail2rels[rel.tail].append(rel)
+
+        id2annotation = {
+            labeled_span_to_id(annotation): annotation
+            for annotation in document.labeled_spans.predictions
+        }
+        annotation = id2annotation.get(annotation_id)
+        # note: we do not need to check if the annotation is different from the reference annotation,
+        # because they come from different documents and we already skip entries from the same document
+        for rel in head2rels.get(annotation, []):
+            result.append(
+                {
+                    "doc_id": doc_id,
+                    "reference_adu": str(annotation),
+                    "sim_score": score,
+                    "rel_score": rel.score,
+                    "relation": rel.label,
+                    "text": str(rel.tail),
+                }
+            )
+
+    # define column order
+    df = pd.DataFrame(
+        result, columns=["text", "relation", "doc_id", "reference_adu", "sim_score", "rel_score"]
+    )
+    return df

rendering_utils.py

@@ -2,7 +2,7 @@ import json
 from collections import defaultdict
 from typing import Dict, List, Optional, Union
 
-from pytorch_ie.annotations import BinaryRelation
+from pytorch_ie.annotations import BinaryRelation, LabeledSpan, Span
 from pytorch_ie.documents import TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
 from rendering_utils_displacy import EntityRenderer
 
@@ -59,6 +59,15 @@ def render_displacy(
     return html
 
 
+def labeled_span_to_id(span: LabeledSpan) -> str:
+    return f"span-{span.start}-{span.end}-{span.label}"
+
+
+def labeled_span_from_id(span_id: str) -> LabeledSpan:
+    parts = span_id.split("-")
+    return LabeledSpan(int(parts[1]), int(parts[2]), parts[3])
+
+
 def inject_relation_data(
     html: str,
     sorted_entities,
@@ -80,7 +89,7 @@ def inject_relation_data(
     entities = soup.find_all(class_="entity")
     for idx, entity in enumerate(entities):
         annotation = sorted_entities[idx]
-        entity["id"] = str(annotation._id)
+        entity["id"] = labeled_span_to_id(annotation)
         original_color = entity["style"].split("background:")[1].split(";")[0].strip()
         entity["data-color-original"] = original_color
         if additional_colors is not None:
@@ -95,13 +104,13 @@ def inject_relation_data(
         entity["data-label"] = entity_annotation.label
         entity["data-relation-tails"] = json.dumps(
             [
-                {"entity-id": str(tail._id), "label": label}
+                {"entity-id": labeled_span_to_id(tail), "label": label}
                 for tail, label in entity2tails.get(entity_annotation, [])
             ]
         )
         entity["data-relation-heads"] = json.dumps(
             [
-                {"entity-id": str(head._id), "label": label}
+                {"entity-id": labeled_span_to_id(head), "label": label}
                 for head, label in entity2heads.get(entity_annotation, [])
             ]
         )

vector_store.py

@@ -1,5 +1,24 @@
+import abc
 from typing import Generic, Hashable, List, Optional, Tuple, TypeVar
 
+T = TypeVar("T", bound=Hashable)
+
+
+class VectorStore(Generic[T], abc.ABC):
+    @abc.abstractmethod
+    def save(self, emb_id: T, embedding: List[float]) -> None:
+        pass
+
+    @abc.abstractmethod
+    def retrieve_similar(
+        self, ref_id: T, top_k: Optional[int] = None, min_similarity: Optional[float] = None
+    ) -> List[Tuple[T, float]]:
+        pass
+
+    @abc.abstractmethod
+    def __len__(self):
+        pass
+
 
 def vector_norm(vector: List[float]) -> float:
     return sum(x**2 for x in vector) ** 0.5
@@ -9,10 +28,7 @@ def cosine_similarity(a: List[float], b: List[float]) -> float:
     return sum(a * b for a, b in zip(a, b)) / (vector_norm(a) * vector_norm(b))
 
 
-T = TypeVar("T", bound=Hashable)
-
-
-class SimpleVectorStore(Generic[T]):
+class SimpleVectorStore(VectorStore[T]):
     def __init__(self):
         self.vectors: dict[T, List[float]] = {}
         self._cache = {}