full pipeline

app.py

@@ -1,9 +1,12 @@
 import json
 import logging
+import os.path
+from collections import defaultdict
 from functools import partial
-from typing import Any, Optional, Tuple
+from typing import Any, 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
@@ -16,6 +19,7 @@ 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
 
 logger = logging.getLogger(__name__)
 
@@ -41,18 +45,26 @@ def embed_text_annotations(
     # 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 = {}
+    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)
 
@@ -80,22 +92,16 @@ def embed_text_annotations(
     return embeddings
 
 
-def predict(
-    text: str,
+def annotate(
+    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
     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)
-    document.labeled_partitions.append(LabeledSpan(start=0, end=len(text), label="text"))
+) -> None:
 
     # execute prediction pipeline
     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,
@@ -105,24 +111,19 @@ def predict(
         )
         # 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
+        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."
+            "No embedding model provided. Skipping embedding extraction. You can load an embedding "
+            "model in the 'Model Configuration' section."
         )
 
-    # 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
-    )
+def render_annotated_document(
+    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+    render_with: str,
+    render_kwargs_json: str,
+) -> str:
     render_kwargs = json.loads(render_kwargs_json)
     if render_with == RENDER_WITH_PRETTY_TABLE:
         html = render_pretty_table(document, **render_kwargs)
@@ -135,26 +136,255 @@ def render(document_txt: str, render_with: str, render_kwargs_json: str) -> str:
 
 
 def add_to_index(
-    output_txt: str, doc_id: str, processed_documents: dict, vector_store: Any
+    document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+    processed_documents: dict,
+    vector_store: SimpleVectorStore,
 ) -> None:
     try:
-        if doc_id in processed_documents:
-            gr.Warning(f"Document {doc_id} already in index. Overwriting.")
-        output = json.loads(output_txt)
-        # get the embeddings from the output and remove them from the output
-        embeddings = output.pop("embeddings")
+        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[doc_id] = output
+        processed_documents[document.id] = document
         # save the embeddings to the vector store
-        for adu_id, embedding in embeddings.items():
-            emb_id = f"{doc_id}:{adu_id}"
-            # TODO: save embedding to vector store at emb_id (embedding is a list of 768 floats)
-
+        for adu_id, embedding in document.metadata["embeddings"].items():
+            vector_store.save((document.id, adu_id), embedding)
         gr.Info(
-            f"Added document {doc_id} to index (index contains {len(processed_documents)} entries). (NOT YET IMPLEMENTED)"
+            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,
+    models: Tuple[Pipeline, Optional[PreTrainedModel], Optional[PreTrainedTokenizer]],
+    processed_documents: dict[
+        str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
+    ],
+    vector_store: SimpleVectorStore,
+) -> Tuple[dict, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions]:
+    document = process_text(
+        text=text,
+        doc_id=doc_id,
+        models=models,
+        processed_documents=processed_documents,
+        vector_store=vector_store,
+    )
+    # Return as dict and document to avoid serialization issues
+    return document.asdict(), document
+
+
+def process_uploaded_file(
+    file_names: List[str],
+    models: Tuple[Pipeline, Optional[PreTrainedModel], Optional[PreTrainedTokenizer]],
+    processed_documents: dict[
+        str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
+    ],
+    vector_store: SimpleVectorStore,
+) -> None:
+    try:
+        for file_name in file_names:
+            if file_name.lower().endswith(".txt"):
+                # read the file content
+                with open(file_name, "r", encoding="utf-8") as f:
+                    text = f.read()
+                base_file_name = os.path.basename(file_name)
+                gr.Info(f"Processing file '{base_file_name}' ...")
+                process_text(text, base_file_name, models, processed_documents, vector_store)
+            else:
+                raise gr.Error(f"Unsupported file format: {file_name}")
+    except Exception as e:
+        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 add document {doc_id} to index: {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():
@@ -202,13 +432,46 @@ def load_models(
     return argumentation_model, embedding_model, embedding_tokenizer
 
 
+def update_processed_documents_df(
+    processed_documents: dict[str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions]
+) -> pd.DataFrame:
+    df = pd.DataFrame(
+        [
+            (
+                doc_id,
+                len(document.labeled_spans.predictions),
+                len(document.binary_relations.predictions),
+            )
+            for doc_id, document in processed_documents.items()
+        ],
+        columns=["doc_id", "num_adus", "num_relations"],
+    )
+    return df
+
+
+def select_processed_document(
+    evt: gr.SelectData,
+    processed_documents_df: pd.DataFrame,
+    processed_documents: Dict[
+        str, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
+    ],
+) -> TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions:
+    row_idx, col_idx = evt.index
+    doc_id = processed_documents_df.iloc[row_idx]["doc_id"]
+    gr.Info(f"Select document: {doc_id}")
+    doc = processed_documents[doc_id]
+    return doc
+
+
 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."
 
-    print("Loading argumentation model ...")
-    argumentation_model = load_argumentation_model(
-        model_name=DEFAULT_MODEL_NAME, revision=DEFAULT_MODEL_REVISION
+    print("Loading models ...")
+    argumentation_model, embedding_model, embedding_tokenizer = load_models(
+        model_name=DEFAULT_MODEL_NAME,
+        revision=DEFAULT_MODEL_REVISION,
+        embedding_model_name=DEFAULT_EMBEDDING_MODEL_NAME,
     )
 
     default_render_kwargs = {
@@ -236,21 +499,18 @@ def main():
         },
     }
 
-    # TODO: setup the vector store
-    vector_store = None
-
     with gr.Blocks() as demo:
         processed_documents_state = gr.State(dict())
-        vector_store_state = gr.State(vector_store)
+        vector_store_state = gr.State(SimpleVectorStore())
         # wrap the pipeline and the embedding model/tokenizer in a tuple to avoid that it gets called
-        models_state = gr.State((argumentation_model, None, None))
+        models_state = gr.State((argumentation_model, embedding_model, embedding_tokenizer))
         with gr.Row():
             with gr.Column(scale=1):
                 doc_id = gr.Textbox(
                     label="Document ID",
                     value="user_input",
                 )
-                text = gr.Textbox(
+                doc_text = gr.Textbox(
                     label="Text",
                     lines=20,
                     value=example_text,
@@ -277,12 +537,12 @@ def main():
 
                 predict_btn = gr.Button("Analyse")
 
-                output_txt = gr.Textbox(visible=False)
+                document_state = gr.State()
 
             with gr.Column(scale=1):
 
                 with gr.Accordion("See plain result ...", open=False) as output_accordion:
-                    output_json = gr.JSON(label="Model Output")
+                    document_json = gr.JSON(label="Model Output")
 
                 with gr.Accordion("Render Options", open=False):
                     render_as = gr.Dropdown(
@@ -299,34 +559,121 @@ def main():
 
                 rendered_output = gr.HTML(label="Rendered Output")
 
-                add_to_index_btn = gr.Button("Add current result to Index")
+                # add_to_index_btn = gr.Button("Add current result to Index")
+                upload_btn = gr.UploadButton(
+                    "Upload & Analyse Documents", file_types=["text"], file_count="multiple"
+                )
 
-        render_button_kwargs = dict(
-            fn=render, inputs=[output_txt, render_as, render_kwargs], outputs=rendered_output
-        )
+            with gr.Column(scale=1):
+                with gr.Accordion("Indexed Documents", open=False):
+                    processed_documents_df = gr.DataFrame(
+                        headers=["id", "num_adus", "num_relations"],
+                        interactive=False,
+                    )
+
+                with gr.Accordion("Reference ADU", open=False):
+                    reference_adu_id = gr.Textbox(label="ID", elem_id="reference_adu_id")
+                    reference_adu_text = gr.Textbox(label="Text")
+
+                with gr.Accordion("Retrieval Configuration", open=False):
+                    min_similarity = gr.Slider(
+                        label="Minimum Similarity",
+                        minimum=0.0,
+                        maximum=1.0,
+                        step=0.01,
+                        value=0.8,
+                    )
+                    retrieve_similar_adus_btn = gr.Button("Retrieve similar ADUs")
+                    similar_adus = gr.DataFrame(headers=["doc_id", "adu_id", "score", "text"])
+
+                # retrieve_relevant_adus_btn = gr.Button("Retrieve relevant ADUs")
+                relevant_adus = gr.DataFrame(
+                    label="Relevant ADUs from other documents",
+                    headers=[
+                        "text",
+                        "relation",
+                        "doc_id",
+                        "reference_adu",
+                        "sim_score",
+                        "rel_score",
+                    ],
+                )
 
-        def _predict(
-            text: str,
-            models: Tuple[Pipeline, Optional[PreTrainedModel], Optional[PreTrainedTokenizer]],
-        ) -> Tuple[dict, str]:
-            return predict(text, *models)
+        render_event_kwargs = dict(
+            fn=render_annotated_document,
+            inputs=[document_state, render_as, render_kwargs],
+            outputs=rendered_output,
+        )
 
-        predict_btn.click(open_accordion, inputs=[], outputs=[output_accordion]).then(
-            fn=_predict,
-            inputs=[text, models_state],
-            outputs=[output_json, output_txt],
+        predict_btn.click(fn=open_accordion, inputs=[], outputs=[output_accordion]).then(
+            fn=wrapped_process_text,
+            inputs=[doc_text, doc_id, models_state, processed_documents_state, vector_store_state],
+            outputs=[document_json, document_state],
             api_name="predict",
-        ).success(**render_button_kwargs).success(
-            close_accordion, inputs=[], outputs=[output_accordion]
+        ).success(
+            fn=update_processed_documents_df,
+            inputs=[processed_documents_state],
+            outputs=[processed_documents_df],
+        )
+        render_btn.click(**render_event_kwargs, api_name="render")
+
+        document_state.change(
+            fn=lambda doc: doc.asdict(),
+            inputs=[document_state],
+            outputs=[document_json],
+        ).success(close_accordion, inputs=[], outputs=[output_accordion]).then(
+            **render_event_kwargs
         )
-        render_btn.click(**render_button_kwargs, api_name="render")
 
-        add_to_index_btn.click(
-            fn=add_to_index,
-            inputs=[output_txt, doc_id, processed_documents_state, vector_store_state],
+        upload_btn.upload(
+            fn=process_uploaded_file,
+            inputs=[upload_btn, models_state, processed_documents_state, vector_store_state],
             outputs=[],
+        ).success(
+            fn=update_processed_documents_df,
+            inputs=[processed_documents_state],
+            outputs=[processed_documents_df],
+        )
+        processed_documents_df.select(
+            select_processed_document,
+            inputs=[processed_documents_df, processed_documents_state],
+            outputs=[document_state],
+        )
+
+        retrieve_relevant_adus_event_kwargs = dict(
+            fn=get_relevant_adus,
+            inputs=[
+                reference_adu_id,
+                document_state,
+                vector_store_state,
+                processed_documents_state,
+                min_similarity,
+            ],
+            outputs=[relevant_adus],
         )
 
+        reference_adu_id.change(
+            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)
+
+        retrieve_similar_adus_btn.click(
+            fn=get_similar_adus,
+            inputs=[
+                reference_adu_id,
+                document_state,
+                vector_store_state,
+                processed_documents_state,
+                min_similarity,
+            ],
+            outputs=[similar_adus],
+        )
+
+        # retrieve_relevant_adus_btn.click(
+        #     **retrieve_relevant_adus_event_kwargs
+        # )
+
         js = """
         () => {
             function maybeSetColor(entity, colorAttributeKey, colorDictKey) {
@@ -337,8 +684,6 @@ def main():
                     color = colors[colorDictKey];
                 } catch (e) {}
                 if (color) {
-                    console.log('setting color', color);
-                    console.log('entity', entity);
                     entity.style.backgroundColor = color;
                     entity.style.color = '#000';
                 }
@@ -391,6 +736,15 @@ def main():
                     });
                 }
             }
+            function setReferenceAduId(entityId) {
+                // get the textarea element that holds the reference adu id
+                let referenceAduIdDiv = document.querySelector('#reference_adu_id textarea');
+                // set the value of the input field
+                referenceAduIdDiv.value = entityId;
+                // trigger an input event to update the state
+                var event = new Event('input');
+                referenceAduIdDiv.dispatchEvent(event);
+            }
 
             const entities = document.querySelectorAll('.entity');
             entities.forEach(entity => {
@@ -400,6 +754,7 @@ def main():
                 }
                 entity.addEventListener('mouseover', () => {
                     highlightRelationArguments(entity.id);
+                    setReferenceAduId(entity.id);
                 });
                 entity.addEventListener('mouseout', () => {
                     highlightRelationArguments(null);

rendering_utils.py

@@ -76,12 +76,11 @@ def inject_relation_data(
         entity2heads[relation.tail].append((relation.head, relation.label))
         entity2tails[relation.head].append((relation.tail, relation.label))
 
-    entity2id = {entity: f"entity-{idx}" for idx, entity in enumerate(sorted_entities)}
-
     # Add unique IDs to each entity
     entities = soup.find_all(class_="entity")
     for idx, entity in enumerate(entities):
-        entity["id"] = f"entity-{idx}"
+        annotation = sorted_entities[idx]
+        entity["id"] = str(annotation._id)
         original_color = entity["style"].split("background:")[1].split(";")[0].strip()
         entity["data-color-original"] = original_color
         if additional_colors is not None:
@@ -96,13 +95,13 @@ def inject_relation_data(
         entity["data-label"] = entity_annotation.label
         entity["data-relation-tails"] = json.dumps(
             [
-                {"entity-id": entity2id[tail], "label": label}
+                {"entity-id": str(tail._id), "label": label}
                 for tail, label in entity2tails.get(entity_annotation, [])
             ]
         )
         entity["data-relation-heads"] = json.dumps(
             [
-                {"entity-id": entity2id[head], "label": label}
+                {"entity-id": str(head._id), "label": label}
                 for head, label in entity2heads.get(entity_annotation, [])
             ]
         )

requirements.txt

@@ -1,4 +1,4 @@
-gradio==4.31.4
+gradio==4.36.0
 prettytable==3.10.0
 pie-modules==0.12.0
 beautifulsoup4==4.12.3

vector_store.py

@@ -0,0 +1,65 @@
+from typing import Generic, Hashable, List, Optional, Tuple, TypeVar
+
+
+def vector_norm(vector: List[float]) -> float:
+    return sum(x**2 for x in vector) ** 0.5
+
+
+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]):
+    def __init__(self):
+        self.vectors: dict[T, List[float]] = {}
+        self._cache = {}
+        self._sim = cosine_similarity
+
+    def save(self, emb_id: T, embedding: List[float]) -> None:
+        self.vectors[emb_id] = embedding
+
+    def get(self, emb_id: T) -> Optional[List[float]]:
+        return self.vectors.get(emb_id)
+
+    def delete(self, emb_id: T) -> None:
+        if emb_id in self.vectors:
+            del self.vectors[emb_id]
+            # remove from cache
+            self._cache = {k: v for k, v in self._cache.items() if emb_id not in k}
+
+    def clear(self) -> None:
+        self.vectors.clear()
+        self._cache.clear()
+
+    def __len__(self):
+        return len(self.vectors)
+
+    def retrieve_similar(
+        self, ref_id: T, top_k: Optional[int] = None, min_similarity: Optional[float] = None
+    ) -> List[Tuple[T, float]]:
+        ref_embedding = self.get(ref_id)
+        if ref_embedding is None:
+            raise ValueError(f"Reference embedding '{ref_id}' not found.")
+
+        # calculate similarity to all embeddings
+        similarities = {}
+        for emb_id, embedding in self.vectors.items():
+            if (emb_id, ref_id) not in self._cache:
+                # use cosine similarity
+                self._cache[(emb_id, ref_id)] = self._sim(ref_embedding, embedding)
+            similarities[emb_id] = self._cache[(emb_id, ref_id)]
+
+        # sort by similarity
+        similar_entries = sorted(similarities.items(), key=lambda x: x[1], reverse=True)
+
+        if min_similarity is not None:
+            similar_entries = [
+                (emb_id, sim) for emb_id, sim in similar_entries if sim >= min_similarity
+            ]
+        if top_k is not None:
+            similar_entries = similar_entries[:top_k]
+
+        return similar_entries