new demo setup with langchain retriever

based on https://github.com/ArneBinder/pie-document-level/pull/298

app.py

@@ -1,32 +1,56 @@
+import pyrootutils
+
+root = pyrootutils.setup_root(
+    search_from=__file__,
+    indicator=[".project-root"],
+    pythonpath=True,
+    dotenv=True,
+)
+
 import json
 import logging
 import os.path
 import re
 import tempfile
-from functools import partial
-from typing import List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import arxiv
 import gradio as gr
 import pandas as pd
 import requests
 import torch
+import yaml
 from bs4 import BeautifulSoup
-from document_store import DocumentStore, get_annotation_from_document
-from embedding import EmbeddingModel
-from model_utils import annotate_document, create_document, load_models
-from pie_modules.taskmodules import PointerNetworkTaskModuleForEnd2EndRE
-from pytorch_ie import Pipeline
-from pytorch_ie.documents import (
-    TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions,
-    TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
+from model_utils import (
+    add_annotated_pie_documents_from_dataset,
+    load_argumentation_model,
+    load_retriever,
+    process_texts,
+    retrieve_all_relevant_spans,
+    retrieve_all_similar_spans,
+    retrieve_relevant_spans,
+    retrieve_similar_spans,
 )
+from pie_modules.taskmodules import PointerNetworkTaskModuleForEnd2EndRE
+from pytorch_ie import Annotation, Pipeline
+from pytorch_ie.annotations import BinaryRelation, LabeledMultiSpan, LabeledSpan
 from rendering_utils import HIGHLIGHT_SPANS_JS, render_displacy, render_pretty_table
-from transformers import PreTrainedModel, PreTrainedTokenizer
-from vector_store import QdrantVectorStore, SimpleVectorStore
+
+from src.langchain_modules import (
+    DocumentAwareSpanRetriever,
+    DocumentAwareSpanRetrieverWithRelations,
+)
 
 logger = logging.getLogger(__name__)
 
+
+def load_retriever_config(path: str) -> str:
+    with open(path, "r") as file:
+        yaml_string = file.read()
+    config = yaml.safe_load(yaml_string)
+    return yaml.dump(config)
+
+
 RENDER_WITH_DISPLACY = "displaCy + highlighted arguments"
 RENDER_WITH_PRETTY_TABLE = "Pretty Table"
 
@@ -35,16 +59,31 @@ DEFAULT_MODEL_REVISION = "76300f8e534e2fcf695f00cb49bba166739b8d8a"
 # local path
 # DEFAULT_MODEL_NAME = "models/dataset-sciarg/task-ner_re/v0.3/2024-05-28_23-33-46"
 # DEFAULT_MODEL_REVISION = None
-DEFAULT_EMBEDDING_MODEL_NAME = "allenai/scibert_scivocab_uncased"
+DEFAULT_RETRIEVER_CONFIG = load_retriever_config(
+    "configs/retriever/related_span_retriever_with_relations_from_other_docs.yaml"
+)
+# 0.943180 from data_dir="predictions/default/2024-10-15_23-40-18"
+DEFAULT_MIN_SIMILARITY = 0.95
 DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-DEFAULT_EMBEDDING_MAX_LENGTH = 512
-DEFAULT_EMBEDDING_BATCH_SIZE = 32
 DEFAULT_SPLIT_REGEX = "\n\n\n+"
 DEFAULT_ARXIV_ID = "1706.03762"
+DEFAULT_LOAD_PIE_DATASET_KWARGS_STR = json.dumps(
+    dict(path="pie/sciarg", name="resolve_parts_of_same", split="train"), indent=2
+)
 
 # Whether to handle segmented entities in the document. If True, labeled_spans are converted
 # to labeled_multi_spans and binary_relations with label "parts_of_same" are used to merge them.
+# This requires the networkx package to be installed.
 HANDLE_PARTS_OF_SAME = True
+LAYER_CAPTIONS = {
+    "labeled_multi_spans": "adus",
+    "binary_relations": "relations",
+    "labeled_partitions": "partitions",
+}
+RELATION_NAME_MAPPING = {
+    "supports_reversed": "supported by",
+    "contradicts_reversed": "contradicts",
+}
 
 
 def escape_regex(regex: str) -> str:
@@ -59,19 +98,38 @@ def unescape_regex(regex: str) -> str:
     return result
 
 
+def get_document_as_dict(retriever: DocumentAwareSpanRetriever, doc_id: str) -> Dict:
+    document = retriever.get_document(doc_id=doc_id)
+    return retriever.docstore.as_dict(document)
+
+
 def render_annotated_document(
-    document: Union[
-        TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-        TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions,
-    ],
+    retriever: DocumentAwareSpanRetrieverWithRelations,
+    document_id: str,
     render_with: str,
     render_kwargs_json: str,
 ) -> str:
+    text, spans, span_id2idx, relations = get_text_spans_and_relations_from_document(
+        retriever=retriever, document_id=document_id
+    )
+
     render_kwargs = json.loads(render_kwargs_json)
     if render_with == RENDER_WITH_PRETTY_TABLE:
-        html = render_pretty_table(document, **render_kwargs)
+        html = render_pretty_table(
+            text=text,
+            spans=spans,
+            span_id2idx=span_id2idx,
+            binary_relations=relations,
+            **render_kwargs,
+        )
     elif render_with == RENDER_WITH_DISPLACY:
-        html = render_displacy(document, **render_kwargs)
+        html = render_displacy(
+            text=text,
+            spans=spans,
+            span_id2idx=span_id2idx,
+            binary_relations=relations,
+            **render_kwargs,
+        )
     else:
         raise ValueError(f"Unknown render_with value: {render_with}")
 
@@ -79,84 +137,47 @@ def render_annotated_document(
 
 
 def wrapped_process_text(
-    text: str,
-    doc_id: str,
-    models: Tuple[Pipeline, Optional[EmbeddingModel]],
-    document_store: DocumentStore,
-    split_regex_escaped: str,
-    handle_parts_of_same: bool = False,
-) -> Tuple[
-    dict,
-    Union[
-        TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-        TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions,
-    ],
-]:
+    doc_id: str, text: str, retriever: DocumentAwareSpanRetriever, **kwargs
+) -> str:
     try:
-        document = create_document(
-            text=text,
-            doc_id=doc_id,
-            split_regex=unescape_regex(split_regex_escaped)
-            if len(split_regex_escaped) > 0
-            else None,
-        )
-        document = annotate_document(
-            document=document,
-            annotation_pipeline=models[0],
-            embedding_model=models[1],
-            handle_parts_of_same=handle_parts_of_same,
-        )
-        document_store.add_document(document)
+        process_texts(doc_ids=[doc_id], texts=[text], retriever=retriever, **kwargs)
     except Exception as e:
         raise gr.Error(f"Failed to process text: {e}")
-    # remove the embeddings because they are very large
-    if document.metadata.get("embeddings"):
-        document.metadata = {k: v for k, v in document.metadata.items() if k != "embeddings"}
     # Return as dict and document to avoid serialization issues
-    return document.asdict(), document
+    return doc_id
 
 
 def process_uploaded_files(
-    file_names: List[str],
-    models: Tuple[Pipeline, Optional[EmbeddingModel]],
-    document_store: DocumentStore,
-    split_regex_escaped: str,
-    show_max_cross_doc_sims: bool = False,
-    min_similarity: float = 0.95,
-    handle_parts_of_same: bool = False,
+    file_names: List[str], retriever: DocumentAwareSpanRetriever, **kwargs
 ) -> pd.DataFrame:
     try:
-        new_documents = []
+        doc_ids = []
+        texts = []
         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}' ...")
-                new_document = create_document(
-                    text=text,
-                    doc_id=base_file_name,
-                    split_regex=unescape_regex(split_regex_escaped)
-                    if len(split_regex_escaped) > 0
-                    else None,
-                )
-                new_document = annotate_document(
-                    document=new_document,
-                    annotation_pipeline=models[0],
-                    embedding_model=models[1],
-                    handle_parts_of_same=handle_parts_of_same,
-                )
-                new_documents.append(new_document)
+                doc_ids.append(base_file_name)
+                texts.append(text)
             else:
                 raise gr.Error(f"Unsupported file format: {file_name}")
-        document_store.add_documents(new_documents)
+        process_texts(texts=texts, doc_ids=doc_ids, retriever=retriever, verbose=True, **kwargs)
     except Exception as e:
         raise gr.Error(f"Failed to process uploaded files: {e}")
 
-    return document_store.overview(
-        with_max_cross_doc_sims=show_max_cross_doc_sims, min_similarity=min_similarity
-    )
+    return retriever.docstore.overview(layer_captions=LAYER_CAPTIONS, use_predictions=True)
+
+
+def wrapped_add_annotated_pie_documents_from_dataset(
+    retriever: DocumentAwareSpanRetriever, verbose: bool, **kwargs
+) -> pd.DataFrame:
+    try:
+        add_annotated_pie_documents_from_dataset(retriever=retriever, verbose=verbose, **kwargs)
+    except Exception as e:
+        raise gr.Error(f"Failed to add annotated PIE documents from dataset: {e}")
+    return retriever.docstore.overview(layer_captions=LAYER_CAPTIONS, use_predictions=True)
 
 
 def open_accordion():
@@ -167,30 +188,34 @@ def close_accordion():
     return gr.Accordion(open=False)
 
 
-def select_processed_document(
+def get_cell_for_fixed_column_from_df(
     evt: gr.SelectData,
-    processed_documents_df: pd.DataFrame,
-    document_store: DocumentStore,
-) -> Union[
-    TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-    TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions,
-]:
+    df: pd.DataFrame,
+    column: str,
+) -> Any:
+    """Get the value of the fixed column for the selected row in the DataFrame.
+    This is required can *not* with a lambda function because that will not get
+    the evt parameter.
+
+    Args:
+        evt: The event object.
+        df: The DataFrame.
+        column: The name of the column.
+
+    Returns:
+        The value of the fixed column for the selected row.
+    """
     row_idx, col_idx = evt.index
-    col_name = processed_documents_df.columns[col_idx]
-    if not col_name.endswith("doc_id"):
-        col_name = "doc_id"
-    doc_id = processed_documents_df.iloc[row_idx][col_name]
-    doc = document_store.get_document(doc_id, with_embeddings=False)
-    return doc
+    doc_id = df.iloc[row_idx][column]
+    return doc_id
 
 
 def set_relation_types(
-    models: Tuple[Pipeline, Optional[PreTrainedModel], Optional[PreTrainedTokenizer]],
+    argumentation_model: Pipeline,
     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"]
+    if isinstance(argumentation_model.taskmodule, PointerNetworkTaskModuleForEnd2EndRE):
+        relation_types = argumentation_model.taskmodule.labels_per_layer["binary_relations"]
     else:
         raise gr.Error("Unsupported taskmodule for relation types")
 
@@ -202,21 +227,64 @@ def set_relation_types(
     )
 
 
+def get_span_annotation(
+    retriever: DocumentAwareSpanRetriever,
+    span_id: str,
+) -> Annotation:
+    if span_id.strip() == "":
+        raise gr.Error("No span selected.")
+    try:
+        return retriever.get_span_by_id(span_id=span_id)
+    except Exception as e:
+        raise gr.Error(f"Failed to retrieve span annotation: {e}")
+
+
+def get_text_spans_and_relations_from_document(
+    retriever: DocumentAwareSpanRetrieverWithRelations, document_id: str
+) -> Tuple[
+    str,
+    Union[Sequence[LabeledSpan], Sequence[LabeledMultiSpan]],
+    Dict[str, int],
+    Sequence[BinaryRelation],
+]:
+    document = retriever.get_document(doc_id=document_id)
+    pie_document = retriever.docstore.unwrap(document)
+    use_predicted_annotations = retriever.use_predicted_annotations(document)
+    spans = retriever.get_base_layer(
+        pie_document=pie_document, use_predicted_annotations=use_predicted_annotations
+    )
+    relations = retriever.get_relation_layer(
+        pie_document=pie_document, use_predicted_annotations=use_predicted_annotations
+    )
+    span_id2idx = retriever.get_span_id2idx_from_doc(document)
+    return pie_document.text, spans, span_id2idx, relations
+
+
 def download_processed_documents(
-    document_store: DocumentStore,
-    file_name: str = "processed_documents.json",
-) -> str:
+    retriever: DocumentAwareSpanRetriever,
+    file_name: str = "retriever_store",
+) -> Optional[str]:
+    if len(retriever.docstore) == 0:
+        gr.Warning("No documents to download.")
+        return None
+
+    # zip the directory
     file_path = os.path.join(tempfile.gettempdir(), file_name)
-    document_store.save_to_file(file_path, indent=2)
-    return file_path
+
+    gr.Info(f"Zipping the retriever store to '{file_name}' ...")
+    result_file_path = retriever.save_to_archive(base_name=file_path, format="zip")
+
+    return result_file_path
 
 
 def upload_processed_documents(
     file_name: str,
-    document_store: DocumentStore,
+    retriever: DocumentAwareSpanRetriever,
 ) -> pd.DataFrame:
-    document_store.add_documents_from_file(file_name)
-    return document_store.overview()
+    # load the documents from the zip file or directory
+    retriever.load_from_disc(file_name)
+    # return the overview of the document store
+    return retriever.docstore.overview(layer_captions=LAYER_CAPTIONS, use_predictions=True)
 
 
 def clean_spaces(text: str) -> str:
@@ -253,21 +321,21 @@ def load_text_from_arxiv(arxiv_id: str, abstract_only: bool = False) -> Tuple[st
     except arxiv.HTTPError as e:
         raise gr.Error(f"Failed to fetch arXiv data: {e}")
     if len(result) == 0:
-        raise gr.Error(f"Could not find any paper with arXive ID '{arxiv_id}'")
+        raise gr.Error(f"Could not find any paper with arXiv ID '{arxiv_id}'")
     first_result = result[0]
     if abstract_only:
         abstract_clean = first_result.summary.replace("\n", " ")
         return abstract_clean, first_result.entry_id
     if "/abs/" not in first_result.entry_id:
         raise gr.Error(
-            f"Could not create the HTML URL for arXive ID '{arxiv_id}' because its entry ID has "
+            f"Could not create the HTML URL for arXiv ID '{arxiv_id}' because its entry ID has "
             f"an unexpected format: {first_result.entry_id}"
         )
     html_url = first_result.entry_id.replace("/abs/", "/html/")
     request_result = requests.get(html_url)
     if request_result.status_code != 200:
         raise gr.Error(
-            f"Could not fetch the HTML content for arXive ID '{arxiv_id}', status code: "
+            f"Could not fetch the HTML content for arXiv ID '{arxiv_id}', status code: "
             f"{request_result.status_code}"
         )
     html_content = request_result.text
@@ -275,19 +343,31 @@ def load_text_from_arxiv(arxiv_id: str, abstract_only: bool = False) -> Tuple[st
     return text_clean, html_url
 
 
+def process_text_from_arxiv(
+    arxiv_id: str, retriever: DocumentAwareSpanRetriever, abstract_only: bool = False, **kwargs
+) -> str:
+    try:
+        text, doc_id = load_text_from_arxiv(arxiv_id=arxiv_id, abstract_only=abstract_only)
+    except Exception as e:
+        raise gr.Error(f"Failed to load text from arXiv: {e}")
+    return wrapped_process_text(doc_id=doc_id, text=text, retriever=retriever, **kwargs)
+
+
 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 models ...")
-    argumentation_model, embedding_model = load_models(
+    print("Loading argumentation model ...")
+    argumentation_model = load_argumentation_model(
         model_name=DEFAULT_MODEL_NAME,
         revision=DEFAULT_MODEL_REVISION,
-        embedding_model_name=DEFAULT_EMBEDDING_MODEL_NAME,
-        embedding_max_length=DEFAULT_EMBEDDING_MAX_LENGTH,
-        embedding_batch_size=DEFAULT_EMBEDDING_BATCH_SIZE,
         device=DEFAULT_DEVICE,
     )
+    print("Loading retriever ...")
+    retriever = load_retriever(
+        DEFAULT_RETRIEVER_CONFIG, device=DEFAULT_DEVICE, config_format="yaml"
+    )
+    print("Models loaded.")
 
     default_render_kwargs = {
         "entity_options": {
@@ -315,21 +395,10 @@ def main():
     }
 
     with gr.Blocks() as demo:
-        document_store_state = gr.State(
-            DocumentStore(
-                span_annotation_caption="adu",
-                relation_annotation_caption="relation",
-                vector_store=QdrantVectorStore(),
-                document_type=TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
-                if not HANDLE_PARTS_OF_SAME
-                else TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions,
-                span_layer_name="labeled_spans"
-                if not HANDLE_PARTS_OF_SAME
-                else "labeled_multi_spans",
-            )
-        )
         # wrap the pipeline and the embedding model/tokenizer in a tuple to avoid that it gets called
-        models_state = gr.State((argumentation_model, embedding_model))
+        # models_state = gr.State((argumentation_model, embedding_model))
+        argumentation_model_state = gr.State((argumentation_model,))
+        retriever_state = gr.State((retriever,))
         with gr.Row():
             with gr.Column(scale=1):
                 doc_id = gr.Textbox(
@@ -341,63 +410,54 @@ def main():
                     lines=20,
                     value=example_text,
                 )
-                with gr.Accordion("Load Text from arXiv", open=False):
-                    arxiv_id = gr.Textbox(
-                        label="arXiv paper ID",
-                        placeholder=f"e.g. {DEFAULT_ARXIV_ID}",
-                        max_lines=1,
-                    )
-                    load_arxiv_only_abstract = gr.Checkbox(label="abstract only", value=False)
-                    load_arxiv_btn = gr.Button("Load Text from arXiv", variant="secondary")
-                    load_arxiv_btn.click(
-                        fn=load_text_from_arxiv,
-                        inputs=[arxiv_id, load_arxiv_only_abstract],
-                        outputs=[doc_text, doc_id],
-                    )
+
                 with gr.Accordion("Model Configuration", open=False):
-                    model_name = gr.Textbox(
-                        label="Model Name",
-                        value=DEFAULT_MODEL_NAME,
-                    )
-                    model_revision = gr.Textbox(
-                        label="Model Revision",
-                        value=DEFAULT_MODEL_REVISION,
-                    )
-                    embedding_model_name = gr.Textbox(
-                        label=f"Embedding Model Name (e.g. {DEFAULT_EMBEDDING_MODEL_NAME})",
-                        value=DEFAULT_EMBEDDING_MODEL_NAME,
-                    )
-                    embedding_max_length = gr.Slider(
-                        label="Embedding Model Max Length",
-                        minimum=16,
-                        maximum=2048,
-                        step=8,
-                        value=DEFAULT_EMBEDDING_MAX_LENGTH,
-                    )
-                    embedding_batch_size = gr.Slider(
-                        label="Embedding Model Batch Size",
-                        minimum=1,
-                        maximum=128,
-                        step=1,
-                        value=DEFAULT_EMBEDDING_BATCH_SIZE,
-                    )
+                    with gr.Accordion("argumentation structure", open=True):
+                        model_name = gr.Textbox(
+                            label="Model Name",
+                            value=DEFAULT_MODEL_NAME,
+                        )
+                        model_revision = gr.Textbox(
+                            label="Model Revision",
+                            value=DEFAULT_MODEL_REVISION,
+                        )
+                        load_arg_model_btn = gr.Button("Load Argumentation Model")
+
+                    with gr.Accordion("retriever", open=True):
+                        retriever_config = gr.Textbox(
+                            label="Retriever Configuration",
+                            placeholder="Configuration for the retriever",
+                            value=DEFAULT_RETRIEVER_CONFIG,
+                            lines=len(DEFAULT_RETRIEVER_CONFIG.split("\n")),
+                        )
+                        load_retriever_btn = gr.Button("Load Retriever")
+
                     device = gr.Textbox(
                         label="Device (e.g. 'cuda' or 'cpu')",
                         value=DEFAULT_DEVICE,
                     )
-                    load_models_btn = gr.Button("Load Models")
-                    load_models_btn.click(
-                        fn=load_models,
-                        inputs=[
-                            model_name,
-                            model_revision,
-                            embedding_model_name,
-                            embedding_max_length,
-                            embedding_batch_size,
-                            device,
-                        ],
-                        outputs=models_state,
+                    load_arg_model_btn.click(
+                        fn=lambda _model_name, _model_revision, _device: (
+                            load_argumentation_model(
+                                model_name=_model_name, revision=_model_revision, device=_device
+                            ),
+                        ),
+                        inputs=[model_name, model_revision, device],
+                        outputs=argumentation_model_state,
+                    )
+                    load_retriever_btn.click(
+                        fn=lambda _retriever_config, _device, _previous_retriever: (
+                            load_retriever(
+                                retriever_config=_retriever_config,
+                                device=_device,
+                                previous_retriever=_previous_retriever[0],
+                                config_format="yaml",
+                            ),
+                        ),
+                        inputs=[retriever_config, device, retriever_state],
+                        outputs=retriever_state,
                     )
+
                     split_regex_escaped = gr.Textbox(
                         label="Regex to partition the text",
                         placeholder="Regular expression pattern to split the text into partitions",
@@ -406,12 +466,12 @@ def main():
 
                 predict_btn = gr.Button("Analyse")
 
-                document_state = gr.State()
-
             with gr.Column(scale=1):
 
-                with gr.Accordion("See plain result ...", open=False) as output_accordion:
-                    document_json = gr.JSON(label="Model Output")
+                selected_document_id = gr.Textbox(
+                    label="Selected Document", max_lines=1, interactive=False
+                )
+                rendered_output = gr.HTML(label="Rendered Output")
 
                 with gr.Accordion("Render Options", open=False):
                     render_as = gr.Dropdown(
@@ -424,9 +484,11 @@ def main():
                         lines=5,
                         value=json.dumps(default_render_kwargs, indent=2),
                     )
-                render_btn = gr.Button("Re-render")
+                    render_btn = gr.Button("Re-render")
 
-                rendered_output = gr.HTML(label="Rendered Output")
+                with gr.Accordion("See plain result ...", open=False) as document_json_accordion:
+                    get_document_json_btn = gr.Button("Fetch annotated document as JSON")
+                    document_json = gr.JSON(label="Model Output")
 
             with gr.Column(scale=1):
                 with gr.Accordion(
@@ -436,14 +498,11 @@ def main():
                         headers=["id", "num_adus", "num_relations"],
                         interactive=False,
                     )
-                    show_max_cross_docu_sims = gr.Checkbox(
-                        label="Show max cross-document similarities", value=False
-                    )
                     gr.Markdown("Data Snapshot:")
                     with gr.Row():
                         download_processed_documents_btn = gr.DownloadButton("Download")
                         upload_processed_documents_btn = gr.UploadButton(
-                            "Upload", file_types=["json"]
+                            "Upload", file_types=["files"]
                         )
 
                 upload_btn = gr.UploadButton(
@@ -452,9 +511,22 @@ def main():
                     file_count="multiple",
                 )
 
-                with gr.Accordion("Selected ADU", open=False):
-                    selected_adu_id = gr.Textbox(label="ID", elem_id="selected_adu_id")
-                    selected_adu_text = gr.Textbox(label="Text")
+                with gr.Accordion("Import text from arXiv", open=False):
+                    arxiv_id = gr.Textbox(
+                        label="arXiv paper ID",
+                        placeholder=f"e.g. {DEFAULT_ARXIV_ID}",
+                        max_lines=1,
+                    )
+                    load_arxiv_only_abstract = gr.Checkbox(label="abstract only", value=False)
+                    load_arxiv_btn = gr.Button("Load & process from arXiv", variant="secondary")
+
+                with gr.Accordion("Import annotated PIE dataset", open=False):
+                    load_pie_dataset_kwargs_str = gr.Textbox(
+                        label="Parameters for Loading the PIE Dataset",
+                        value=DEFAULT_LOAD_PIE_DATASET_KWARGS_STR,
+                        lines=len(DEFAULT_LOAD_PIE_DATASET_KWARGS_STR.split("\n")),
+                    )
+                    load_pie_dataset_btn = gr.Button("Load & Embed PIE Dataset")
 
                 with gr.Accordion("Retrieval Configuration", open=False):
                     min_similarity = gr.Slider(
@@ -462,174 +534,294 @@ def main():
                         minimum=0.0,
                         maximum=1.0,
                         step=0.01,
-                        value=0.95,
+                        value=DEFAULT_MIN_SIMILARITY,
                     )
                     top_k = gr.Slider(
                         label="Top K",
                         minimum=2,
                         maximum=50,
                         step=1,
-                        value=20,
+                        value=10,
                     )
-                    retrieve_similar_adus_btn = gr.Button("Retrieve similar ADUs")
-                    similar_adus = gr.DataFrame(headers=["doc_id", "adu_id", "score", "text"])
-
-                    all2all_adu_similarities_button = gr.Button(
-                        "Compute all ADU-to-ADU similarities"
+                    retrieve_similar_adus_btn = gr.Button(
+                        "Retrieve *similar* ADUs for *selected* ADU"
                     )
-                    all2all_adu_similarities = gr.DataFrame(
-                        headers=["sim_score", "doc_id", "other_doc_id", "text", "other_text"]
+                    similar_adus_df = gr.DataFrame(
+                        headers=["doc_id", "adu_id", "score", "text"], interactive=False
                     )
-
-                    relation_types = set_relation_types(
-                        models_state.value, default=["supports", "contradicts"]
+                    retrieve_all_similar_adus_btn = gr.Button(
+                        "Retrieve *similar* ADUs for *all* ADUs in the document"
+                    )
+                    all_similar_adus_df = gr.DataFrame(
+                        headers=["doc_id", "query_adu_id", "adu_id", "score", "text"],
+                        interactive=False,
+                    )
+                    retrieve_all_relevant_adus_btn = gr.Button(
+                        "Retrieve *relevant* ADUs for *all* ADUs in the document"
+                    )
+                    all_relevant_adus_df = gr.DataFrame(
+                        headers=["doc_id", "adu_id", "score", "text"], interactive=False
                     )
 
-                # retrieve_relevant_adus_btn = gr.Button("Retrieve relevant ADUs")
-                relevant_adus = gr.DataFrame(
-                    label="Relevant ADUs from other documents",
-                    headers=[
-                        "relation",
-                        "adu",
-                        "reference_adu",
-                        "doc_id",
-                        "sim_score",
-                        "rel_score",
-                    ],
+                # currently not used
+                # relation_types = set_relation_types(
+                #    argumentation_model_state.value[0], default=["supports_reversed", "contradicts_reversed"]
+                # )
+
+                # Dummy textbox to hold the hover adu id. On click on the rendered output,
+                # its content will be copied to selected_adu_id which will trigger the retrieval.
+                hover_adu_id = gr.Textbox(
+                    label="ID (hover)",
+                    elem_id="hover_adu_id",
+                    interactive=False,
+                    visible=False,
+                )
+                selected_adu_id = gr.Textbox(
+                    label="ID (selected)",
+                    elem_id="selected_adu_id",
                     interactive=False,
+                    visible=False,
                 )
+                selected_adu_text = gr.Textbox(label="Selected ADU", interactive=False)
+
+                with gr.Accordion("Relevant ADUs from other documents", open=True):
+
+                    relevant_adus_df = gr.DataFrame(
+                        headers=[
+                            "relation",
+                            "adu",
+                            "reference_adu",
+                            "doc_id",
+                            "sim_score",
+                            "rel_score",
+                        ],
+                        interactive=False,
+                    )
 
         render_event_kwargs = dict(
-            fn=render_annotated_document,
-            inputs=[document_state, render_as, render_kwargs],
+            fn=lambda _retriever, _document_id, _render_as, _render_kwargs: render_annotated_document(
+                retriever=_retriever[0],
+                document_id=_document_id,
+                render_with=_render_as,
+                render_kwargs_json=_render_kwargs,
+            ),
+            inputs=[retriever_state, selected_document_id, render_as, render_kwargs],
             outputs=rendered_output,
         )
 
         show_overview_kwargs = dict(
-            fn=lambda document_store, show_max_sims, min_sim: document_store.overview(
-                with_max_cross_doc_sims=show_max_sims, min_similarity=min_sim
+            fn=lambda _retriever: _retriever[0].docstore.overview(
+                layer_captions=LAYER_CAPTIONS, use_predictions=True
             ),
-            inputs=[document_store_state, show_max_cross_docu_sims, min_similarity],
+            inputs=[retriever_state],
             outputs=[processed_documents_df],
         )
-        predict_btn.click(fn=open_accordion, inputs=[], outputs=[output_accordion]).then(
-            fn=partial(wrapped_process_text, handle_parts_of_same=HANDLE_PARTS_OF_SAME),
+        predict_btn.click(
+            fn=lambda _doc_text, _doc_id, _argumentation_model, _retriever, _split_regex_escaped: wrapped_process_text(
+                text=_doc_text,
+                doc_id=_doc_id,
+                argumentation_model=_argumentation_model[0],
+                retriever=_retriever[0],
+                split_regex_escaped=(
+                    unescape_regex(_split_regex_escaped) if _split_regex_escaped else None
+                ),
+                handle_parts_of_same=HANDLE_PARTS_OF_SAME,
+            ),
             inputs=[
                 doc_text,
                 doc_id,
-                models_state,
-                document_store_state,
+                argumentation_model_state,
+                retriever_state,
                 split_regex_escaped,
             ],
-            outputs=[document_json, document_state],
+            outputs=[selected_document_id],
             api_name="predict",
-        ).success(**show_overview_kwargs)
+        ).success(**show_overview_kwargs).success(**render_event_kwargs)
         render_btn.click(**render_event_kwargs, api_name="render")
 
-        document_state.change(
-            fn=lambda doc: doc.asdict(),
-            inputs=[document_state],
+        load_arxiv_btn.click(
+            fn=lambda _arxiv_id, _load_arxiv_only_abstract, _argumentation_model, _retriever, _split_regex_escaped: process_text_from_arxiv(
+                arxiv_id=_arxiv_id,
+                abstract_only=_load_arxiv_only_abstract,
+                argumentation_model=_argumentation_model[0],
+                retriever=_retriever[0],
+                split_regex_escaped=(
+                    unescape_regex(_split_regex_escaped) if _split_regex_escaped else None
+                ),
+                handle_parts_of_same=HANDLE_PARTS_OF_SAME,
+            ),
+            inputs=[
+                arxiv_id,
+                load_arxiv_only_abstract,
+                argumentation_model_state,
+                retriever_state,
+                split_regex_escaped,
+            ],
+            outputs=[selected_document_id],
+            api_name="predict",
+        ).success(**show_overview_kwargs)
+
+        load_pie_dataset_btn.click(
+            fn=open_accordion, inputs=[], outputs=[processed_documents_accordion]
+        ).then(
+            fn=lambda _retriever, _load_pie_dataset_kwargs_str: wrapped_add_annotated_pie_documents_from_dataset(
+                retriever=_retriever[0], verbose=True, **json.loads(_load_pie_dataset_kwargs_str)
+            ),
+            inputs=[retriever_state, load_pie_dataset_kwargs_str],
+            outputs=[processed_documents_df],
+        )
+
+        selected_document_id.change(**render_event_kwargs)
+
+        get_document_json_btn.click(
+            fn=lambda _retriever, _document_id: get_document_as_dict(
+                retriever=_retriever[0], doc_id=_document_id
+            ),
+            inputs=[retriever_state, selected_document_id],
             outputs=[document_json],
-        ).success(close_accordion, inputs=[], outputs=[output_accordion]).then(
-            **render_event_kwargs
         )
 
         upload_btn.upload(
             fn=open_accordion, inputs=[], outputs=[processed_documents_accordion]
         ).then(
-            fn=partial(process_uploaded_files, handle_parts_of_same=HANDLE_PARTS_OF_SAME),
+            fn=lambda _file_names, _argumentation_model, _retriever, _split_regex_escaped: process_uploaded_files(
+                file_names=_file_names,
+                argumentation_model=_argumentation_model[0],
+                retriever=_retriever[0],
+                split_regex_escaped=unescape_regex(_split_regex_escaped),
+                handle_parts_of_same=HANDLE_PARTS_OF_SAME,
+            ),
             inputs=[
                 upload_btn,
-                models_state,
-                document_store_state,
+                argumentation_model_state,
+                retriever_state,
                 split_regex_escaped,
-                show_max_cross_docu_sims,
-                min_similarity,
             ],
             outputs=[processed_documents_df],
         )
         processed_documents_df.select(
-            select_processed_document,
-            inputs=[processed_documents_df, document_store_state],
-            outputs=[document_state],
+            fn=get_cell_for_fixed_column_from_df,
+            inputs=[processed_documents_df, gr.State("doc_id")],
+            outputs=[selected_document_id],
         )
-        show_max_cross_docu_sims.change(**show_overview_kwargs)
 
         download_processed_documents_btn.click(
-            fn=partial(download_processed_documents, file_name="processed_documents.zip"),
-            inputs=[document_store_state],
+            fn=lambda _retriever: download_processed_documents(
+                _retriever[0], file_name="processed_documents"
+            ),
+            inputs=[retriever_state],
             outputs=[download_processed_documents_btn],
         )
         upload_processed_documents_btn.upload(
-            fn=upload_processed_documents,
-            inputs=[upload_processed_documents_btn, document_store_state],
+            fn=lambda file_name, _retriever: upload_processed_documents(
+                file_name, retriever=_retriever[0]
+            ),
+            inputs=[upload_processed_documents_btn, retriever_state],
             outputs=[processed_documents_df],
         )
 
         retrieve_relevant_adus_event_kwargs = dict(
-            fn=partial(
-                DocumentStore.get_related_annotations_from_other_documents_df,
-                columns=relevant_adus.headers,
+            fn=lambda _retriever, _selected_adu_id, _min_similarity, _top_k: retrieve_relevant_spans(
+                retriever=_retriever[0],
+                query_span_id=_selected_adu_id,
+                k=_top_k,
+                score_threshold=_min_similarity,
+                relation_label_mapping=RELATION_NAME_MAPPING,
+                # columns=relevant_adus.headers
             ),
             inputs=[
-                document_store_state,
+                retriever_state,
                 selected_adu_id,
-                document_state,
                 min_similarity,
                 top_k,
-                relation_types,
             ],
-            outputs=[relevant_adus],
+            outputs=[relevant_adus_df],
+        )
+        relevant_adus_df.select(
+            fn=get_cell_for_fixed_column_from_df,
+            inputs=[relevant_adus_df, gr.State("doc_id")],
+            outputs=[selected_document_id],
         )
 
         selected_adu_id.change(
-            fn=partial(
-                get_annotation_from_document,
-                annotation_layer="labeled_spans"
-                if not HANDLE_PARTS_OF_SAME
-                else "labeled_multi_spans",
-                use_predictions=True,
+            fn=lambda _retriever, _selected_adu_id: get_span_annotation(
+                retriever=_retriever[0], span_id=_selected_adu_id
             ),
-            inputs=[document_state, selected_adu_id],
+            inputs=[retriever_state, selected_adu_id],
             outputs=[selected_adu_text],
         ).success(**retrieve_relevant_adus_event_kwargs)
 
         retrieve_similar_adus_btn.click(
-            fn=lambda document_store, ann_id, document, min_sim, k: document_store.get_similar_annotations_df(
-                ref_annotation_id=ann_id,
-                ref_document=document,
-                min_similarity=min_sim,
-                top_k=k,
-                annotation_layer="labeled_spans"
-                if not HANDLE_PARTS_OF_SAME
-                else "labeled_multi_spans",
+            fn=lambda _retriever, _selected_adu_id, _min_similarity, _tok_k: retrieve_similar_spans(
+                retriever=_retriever[0],
+                query_span_id=_selected_adu_id,
+                k=_tok_k,
+                score_threshold=_min_similarity,
             ),
             inputs=[
-                document_store_state,
+                retriever_state,
                 selected_adu_id,
-                document_state,
                 min_similarity,
                 top_k,
             ],
-            outputs=[similar_adus],
+            outputs=[similar_adus_df],
+        )
+        similar_adus_df.select(
+            fn=get_cell_for_fixed_column_from_df,
+            inputs=[similar_adus_df, gr.State("doc_id")],
+            outputs=[selected_document_id],
         )
 
-        models_state.change(
-            fn=set_relation_types,
-            inputs=[models_state],
-            outputs=[relation_types],
+        retrieve_all_similar_adus_btn.click(
+            fn=lambda _retriever, _document_id, _min_similarity, _tok_k: retrieve_all_similar_spans(
+                retriever=_retriever[0],
+                query_doc_id=_document_id,
+                k=_tok_k,
+                score_threshold=_min_similarity,
+                query_span_id_column="query_span_id",
+            ),
+            inputs=[
+                retriever_state,
+                selected_document_id,
+                min_similarity,
+                top_k,
+            ],
+            outputs=[all_similar_adus_df],
         )
-        all2all_adu_similarities_button.click(
-            fn=partial(
-                DocumentStore.get_all2all_adu_similarities,
-                columns=all2all_adu_similarities.headers,
+
+        retrieve_all_relevant_adus_btn.click(
+            fn=lambda _retriever, _document_id, _min_similarity, _tok_k: retrieve_all_relevant_spans(
+                retriever=_retriever[0],
+                query_doc_id=_document_id,
+                k=_tok_k,
+                score_threshold=_min_similarity,
+                query_span_id_column="query_span_id",
             ),
-            inputs=[document_store_state, min_similarity],
-            outputs=[all2all_adu_similarities],
+            inputs=[
+                retriever_state,
+                selected_document_id,
+                min_similarity,
+                top_k,
+            ],
+            outputs=[all_relevant_adus_df],
+        )
+
+        # select query span id from the "retrieve all" result data frames
+        all_similar_adus_df.select(
+            fn=get_cell_for_fixed_column_from_df,
+            inputs=[all_similar_adus_df, gr.State("query_span_id")],
+            outputs=[selected_adu_id],
+        )
+        all_relevant_adus_df.select(
+            fn=get_cell_for_fixed_column_from_df,
+            inputs=[all_relevant_adus_df, gr.State("query_span_id")],
+            outputs=[selected_adu_id],
         )
 
-        # retrieve_relevant_adus_btn.click(
-        #     **retrieve_relevant_adus_event_kwargs
+        # argumentation_model_state.change(
+        #    fn=lambda _argumentation_model: set_relation_types(_argumentation_model[0]),
+        #    inputs=[argumentation_model_state],
+        #    outputs=[relation_types],
         # )
 
         rendered_output.change(fn=None, js=HIGHLIGHT_SPANS_JS, inputs=[], outputs=[])

model_utils.py

@@ -1,10 +1,10 @@
+import json
 import logging
-from typing import Optional, Tuple, Union
+from typing import Iterable, Optional, Sequence, Union
 
 import gradio as gr
-import torch
-from annotation_utils import labeled_span_to_id
-from embedding import EmbeddingModel, HuggingfaceEmbeddingModel
+import pandas as pd
+from pie_datasets import Dataset, IterableDataset, load_dataset
 from pie_modules.document.processing import RegexPartitioner, SpansViaRelationMerger
 from pytorch_ie import Pipeline
 from pytorch_ie.annotations import LabeledSpan
@@ -13,31 +13,35 @@ from pytorch_ie.documents import (
     TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions,
     TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
 )
+from typing_extensions import Protocol
+
+from src.langchain_modules import DocumentAwareSpanRetriever
+from src.langchain_modules.span_retriever import (
+    DocumentAwareSpanRetrieverWithRelations,
+    _parse_config,
+)
 
 logger = logging.getLogger(__name__)
 
 
 def annotate_document(
     document: TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-    annotation_pipeline: Pipeline,
-    embedding_model: Optional[EmbeddingModel] = None,
+    argumentation_model: Pipeline,
     handle_parts_of_same: bool = False,
 ) -> Union[
     TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
     TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions,
 ]:
-    """Annotate a document with the provided pipeline. If an embedding model is provided, also
-    extract embeddings for the labeled spans.
+    """Annotate a document with the provided pipeline.
 
     Args:
         document: The document to annotate.
-        annotation_pipeline: The pipeline to use for annotation.
-        embedding_model: The embedding model to use for extracting text span embeddings.
+        argumentation_model: The pipeline to use for annotation.
         handle_parts_of_same: Whether to merge spans that are part of the same entity into a single multi span.
     """
 
     # execute prediction pipeline
-    annotation_pipeline(document)
+    argumentation_model(document)
 
     if handle_parts_of_same:
         merger = SpansViaRelationMerger(
@@ -53,22 +57,6 @@ def annotate_document(
         )
         document = merger(document)
 
-    if embedding_model is not None:
-        text_span_embeddings = embedding_model(
-            document=document,
-            span_layer_name="labeled_spans" if not handle_parts_of_same else "labeled_multi_spans",
-        )
-        # convert keys to str because JSON keys must be strings
-        text_span_embeddings_dict = {
-            labeled_span_to_id(k): v.tolist() for k, v in text_span_embeddings.items()
-        }
-        document.metadata["embeddings"] = text_span_embeddings_dict
-    else:
-        gr.Warning(
-            "No embedding model provided. Skipping embedding extraction. You can load an embedding "
-            "model in the 'Model Configuration' section."
-        )
-
     return document
 
 
@@ -101,6 +89,80 @@ def create_document(
     return document
 
 
+def add_annotated_pie_documents(
+    retriever: DocumentAwareSpanRetriever,
+    pie_documents: Sequence[TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions],
+    use_predicted_annotations: bool,
+    verbose: bool = False,
+) -> None:
+    if verbose:
+        gr.Info(f"Create span embeddings for {len(pie_documents)} documents...")
+    num_docs_before = len(retriever.docstore)
+    retriever.add_pie_documents(pie_documents, use_predicted_annotations=use_predicted_annotations)
+    # number of documents that were overwritten
+    num_overwritten_docs = num_docs_before + len(pie_documents) - len(retriever.docstore)
+    # warn if documents were overwritten
+    if num_overwritten_docs > 0:
+        gr.Warning(f"{num_overwritten_docs} documents were overwritten.")
+
+
+def process_texts(
+    texts: Iterable[str],
+    doc_ids: Iterable[str],
+    argumentation_model: Pipeline,
+    retriever: DocumentAwareSpanRetriever,
+    split_regex_escaped: Optional[str],
+    handle_parts_of_same: bool = False,
+    verbose: bool = False,
+) -> None:
+    # check that doc_ids are unique
+    if len(set(doc_ids)) != len(list(doc_ids)):
+        raise gr.Error("Document IDs must be unique.")
+    pie_documents = [
+        create_document(text=text, doc_id=doc_id, split_regex=split_regex_escaped)
+        for text, doc_id in zip(texts, doc_ids)
+    ]
+    if verbose:
+        gr.Info(f"Annotate {len(pie_documents)} documents...")
+    pie_documents = [
+        annotate_document(
+            document=pie_document,
+            argumentation_model=argumentation_model,
+            handle_parts_of_same=handle_parts_of_same,
+        )
+        for pie_document in pie_documents
+    ]
+    add_annotated_pie_documents(
+        retriever=retriever,
+        pie_documents=pie_documents,
+        use_predicted_annotations=True,
+        verbose=verbose,
+    )
+
+
+def add_annotated_pie_documents_from_dataset(
+    retriever: DocumentAwareSpanRetriever, verbose: bool = False, **load_dataset_kwargs
+) -> None:
+    try:
+        gr.Info(
+            "Loading PIE dataset with parameters:\n" + json.dumps(load_dataset_kwargs, indent=2)
+        )
+        dataset = load_dataset(**load_dataset_kwargs)
+        if not isinstance(dataset, (Dataset, IterableDataset)):
+            raise gr.Error("Loaded dataset is not of type PIE (Iterable)Dataset.")
+        dataset_converted = dataset.to_document_type(
+            TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions
+        )
+        add_annotated_pie_documents(
+            retriever=retriever,
+            pie_documents=dataset_converted,
+            use_predicted_annotations=False,
+            verbose=verbose,
+        )
+    except Exception as e:
+        raise gr.Error(f"Failed to load dataset: {e}")
+
+
 def load_argumentation_model(
     model_name: str,
     revision: Optional[str] = None,
@@ -133,48 +195,260 @@ def load_argumentation_model(
     return model
 
 
-def load_embedding_model(
-    embedding_model_name: Optional[str] = None,
-    # embedding_model_revision: Optional[str] = None,
-    embedding_max_length: int = 512,
-    embedding_batch_size: int = 16,
+def load_retriever(
+    retriever_config: str,
+    config_format: str,
     device: str = "cpu",
-) -> Optional[EmbeddingModel]:
-    if embedding_model_name is not None and embedding_model_name.strip():
-        try:
-            embedding_model = HuggingfaceEmbeddingModel(
-                embedding_model_name.strip(),
-                # revision=embedding_model_revision,
-                device=device,
-                max_length=embedding_max_length,
-                batch_size=embedding_batch_size,
+    previous_retriever: Optional[DocumentAwareSpanRetrieverWithRelations] = None,
+) -> DocumentAwareSpanRetrieverWithRelations:
+    try:
+        retriever_config = _parse_config(retriever_config, format=config_format)
+        # set device for the embeddings pipeline
+        retriever_config["vectorstore"]["embedding"]["pipeline_kwargs"]["device"] = device
+        result = DocumentAwareSpanRetrieverWithRelations.instantiate_from_config(retriever_config)
+        # if a previous retriever is provided, load all documents and vectors from the previous retriever
+        if previous_retriever is not None:
+            # documents
+            all_doc_ids = list(previous_retriever.docstore.yield_keys())
+            gr.Info(f"Storing {len(all_doc_ids)} documents from previous retriever...")
+            all_docs = previous_retriever.docstore.mget(all_doc_ids)
+            result.docstore.mset([(doc.id, doc) for doc in all_docs])
+            # spans (with vectors)
+            all_span_ids = list(previous_retriever.vectorstore.yield_keys())
+            all_spans = previous_retriever.vectorstore.mget(all_span_ids)
+            result.vectorstore.mset([(span.id, span) for span in all_spans])
+
+        gr.Info("Retriever loaded successfully.")
+        return result
+    except Exception as e:
+        raise gr.Error(f"Failed to load retriever: {e}")
+
+
+def retrieve_similar_spans(
+    retriever: DocumentAwareSpanRetriever,
+    query_span_id: str,
+    **kwargs,
+) -> pd.DataFrame:
+    if not query_span_id.strip():
+        raise gr.Error("No query span selected.")
+    try:
+        retrieval_result = retriever.invoke(input=query_span_id, **kwargs)
+        records = []
+        for similar_span_doc in retrieval_result:
+            pie_doc, metadata = retriever.docstore.unwrap_with_metadata(similar_span_doc)
+            span_ann = metadata["attached_span"]
+            records.append(
+                {
+                    "doc_id": pie_doc.id,
+                    "span_id": similar_span_doc.id,
+                    "score": metadata["relevance_score"],
+                    "label": span_ann.label,
+                    "text": str(span_ann),
+                }
             )
-            gr.Info(f"Loaded embedding model: model_name={embedding_model_name}, device={device}")
-        except Exception as e:
-            raise gr.Error(f"Failed to load embedding model: {e}")
-    else:
-        embedding_model = None
+        return (
+            pd.DataFrame(records, columns=["doc_id", "score", "label", "text", "span_id"])
+            .sort_values(by="score", ascending=False)
+            .round(3)
+        )
+    except Exception as e:
+        raise gr.Error(f"Failed to retrieve similar ADUs: {e}")
+
+
+def retrieve_relevant_spans(
+    retriever: DocumentAwareSpanRetriever,
+    query_span_id: str,
+    relation_label_mapping: Optional[dict[str, str]] = None,
+    **kwargs,
+) -> pd.DataFrame:
+    if not query_span_id.strip():
+        raise gr.Error("No query span selected.")
+    try:
+        relation_label_mapping = relation_label_mapping or {}
+        retrieval_result = retriever.invoke(input=query_span_id, return_related=True, **kwargs)
+        records = []
+        for relevant_span_doc in retrieval_result:
+            pie_doc, metadata = retriever.docstore.unwrap_with_metadata(relevant_span_doc)
+            span_ann = metadata["attached_span"]
+            tail_span_ann = metadata["attached_tail_span"]
+            mapped_relation_label = relation_label_mapping.get(
+                metadata["relation_label"], metadata["relation_label"]
+            )
+            records.append(
+                {
+                    "doc_id": pie_doc.id,
+                    "type": mapped_relation_label,
+                    "rel_score": metadata["relation_score"],
+                    "text": str(tail_span_ann),
+                    "span_id": relevant_span_doc.id,
+                    "label": tail_span_ann.label,
+                    "ref_score": metadata["relevance_score"],
+                    "ref_label": span_ann.label,
+                    "ref_text": str(span_ann),
+                    "ref_span_id": metadata["head_id"],
+                }
+            )
+        return (
+            pd.DataFrame(
+                records,
+                columns=[
+                    "type",
+                    # omitted for now, we get no valid relation scores for the generative model
+                    # "rel_score",
+                    "ref_score",
+                    "label",
+                    "text",
+                    "ref_label",
+                    "ref_text",
+                    "doc_id",
+                    "span_id",
+                    "ref_span_id",
+                ],
+            )
+            .sort_values(by=["ref_score"], ascending=False)
+            .round(3)
+        )
+    except Exception as e:
+        raise gr.Error(f"Failed to retrieve relevant ADUs: {e}")
 
-    return embedding_model
 
+class RetrieverCallable(Protocol):
+    def __call__(
+        self,
+        retriever: DocumentAwareSpanRetriever,
+        query_span_id: str,
+        **kwargs,
+    ) -> Optional[pd.DataFrame]:
+        pass
 
-def load_models(
-    model_name: str,
-    revision: Optional[str] = None,
-    embedding_model_name: Optional[str] = None,
-    # embedding_model_revision: Optional[str] = None,
-    embedding_max_length: int = 512,
-    embedding_batch_size: int = 16,
-    device: str = "cpu",
-) -> Tuple[Pipeline, Optional[EmbeddingModel]]:
-    torch.cuda.empty_cache()
-    argumentation_model = load_argumentation_model(model_name, revision=revision, device=device)
-    embedding_model = load_embedding_model(
-        embedding_model_name=embedding_model_name,
-        # embedding_model_revision=embedding_model_revision,
-        embedding_max_length=embedding_max_length,
-        embedding_batch_size=embedding_batch_size,
-        device=device,
+
+def _retrieve_for_all_spans(
+    retriever: DocumentAwareSpanRetriever,
+    query_doc_id: str,
+    retrieve_func: RetrieverCallable,
+    query_span_id_column: str = "query_span_id",
+    **kwargs,
+) -> Optional[pd.DataFrame]:
+    if not query_doc_id.strip():
+        raise gr.Error("No query document selected.")
+    try:
+        span_id2idx = retriever.get_span_id2idx_from_doc(query_doc_id)
+        gr.Info(f"Retrieving results for {len(span_id2idx)} ADUs in document {query_doc_id}...")
+        span_results = {
+            query_span_id: retrieve_func(
+                retriever=retriever,
+                query_span_id=query_span_id,
+                **kwargs,
+            )
+            for query_span_id in span_id2idx.keys()
+        }
+        span_results_not_empty = {
+            query_span_id: df
+            for query_span_id, df in span_results.items()
+            if df is not None and not df.empty
+        }
+
+        # add column with query_span_id
+        for query_span_id, query_span_result in span_results_not_empty.items():
+            query_span_result[query_span_id_column] = query_span_id
+
+        if len(span_results_not_empty) == 0:
+            gr.Info(f"No results found for any ADU in document {query_doc_id}.")
+            return None
+        else:
+            result = pd.concat(span_results_not_empty.values(), ignore_index=True)
+            gr.Info(f"Retrieved {len(result)} ADUs for document {query_doc_id}.")
+            return result
+    except Exception as e:
+        raise gr.Error(
+            f'Failed to retrieve results for all ADUs in document "{query_doc_id}": {e}'
+        )
+
+
+def retrieve_all_similar_spans(
+    retriever: DocumentAwareSpanRetriever,
+    query_doc_id: str,
+    **kwargs,
+) -> Optional[pd.DataFrame]:
+    return _retrieve_for_all_spans(
+        retriever=retriever,
+        query_doc_id=query_doc_id,
+        retrieve_func=retrieve_similar_spans,
+        **kwargs,
     )
 
-    return argumentation_model, embedding_model
+
+def retrieve_all_relevant_spans(
+    retriever: DocumentAwareSpanRetriever,
+    query_doc_id: str,
+    **kwargs,
+) -> Optional[pd.DataFrame]:
+    return _retrieve_for_all_spans(
+        retriever=retriever,
+        query_doc_id=query_doc_id,
+        retrieve_func=retrieve_relevant_spans,
+        **kwargs,
+    )
+
+
+class RetrieverForAllSpansCallable(Protocol):
+    def __call__(
+        self,
+        retriever: DocumentAwareSpanRetriever,
+        query_doc_id: str,
+        **kwargs,
+    ) -> Optional[pd.DataFrame]:
+        pass
+
+
+def _retrieve_for_all_documents(
+    retriever: DocumentAwareSpanRetriever,
+    retrieve_func: RetrieverForAllSpansCallable,
+    query_doc_id_column: str = "query_doc_id",
+    **kwargs,
+) -> Optional[pd.DataFrame]:
+    try:
+        all_doc_ids = list(retriever.docstore.yield_keys())
+        gr.Info(f"Retrieving results for {len(all_doc_ids)} documents...")
+        doc_results = {
+            doc_id: retrieve_func(retriever=retriever, query_doc_id=doc_id, **kwargs)
+            for doc_id in all_doc_ids
+        }
+        doc_results_not_empty = {
+            doc_id: df for doc_id, df in doc_results.items() if df is not None and not df.empty
+        }
+        # add column with query_doc_id
+        for doc_id, doc_result in doc_results_not_empty.items():
+            doc_result[query_doc_id_column] = doc_id
+
+        if len(doc_results_not_empty) == 0:
+            gr.Info("No results found for any document.")
+            return None
+        else:
+            result = pd.concat(doc_results_not_empty, ignore_index=True)
+            gr.Info(f"Retrieved {len(result)} ADUs for all documents.")
+            return result
+    except Exception as e:
+        raise gr.Error(f"Failed to retrieve results for all documents: {e}")
+
+
+def retrieve_all_similar_spans_for_all_documents(
+    retriever: DocumentAwareSpanRetriever,
+    **kwargs,
+) -> Optional[pd.DataFrame]:
+    return _retrieve_for_all_documents(
+        retriever=retriever,
+        retrieve_func=retrieve_all_similar_spans,
+        **kwargs,
+    )
+
+
+def retrieve_all_relevant_spans_for_all_documents(
+    retriever: DocumentAwareSpanRetriever,
+    **kwargs,
+) -> Optional[pd.DataFrame]:
+    return _retrieve_for_all_documents(
+        retriever=retriever,
+        retrieve_func=retrieve_all_relevant_spans,
+        **kwargs,
+    )

rendering_utils.py

@@ -1,14 +1,9 @@
 import json
 import logging
 from collections import defaultdict
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Optional, Sequence, Union
 
-from annotation_utils import labeled_span_to_id
 from pytorch_ie.annotations import BinaryRelation, LabeledMultiSpan, LabeledSpan
-from pytorch_ie.documents import (
-    TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions,
-    TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-)
 from rendering_utils_displacy import EntityRenderer
 
 logger = logging.getLogger(__name__)
@@ -92,7 +87,20 @@ HIGHLIGHT_SPANS_JS = """
             });
         }
     }
-    function setReferenceAduId(entityId) {
+    function setHoverAduId(entityId) {
+        // get the textarea element that holds the reference adu id
+        let hoverAduIdDiv = document.querySelector('#hover_adu_id textarea');
+        // set the value of the input field
+        hoverAduIdDiv.value = entityId;
+        // trigger an input event to update the state
+        var event = new Event('input');
+        hoverAduIdDiv.dispatchEvent(event);
+    }
+    function setReferenceAduIdFromHover() {
+        // get the hover adu id
+        const hoverAduIdDiv = document.querySelector('#hover_adu_id textarea');
+        // get the value of the input field
+        const entityId = hoverAduIdDiv.value;
         // get the textarea element that holds the reference adu id
         let referenceAduIdDiv = document.querySelector('#selected_adu_id textarea');
         // set the value of the input field
@@ -104,6 +112,8 @@ HIGHLIGHT_SPANS_JS = """
 
     const entities = document.querySelectorAll('.entity');
     entities.forEach(entity => {
+        // make the cursor a pointer
+        entity.style.cursor = 'pointer';
         const alreadyHasListener = entity.getAttribute('data-has-listener');
         if (alreadyHasListener) {
             return;
@@ -111,22 +121,30 @@ HIGHLIGHT_SPANS_JS = """
         entity.addEventListener('mouseover', () => {
             const entityId = entity.getAttribute('data-entity-id');
             highlightRelationArguments(entityId);
-            setReferenceAduId(entityId);
+            setHoverAduId(entityId);
         });
         entity.addEventListener('mouseout', () => {
             highlightRelationArguments(null);
         });
         entity.setAttribute('data-has-listener', 'true');
     });
+    const entityContainer = document.querySelector('.entities');
+    if (entityContainer) {
+        entityContainer.addEventListener('click', () => {
+            setReferenceAduIdFromHover();
+        });
+        // make the cursor a pointer
+        // entityContainer.style.cursor = 'pointer';
+    }
 }
 """
 
 
 def render_pretty_table(
-    document: Union[
-        TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-        TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions,
-    ],
+    text: str,
+    spans: Union[Sequence[LabeledSpan], Sequence[LabeledMultiSpan]],
+    span_id2idx: Dict[str, int],
+    binary_relations: Sequence[BinaryRelation],
     **render_kwargs,
 ):
     from prettytable import PrettyTable
@@ -134,7 +152,7 @@ def render_pretty_table(
     t = PrettyTable()
     t.field_names = ["head", "tail", "relation"]
     t.align = "l"
-    for relation in list(document.binary_relations) + list(document.binary_relations.predictions):
+    for relation in list(binary_relations) + list(binary_relations):
         t.add_row([str(relation.head), str(relation.tail), relation.label])
 
     html = t.get_html_string(format=True)
@@ -144,29 +162,19 @@ def render_pretty_table(
 
 
 def render_displacy(
-    document: Union[
-        TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
-        TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions,
-    ],
+    text: str,
+    spans: Union[Sequence[LabeledSpan], Sequence[LabeledMultiSpan]],
+    span_id2idx: Dict[str, int],
+    binary_relations: Sequence[BinaryRelation],
     inject_relations=True,
     colors_hover=None,
     entity_options={},
     **render_kwargs,
 ):
 
-    if isinstance(document, TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions):
-        span_layer = document.labeled_spans
-    elif isinstance(
-        document, TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions
-    ):
-        span_layer = document.labeled_multi_spans
-    else:
-        raise ValueError(f"Unsupported document type: {type(document)}")
-
-    span_annotations = list(span_layer) + list(span_layer.predictions)
     ents = []
-    for labeled_span in span_annotations:
-        entity_id = labeled_span_to_id(labeled_span)
+    for entity_id, idx in span_id2idx.items():
+        labeled_span = spans[idx]
         # pass the ID as a parameter to the entity. The id is required to fetch the entity annotations
         # on hover and to inject the relation data.
         if isinstance(labeled_span, LabeledSpan):
@@ -192,7 +200,7 @@ def render_displacy(
             raise ValueError(f"Unsupported labeled span type: {type(labeled_span)}")
 
     spacy_doc = {
-        "text": document.text,
+        "text": text,
         # the ents MUST be sorted by start and end
         "ents": sorted(ents, key=lambda x: (x["start"], x["end"])),
         "title": None,
@@ -207,12 +215,10 @@ def render_displacy(
 
     html = "<div style='max-width:100%; max-height:360px; overflow:auto'>" + html + "</div>"
     if inject_relations:
-        binary_relations = list(document.binary_relations) + list(
-            document.binary_relations.predictions
-        )
         html = inject_relation_data(
             html,
-            span_annotations=span_annotations,
+            spans=spans,
+            span_id2idx=span_id2idx,
             binary_relations=binary_relations,
             additional_colors=colors_hover,
         )
@@ -221,8 +227,9 @@ def render_displacy(
 
 def inject_relation_data(
     html: str,
-    span_annotations: Union[List[LabeledSpan], List[LabeledMultiSpan]],
-    binary_relations: List[BinaryRelation],
+    spans: Union[Sequence[LabeledSpan], Sequence[LabeledMultiSpan]],
+    span_id2idx: Dict[str, int],
+    binary_relations: Sequence[BinaryRelation],
     additional_colors: Optional[Dict[str, Union[str, dict]]] = None,
 ) -> str:
     from bs4 import BeautifulSoup
@@ -236,7 +243,7 @@ def inject_relation_data(
         entity2heads[relation.tail].append((relation.head, relation.label))
         entity2tails[relation.head].append((relation.tail, relation.label))
 
-    ann_id2annotation = {labeled_span_to_id(entity): entity for entity in span_annotations}
+    annotation2id = {spans[span_idx]: span_id for span_id, span_idx in span_id2idx.items()}
     # Add unique IDs to each entity
     entities = soup.find_all(class_="entity")
     for entity in entities:
@@ -247,7 +254,8 @@ def inject_relation_data(
                 entity[f"data-color-{key}"] = (
                     json.dumps(color) if isinstance(color, dict) else color
                 )
-        entity_annotation = ann_id2annotation[entity["data-entity-id"]]
+
+        entity_annotation = spans[span_id2idx[entity["data-entity-id"]]]
 
         # sanity check.
         if isinstance(entity_annotation, LabeledSpan):
@@ -265,14 +273,16 @@ def inject_relation_data(
         entity["data-label"] = entity_annotation.label
         entity["data-relation-tails"] = json.dumps(
             [
-                {"entity-id": labeled_span_to_id(tail), "label": label}
+                {"entity-id": annotation2id[tail], "label": label}
                 for tail, label in entity2tails.get(entity_annotation, [])
+                if tail in annotation2id
             ]
         )
         entity["data-relation-heads"] = json.dumps(
             [
-                {"entity-id": labeled_span_to_id(head), "label": label}
+                {"entity-id": annotation2id[head], "label": label}
                 for head, label in entity2heads.get(entity_annotation, [])
+                if head in annotation2id
             ]
         )
 

requirements.txt

@@ -1,11 +1,34 @@
-pytorch-ie==0.31.1
-gradio==4.36.0
+gradio==4.44.1
 prettytable==3.10.0
-pie-modules==0.12.0
 beautifulsoup4==4.12.3
-datasets==2.14.4
 # numpy 2.0.0 breaks the code
 numpy==1.25.2
-qdrant-client==1.9.1
 scipy==1.13.0
 arxiv==2.1.3
+pyrootutils>=1.0.0,<1.1.0
+
+########## from root requirements ##########
+# --------- pytorch-ie --------- #
+pytorch-ie>=0.29.6,<0.32.0
+pie-datasets>=0.10.5,<0.11.0
+pie-modules>=0.14.0,<0.15.0
+
+# --------- models -------- #
+adapters>=0.1.2,<0.2.0
+# ADU retrieval (and demo, in future):
+langchain>=0.3.0,<0.4.0
+langchain-core>=0.3.0,<0.4.0
+langchain-community>=0.3.0,<0.4.0
+# we use QDrant as vectorstore backend
+langchain-qdrant>=0.1.0,<0.2.0
+qdrant-client>=1.12.0,<2.0.0
+# 0.26 seems to be broken when used with adapters, see https://github.com/adapter-hub/adapters/issues/748
+huggingface_hub<0.26.0  # 0.26 seems to be broken
+# to to handle segmented entities (if HANDLE_PARTS_OF_SAME=True)
+networkx>=3.0.0,<4.0.0
+
+# --------- config --------- #
+hydra-core>=1.3.0
+
+# --------- dev --------- #
+pre-commit      # hooks for applying linters on commit

retrieve_and_dump_all_relevant.py

@@ -0,0 +1,101 @@
+import pyrootutils
+
+root = pyrootutils.setup_root(
+    search_from=__file__,
+    indicator=[".project-root"],
+    pythonpath=True,
+    dotenv=True,
+)
+
+import argparse
+import logging
+
+from demo.model_utils import (
+    retrieve_all_relevant_spans,
+    retrieve_all_relevant_spans_for_all_documents,
+    retrieve_relevant_spans,
+)
+from src.langchain_modules import DocumentAwareSpanRetrieverWithRelations
+
+logger = logging.getLogger(__name__)
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-c",
+        "--config_path",
+        type=str,
+        default="configs/retriever/related_span_retriever_with_relations_from_other_docs.yaml",
+    )
+    parser.add_argument(
+        "--data_path",
+        type=str,
+        required=True,
+        help="Path to a zip or directory containing a retriever dump.",
+    )
+    parser.add_argument("-k", "--top_k", type=int, default=10)
+    parser.add_argument("-t", "--threshold", type=float, default=0.95)
+    parser.add_argument(
+        "-o",
+        "--output_path",
+        type=str,
+        required=True,
+    )
+    parser.add_argument(
+        "--query_doc_id",
+        type=str,
+        default=None,
+        help="If provided, retrieve all spans for only this query document.",
+    )
+    parser.add_argument(
+        "--query_span_id",
+        type=str,
+        default=None,
+        help="If provided, retrieve all spans for only this query span.",
+    )
+    args = parser.parse_args()
+
+    logging.basicConfig(
+        format="%(asctime)s %(levelname)-8s %(message)s",
+        level=logging.INFO,
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    if not args.output_path.endswith(".json"):
+        raise ValueError("only support json output")
+
+    logger.info(f"instantiating retriever from {args.config_path}...")
+    retriever = DocumentAwareSpanRetrieverWithRelations.instantiate_from_config_file(
+        args.config_path
+    )
+    logger.info(f"loading data from {args.data_path}...")
+    retriever.load_from_disc(args.data_path)
+
+    search_kwargs = {"k": args.top_k, "score_threshold": args.threshold}
+    logger.info(f"use search_kwargs: {search_kwargs}")
+
+    if args.query_span_id is not None:
+        logger.warning(f"retrieving results for single span: {args.query_span_id}")
+        all_spans_for_all_documents = retrieve_relevant_spans(
+            retriever=retriever, query_span_id=args.query_span_id, **search_kwargs
+        )
+    elif args.query_doc_id is not None:
+        logger.warning(f"retrieving results for single document: {args.query_doc_id}")
+        all_spans_for_all_documents = retrieve_all_relevant_spans(
+            retriever=retriever, query_doc_id=args.query_doc_id, **search_kwargs
+        )
+    else:
+        all_spans_for_all_documents = retrieve_all_relevant_spans_for_all_documents(
+            retriever=retriever, **search_kwargs
+        )
+
+    if all_spans_for_all_documents is None:
+        logger.warning("no relevant spans found in any document")
+        exit(0)
+
+    logger.info(f"dumping results to {args.output_path}...")
+    all_spans_for_all_documents.to_json(args.output_path)
+
+    logger.info("done")

retriever/related_span_retriever_with_relations_from_other_docs.yaml

@@ -0,0 +1,47 @@
+_target_: src.langchain_modules.DocumentAwareSpanRetrieverWithRelations
+reversed_relations_suffix: _reversed
+relation_labels:
+  - supports_reversed
+  - contradicts_reversed
+retrieve_from_same_document: false
+retrieve_from_different_documents: true
+pie_document_type:
+  _target_: pie_modules.utils.resolve_type
+  type_or_str: pytorch_ie.documents.TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions
+docstore:
+  _target_: src.langchain_modules.DatasetsPieDocumentStore
+search_kwargs:
+  k: 10
+search_type: similarity_score_threshold
+vectorstore:
+  _target_: src.langchain_modules.QdrantSpanVectorStore
+  embedding:
+    _target_: src.langchain_modules.HuggingFaceSpanEmbeddings
+    model:
+      _target_: src.models.utils.adapters.load_model_with_adapter
+      model_kwargs:
+        pretrained_model_name_or_path: allenai/specter2_base
+      adapter_kwargs:
+        adapter_name_or_path: allenai/specter2
+        load_as: proximity
+        source: hf
+    pipeline_kwargs:
+      tokenizer: allenai/specter2_base
+      stride: 64
+      batch_size: 32
+    model_max_length: 512
+  client:
+    _target_: qdrant_client.QdrantClient
+    location: ":memory:"
+  collection_name: adus
+  vector_params:
+    distance:
+      _target_: qdrant_client.http.models.Distance
+      value: Cosine
+  label_mapping:
+    background_claim:
+      - background_claim
+      - own_claim
+    own_claim:
+      - background_claim
+      - own_claim

src/hf_pipeline/__init__.py

@@ -0,0 +1 @@
+from .feature_extraction import FeatureExtractionPipelineWithStriding

src/hf_pipeline/feature_extraction.py

@@ -0,0 +1,317 @@
+import types
+from collections import defaultdict
+from typing import List, Optional, Tuple, Union
+
+import torch
+from transformers.pipelines.base import ArgumentHandler, ChunkPipeline, Dataset
+from transformers.utils import is_tf_available, is_torch_available
+
+if is_tf_available():
+    import tensorflow as tf
+    from transformers.models.auto.modeling_tf_auto import (
+        TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES,
+    )
+if is_torch_available():
+    from transformers.models.auto.modeling_auto import MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES
+
+
+def list_of_dicts2dict_of_lists(list_of_dicts: list[dict]) -> dict[str, list]:
+    return {k: [d[k] for d in list_of_dicts] for k in list_of_dicts[0].keys()}
+
+
+class FeatureExtractionArgumentHandler(ArgumentHandler):
+    """Handles arguments for feature extraction."""
+
+    def __call__(self, inputs: Union[str, List[str]], **kwargs):
+        if inputs is not None and isinstance(inputs, (list, tuple)) and len(inputs) > 0:
+            inputs = list(inputs)
+            batch_size = len(inputs)
+        elif isinstance(inputs, str):
+            inputs = [inputs]
+            batch_size = 1
+        elif (
+            Dataset is not None
+            and isinstance(inputs, Dataset)
+            or isinstance(inputs, types.GeneratorType)
+        ):
+            return inputs, None
+        else:
+            raise ValueError("At least one input is required.")
+
+        offset_mapping = kwargs.get("offset_mapping")
+        if offset_mapping:
+            if isinstance(offset_mapping, list) and isinstance(offset_mapping[0], tuple):
+                offset_mapping = [offset_mapping]
+            if len(offset_mapping) != batch_size:
+                raise ValueError("offset_mapping should have the same batch size as the input")
+        return inputs, offset_mapping
+
+
+class FeatureExtractionPipelineWithStriding(ChunkPipeline):
+    """Same as transformers.FeatureExtractionPipeline, but with long input handling. Inspired by
+    transformers.TokenClassificationPipeline. The functionality is triggered when providing the
+    "stride" parameter (can be 0). When passing "create_unique_embeddings_per_token=True", only one
+    embedding (and other data, see flags below) per token will be returned (this makes use of
+    min_distance_to_border, see "return_min_distance_to_border" below for details). Note that this
+    removes data for special token positions!
+
+    Per default, it will return just the embeddings. If any of the return_ADDITIONAL_RESULT is
+    enabled (see below), it will return dictionaries with "last_hidden_state" and all
+    ADDITIONAL_RESULT depending on the flags.
+
+    Flags to return additional results:
+        return_offset_mapping: If enabled, return the offset mapping.
+        return_special_tokens_mask: If enabled, return the special tokens mask.
+        return_sequence_indices: If enabled, return the sequence indices.
+        return_position_ids: If enabled, return the position ids from, values are in [0, model_max_length).
+        return_min_distance_to_border: If enabled, return the minimum distance to the "border" of
+            the input that gets passed into the model. This is useful when striding is used which may
+            produce multiple embeddings for a token (compare values in offset_mapping). In this case,
+            min_distance_to_border can be used to select the embedding that is more in the center
+            of the input by choosing the entry with the *higher* min_distance_to_border.
+    """
+
+    default_input_names = "sequences"
+
+    def __init__(self, args_parser=FeatureExtractionArgumentHandler(), *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.check_model_type(
+            TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES
+            if self.framework == "tf"
+            else MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES
+        )
+
+        self._args_parser = args_parser
+
+    def _sanitize_parameters(
+        self,
+        offset_mapping: Optional[List[Tuple[int, int]]] = None,
+        stride: Optional[int] = None,
+        create_unique_embeddings_per_token: Optional[bool] = False,
+        return_offset_mapping: Optional[bool] = None,
+        return_special_tokens_mask: Optional[bool] = None,
+        return_sequence_indices: Optional[bool] = None,
+        return_position_ids: Optional[bool] = None,
+        return_min_distance_to_border: Optional[bool] = None,
+        return_tensors=None,
+    ):
+        preprocess_params = {}
+        if offset_mapping is not None:
+            preprocess_params["offset_mapping"] = offset_mapping
+
+        if stride is not None:
+            if stride >= self.tokenizer.model_max_length:
+                raise ValueError(
+                    "`stride` must be less than `tokenizer.model_max_length` (or even lower if the tokenizer adds special tokens)"
+                )
+
+            if self.tokenizer.is_fast:
+                tokenizer_params = {
+                    "return_overflowing_tokens": True,
+                    "padding": True,
+                    "stride": stride,
+                }
+                preprocess_params["tokenizer_params"] = tokenizer_params  # type: ignore
+            else:
+                raise ValueError(
+                    "`stride` was provided to process all the text but you're using a slow tokenizer."
+                    " Please use a fast tokenizer."
+                )
+        postprocess_params = {}
+        if create_unique_embeddings_per_token is not None:
+            postprocess_params["create_unique_embeddings_per_token"] = (
+                create_unique_embeddings_per_token
+            )
+        if return_offset_mapping is not None:
+            postprocess_params["return_offset_mapping"] = return_offset_mapping
+        if return_special_tokens_mask is not None:
+            postprocess_params["return_special_tokens_mask"] = return_special_tokens_mask
+        if return_sequence_indices is not None:
+            postprocess_params["return_sequence_indices"] = return_sequence_indices
+        if return_position_ids is not None:
+            postprocess_params["return_position_ids"] = return_position_ids
+        if return_min_distance_to_border is not None:
+            postprocess_params["return_min_distance_to_border"] = return_min_distance_to_border
+        if return_tensors is not None:
+            postprocess_params["return_tensors"] = return_tensors
+        return preprocess_params, {}, postprocess_params
+
+    def __call__(self, inputs: Union[str, List[str]], **kwargs):
+
+        _inputs, offset_mapping = self._args_parser(inputs, **kwargs)
+        if offset_mapping:
+            kwargs["offset_mapping"] = offset_mapping
+
+        return super().__call__(inputs, **kwargs)
+
+    def preprocess(self, sentence, offset_mapping=None, **preprocess_params):
+        tokenizer_params = preprocess_params.pop("tokenizer_params", {})
+        truncation = (
+            True
+            if self.tokenizer.model_max_length and self.tokenizer.model_max_length > 0
+            else False
+        )
+        inputs = self.tokenizer(
+            sentence,
+            return_tensors=self.framework,
+            truncation=truncation,
+            return_special_tokens_mask=True,
+            return_offsets_mapping=self.tokenizer.is_fast,
+            **tokenizer_params,
+        )
+        inputs.pop("overflow_to_sample_mapping", None)
+        num_chunks = len(inputs["input_ids"])
+
+        for i in range(num_chunks):
+            if self.framework == "tf":
+                model_inputs = {k: tf.expand_dims(v[i], 0) for k, v in inputs.items()}
+            else:
+                model_inputs = {k: v[i].unsqueeze(0) for k, v in inputs.items()}
+            if offset_mapping is not None:
+                model_inputs["offset_mapping"] = offset_mapping
+            model_inputs["sentence"] = sentence if i == 0 else None
+            model_inputs["is_last"] = i == num_chunks - 1
+
+            yield model_inputs
+
+    def _forward(self, model_inputs, **kwargs):
+        # Forward
+        special_tokens_mask = model_inputs.pop("special_tokens_mask")
+        offset_mapping = model_inputs.pop("offset_mapping", None)
+        sentence = model_inputs.pop("sentence")
+        is_last = model_inputs.pop("is_last")
+        if self.framework == "tf":
+            last_hidden_state = self.model(**model_inputs)[0]
+        else:
+            output = self.model(**model_inputs)
+            last_hidden_state = (
+                output["last_hidden_state"] if isinstance(output, dict) else output[0]
+            )
+
+        return {
+            "last_hidden_state": last_hidden_state,
+            "special_tokens_mask": special_tokens_mask,
+            "offset_mapping": offset_mapping,
+            "sentence": sentence,
+            "is_last": is_last,
+            **model_inputs,
+        }
+
+    def postprocess_tensor(self, data, return_tensors=False):
+        if return_tensors:
+            return data
+        if self.framework == "pt":
+            return data.tolist()
+        elif self.framework == "tf":
+            return data.numpy().tolist()
+        else:
+            raise ValueError(f"unknown framework: {self.framework}")
+
+    def make_embeddings_unique_per_token(
+        self, data, offset_mapping, special_tokens_mask, min_distance_to_border
+    ):
+        char_offsets2token_pos = defaultdict(list)
+        bs, seq_len = offset_mapping.shape[:2]
+        if bs != 1:
+            raise ValueError(f"expected result batch size 1, but it is: {bs}")
+        for token_idx, ((char_start, shar_end), is_special_token, min_dist) in enumerate(
+            zip(
+                offset_mapping[0].tolist(),
+                special_tokens_mask[0].tolist(),
+                min_distance_to_border[0].tolist(),
+            )
+        ):
+
+            if not is_special_token:
+                char_offsets2token_pos[(char_start, shar_end)].append((token_idx, min_dist))
+
+        # tokens_with_multiple_embeddings = {k: v for k, v in char_offsets2token_pos.items() if len(v) > 1}
+        char_offsets2best_token_pos = {
+            k: max(v, key=lambda pos_dist: pos_dist[1])[0]
+            for k, v in char_offsets2token_pos.items()
+        }
+        # sort by char offsets (start and end)
+        sorted_char_offsets_token_positions = sorted(
+            char_offsets2best_token_pos.items(),
+            key=lambda char_offsets_tok_pos: (
+                char_offsets_tok_pos[0][0],
+                char_offsets_tok_pos[0][1],
+            ),
+        )
+        best_indices = [tok_pos for char_offsets, tok_pos in sorted_char_offsets_token_positions]
+
+        result = {k: v[0][best_indices].unsqueeze(0) for k, v in data.items()}
+        return result
+
+    def postprocess(
+        self,
+        all_outputs,
+        create_unique_embeddings_per_token: bool = False,
+        return_offset_mapping: bool = False,
+        return_special_tokens_mask: bool = False,
+        return_sequence_indices: bool = False,
+        return_position_ids: bool = False,
+        return_min_distance_to_border: bool = False,
+        return_tensors: bool = False,
+    ):
+
+        all_outputs_dict = list_of_dicts2dict_of_lists(all_outputs)
+        if self.framework == "pt":
+            result = {
+                "last_hidden_state": torch.concat(all_outputs_dict["last_hidden_state"], axis=1)
+            }
+            if return_offset_mapping or create_unique_embeddings_per_token:
+                result["offset_mapping"] = torch.concat(all_outputs_dict["offset_mapping"], axis=1)
+            if return_special_tokens_mask or create_unique_embeddings_per_token:
+                result["special_tokens_mask"] = torch.concat(
+                    all_outputs_dict["special_tokens_mask"], axis=1
+                )
+            if return_sequence_indices:
+                sequence_indices = []
+                for seq_idx, model_outputs in enumerate(all_outputs):
+                    sequence_indices.append(torch.ones_like(model_outputs["input_ids"]) * seq_idx)
+                result["sequence_indices"] = torch.concat(sequence_indices, axis=1)
+            if return_position_ids:
+                position_ids = []
+                for seq_idx, model_outputs in enumerate(all_outputs):
+                    seq_len = model_outputs["input_ids"].size(1)
+                    position_ids.append(torch.arange(seq_len).unsqueeze(0))
+                result["indices"] = torch.concat(position_ids, axis=1)
+            if return_min_distance_to_border or create_unique_embeddings_per_token:
+                min_distance_to_border = []
+                for seq_idx, model_outputs in enumerate(all_outputs):
+                    seq_len = model_outputs["input_ids"].size(1)
+                    current_indices = torch.arange(seq_len).unsqueeze(0)
+                    min_distance_to_border.append(
+                        torch.minimum(current_indices, seq_len - current_indices)
+                    )
+                result["min_distance_to_border"] = torch.concat(min_distance_to_border, axis=1)
+        elif self.framework == "tf":
+            raise NotImplementedError()
+        else:
+            raise ValueError(f"unknown framework: {self.framework}")
+
+        if create_unique_embeddings_per_token:
+            offset_mapping = result["offset_mapping"]
+            if not return_offset_mapping:
+                del result["offset_mapping"]
+            special_tokens_mask = result["special_tokens_mask"]
+            if not return_special_tokens_mask:
+                del result["special_tokens_mask"]
+            min_distance_to_border = result["min_distance_to_border"]
+            if not return_min_distance_to_border:
+                del result["min_distance_to_border"]
+            result = self.make_embeddings_unique_per_token(
+                data=result,
+                offset_mapping=offset_mapping,
+                special_tokens_mask=special_tokens_mask,
+                min_distance_to_border=min_distance_to_border,
+            )
+
+        result = {
+            k: self.postprocess_tensor(v, return_tensors=return_tensors) for k, v in result.items()
+        }
+        if set(result) == {"last_hidden_state"}:
+            return result["last_hidden_state"]
+        else:
+            return result

src/langchain_modules/__init__.py

@@ -0,0 +1,9 @@
+from .basic_pie_document_store import BasicPieDocumentStore
+from .datasets_pie_document_store import DatasetsPieDocumentStore
+from .huggingface_span_embeddings import HuggingFaceSpanEmbeddings
+from .pie_document_store import PieDocumentStore
+from .qdrant_span_vectorstore import QdrantSpanVectorStore
+from .serializable_store import SerializableStore
+from .span_embeddings import SpanEmbeddings
+from .span_retriever import DocumentAwareSpanRetriever, DocumentAwareSpanRetrieverWithRelations
+from .span_vectorstore import SpanVectorStore

src/langchain_modules/basic_pie_document_store.py

@@ -0,0 +1,103 @@
+import json
+import logging
+import os
+import shutil
+from itertools import islice
+from typing import Iterator, List, Optional, Sequence, Tuple
+
+from langchain.storage import create_kv_docstore
+from langchain_core.documents import Document as LCDocument
+from langchain_core.stores import BaseStore, ByteStore
+from pie_datasets import Dataset, DatasetDict
+
+from .pie_document_store import PieDocumentStore
+
+logger = logging.getLogger(__name__)
+
+
+class BasicPieDocumentStore(PieDocumentStore):
+    """PIE Document store that uses a client to store and retrieve documents."""
+
+    def __init__(
+        self,
+        client: Optional[BaseStore[str, LCDocument]] = None,
+        byte_store: Optional[ByteStore] = None,
+    ):
+        if byte_store is not None:
+            client = create_kv_docstore(byte_store)
+        elif client is None:
+            raise Exception("You must pass a `byte_store` parameter.")
+
+        self.client = client
+
+    def mget(self, keys: Sequence[str]) -> List[LCDocument]:
+        return self.client.mget(keys)
+
+    def mset(self, items: Sequence[Tuple[str, LCDocument]]) -> None:
+        self.client.mset(items)
+
+    def mdelete(self, keys: Sequence[str]) -> None:
+        self.client.mdelete(keys)
+
+    def yield_keys(self, *, prefix: Optional[str] = None) -> Iterator[str]:
+        return self.client.yield_keys(prefix=prefix)
+
+    def _save_to_directory(self, path: str, batch_size: Optional[int] = None, **kwargs) -> None:
+        all_doc_ids = []
+        all_metadata = []
+        pie_documents_path = os.path.join(path, "pie_documents")
+        if os.path.exists(pie_documents_path):
+            # remove existing directory
+            logger.warning(f"Removing existing directory: {pie_documents_path}")
+            shutil.rmtree(pie_documents_path)
+        os.makedirs(pie_documents_path, exist_ok=True)
+        doc_ids_iter = iter(self.client.yield_keys())
+        while batch_doc_ids := list(islice(doc_ids_iter, batch_size or 1000)):
+            all_doc_ids.extend(batch_doc_ids)
+            docs = self.client.mget(batch_doc_ids)
+            pie_docs = []
+            for doc in docs:
+                pie_doc = doc.metadata[self.METADATA_KEY_PIE_DOCUMENT]
+                pie_docs.append(pie_doc)
+                all_metadata.append(
+                    {k: v for k, v in doc.metadata.items() if k != self.METADATA_KEY_PIE_DOCUMENT}
+                )
+            pie_dataset = Dataset.from_documents(pie_docs)
+            DatasetDict({"train": pie_dataset}).to_json(path=pie_documents_path)
+        if len(all_doc_ids) > 0:
+            doc_ids_path = os.path.join(path, "doc_ids.json")
+            with open(doc_ids_path, "w") as f:
+                json.dump(all_doc_ids, f)
+        if len(all_metadata) > 0:
+            metadata_path = os.path.join(path, "metadata.json")
+            with open(metadata_path, "w") as f:
+                json.dump(all_metadata, f)
+
+    def _load_from_directory(self, path: str, **kwargs) -> None:
+        pie_documents_path = os.path.join(path, "pie_documents")
+        if not os.path.exists(pie_documents_path):
+            logger.warning(
+                f"Directory {pie_documents_path} does not exist, don't load any documents."
+            )
+            return None
+        pie_dataset = DatasetDict.from_json(data_dir=pie_documents_path)
+        pie_docs = pie_dataset["train"]
+        metadata_path = os.path.join(path, "metadata.json")
+        if os.path.exists(metadata_path):
+            with open(metadata_path, "r") as f:
+                all_metadata = json.load(f)
+        else:
+            logger.warning(f"File {metadata_path} does not exist, don't load any metadata.")
+            all_metadata = [{} for _ in pie_docs]
+        docs = [
+            self.wrap(pie_doc, **metadata) for pie_doc, metadata in zip(pie_docs, all_metadata)
+        ]
+        doc_ids_path = os.path.join(path, "doc_ids.json")
+        if os.path.exists(doc_ids_path):
+            with open(doc_ids_path, "r") as f:
+                all_doc_ids = json.load(f)
+        else:
+            logger.warning(f"File {doc_ids_path} does not exist, don't load any document ids.")
+            all_doc_ids = [doc.id for doc in pie_docs]
+        self.client.mset(zip(all_doc_ids, docs))
+        logger.info(f"Loaded {len(docs)} documents from {path} into docstore")

src/langchain_modules/datasets_pie_document_store.py

@@ -0,0 +1,156 @@
+import json
+import logging
+import os
+import shutil
+from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple
+
+from datasets import Dataset as HFDataset
+from langchain_core.documents import Document as LCDocument
+from pie_datasets import Dataset, DatasetDict, concatenate_datasets
+from pytorch_ie.documents import TextBasedDocument
+
+from .pie_document_store import PieDocumentStore
+
+logger = logging.getLogger(__name__)
+
+
+class DatasetsPieDocumentStore(PieDocumentStore):
+    """PIE Document store that uses Huggingface Datasets as the backend."""
+
+    def __init__(self) -> None:
+        self._data: Optional[Dataset] = None
+        # keys map to indices in the dataset
+        self._keys: Dict[str, int] = {}
+        self._metadata: Dict[str, Any] = {}
+
+    def __len__(self):
+        return len(self._keys)
+
+    def _get_pie_docs_by_indices(self, indices: Iterable[int]) -> Sequence[TextBasedDocument]:
+        if self._data is None:
+            return []
+        return self._data.apply_hf_func(func=HFDataset.select, indices=indices)
+
+    def mget(self, keys: Sequence[str]) -> List[LCDocument]:
+        if self._data is None or len(keys) == 0:
+            return []
+        keys_in_data = [key for key in keys if key in self._keys]
+        indices = [self._keys[key] for key in keys_in_data]
+        dataset = self._get_pie_docs_by_indices(indices)
+        metadatas = [self._metadata.get(key, {}) for key in keys_in_data]
+        return [self.wrap(pie_doc, **metadata) for pie_doc, metadata in zip(dataset, metadatas)]
+
+    def mset(self, items: Sequence[Tuple[str, LCDocument]]) -> None:
+        if len(items) == 0:
+            return
+        keys, new_docs = zip(*items)
+        pie_docs, metadatas = zip(*[self.unwrap_with_metadata(doc) for doc in new_docs])
+        if self._data is None:
+            idx_start = 0
+            self._data = Dataset.from_documents(pie_docs)
+        else:
+            # we pass the features to the new dataset to mitigate issues caused by
+            # slightly different inferred features
+            dataset = Dataset.from_documents(pie_docs, features=self._data.features)
+            idx_start = len(self._data)
+            self._data = concatenate_datasets([self._data, dataset], clear_metadata=False)
+        keys_dict = {key: idx for idx, key in zip(range(idx_start, len(self._data)), keys)}
+        self._keys.update(keys_dict)
+        self._metadata.update(
+            {key: metadata for key, metadata in zip(keys, metadatas) if metadata}
+        )
+
+    def add_pie_dataset(
+        self,
+        dataset: Dataset,
+        keys: Optional[List[str]] = None,
+        metadatas: Optional[List[Dict[str, Any]]] = None,
+    ) -> None:
+        if len(dataset) == 0:
+            return
+        if keys is None:
+            keys = [doc.id for doc in dataset]
+        if len(keys) != len(set(keys)):
+            raise ValueError("Keys must be unique.")
+        if None in keys:
+            raise ValueError("Keys must not be None.")
+        if metadatas is None:
+            metadatas = [{} for _ in range(len(dataset))]
+        if len(keys) != len(dataset) or len(keys) != len(metadatas):
+            raise ValueError("Keys, dataset and metadatas must have the same length.")
+
+        if self._data is None:
+            idx_start = 0
+            self._data = dataset
+        else:
+            idx_start = len(self._data)
+            self._data = concatenate_datasets([self._data, dataset], clear_metadata=False)
+        keys_dict = {key: idx for idx, key in zip(range(idx_start, len(self._data)), keys)}
+        self._keys.update(keys_dict)
+        metadatas_dict = {key: metadata for key, metadata in zip(keys, metadatas) if metadata}
+        self._metadata.update(metadatas_dict)
+
+    def mdelete(self, keys: Sequence[str]) -> None:
+        for key in keys:
+            idx = self._keys.pop(key, None)
+            if idx is not None:
+                self._metadata.pop(key, None)
+
+    def yield_keys(self, *, prefix: Optional[str] = None) -> Iterator[str]:
+        return (key for key in self._keys if prefix is None or key.startswith(prefix))
+
+    def _purge_invalid_entries(self):
+        if self._data is None or len(self._keys) == len(self._data):
+            return
+        self._data = self._get_pie_docs_by_indices(self._keys.values())
+
+    def _save_to_directory(self, path: str, batch_size: Optional[int] = None, **kwargs) -> None:
+        self._purge_invalid_entries()
+        if len(self) == 0:
+            logger.warning("No documents to save.")
+            return
+
+        all_doc_ids = list(self._keys)
+        all_metadatas: List[Dict[str, Any]] = [self._metadata.get(key, {}) for key in all_doc_ids]
+        pie_documents_path = os.path.join(path, "pie_documents")
+        if os.path.exists(pie_documents_path):
+            # remove existing directory
+            logger.warning(f"Removing existing directory: {pie_documents_path}")
+            shutil.rmtree(pie_documents_path)
+        os.makedirs(pie_documents_path, exist_ok=True)
+        DatasetDict({"train": self._data}).to_json(pie_documents_path)
+        doc_ids_path = os.path.join(path, "doc_ids.json")
+        with open(doc_ids_path, "w") as f:
+            json.dump(all_doc_ids, f)
+        metadata_path = os.path.join(path, "metadata.json")
+        with open(metadata_path, "w") as f:
+            json.dump(all_metadatas, f)
+
+    def _load_from_directory(self, path: str, **kwargs) -> None:
+        doc_ids_path = os.path.join(path, "doc_ids.json")
+        if os.path.exists(doc_ids_path):
+            with open(doc_ids_path, "r") as f:
+                all_doc_ids = json.load(f)
+        else:
+            logger.warning(f"File {doc_ids_path} does not exist, don't load any document ids.")
+            all_doc_ids = None
+        metadata_path = os.path.join(path, "metadata.json")
+        if os.path.exists(metadata_path):
+            with open(metadata_path, "r") as f:
+                all_metadata = json.load(f)
+        else:
+            logger.warning(f"File {metadata_path} does not exist, don't load any metadata.")
+            all_metadata = None
+        pie_documents_path = os.path.join(path, "pie_documents")
+        if not os.path.exists(pie_documents_path):
+            logger.warning(
+                f"Directory {pie_documents_path} does not exist, don't load any documents."
+            )
+            return None
+        # If we have a dataset already loaded, we use its features to load the new dataset
+        # This is to mitigate issues caused by slightly different inferred features.
+        features = self._data.features if self._data is not None else None
+        pie_dataset = DatasetDict.from_json(data_dir=pie_documents_path, features=features)
+        pie_docs = pie_dataset["train"]
+        self.add_pie_dataset(pie_docs, keys=all_doc_ids, metadatas=all_metadata)
+        logger.info(f"Loaded {len(pie_docs)} documents from {path} into docstore")

src/langchain_modules/huggingface_span_embeddings.py

@@ -0,0 +1,192 @@
+import logging
+from typing import Any, Dict, List, Optional, Union  # type: ignore[import-not-found]
+
+import torch
+from pydantic import BaseModel, ConfigDict, Field
+from transformers import pipeline
+
+from ..hf_pipeline import FeatureExtractionPipelineWithStriding
+from .span_embeddings import SpanEmbeddings
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL_NAME = "allenai/specter2_base"
+
+
+class HuggingFaceSpanEmbeddings(BaseModel, SpanEmbeddings):
+    """An implementation of SpanEmbeddings using a modified HuggingFace Transformers
+    feature-extraction pipeline, adapted for long text inputs by chunking with optional stride
+    (see src.hf_pipeline.FeatureExtractionPipelineWithStriding).
+
+    Note that calculating embeddings for multiple spans is efficient when all spans for a
+    text are passed in a single call to embed_document_spans, as the text embedding is computed
+    only once per unique text, and the span embeddings are simply pooled from these text embeddings.
+
+    It accepts any model that can be used with the HuggingFace feature-extraction pipeline, also
+    models with adapters such as SPECTER2 (see https://huggingface.co/allenai/specter2). In this case,
+    the model should be loaded beforehand and passed as parameter 'model' instead of the model identifier.
+    See https://huggingface.co/docs/transformers/main_classes/pipelines for further information.
+
+    To use, you should have the ``transformers`` python package installed.
+
+    Example:
+        .. code-block:: python
+            from transformers import AutoModel
+
+            model = "allenai/specter2_base"
+            pipeline_kwargs = {'device': 'cpu', 'stride': 64, 'batch_size': 8}
+            encode_kwargs = {'normalize_embeddings': False}
+            hf = HuggingFaceSpanEmbeddings(
+                model=model,
+                pipeline_kwargs=pipeline_kwargs,
+            )
+
+            text = "This is a test sentence."
+
+            # calculate embeddings for text[0:4]="This" and text[15:23]="sentence"
+            embeddings = hf.embed_document_spans(texts=[text, text], starts=[0, 11], ends=[4, 19])
+    """
+
+    client: Any = None  #: :meta private:
+    model: Optional[Any] = DEFAULT_MODEL_NAME
+    pooling_strategy: str = "mean"
+    """Model name to use."""
+    pipeline_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    """Keyword arguments to pass to the Huggingface pipeline constructor."""
+    encode_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    """Keyword arguments to pass when calling the pipeline."""
+    # show_progress: bool = False
+    """Whether to show a progress bar."""
+    model_max_length: Optional[int] = None
+    """The maximum input length of the model. Required for some model checkpoints with outdated configs."""
+
+    def __init__(self, **kwargs: Any):
+        """Initialize the sentence_transformer."""
+        super().__init__(**kwargs)
+
+        self.client = pipeline(
+            "feature-extraction",
+            model=self.model,
+            pipeline_class=FeatureExtractionPipelineWithStriding,
+            trust_remote_code=True,
+            **self.pipeline_kwargs,
+        )
+
+        # The Transformers library is buggy since 4.40.0,
+        # see https://github.com/huggingface/transformers/issues/30643,
+        # so we need to set the max_length to e.g. 512 manually
+        if self.model_max_length is not None:
+            self.client.tokenizer.model_max_length = self.model_max_length
+
+        # Check if the model has a valid max length
+        max_input_size = self.client.tokenizer.model_max_length
+        if max_input_size > 1e5:  # A high threshold to catch "unlimited" values
+            raise ValueError(
+                "The tokenizer does not specify a valid `model_max_length` attribute. "
+                "Consider setting it manually by passing `model_max_length` to the "
+                "HuggingFaceSpanEmbeddings constructor."
+            )
+
+    model_config = ConfigDict(
+        extra="forbid",
+        protected_namespaces=(),
+    )
+
+    def get_span_embedding(
+        self,
+        last_hidden_state: torch.Tensor,
+        offset_mapping: torch.Tensor,
+        start: Union[int, List[int]],
+        end: Union[int, List[int]],
+        **unused_kwargs,
+    ) -> Optional[List[float]]:
+        """Pool the span embeddings."""
+        if isinstance(start, int):
+            start = [start]
+        if isinstance(end, int):
+            end = [end]
+        if len(start) != len(end):
+            raise ValueError("start and end should have the same length.")
+        if len(start) == 0:
+            raise ValueError("start and end should not be empty.")
+        if last_hidden_state.shape[0] != 1:
+            raise ValueError("last_hidden_state should have a batch size of 1.")
+        if last_hidden_state.shape[0] != offset_mapping.shape[0]:
+            raise ValueError(
+                "last_hidden_state and offset_mapping should have the same batch size."
+            )
+        offset_mapping = offset_mapping[0]
+        last_hidden_state = last_hidden_state[0]
+
+        mask = (start[0] <= offset_mapping[:, 0]) & (offset_mapping[:, 1] <= end[0])
+        for s, e in zip(start[1:], end[1:]):
+            mask = mask | ((s <= offset_mapping[:, 0]) & (offset_mapping[:, 1] <= e))
+        span_embeddings = last_hidden_state[mask]
+        if span_embeddings.shape[0] == 0:
+            return None
+        if self.pooling_strategy == "mean":
+            return span_embeddings.mean(dim=0).tolist()
+        elif self.pooling_strategy == "max":
+            return span_embeddings.max(dim=0).values.tolist()
+        else:
+            raise ValueError(f"Unknown pool strategy: {self.pooling_strategy}")
+
+    def embed_document_spans(
+        self,
+        texts: List[str],
+        starts: Union[List[int], List[List[int]]],
+        ends: Union[List[int], List[List[int]]],
+    ) -> List[Optional[List[float]]]:
+        """Compute doc embeddings using a HuggingFace transformer model.
+
+        Args:
+            texts: The list of texts to embed.
+            starts: The list of start indices or list of lists of start indices (multi-span).
+            ends: The list of end indices or list of lists of end indices (multi-span).
+
+        Returns:
+            List of embeddings, one for each text.
+        """
+        pipeline_kwargs = self.encode_kwargs.copy()
+        pipeline_kwargs["return_offset_mapping"] = True
+        # we enable long text handling by providing the stride parameter
+        if pipeline_kwargs.get("stride", None) is None:
+            pipeline_kwargs["stride"] = 0
+        # when stride is positive, we need to create unique embeddings per token
+        if pipeline_kwargs["stride"] > 0:
+            pipeline_kwargs["create_unique_embeddings_per_token"] = True
+        # we ask for tensors to efficiently compute the span embeddings
+        pipeline_kwargs["return_tensors"] = True
+
+        unique_texts = sorted(set(texts))
+        idx2unique_idx = {i: unique_texts.index(text) for i, text in enumerate(texts)}
+        pipeline_results = self.client(unique_texts, **pipeline_kwargs)
+        embeddings = [
+            self.get_span_embedding(
+                start=starts[idx], end=ends[idx], **pipeline_results[idx2unique_idx[idx]]
+            )
+            for idx in range(len(texts))
+        ]
+        return embeddings
+
+    def embed_query_span(
+        self, text: str, start: Union[int, List[int]], end: Union[int, List[int]]
+    ) -> Optional[List[float]]:
+        """Compute query embeddings using a HuggingFace transformer model.
+
+        Args:
+            text: The text to embed.
+            start: The start index or list of start indices (multi-span).
+            end: The end index or list of end indices (multi-span).
+
+        Returns:
+            Embeddings for the text.
+        """
+        starts: Union[List[int], List[List[int]]] = [start]  # type: ignore[assignment]
+        ends: Union[List[int], List[List[int]]] = [end]  # type: ignore[assignment]
+        return self.embed_document_spans([text], starts=starts, ends=ends)[0]
+
+    @property
+    def embedding_dim(self) -> int:
+        """Get the embedding dimension."""
+        return self.client.model.config.hidden_size

src/langchain_modules/pie_document_store.py

@@ -0,0 +1,88 @@
+import abc
+import logging
+from copy import copy
+from typing import Iterator, List, Optional, Sequence, Tuple
+
+import pandas as pd
+from langchain_core.documents import Document as LCDocument
+from langchain_core.stores import BaseStore
+from pytorch_ie.documents import TextBasedDocument
+
+from .serializable_store import SerializableStore
+
+logger = logging.getLogger(__name__)
+
+
+class PieDocumentStore(SerializableStore, BaseStore[str, LCDocument], abc.ABC):
+    """Abstract base class for document stores specialized in storing and retrieving pie documents."""
+
+    METADATA_KEY_PIE_DOCUMENT: str = "pie_document"
+    """Key for the pie document in the (langchain) document metadata."""
+
+    def wrap(self, pie_document: TextBasedDocument, **metadata) -> LCDocument:
+        """Wrap the pie document in an LCDocument."""
+        return LCDocument(
+            id=pie_document.id,
+            page_content="",
+            metadata={self.METADATA_KEY_PIE_DOCUMENT: pie_document, **metadata},
+        )
+
+    def unwrap(self, document: LCDocument) -> TextBasedDocument:
+        """Get the pie document from the langchain document."""
+        return document.metadata[self.METADATA_KEY_PIE_DOCUMENT]
+
+    def unwrap_with_metadata(self, document: LCDocument) -> Tuple[TextBasedDocument, dict]:
+        """Get the pie document and metadata from the langchain document."""
+        metadata = copy(document.metadata)
+        pie_document = metadata.pop(self.METADATA_KEY_PIE_DOCUMENT)
+        return pie_document, metadata
+
+    @abc.abstractmethod
+    def mget(self, keys: Sequence[str]) -> List[LCDocument]:
+        pass
+
+    @abc.abstractmethod
+    def mset(self, items: Sequence[Tuple[str, LCDocument]]) -> None:
+        pass
+
+    @abc.abstractmethod
+    def mdelete(self, keys: Sequence[str]) -> None:
+        pass
+
+    @abc.abstractmethod
+    def yield_keys(self, *, prefix: Optional[str] = None) -> Iterator[str]:
+        pass
+
+    def __len__(self):
+        return len(list(self.yield_keys()))
+
+    def overview(self, layer_captions: dict, use_predictions: bool = False) -> pd.DataFrame:
+        """Get an overview of the document store, including the number of items in each layer for each document
+        in the store.
+
+        Args:
+            layer_captions: A dictionary mapping layer names to captions.
+            use_predictions: Whether to use predictions instead of the actual layers.
+
+        Returns:
+            DataFrame: A pandas DataFrame containing the overview.
+        """
+        rows = []
+        for doc_id in self.yield_keys():
+            document = self.mget([doc_id])[0]
+            pie_document = self.unwrap(document)
+            layers = {
+                caption: pie_document[layer_name] for layer_name, caption in layer_captions.items()
+            }
+            layer_sizes = {
+                f"num_{caption}s": len(layer) + (len(layer.predictions) if use_predictions else 0)
+                for caption, layer in layers.items()
+            }
+            rows.append({"doc_id": doc_id, **layer_sizes})
+        df = pd.DataFrame(rows)
+        return df
+
+    def as_dict(self, document: LCDocument) -> dict:
+        """Convert the langchain document to a dictionary."""
+        pie_document, metadata = self.unwrap_with_metadata(document)
+        return {self.METADATA_KEY_PIE_DOCUMENT: pie_document.asdict(), "metadata": metadata}

src/langchain_modules/qdrant_span_vectorstore.py

@@ -0,0 +1,349 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+import uuid
+from collections import defaultdict
+from itertools import islice
+from typing import (  # type: ignore[import-not-found]
+    Any,
+    Dict,
+    Generator,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+)
+
+import numpy as np
+from langchain_core.documents import Document as LCDocument
+from langchain_qdrant import QdrantVectorStore, RetrievalMode
+from pytorch_ie.annotations import LabeledMultiSpan, LabeledSpan, MultiSpan, Span
+from qdrant_client import QdrantClient, models
+from qdrant_client.http.models import Record
+
+from .span_embeddings import SpanEmbeddings
+from .span_vectorstore import SpanVectorStore
+
+logger = logging.getLogger(__name__)
+
+
+class QdrantSpanVectorStore(SpanVectorStore, QdrantVectorStore):
+    """An implementation of the SpanVectorStore interface that uses Qdrant
+    as backend for storing and retrieving span embeddings."""
+
+    EMBEDDINGS_FILE = "embeddings.npy"
+    PAYLOADS_FILE = "payloads.json"
+    INDEX_FILE = "index.json"
+
+    def __init__(
+        self,
+        client: QdrantClient,
+        collection_name: str,
+        embedding: SpanEmbeddings,
+        vector_params: Optional[Dict[str, Any]] = None,
+        **kwargs,
+    ):
+        if not client.collection_exists(collection_name):
+            logger.info(f'Collection "{collection_name}" does not exist. Creating it now.')
+            client.create_collection(
+                collection_name=collection_name,
+                vectors_config=models.VectorParams(size=embedding.embedding_dim, **vector_params),
+            )
+        else:
+            logger.info(f'Collection "{collection_name}" already exists.')
+        super().__init__(
+            client=client, collection_name=collection_name, embedding=embedding, **kwargs
+        )
+
+    def __len__(self):
+        return self.client.get_collection(collection_name=self.collection_name).points_count
+
+    def get_by_ids_with_vectors(self, ids: Sequence[str | int], /) -> List[LCDocument]:
+        results = self.client.retrieve(
+            self.collection_name, ids, with_payload=True, with_vectors=True
+        )
+
+        return [
+            self._document_from_point(
+                result,
+                self.collection_name,
+                self.content_payload_key,
+                self.metadata_payload_key,
+            )
+            for result in results
+        ]
+
+    def construct_filter(
+        self,
+        query_span: Union[Span, MultiSpan],
+        metadata_doc_id_key: str,
+        doc_id_whitelist: Optional[Sequence[str]] = None,
+        doc_id_blacklist: Optional[Sequence[str]] = None,
+    ) -> Optional[models.Filter]:
+        """Construct a filter for the retrieval. It should enforce that:
+        - if the span is labeled, the retrieved span has the same label, or
+        - if, in addition, a label mapping is provided, the retrieved span has a label that is in the mapping for the query span's label
+        - if `doc_id_whitelist` is provided, the retrieved span is from a document in the whitelist
+        - if `doc_id_blacklist` is provided, the retrieved span is not from a document in the blacklist
+
+        Args:
+            query_span: The query span.
+            metadata_doc_id_key: The key in the metadata that holds the document id.
+            doc_id_whitelist: A list of document ids to restrict the retrieval to.
+            doc_id_blacklist: A list of document ids to exclude from the retrieval.
+
+        Returns:
+            A filter object.
+        """
+        filter_kwargs = defaultdict(list)
+        # if the span is labeled, enforce that the retrieved span has the same label
+        if isinstance(query_span, (LabeledSpan, LabeledMultiSpan)):
+            if self.label_mapping is not None:
+                target_labels = self.label_mapping.get(query_span.label, [])
+            else:
+                target_labels = [query_span.label]
+            filter_kwargs["must"].append(
+                models.FieldCondition(
+                    key=f"metadata.{self.METADATA_SPAN_KEY}.label",
+                    match=models.MatchAny(any=target_labels),
+                )
+            )
+        elif self.label_mapping is not None:
+            raise TypeError("Label mapping is only supported for labeled spans")
+
+        if doc_id_blacklist is not None and doc_id_whitelist is not None:
+            overlap = set(doc_id_whitelist) & set(doc_id_blacklist)
+            if len(overlap) > 0:
+                raise ValueError(
+                    f"Overlap between doc_id_whitelist and doc_id_blacklist: {overlap}"
+                )
+
+        if doc_id_whitelist is not None:
+            filter_kwargs["must"].append(
+                models.FieldCondition(
+                    key=f"metadata.{metadata_doc_id_key}",
+                    match=(
+                        models.MatchValue(value=doc_id_whitelist[0])
+                        if len(doc_id_whitelist) == 1
+                        else models.MatchAny(any=doc_id_whitelist)
+                    ),
+                )
+            )
+        if doc_id_blacklist is not None:
+            filter_kwargs["must_not"].append(
+                models.FieldCondition(
+                    key=f"metadata.{metadata_doc_id_key}",
+                    match=(
+                        models.MatchValue(value=doc_id_blacklist[0])
+                        if len(doc_id_blacklist) == 1
+                        else models.MatchAny(any=doc_id_blacklist)
+                    ),
+                )
+            )
+        if len(filter_kwargs) > 0:
+            return models.Filter(**filter_kwargs)
+        return None
+
+    @classmethod
+    def _document_from_point(
+        cls,
+        scored_point: Any,
+        collection_name: str,
+        content_payload_key: str,
+        metadata_payload_key: str,
+    ) -> LCDocument:
+        metadata = scored_point.payload.get(metadata_payload_key) or {}
+        metadata["_collection_name"] = collection_name
+        if hasattr(scored_point, "score"):
+            metadata[cls.RELEVANCE_SCORE_KEY] = scored_point.score
+        if hasattr(scored_point, "vector"):
+            metadata[cls.METADATA_VECTOR_KEY] = scored_point.vector
+        return LCDocument(
+            id=scored_point.id,
+            page_content=scored_point.payload.get(content_payload_key, ""),
+            metadata=metadata,
+        )
+
+    def _build_vectors_with_metadata(
+        self,
+        texts: Iterable[str],
+        metadatas: List[dict],
+    ) -> List[models.VectorStruct]:
+        starts = [metadata[self.METADATA_SPAN_KEY][self.SPAN_START_KEY] for metadata in metadatas]
+        ends = [metadata[self.METADATA_SPAN_KEY][self.SPAN_END_KEY] for metadata in metadatas]
+        if self.retrieval_mode == RetrievalMode.DENSE:
+            batch_embeddings = self.embeddings.embed_document_spans(list(texts), starts, ends)
+            return [
+                {
+                    self.vector_name: vector,
+                }
+                for vector in batch_embeddings
+            ]
+
+        elif self.retrieval_mode == RetrievalMode.SPARSE:
+            raise ValueError("Sparse retrieval mode is not yet implemented.")
+
+        elif self.retrieval_mode == RetrievalMode.HYBRID:
+            raise NotImplementedError("Hybrid retrieval mode is not yet implemented.")
+        else:
+            raise ValueError(f"Unknown retrieval mode. {self.retrieval_mode} to build vectors.")
+
+    def _build_payloads_from_metadata(
+        self,
+        metadatas: Iterable[dict],
+        metadata_payload_key: str,
+    ) -> List[dict]:
+        payloads = [{metadata_payload_key: metadata} for metadata in metadatas]
+
+        return payloads
+
+    def _generate_batches(
+        self,
+        texts: Iterable[str],
+        metadatas: Optional[List[dict]] = None,
+        ids: Optional[Sequence[str | int]] = None,
+        batch_size: int = 64,
+    ) -> Generator[tuple[list[str | int], list[models.PointStruct]], Any, None]:
+        """Generate batches of points to index. Same as in `QdrantVectorStore` but metadata is used
+        to build vectors and payloads."""
+
+        texts_iterator = iter(texts)
+        if metadatas is None:
+            raise ValueError("Metadata must be provided to generate batches.")
+        metadatas_iterator = iter(metadatas)
+        ids_iterator = iter(ids or [uuid.uuid4().hex for _ in iter(texts)])
+
+        while batch_texts := list(islice(texts_iterator, batch_size)):
+            batch_metadatas = list(islice(metadatas_iterator, batch_size))
+            batch_ids = list(islice(ids_iterator, batch_size))
+            points = [
+                models.PointStruct(
+                    id=point_id,
+                    vector=vector,
+                    payload=payload,
+                )
+                for point_id, vector, payload in zip(
+                    batch_ids,
+                    self._build_vectors_with_metadata(batch_texts, metadatas=batch_metadatas),
+                    # we do not save the text in the payload because the text is the full
+                    # document which is usually already saved in the docstore
+                    self._build_payloads_from_metadata(
+                        metadatas=batch_metadatas,
+                        metadata_payload_key=self.metadata_payload_key,
+                    ),
+                )
+                if vector[self.vector_name] is not None
+            ]
+
+            yield [point.id for point in points], points
+
+    def similarity_search_with_score_by_vector(
+        self,
+        embedding: List[float],
+        k: int = 4,
+        filter: Optional[models.Filter] = None,
+        search_params: Optional[models.SearchParams] = None,
+        offset: int = 0,
+        score_threshold: Optional[float] = None,
+        consistency: Optional[models.ReadConsistency] = None,
+        **kwargs: Any,
+    ) -> List[Tuple[LCDocument, float]]:
+        """Return docs most similar to query vector.
+
+        Returns:
+            List of documents most similar to the query text and distance for each.
+        """
+        query_options = {
+            "collection_name": self.collection_name,
+            "query_filter": filter,
+            "search_params": search_params,
+            "limit": k,
+            "offset": offset,
+            "with_payload": True,
+            "with_vectors": False,
+            "score_threshold": score_threshold,
+            "consistency": consistency,
+            **kwargs,
+        }
+
+        results = self.client.query_points(
+            query=embedding,
+            using=self.vector_name,
+            **query_options,
+        ).points
+
+        return [
+            (
+                self._document_from_point(
+                    result,
+                    self.collection_name,
+                    self.content_payload_key,
+                    self.metadata_payload_key,
+                ),
+                result.score,
+            )
+            for result in results
+        ]
+
+    def _as_indices_vectors_payloads(self) -> Tuple[List[str], np.ndarray, List[Any]]:
+        data, _ = self.client.scroll(
+            collection_name=self.collection_name, with_vectors=True, limit=len(self)
+        )
+        vectors_np = np.array([point.vector for point in data])
+        payloads = [point.payload for point in data]
+        emb_ids = [point.id for point in data]
+        return emb_ids, vectors_np, payloads
+
+    # TODO: or use create_snapshot and restore_snapshot?
+    def _save_to_directory(self, path: str, **kwargs) -> None:
+        indices, vectors, payloads = self._as_indices_vectors_payloads()
+        np.save(os.path.join(path, self.EMBEDDINGS_FILE), vectors)
+        with open(os.path.join(path, self.PAYLOADS_FILE), "w") as f:
+            json.dump(payloads, f, indent=2)
+        with open(os.path.join(path, self.INDEX_FILE), "w") as f:
+            json.dump(indices, f)
+
+    def _load_from_directory(self, path: str, **kwargs) -> None:
+        with open(os.path.join(path, self.INDEX_FILE), "r") as f:
+            index = json.load(f)
+        embeddings_np: np.ndarray = np.load(os.path.join(path, self.EMBEDDINGS_FILE))
+        with open(os.path.join(path, self.PAYLOADS_FILE), "r") as f:
+            payloads = json.load(f)
+        points = models.Batch(ids=index, vectors=embeddings_np.tolist(), payloads=payloads)
+        self.client.upsert(
+            collection_name=self.collection_name,
+            points=points,
+        )
+        logger.info(f"Loaded {len(index)} points into collection {self.collection_name}.")
+
+    def mget(self, keys: Sequence[str]) -> list[Optional[Record]]:
+        return self.client.retrieve(
+            self.collection_name, ids=keys, with_payload=True, with_vectors=True
+        )
+
+    def mset(self, key_value_pairs: Sequence[tuple[str, Record]]) -> None:
+        self.client.upsert(
+            collection_name=self.collection_name,
+            points=models.Batch(
+                ids=[key for key, _ in key_value_pairs],
+                vectors=[value.vector for _, value in key_value_pairs],
+                payloads=[value.payload for _, value in key_value_pairs],
+            ),
+        )
+
+    def mdelete(self, keys: Sequence[str]) -> None:
+        self.client.delete(collection_name=self.collection_name, points_selector=keys)
+
+    def yield_keys(self, *, prefix: Optional[str] = None) -> Iterator[str]:
+        for point in self.client.scroll(
+            collection_name=self.collection_name,
+            with_vectors=False,
+            with_payload=False,
+            limit=len(self),
+        )[0]:
+            yield point.id

src/langchain_modules/serializable_store.py

@@ -0,0 +1,137 @@
+import logging
+import os
+import shutil
+import tempfile
+from abc import ABC, abstractmethod
+
+logger = logging.getLogger(__name__)
+
+
+class SerializableStore(ABC):
+    """Abstract base class for serializable stores."""
+
+    @abstractmethod
+    def _save_to_directory(self, path: str, **kwargs) -> None:
+        """Save the store to a directory.
+
+        Args:
+            path (str): The path to a directory.
+        """
+
+    @abstractmethod
+    def _load_from_directory(self, path: str, **kwargs) -> None:
+        """Load the store from a directory.
+
+        Args:
+            path (str): The path to the file.
+        """
+
+    def save_to_directory(self, path: str, **kwargs) -> None:
+        """Save the store to a directory.
+
+        Args:
+            path (str): The path to a directory.
+        """
+        os.makedirs(path, exist_ok=True)
+        self._save_to_directory(path, **kwargs)
+
+    def load_from_directory(self, path: str, **kwargs) -> None:
+        """Load the store from a directory.
+
+        Args:
+            path (str): The path to the file.
+        """
+        if not os.path.exists(path):
+            raise FileNotFoundError(f'Directory "{path}" not found.')
+
+        self._load_from_directory(path, **kwargs)
+
+    def save_to_archive(
+        self,
+        base_name: str,
+        format: str,
+    ) -> str:
+        """Save the store to an archive.
+
+        Args:
+            base_name (str): The base name of the archive.
+            format (str): The format of the archive.
+
+        Returns:
+            str: The path to the archive.
+        """
+        temp_dir = os.path.join(tempfile.gettempdir(), "retriever_store")
+        # remove the temporary directory if it already exists
+        if os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+        # save the documents to the directory
+        self.save_to_directory(temp_dir)
+        # zip the directory
+        result_file_path = shutil.make_archive(
+            base_name=base_name, root_dir=temp_dir, format=format
+        )
+        # remove the temporary directory
+        shutil.rmtree(temp_dir)
+        return result_file_path
+
+    def load_from_archive(
+        self,
+        file_name: str,
+    ) -> None:
+        """Load the store from an archive.
+
+        Args:
+            file_name (str): The path to the archive.
+        """
+        if not os.path.exists(file_name):
+            raise FileNotFoundError(f'Archive file "{file_name}" not found.')
+
+        temp_dir = os.path.join(tempfile.gettempdir(), "retriever_store")
+        # remove the temporary directory if it already exists
+        if os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+
+        # unzip the file
+        shutil.unpack_archive(file_name, temp_dir)
+        # load the documents from the directory
+        self.load_from_directory(temp_dir)
+        # remove the temporary directory
+        shutil.rmtree(temp_dir)
+
+    def save_to_disc(self, path: str) -> None:
+        """Save the store to disc. Depending on the path, this can be a directory or an archive.
+
+        Args:
+            path (str): The path to a directory or an archive.
+        """
+        # if it is a zip file, save to archive
+        if path.lower().endswith(".zip"):
+            logger.info(f"Saving to archive at {path} ...")
+            # get base name without extension
+            base_name = os.path.splitext(path)[0]
+            result_path = self.save_to_archive(base_name, format="zip")
+            if not result_path.endswith(path):
+                logger.warning(f"Saved to {result_path} instead of {path}.")
+        # if it does not have an extension, save to directory
+        elif not os.path.splitext(path)[1]:
+            logger.info(f"Saving to directory at {path} ...")
+            self.save_to_directory(path)
+        else:
+            raise ValueError("Unsupported file format. Only .zip and directories are supported.")
+
+    def load_from_disc(self, path: str) -> None:
+        """Load the store from disc. Depending on the path, this can be a directory or an archive.
+
+        Args:
+            path (str): The path to a directory or an archive.
+        """
+        # if it is a zip file, load from archive
+        if path.lower().endswith(".zip"):
+            logger.info(f"Loading from archive at {path} ...")
+            self.load_from_archive(path)
+        # if it is a directory, load from directory
+        elif os.path.isdir(path):
+            logger.info(f"Loading from directory at {path} ...")
+            self.load_from_directory(path)
+        else:
+            raise ValueError("Unsupported file format. Only .zip and directories are supported.")

src/langchain_modules/span_embeddings.py

@@ -0,0 +1,103 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Union  # type: ignore[import-not-found]
+
+from langchain_core.embeddings import Embeddings
+from langchain_core.runnables.config import run_in_executor
+
+
+class SpanEmbeddings(Embeddings, ABC):
+    """Interface for models that embed text spans within documents."""
+
+    @abstractmethod
+    def embed_document_spans(
+        self,
+        texts: list[str],
+        starts: Union[list[int], List[List[int]]],
+        ends: Union[list[int], List[List[int]]],
+    ) -> list[Optional[list[float]]]:
+        """Embed search docs.
+
+        Args:
+            texts: List of text to embed.
+            starts: List of start indices or list of lists of start indices (multi-span).
+            ends: List of end indices or list of lists of end indices (multi-span).
+
+        Returns:
+            List of embeddings.
+        """
+
+    @abstractmethod
+    def embed_query_span(
+        self, text: str, start: Union[int, list[int]], end: Union[int, list[int]]
+    ) -> Optional[list[float]]:
+        """Embed query text.
+
+        Args:
+            text: Text to embed.
+            start: Start index or list of start indices (multi-span).
+            end: End index or list of end indices (multi-span).
+
+        Returns:
+            Embedding.
+        """
+
+    def embed_documents(self, texts: list[str]) -> list[Optional[list[float]]]:
+        """Embed search docs.
+
+        Args:
+            texts: List of text to embed.
+
+        Returns:
+            List of embeddings.
+        """
+        return self.embed_document_spans(texts, [0] * len(texts), [len(text) for text in texts])
+
+    def embed_query(self, text: str) -> Optional[list[float]]:
+        """Embed query text.
+
+        Args:
+            text: Text to embed.
+
+        Returns:
+            Embedding.
+        """
+        return self.embed_query_span(text, 0, len(text))
+
+    async def aembed_document_spans(
+        self,
+        texts: list[str],
+        starts: Union[list[int], list[list[int]]],
+        ends: Union[list[int], list[list[int]]],
+    ) -> list[Optional[list[float]]]:
+        """Asynchronous Embed search docs.
+
+        Args:
+            texts: List of text to embed.
+            starts: List of start indices or list of lists of start indices (multi-span).
+            ends: List of end indices or list of lists of end indices (multi-span).
+
+        Returns:
+            List of embeddings.
+        """
+        return await run_in_executor(None, self.embed_document_spans, texts, starts, ends)
+
+    async def aembed_query_spans(
+        self, text: str, start: Union[int, list[int]], end: Union[int, list[int]]
+    ) -> Optional[list[float]]:
+        """Asynchronous Embed query text.
+
+        Args:
+            text: Text to embed.
+            start: Start index or list of start indices (multi-span).
+            end: End index or list of end indices (multi-span).
+
+        Returns:
+            Embedding.
+        """
+        return await run_in_executor(None, self.embed_query_span, text, start, end)
+
+    @property
+    @abstractmethod
+    def embedding_dim(self) -> int:
+        """Get the embedding dimension."""
+        ...

src/langchain_modules/span_retriever.py

@@ -0,0 +1,860 @@
+import logging
+import os
+import uuid
+from collections import defaultdict
+from copy import copy
+from enum import Enum
+from typing import Any, Dict, List, Optional, Sequence, Set, Tuple, Type, Union
+
+from langchain_core.callbacks import (
+    AsyncCallbackManagerForRetrieverRun,
+    CallbackManagerForRetrieverRun,
+)
+from langchain_core.documents import BaseDocumentCompressor
+from langchain_core.documents import Document as LCDocument
+from langchain_core.retrievers import BaseRetriever
+from pydantic import Field
+from pytorch_ie.annotations import LabeledMultiSpan, LabeledSpan, MultiSpan, Span
+from pytorch_ie.core.document import BaseAnnotationList
+from pytorch_ie.documents import (
+    TextBasedDocument,
+    TextDocumentWithLabeledMultiSpans,
+    TextDocumentWithLabeledSpans,
+    TextDocumentWithSpans,
+)
+
+from .pie_document_store import PieDocumentStore
+from .serializable_store import SerializableStore
+from .span_vectorstore import SpanVectorStore
+
+logger = logging.getLogger(__name__)
+
+
+def _parse_config(config_string: str, format: str) -> Dict[str, Any]:
+    """Parse a configuration string."""
+    if format == "json":
+        import json
+
+        return json.loads(config_string)
+    elif format == "yaml":
+        import yaml
+
+        return yaml.safe_load(config_string)
+    else:
+        raise ValueError(f"Unsupported format: {format}. Use 'json' or 'yaml'.")
+
+
+METADATA_KEY_CHILD_ID2IDX = "child_id2idx"
+
+
+class SpanNotFoundError(ValueError):
+    def __init__(self, span_id: str, doc_id: Optional[str] = None, msg: Optional[str] = None):
+        if msg is None:
+            if doc_id is not None:
+                msg = f"Span with id [{span_id}] not found in document [{doc_id}]"
+            else:
+                msg = f"Span with id [{span_id}] not found in the vectorstore"
+        super().__init__(msg)
+        self.span_id = span_id
+        self.doc_id = doc_id
+
+
+class DocumentNotFoundError(ValueError):
+    def __init__(self, doc_id: str, msg: Optional[str] = None):
+        msg = msg or f"Document with id [{doc_id}] not found in the docstore"
+        super().__init__(msg)
+        self.doc_id = doc_id
+
+
+class SearchType(str, Enum):
+    """Enumerator of the types of search to perform."""
+
+    similarity = "similarity"
+    """Similarity search."""
+    similarity_score_threshold = "similarity_score_threshold"
+    """Similarity search with a score threshold."""
+    mmr = "mmr"
+    """Maximal Marginal Relevance reranking of similarity search."""
+
+
+class DocumentAwareSpanRetriever(BaseRetriever, SerializableStore):
+    """Retriever for contextualized text spans, i.e. spans within text documents.
+    It accepts spans as queries and retrieves spans with their containing document.
+    Note that the query span (and its document) must already be in the retriever's
+    store."""
+
+    pie_document_type: Type[TextBasedDocument]
+    """The name of the span annotation layer in the pie document."""
+    use_predicted_annotations_key: str = "use_predicted_annotations"
+    """Whether to use the predicted annotations or the gold annotations."""
+    retrieve_from_same_document: bool = False
+    """Whether to retrieve spans exclusively from the same document as the query span."""
+    retrieve_from_different_documents: bool = False
+    """Whether to retrieve spans exclusively from different documents than the query span."""
+
+    # content from langchain_core.retrievers.MultiVectorRetriever
+    vectorstore: SpanVectorStore
+    """The underlying vectorstore to use to store small chunks
+    and their embedding vectors"""
+    docstore: PieDocumentStore
+    """The storage interface for the parent documents"""
+    id_key: str = "doc_id"
+    """The key to use to track the parent id. This will be stored in the
+    metadata of child documents."""
+    search_kwargs: dict = Field(default_factory=dict)
+    """Keyword arguments to pass to the search function."""
+    search_type: SearchType = SearchType.similarity
+    """Type of search to perform (similarity / mmr)"""
+
+    # content taken from langchain_core.retrievers.ParentDocumentRetriever
+    child_metadata_fields: Optional[Sequence[str]] = None
+    """Metadata fields to leave in child documents. If None, leave all parent document
+        metadata.
+    """
+
+    # re-ranking
+    compressor: Optional[BaseDocumentCompressor] = None
+    """Compressor for compressing retrieved documents."""
+    compressor_context_size: int = 50
+    """Size of the context to use around the query and retrieved spans when compressing."""
+    compressor_query_context_size: Optional[int] = 10
+    """Size of the context to use around the query when compressing. If None, will use the
+    same value as `compressor_context_size`."""
+
+    @classmethod
+    def instantiate_from_config(
+        cls, config: Dict[str, Any], overwrites: Optional[Dict[str, Any]] = None
+    ) -> "DocumentAwareSpanRetriever":
+        """Instantiate a retriever from a configuration dictionary."""
+        from hydra.utils import instantiate
+
+        return instantiate(config, **(overwrites or {}))
+
+    @classmethod
+    def instantiate_from_config_string(
+        cls, config_string: str, format: str, overwrites: Optional[Dict[str, Any]] = None
+    ) -> "DocumentAwareSpanRetriever":
+        """Instantiate a retriever from a configuration string."""
+        return cls.instantiate_from_config(
+            _parse_config(config_string, format=format), overwrites=overwrites
+        )
+
+    @classmethod
+    def instantiate_from_config_file(
+        cls, config_path: str, overwrites: Optional[Dict[str, Any]] = None
+    ) -> "DocumentAwareSpanRetriever":
+        """Instantiate a retriever from a configuration file."""
+        with open(config_path, "r") as file:
+            config_string = file.read()
+        if config_path.endswith(".json"):
+            return cls.instantiate_from_config_string(
+                config_string, format="json", overwrites=overwrites
+            )
+        elif config_path.endswith(".yaml"):
+            return cls.instantiate_from_config_string(
+                config_string, format="yaml", overwrites=overwrites
+            )
+        else:
+            raise ValueError(f"Unsupported file extension: {config_path}")
+
+    @property
+    def pie_annotation_layer_name(self) -> str:
+        if issubclass(self.pie_document_type, TextDocumentWithSpans):
+            return "spans"
+        elif issubclass(self.pie_document_type, TextDocumentWithLabeledSpans):
+            return "labeled_spans"
+        elif issubclass(self.pie_document_type, TextDocumentWithLabeledMultiSpans):
+            return "labeled_multi_spans"
+        else:
+            raise ValueError(
+                f"Unsupported pie document type: {self.pie_document_type}. "
+                "Must be one of TextDocumentWithSpans, TextDocumentWithLabeledSpans, "
+                "or TextDocumentWithLabeledMultiSpans."
+            )
+
+    def _span_to_dict(self, span: Union[Span, MultiSpan]) -> dict:
+        span_dict = {}
+        if isinstance(span, Span):
+            span_dict[self.vectorstore.SPAN_START_KEY] = span.start
+            span_dict[self.vectorstore.SPAN_END_KEY] = span.end
+            span_dict["type"] = "Span"
+        elif isinstance(span, MultiSpan):
+            starts, ends = zip(*span.slices)
+            span_dict[self.vectorstore.SPAN_START_KEY] = starts
+            span_dict[self.vectorstore.SPAN_END_KEY] = ends
+            span_dict["type"] = "MultiSpan"
+        else:
+            raise ValueError(f"Unsupported span type: {type(span)}")
+        if isinstance(span, (LabeledSpan, LabeledMultiSpan)):
+            span_dict["label"] = span.label
+            span_dict["score"] = span.score
+        return span_dict
+
+    def _dict_to_span(self, span_dict: dict) -> Union[Span, MultiSpan]:
+
+        if span_dict["type"] == "Span":
+            kwargs = dict(
+                start=span_dict[self.vectorstore.SPAN_START_KEY],
+                end=span_dict[self.vectorstore.SPAN_END_KEY],
+            )
+            if "label" in span_dict:
+                kwargs["label"] = span_dict["label"]
+                kwargs["score"] = span_dict["score"]
+                return LabeledSpan(**kwargs)
+            else:
+                return Span(**kwargs)
+        elif span_dict["type"] == "MultiSpan":
+            starts = span_dict[self.vectorstore.SPAN_START_KEY]
+            ends = span_dict[self.vectorstore.SPAN_END_KEY]
+            slices = tuple((start, end) for start, end in zip(starts, ends))
+            kwargs = dict(slices=slices)
+            if "label" in span_dict:
+                kwargs["label"] = span_dict["label"]
+                kwargs["score"] = span_dict["score"]
+                return LabeledMultiSpan(**kwargs)
+            else:
+                return MultiSpan(**kwargs)
+        else:
+            raise ValueError(f"Unsupported span type: {span_dict['type']}")
+
+    def use_predicted_annotations(self, doc: LCDocument) -> bool:
+        """Check if the document uses predicted spans."""
+        return doc.metadata.get(self.use_predicted_annotations_key, True)
+
+    def get_document(self, doc_id: str) -> LCDocument:
+        """Get a document by its id."""
+        documents = self.docstore.mget([doc_id])
+        if len(documents) == 0 or documents[0] is None:
+            raise DocumentNotFoundError(doc_id=doc_id)
+        if len(documents) > 1:
+            raise ValueError(f"Multiple documents found with id: {doc_id}")
+        return documents[0]
+
+    def get_span_document(self, span_id: str, with_vector: bool = False) -> LCDocument:
+        """Get a span document by its id."""
+        if with_vector:
+            span_docs = self.vectorstore.get_by_ids_with_vectors([span_id])
+        else:
+            span_docs = self.vectorstore.get_by_ids([span_id])
+        if len(span_docs) == 0 or span_docs[0] is None:
+            raise SpanNotFoundError(span_id=span_id)
+        if len(span_docs) > 1:
+            raise ValueError(f"Multiple span documents found with id: {span_id}")
+        return span_docs[0]
+
+    def get_base_layer(
+        self, pie_document: TextBasedDocument, use_predicted_annotations: bool
+    ) -> BaseAnnotationList:
+        """Get the base layer of the pie document."""
+
+        if self.pie_annotation_layer_name not in pie_document:
+            raise ValueError(
+                f'The pie document must contain the annotation layer "{self.pie_annotation_layer_name}"'
+            )
+        layer = pie_document[self.pie_annotation_layer_name]
+        return layer.predictions if use_predicted_annotations else layer
+
+    def get_span_by_id(self, span_id: str) -> Union[Span, MultiSpan]:
+        """Get a span annotation by its id."""
+        span_doc = self.get_span_document(span_id)
+        doc_id = span_doc.metadata[self.id_key]
+        doc = self.get_document(doc_id)
+        return self.get_span_from_doc_by_id(doc=doc, span_id=span_id)
+
+    def get_span_from_doc_by_id(self, doc: LCDocument, span_id: str) -> Union[Span, MultiSpan]:
+        """Get the span of a query."""
+        base_layer = self.get_base_layer(
+            self.docstore.unwrap(doc),
+            use_predicted_annotations=self.use_predicted_annotations(doc),
+        )
+        span_idx = doc.metadata[METADATA_KEY_CHILD_ID2IDX].get(span_id)
+        if span_idx is None:
+            raise SpanNotFoundError(span_id=span_id, doc_id=doc.id)
+        return base_layer[span_idx]
+
+    def get_span_id2idx_from_doc(self, doc: Union[LCDocument, str]) -> Dict[str, int]:
+        """Get all span ids from a document.
+
+        Args:
+            doc: Document or document id
+
+        Returns:
+            Dictionary mapping span ids to their index in the base layer.
+        """
+
+        if isinstance(doc, str):
+            doc = self.get_document(doc)
+        return doc.metadata[METADATA_KEY_CHILD_ID2IDX]
+
+    def prepare_search_kwargs(
+        self,
+        span_id: str,
+        doc_id_whitelist: Optional[List[str]] = None,
+        doc_id_blacklist: Optional[List[str]] = None,
+        kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[dict, LCDocument]:
+        # get the span document
+        query_span_doc = self.get_span_document(span_id, with_vector=True)
+        query_doc_id = query_span_doc.metadata[self.id_key]
+        query_doc = self.get_document(query_doc_id)
+
+        # TODO: why do we do this? Just to be the same as the result of the search when doing compression?
+        # add "pie_document" to the metadata
+        query_span_doc.metadata[self.docstore.METADATA_KEY_PIE_DOCUMENT] = self.docstore.unwrap(
+            query_doc
+        )
+
+        search_kwargs = copy(self.search_kwargs)
+        search_kwargs.update(kwargs or {})
+
+        query_span = self.get_span_from_doc_by_id(doc=query_doc, span_id=span_id)
+
+        if self.retrieve_from_different_documents and self.retrieve_from_same_document:
+            raise ValueError("Cannot retrieve from both same and different documents")
+
+        if self.retrieve_from_same_document:
+            if doc_id_whitelist is None:
+                doc_id_whitelist = [query_doc_id]
+            elif query_doc_id not in doc_id_whitelist:
+                doc_id_whitelist.append(query_doc_id)
+
+        if self.retrieve_from_different_documents:
+            if doc_id_blacklist is None:
+                doc_id_blacklist = [query_doc_id]
+            elif query_doc_id not in doc_id_blacklist:
+                doc_id_blacklist.append(query_doc_id)
+
+        query_filter = self.vectorstore.construct_filter(
+            query_span=query_span,
+            metadata_doc_id_key=self.id_key,
+            doc_id_whitelist=doc_id_whitelist,
+            doc_id_blacklist=doc_id_blacklist,
+        )
+        if query_filter is not None:
+            search_kwargs["filter"] = query_filter
+
+        # get the vector of the reference span
+        search_kwargs["embedding"] = query_span_doc.metadata[self.vectorstore.METADATA_VECTOR_KEY]
+        return search_kwargs, query_span_doc
+
+    def _prepare_query_for_compression(self, query_doc: LCDocument) -> str:
+        return self._prepare_doc_for_compression(
+            query_doc, context_size=self.compressor_query_context_size
+        ).page_content
+
+    def _prepare_doc_for_compression(
+        self, doc: LCDocument, context_size: Optional[int] = None
+    ) -> LCDocument:
+        if context_size is None:
+            context_size = self.compressor_context_size
+        pie_doc: TextBasedDocument = self.docstore.unwrap(doc)
+        text = pie_doc.text
+        span_dict = doc.metadata[self.vectorstore.METADATA_SPAN_KEY]
+        span_start = span_dict[self.vectorstore.SPAN_START_KEY]
+        span_end = span_dict[self.vectorstore.SPAN_END_KEY]
+        if isinstance(span_start, list):
+            span_start = span_start[0]
+        if isinstance(span_end, list):
+            span_end = span_end[0]
+        context_start = span_start - context_size
+        context_end = span_end + context_size
+        doc.page_content = text[max(0, context_start) : min(context_end, len(text))]
+        # save the original relevance score and remove it because otherwise we will not be able to get
+        # the reranking relevance score
+        if "relevance_score" in doc.metadata:
+            doc.metadata["relevance_score_without_reranking"] = doc.metadata.pop("relevance_score")
+        return doc
+
+    def _get_relevant_documents(
+        self,
+        query: str,
+        *,
+        run_manager: CallbackManagerForRetrieverRun,
+        doc_id_whitelist: Optional[List[str]] = None,
+        doc_id_blacklist: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> List[LCDocument]:
+        """Get span documents relevant to a query span
+        Args:
+            query: The span id to find relevant spans for
+            run_manager: The callbacks handler to use
+        Returns:
+            List of relevant span documents with metadata from the parent document
+        """
+
+        search_kwargs, query_span_doc = self.prepare_search_kwargs(
+            span_id=query,
+            kwargs=kwargs,
+            doc_id_whitelist=doc_id_whitelist,
+            doc_id_blacklist=doc_id_blacklist,
+        )
+        if self.search_type == SearchType.mmr:
+            span_docs = self.vectorstore.max_marginal_relevance_search_by_vector(**search_kwargs)
+        elif self.search_type == SearchType.similarity_score_threshold:
+            sub_docs_and_similarities = self.vectorstore.similarity_search_with_score_by_vector(
+                **search_kwargs
+            )
+            span_docs = [sub_doc for sub_doc, _ in sub_docs_and_similarities]
+        else:
+            span_docs = self.vectorstore.similarity_search_by_vector(**search_kwargs)
+
+        # We do this to maintain the order of the ids that are returned
+        doc_ids = []
+        for span_doc in span_docs:
+            if self.id_key not in span_doc.metadata:
+                raise ValueError(f"Metadata must contain the key {self.id_key}")
+            if span_doc.metadata[self.id_key] not in doc_ids:
+                doc_ids.append(span_doc.metadata[self.id_key])
+        docs = self.docstore.mget(doc_ids)
+        doc_id2doc = dict(zip(doc_ids, docs))
+        for span_doc in span_docs:
+            doc = doc_id2doc[span_doc.metadata[self.id_key]]
+            span_doc.metadata.update(doc.metadata)
+            span_doc.metadata["attached_span"] = self.get_span_from_doc_by_id(
+                doc=doc, span_id=span_doc.id
+            )
+            span_doc.metadata["query_span_id"] = query
+        # filter out the query span doc
+        span_docs_filtered = [
+            span_doc for span_doc in span_docs if span_doc.id != query_span_doc.id
+        ]
+        if self.compressor is None:
+            return span_docs_filtered
+        if span_docs_filtered:
+            prepared_docs = [
+                self._prepare_doc_for_compression(sub_doc) for sub_doc in span_docs_filtered
+            ]
+            prepared_query = self._prepare_query_for_compression(query_span_doc)
+            compressed_docs = self.compressor.compress_documents(
+                documents=prepared_docs, query=prepared_query, callbacks=run_manager.get_child()
+            )
+            return list(compressed_docs)
+        else:
+            return []
+
+    async def _aget_relevant_documents(
+        self,
+        query: str,
+        *,
+        run_manager: AsyncCallbackManagerForRetrieverRun,
+        doc_id_whitelist: Optional[List[str]] = None,
+        doc_id_blacklist: Optional[List[str]] = None,
+        **kwargs: Any,
+    ) -> List[LCDocument]:
+        """Asynchronously get span documents relevant to a query span
+        Args:
+            query: The span id to find relevant spans for
+            run_manager: The callbacks handler to use
+        Returns:
+            List of relevant span documents with metadata from the parent document
+        """
+        search_kwargs, query_span_doc = self.prepare_search_kwargs(
+            span_id=query,
+            kwargs=kwargs,
+            doc_id_whitelist=doc_id_whitelist,
+            doc_id_blacklist=doc_id_blacklist,
+        )
+        if self.search_type == SearchType.mmr:
+            span_docs = await self.vectorstore.amax_marginal_relevance_search_by_vector(
+                **search_kwargs
+            )
+        elif self.search_type == SearchType.similarity_score_threshold:
+            sub_docs_and_similarities = (
+                await self.vectorstore.asimilarity_search_with_score_by_vector(**search_kwargs)
+            )
+            span_docs = [sub_doc for sub_doc, _ in sub_docs_and_similarities]
+        else:
+            span_docs = await self.vectorstore.asimilarity_search_by_vector(**search_kwargs)
+
+        # We do this to maintain the order of the ids that are returned
+        ids = []
+        for span_doc in span_docs:
+            if self.id_key not in span_doc.metadata:
+                raise ValueError(f"Metadata must contain the key {self.id_key}")
+            if span_doc.metadata[self.id_key] not in ids:
+                ids.append(span_doc.metadata[self.id_key])
+        docs = await self.docstore.amget(ids)
+        doc_id2doc = dict(zip(ids, docs))
+        for span_doc in span_docs:
+            doc = doc_id2doc[span_doc.metadata[self.id_key]]
+            span_doc.metadata.update(doc.metadata)
+            span_doc.metadata["attached_span"] = self.get_span_from_doc_by_id(
+                doc=doc, span_id=span_doc.id
+            )
+            span_doc.metadata["query_span_id"] = query
+        # filter out the query span doc
+        span_docs_filtered = [
+            span_doc for span_doc in span_docs if span_doc.id != query_span_doc.id
+        ]
+
+        if self.compressor is None:
+            return span_docs_filtered
+        if docs:
+            prepared_docs = [
+                self._prepare_doc_for_compression(sub_doc) for sub_doc in span_docs_filtered
+            ]
+            prepared_query = self._prepare_query_for_compression(query_span_doc)
+            compressed_docs = await self.base_compressor.acompress_documents(
+                prepared_docs, query=prepared_query, callbacks=run_manager.get_child()
+            )
+            return list(compressed_docs)
+        else:
+            return []
+
+    def create_span_documents(
+        self, documents: List[LCDocument]
+    ) -> Tuple[List[LCDocument], Dict[str, int]]:
+        span_docs = []
+        id2idx = {}
+        for i, doc in enumerate(documents):
+            pie_doc, metadata = self.docstore.unwrap_with_metadata(doc)
+            base_layer = self.get_base_layer(
+                pie_doc, use_predicted_annotations=self.use_predicted_annotations(doc)
+            )
+            if len(base_layer) == 0:
+                logger.warning(f"No spans found in document {i} (id: {doc.id})")
+            for idx, labeled_span in enumerate(base_layer):
+                _metadata = {k: v for k, v in metadata.items() if k != METADATA_KEY_CHILD_ID2IDX}
+                # save as dict to avoid serialization issues
+                _metadata[self.vectorstore.METADATA_SPAN_KEY] = self._span_to_dict(labeled_span)
+                new_doc = LCDocument(
+                    id=str(uuid.uuid4()), page_content=pie_doc.text, metadata=_metadata
+                )
+                span_docs.append(new_doc)
+                id2idx[new_doc.id] = idx
+        return span_docs, id2idx
+
+    def _split_docs_for_adding(
+        self,
+        documents: List[LCDocument],
+        ids: Optional[List[str]] = None,
+        add_to_docstore: bool = True,
+    ) -> Tuple[List[LCDocument], List[Tuple[str, LCDocument]]]:
+        if ids is None:
+            doc_ids = [doc.id for doc in documents]
+            if not add_to_docstore:
+                raise ValueError("If ids are not passed in, `add_to_docstore` MUST be True")
+        else:
+            if len(documents) != len(ids):
+                raise ValueError(
+                    "Got uneven list of documents and ids. "
+                    "If `ids` is provided, should be same length as `documents`."
+                )
+            doc_ids = ids
+
+        if len(set(doc_ids)) != len(doc_ids):
+            raise ValueError("IDs must be unique")
+
+        docs = []
+        full_docs = []
+        for i, doc in enumerate(documents):
+            _id = doc_ids[i]
+            sub_docs, sub_doc_id2idx = self.create_span_documents([doc])
+            if self.child_metadata_fields is not None:
+                for sub_doc in sub_docs:
+                    sub_doc.metadata = {k: sub_doc.metadata[k] for k in self.child_metadata_fields}
+            for sub_doc in sub_docs:
+                # Add the parent id to the child document id
+                sub_doc.metadata[self.id_key] = _id
+            docs.extend(sub_docs)
+            doc.metadata[METADATA_KEY_CHILD_ID2IDX] = sub_doc_id2idx
+            full_docs.append((_id, doc))
+
+        return docs, full_docs
+
+    def remove_missing_span_ids_from_document(
+        self, document: LCDocument, span_ids: Set[str]
+    ) -> LCDocument:
+        """Remove invalid span ids from the span to idx mapping
+        of the document.
+
+        Args:
+            document: Document to remove invalid span ids from
+            span_ids: Set of valid span ids
+
+        Returns:
+            Document with invalid span ids removed
+        """
+        span_id2idx = document.metadata[METADATA_KEY_CHILD_ID2IDX]
+        new_doc = copy(document)
+        filtered_span_id2idx = {
+            span_id: idx for span_id, idx in span_id2idx.items() if span_id in span_ids
+        }
+        new_doc.metadata[METADATA_KEY_CHILD_ID2IDX] = filtered_span_id2idx
+        missed_span_ids = set(span_id2idx.keys()) - span_ids
+        if len(missed_span_ids) > 0:
+            layer = self.get_base_layer(
+                self.docstore.unwrap(document),
+                use_predicted_annotations=self.use_predicted_annotations(document),
+            )
+            resolved_missed_spans = [
+                layer[span_id2idx[span_id]].resolve() for span_id in missed_span_ids
+            ]
+            logger.warning(
+                f"Document {document.id} contains spans that can not be added to the "
+                f"vectorstore because no vector could be calculated:\n{resolved_missed_spans}.\n"
+                "These spans will be not queryable."
+            )
+        return document
+
+    def add_documents(
+        self,
+        documents: List[LCDocument],
+        ids: Optional[List[str]] = None,
+        add_to_docstore: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        """Adds documents to the docstore and vectorstores.
+
+        Args:
+            documents: List of documents to add
+            ids: Optional list of ids for documents. If provided should be the same
+                length as the list of documents. Can be provided if parent documents
+                are already in the document store and you don't want to re-add
+                to the docstore. If not provided, random UUIDs will be used as
+                ids.
+            add_to_docstore: Boolean of whether to add documents to docstore.
+                This can be false if and only if `ids` are provided. You may want
+                to set this to False if the documents are already in the docstore
+                and you don't want to re-add them.
+        """
+        docs, full_docs = self._split_docs_for_adding(documents, ids, add_to_docstore)
+        added_span_ids = self.vectorstore.add_documents(docs, **kwargs)
+        full_docs = [
+            (doc_id, self.remove_missing_span_ids_from_document(doc, set(added_span_ids)))
+            for doc_id, doc in full_docs
+        ]
+        if add_to_docstore:
+            self.docstore.mset(full_docs)
+
+    async def aadd_documents(
+        self,
+        documents: List[LCDocument],
+        ids: Optional[List[str]] = None,
+        add_to_docstore: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        docs, full_docs = self._split_docs_for_adding(documents, ids, add_to_docstore)
+        added_span_ids = await self.vectorstore.aadd_documents(docs, **kwargs)
+        full_docs = [
+            (doc_id, self.remove_missing_span_ids_from_document(doc, set(added_span_ids)))
+            for doc_id, doc in full_docs
+        ]
+        if add_to_docstore:
+            await self.docstore.amset(full_docs)
+
+    def delete_documents(self, ids: List[str]) -> None:
+        """Remove documents from the docstore and vectorstores.
+
+        Args:
+            ids: List of ids to remove
+        """
+        # get all child ids
+        child_ids = []
+        for doc in self.docstore.mget(ids):
+            child_ids.extend(doc.metadata[METADATA_KEY_CHILD_ID2IDX])
+
+        self.vectorstore.delete(child_ids)
+        self.docstore.mdelete(ids)
+
+    async def adelete_documents(self, ids: List[str]) -> None:
+        """Asynchronously remove documents from the docstore and vectorstores.
+
+        Args:
+            ids: List of ids to remove
+        """
+        # get all child ids
+        child_ids = []
+        docs: List[LCDocument] = await self.docstore.amget(ids)
+        for doc in docs:
+            child_ids.extend(doc.metadata[METADATA_KEY_CHILD_ID2IDX])
+
+        await self.vectorstore.adelete(child_ids)
+        await self.docstore.amdelete(ids)
+
+    def add_pie_documents(
+        self,
+        documents: List[TextBasedDocument],
+        use_predicted_annotations: bool,
+        metadata: Optional[Dict[str, Any]] = None,
+    ) -> None:
+        """Add pie documents to the retriever.
+
+        Args:
+            documents: List of pie documents to add
+            use_predicted_annotations: Whether to use the predicted annotations or the gold annotations
+            metadata: Optional metadata to add to each document
+        """
+        metadata = metadata or {}
+        metadata = copy(metadata)
+        metadata[self.use_predicted_annotations_key] = use_predicted_annotations
+        docs = [self.docstore.wrap(doc, **metadata) for doc in documents]
+
+        # delete any existing documents with the same ids (simply overwriting would keep the spans)
+        new_docs_ids = [doc.id for doc in docs]
+        existing_docs = self.docstore.mget(new_docs_ids)
+        existing_doc_ids = [doc.id for doc in existing_docs]
+        self.delete_documents(existing_doc_ids)
+
+        self.add_documents(docs)
+
+    def _save_to_directory(self, path: str, **kwargs) -> None:
+        logger.info(f'Saving docstore and vectorstore to "{path}" ...')
+        self.docstore.save_to_directory(os.path.join(path, "docstore"))
+        self.vectorstore.save_to_directory(os.path.join(path, "vectorstore"))
+
+    def _load_from_directory(self, path: str, **kwargs) -> None:
+        logger.info(f'Loading docstore and vectorstore from "{path}" ...')
+        self.docstore.load_from_directory(os.path.join(path, "docstore"))
+        self.vectorstore.load_from_directory(os.path.join(path, "vectorstore"))
+
+
+METADATA_KEY_RELATION_LABEL2TAILS_WITH_SCORES = "relation_label2tails_with_scores"
+
+
+class DocumentAwareSpanRetrieverWithRelations(DocumentAwareSpanRetriever):
+    """Retriever for related contextualized text spans, i.e. spans linked by relations
+    to reference spans that are similar to the query span. It accepts spans as queries and
+    retrieves spans with their containing document and the reference span."""
+
+    relation_layer_name: str = "binary_relations"
+    """The name of the relation annotation layer in the pie document."""
+    relation_labels: Optional[List[str]] = None
+    """The list of relation labels to consider."""
+    span_labels: Optional[List[str]] = None
+    """The list of span labels to consider."""
+    reversed_relations_suffix: Optional[str] = None
+    """Whether to consider reverse relations as well."""
+
+    def get_relation_layer(
+        self, pie_document: TextBasedDocument, use_predicted_annotations: bool
+    ) -> BaseAnnotationList:
+        """Get the relation layer of the pie document."""
+        if self.relation_layer_name not in pie_document:
+            raise ValueError(
+                f'The pie document must contain the annotation layer "{self.relation_layer_name}"'
+            )
+        layer = pie_document[self.relation_layer_name]
+        return layer.predictions if use_predicted_annotations else layer
+
+    def create_span_documents(
+        self, documents: List[LCDocument]
+    ) -> Tuple[List[LCDocument], Dict[str, int]]:
+        span_docs = []
+        id2idx = {}
+        for i, doc in enumerate(documents):
+            pie_doc, metadata = self.docstore.unwrap_with_metadata(doc)
+            use_predicted_annotations = self.use_predicted_annotations(doc)
+            base_layer = self.get_base_layer(
+                pie_doc, use_predicted_annotations=use_predicted_annotations
+            )
+            if len(base_layer) == 0:
+                logger.warning(f"No spans found in document {i} (id: {doc.id})")
+            id2span = {str(uuid.uuid4()): span for span in base_layer}
+            span2id = {span: span_id for span_id, span in id2span.items()}
+            if len(id2span) != len(span2id):
+                raise ValueError("Span ids and spans must be unique")
+            relations = self.get_relation_layer(
+                pie_doc, use_predicted_annotations=use_predicted_annotations
+            )
+            head2label2tails_with_scores: Dict[str, Dict[str, List[Tuple[str, float]]]] = (
+                defaultdict(lambda: defaultdict(list))
+            )
+
+            for relation in relations:
+                if self.relation_labels is None or relation.label in self.relation_labels:
+                    head2label2tails_with_scores[span2id[relation.head]][relation.label].append(
+                        (span2id[relation.tail], relation.score)
+                    )
+                if self.reversed_relations_suffix is not None:
+                    reversed_label = f"{relation.label}{self.reversed_relations_suffix}"
+                    if self.relation_labels is None or reversed_label in self.relation_labels:
+                        head2label2tails_with_scores[span2id[relation.tail]][
+                            reversed_label
+                        ].append((span2id[relation.head], relation.score))
+
+            for idx, span in enumerate(base_layer):
+                span_id = span2id[span]
+                _metadata = {k: v for k, v in metadata.items() if k != METADATA_KEY_CHILD_ID2IDX}
+                # save as dict to avoid serialization issues
+                _metadata[self.vectorstore.METADATA_SPAN_KEY] = self._span_to_dict(span)
+                relation_label2tails_with_scores = head2label2tails_with_scores[span_id]
+                _metadata[METADATA_KEY_RELATION_LABEL2TAILS_WITH_SCORES] = dict(
+                    relation_label2tails_with_scores
+                )
+                new_doc = LCDocument(id=span_id, page_content=pie_doc.text, metadata=_metadata)
+                span_docs.append(new_doc)
+                id2idx[span_id] = idx
+        return span_docs, id2idx
+
+    def _get_relevant_documents(
+        self,
+        query: str,
+        return_related: bool = False,
+        *,
+        run_manager: CallbackManagerForRetrieverRun,
+        **kwargs: Any,
+    ) -> List[LCDocument]:
+        """Get span documents relevant to a query span. We follow one hop of relations.
+
+        Args:
+            query: The span id to find relevant spans for
+            return_related: Whether to return related spans
+            run_manager: The callbacks handler to use
+        Returns:
+            List of relevant span documents with metadata from the parent document
+        """
+        similar_span_docs = super()._get_relevant_documents(
+            query=query, run_manager=run_manager, **kwargs
+        )
+        if not return_related:
+            return similar_span_docs
+
+        related_docs = []
+        for head_span_doc in similar_span_docs:
+            doc_id = head_span_doc.metadata[self.id_key]
+            doc = self.get_document(doc_id)
+            query_span_id = head_span_doc.metadata["query_span_id"]
+
+            for relation_label, tails_with_score in head_span_doc.metadata[
+                METADATA_KEY_RELATION_LABEL2TAILS_WITH_SCORES
+            ].items():
+                for tail_id, relation_score in tails_with_score:
+                    # in the case that we query against the same document,
+                    # we don't want to return the same span as the query span
+                    if tail_id == query_span_id:
+                        continue
+
+                    try:
+                        attached_tail_span = self.get_span_from_doc_by_id(doc=doc, span_id=tail_id)
+                    # this may happen if the tail span could not be added to the vectorstore, e.g. because
+                    # the token span length is zero and no vector could be calculated
+                    except SpanNotFoundError:
+                        logger.warning(
+                            f"Tail span with id [{tail_id}] not found in the vectorstore. Skipping."
+                        )
+                        continue
+
+                    # TODO: handle via filter? see vectorstore.construct_filter
+                    if self.span_labels is not None:
+                        if not isinstance(attached_tail_span, (LabeledSpan, LabeledMultiSpan)):
+                            raise ValueError(
+                                "Span must must be a labeled span if span_labels is provided"
+                            )
+                        if attached_tail_span.label not in self.span_labels:
+                            continue
+
+                    related_docs.append(
+                        LCDocument(
+                            id=tail_id,
+                            page_content="",
+                            metadata={
+                                "relation_score": relation_score,
+                                "head_id": head_span_doc.id,
+                                "relation_label": relation_label,
+                                "attached_tail_span": attached_tail_span,
+                                **head_span_doc.metadata,
+                            },
+                        )
+                    )
+        return related_docs

src/langchain_modules/span_vectorstore.py

@@ -0,0 +1,131 @@
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Sequence, Union
+
+from langchain_core.documents import Document as LCDocument
+from langchain_core.runnables import run_in_executor
+from langchain_core.stores import BaseStore
+from langchain_core.vectorstores import VectorStore
+from pytorch_ie.annotations import MultiSpan, Span
+
+from .serializable_store import SerializableStore
+from .span_embeddings import SpanEmbeddings
+
+
+class SpanVectorStore(VectorStore, BaseStore, SerializableStore, ABC):
+    """Abstract base class for vector stores specialized in storing
+    and retrieving embeddings for text spans within documents."""
+
+    METADATA_SPAN_KEY: str = "pie_labeled_span"
+    """Key for the span data in the (langchain) document metadata."""
+    SPAN_START_KEY: str = "start"
+    """Key for the start of the span in the span data."""
+    SPAN_END_KEY: str = "end"
+    """Key for the end of the span in the span data."""
+    METADATA_VECTOR_KEY: str = "vector"
+    """Key for the vector in the (langchain) document metadata."""
+    RELEVANCE_SCORE_KEY: str = "relevance_score"
+    """Key for the relevance score in the (langchain) document metadata."""
+
+    def __init__(
+        self,
+        label_mapping: Optional[Dict[str, List[str]]] = None,
+        **kwargs: Any,
+    ):
+        """Initialize the SpanVectorStore.
+
+        Args:
+            label_mapping: Mapping from query span labels to target span labels. If provided,
+                only spans with a label in the mapping for the query span's label are retrieved.
+            **kwargs: Additional arguments.
+        """
+        self.label_mapping = label_mapping
+        super().__init__(**kwargs)
+
+    @property
+    def embeddings(self) -> SpanEmbeddings:
+        """Get the dense embeddings instance that is being used.
+
+        Raises:
+            ValueError: If embeddings are `None`.
+
+        Returns:
+            Embeddings: An instance of `Embeddings`.
+        """
+        result = super().embeddings
+        if not isinstance(result, SpanEmbeddings):
+            raise ValueError(f"Embeddings must be of type SpanEmbeddings, but got: {result}")
+        return result
+
+    @abstractmethod
+    def get_by_ids_with_vectors(self, ids: Sequence[Union[str, int]], /) -> List[LCDocument]:
+        """Get documents by their ids.
+
+        Args:
+            ids: List of document ids.
+
+        Returns:
+            List of documents including their vectors in the metadata at key `metadata_vector_key`.
+        """
+        ...
+
+    @abstractmethod
+    def construct_filter(
+        self,
+        query_span: Union[Span, MultiSpan],
+        metadata_doc_id_key: str,
+        doc_id_whitelist: Optional[Sequence[str]] = None,
+        doc_id_blacklist: Optional[Sequence[str]] = None,
+    ) -> Any:
+        """Construct a filter for the retrieval. It should enforce that:
+        - if the span is labeled, the retrieved span has the same label, or
+        - if, in addition, a label mapping is provided, the retrieved span has a label that is in the mapping for the query span's label
+        - if `doc_id_whitelist` is provided, the retrieved span is from a document in the whitelist
+        - if `doc_id_blacklist` is provided, the retrieved span is not from a document in the blacklist
+
+        Args:
+            query_span: The query span.
+            metadata_doc_id_key: The key in the metadata that holds the document id.
+            doc_id_whitelist: A list of document ids to restrict the retrieval to.
+            doc_id_blacklist: A list of document ids to exclude from the retrieval.
+
+        Returns:
+            A filter object.
+        """
+        ...
+
+    @abstractmethod
+    def similarity_search_with_score_by_vector(
+        self, embedding: list[float], k: int = 4, **kwargs: Any
+    ) -> list[LCDocument]:
+        """Return docs most similar to embedding vector.
+
+        Args:
+            embedding: Embedding to look up documents similar to.
+            k: Number of Documents to return. Defaults to 4.
+            **kwargs: Arguments to pass to the search method.
+
+        Returns:
+            List of Documents most similar to the query vector.
+        """
+        ...
+
+    async def asimilarity_search_with_score_by_vector(
+        self, embedding: list[float], k: int = 4, **kwargs: Any
+    ) -> list[LCDocument]:
+        """Async return docs most similar to embedding vector.
+
+        Args:
+            embedding: Embedding to look up documents similar to.
+            k: Number of Documents to return. Defaults to 4.
+            **kwargs: Arguments to pass to the search method.
+
+        Returns:
+            List of Documents most similar to the query vector.
+        """
+
+        # This is a temporary workaround to make the similarity search
+        # asynchronous. The proper solution is to make the similarity search
+        # asynchronous in the vector store implementations.
+        return await run_in_executor(
+            None, self.similarity_search_with_score_by_vector, embedding, k=k, **kwargs
+        )