document_store.py

import json
import logging
import os
import shutil
import tempfile
import zipfile
from collections import defaultdict
from typing import Any, Dict, List, Optional

import gradio as gr
import numpy as np
import pandas as pd
from annotation_utils import labeled_span_to_id
from pytorch_ie import Annotation
from pytorch_ie.documents import (
    TextBasedDocument,
    TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
)
from scipy.sparse import csr_matrix
from vector_store import VectorStore

logger = logging.getLogger(__name__)


def get_annotation_from_document(
    document: TextBasedDocument,
    annotation_id: str,
    annotation_layer: str,
    use_predictions: bool,
) -> Annotation:
    """Get an annotation from a document by its id. Note that the annotation id is constructed from
    the annotation itself, so it is unique within the document.

    Args:
        document: The document to get the annotation from.
        annotation_id: The id of the annotation.
        annotation_layer: The name of the annotation layer.
        use_predictions: Whether to use the predictions of the annotation layer.

    Returns:
        The annotation with the given id.
    """

    annotations = document[annotation_layer]
    if use_predictions:
        annotations = annotations.predictions

    if annotation_layer == "labeled_spans":
        annotation_to_id_func = labeled_span_to_id
    else:
        raise gr.Error(f"Unknown annotation layer '{annotation_layer}'.")

    id2annotation = {annotation_to_id_func(annotation): annotation for annotation in annotations}
    annotation = id2annotation.get(annotation_id)
    if annotation is None:
        raise gr.Error(
            f"annotation '{annotation_id}' not found in document '{document.id}'. Available "
            f"annotations: {id2annotation}"
        )
    return annotation


def get_related_annotation_records_from_document(
    document: TextBasedDocument,
    reference_annotation: Annotation,
    relation_layer_name: str,
    use_predictions: bool,
    annotation_caption: str,
    relation_types: Optional[List[str]] = None,
    additional_static_columns: Optional[Dict[str, str]] = None,
) -> List[Dict[str, str]]:
    """Get related annotations from a document for a given reference annotation. The related
    annotations are all annotations that are targets (tails) of relations with the reference
    annotation as source (head).

    Args:
        document: The document to get the related annotations from.
        reference_annotation: The reference annotation. Should be an annotation from the document.
        relation_layer_name: The name of the relation layer.
        use_predictions: Whether to use the predictions of the relation layer.
        annotation_caption: The caption for the related annotations in the result.
        relation_types: The types of relations to consider. If None, all relation types are considered.
        additional_static_columns: Additional static columns to add to the result.

    Returns:
        A list of dictionaries with the related annotations and additional columns.
    """

    result = []

    # get the relation layer
    relation_layer = document[relation_layer_name]
    if use_predictions:
        relation_layer = relation_layer.predictions

    # create helper dictionaries to quickly find related annotations
    tail2rels = defaultdict(list)
    head2rels = defaultdict(list)
    for rel in relation_layer:
        # skip non-argumentative relations
        if relation_types is not None and rel.label not in relation_types:
            continue
        head2rels[rel.head].append(rel)
        tail2rels[rel.tail].append(rel)

    # get the related annotations: all annotations that are targets (tails) of relations with the reference
    # annotation as source (head)
    for rel in head2rels.get(reference_annotation, []):
        result.append(
            {
                "doc_id": document.id,
                f"reference_{annotation_caption}": str(reference_annotation),
                "rel_score": rel.score,
                "relation": rel.label,
                annotation_caption: str(rel.tail),
                **(additional_static_columns or {}),
            }
        )
    return result


class DocumentStore:
    """A document store that allows to add, retrieve, and search for documents and annotations.

    The store keeps the documents in memory and stores the embeddings of the labeled spans in a vector
    store to efficiently retrieve similar or related spans.

    Args:
        vector_store: The vector store to use. If None, a new SimpleVectorStore is created.
        document_type: The type of the documents to store. Should be a subclass of TextBasedDocument with
            a span and a relation layer (see below).
        span_layer_name: The name of the span annotation layer. This should be a valid annotation layer
            of type LabelSpan in the document type.
        relation_layer_name: The name of the argumentative relation annotation layer. This should be a
            valid annotation layer of type BinaryRelation in the document type.
        span_annotation_caption: The caption for the span annotations (e.g. in the statistical overview)
        relation_annotation_caption: The caption for the relation annotations (e.g. in the statistical
            overview)
        use_predictions: Whether to use the predictions of the annotation layers. If True, the predictions
            are used, otherwise the gold annotations are used.
    """

    JSON_FILE_NAME = "documents.json"

    def __init__(
        self,
        vector_store: VectorStore[Dict[str, Any], List[float]],
        document_type: type[
            TextBasedDocument
        ] = TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
        span_layer_name: str = "labeled_spans",
        relation_layer_name: str = "binary_relations",
        span_annotation_caption: str = "span",
        relation_annotation_caption: str = "relation",
        use_predictions: bool = True,
    ):
        # The annotated documents. As key, we use the document id. All documents keep the embeddings
        # of the spans in the metadata.
        self.documents: Dict[str, TextBasedDocument] = {}
        # The vector store to efficiently retrieve similar spans. Can be constructed from the
        # documents.
        self.vector_store = vector_store
        # the document type (to create new documents from dicts)
        self.document_type = document_type
        self.span_layer_name = span_layer_name
        self.relation_layer_name = relation_layer_name
        self.use_predictions = use_predictions
        self.layer_captions = {
            self.span_layer_name: span_annotation_caption,
            self.relation_layer_name: relation_annotation_caption,
        }

    def get_annotation(
        self,
        doc_id: str,
        annotation_id: str,
        annotation_layer: str,
        use_predictions: bool,
    ) -> Annotation:
        document = self.documents.get(doc_id)
        if document is None:
            raise gr.Error(
                f"Document '{doc_id}' not found in index. Available documents: {list(self.documents)}"
            )
        return get_annotation_from_document(
            document, annotation_id, annotation_layer, use_predictions=use_predictions
        )

    def construct_embedding_payload(self, document: TextBasedDocument, annotation_id: str) -> dict:
        payload = {"doc_id": document.id, "annotation_id": annotation_id}
        return payload

    def get_similar_annotations_df(
        self,
        ref_annotation_id: str,
        ref_document: TextBasedDocument,
        annotation_layer: str,
        **similarity_kwargs,
    ) -> pd.DataFrame:
        """Get similar annotations from documents in the store sorted by similarity. Usually, the
        reference annotation is returned as the most similar annotation.

        Args:
            ref_annotation_id: The id of the reference annotation.
            ref_document: The document of the reference annotation.
            annotation_layer: The name of the annotation layer to consider.
            **similarity_kwargs: Additional keyword arguments that will be passed to the vector
                store to retrieve similar entries (see VectorStore.retrieve_similar()).

        Returns:
            A DataFrame with the similar annotations with columns: doc_id, annotation_id, sim_score,
            and text.
        """

        similar_entries = self.vector_store.retrieve_similar(
            ref_payload=self.construct_embedding_payload(ref_document, ref_annotation_id),
            **similarity_kwargs,
        )

        similar_annotations = [
            self.get_annotation(
                doc_id=payload["doc_id"],
                annotation_id=payload["annotation_id"],
                annotation_layer=annotation_layer,
                use_predictions=self.use_predictions,
            )
            for _, payload, _ in similar_entries
        ]
        df = pd.DataFrame(
            [
                # unpack the tuple (doc_id, annotation_id) to separate columns
                # and add the similarity score and the text of the annotation
                (payload["doc_id"], payload["annotation_id"], score, str(annotation))
                for (_, payload, score), annotation in zip(similar_entries, similar_annotations)
            ],
            columns=["doc_id", "annotation_id", "sim_score", "text"],
        )

        return df

    def get_related_annotations_from_other_documents_df(
        self,
        ref_annotation_id: str,
        ref_document: TextBasedDocument,
        min_similarity: float,
        top_k: int,
        relation_types: List[str],
        columns: List[str],
    ) -> pd.DataFrame:
        """Get related annotations from documents in the store for a given reference annotation.
        First, similar annotations are retrieved from the vector store. Then, annotations that are
        linked to them via relations are returned. Only annotations from other documents are
        considered.

        Args:
            ref_annotation_id: The id of the reference annotation.
            ref_document: The document of the reference annotation.
            min_similarity: The minimum similarity score to consider.
            top_k: The number of related annotations to return.
            relation_types: The types of relations to consider.
            columns: The columns to include in the result DataFrame.

        Returns:
            A DataFrame with the columns that contain: the related annotation, the relation type,
            the similar annotation, the similarity score, and the relation score.
        """

        similar_entries = self.vector_store.retrieve_similar(
            ref_payload=self.construct_embedding_payload(ref_document, ref_annotation_id),
            min_similarity=min_similarity,
            top_k=top_k,
        )
        result = []
        for _, payload, score in similar_entries:
            doc_id = payload["doc_id"]
            # skip entries from the same document
            if doc_id == ref_document.id:
                continue
            document = self.documents[doc_id]
            reference_annotation = get_annotation_from_document(
                document=document,
                annotation_id=payload["annotation_id"],
                annotation_layer=self.span_layer_name,
                use_predictions=self.use_predictions,
            )

            new_entries = get_related_annotation_records_from_document(
                document=document,
                reference_annotation=reference_annotation,
                relation_types=relation_types,
                relation_layer_name=self.relation_layer_name,
                use_predictions=self.use_predictions,
                annotation_caption=self.layer_captions[self.span_layer_name],
                additional_static_columns={"sim_score": str(score)},
            )
            result.extend(new_entries)

        # define column order
        df = pd.DataFrame(result, columns=columns)
        return df

    def add_document(self, document: TextBasedDocument) -> None:
        try:
            if document.id in self.documents:
                gr.Warning(f"Document '{document.id}' already in index. Overwriting.")

            # copy the document to avoid side effects
            document = document.copy()

            # save the processed document to the index
            self.documents[document.id] = document

            # save the embeddings to the vector store, if available
            if "embeddings" in document.metadata:
                for annotation_id, embedding in document.metadata["embeddings"].items():
                    payload = self.construct_embedding_payload(document, annotation_id)
                    self.vector_store.add(payload=payload, embedding=embedding)
                # remove the embeddings from the document metadata
                document.metadata = {
                    k: v for k, v in document.metadata.items() if k != "embeddings"
                }

        except Exception as e:
            raise gr.Error(f"Failed to add document {document.id} to index: {e}")

    def add_document_from_dict(self, document_dict: dict) -> None:
        document = self.document_type.fromdict(document_dict)
        self.add_document(document)

    def add_documents(self, documents: List[TextBasedDocument]) -> None:
        for document in documents:
            self.add_document(document)
        gr.Info(
            f"Added {len(documents)} documents to the index ({len(self.documents)} documents in total)."
        )

    def add_documents_from_json(self, file_path: str) -> None:
        with open(file_path, "r", encoding="utf-8") as f:
            documents_json = json.load(f)
        for _, document_json in documents_json.items():
            self.add_document_from_dict(document_dict=document_json)
        gr.Info(
            f"Added {len(documents_json)} documents to the index ({len(self.documents)} documents in total)."
        )

    def get_payloads_for_missing_and_unexpected_embeddings(self) -> dict[str, dict[str, Any]]:
        """Get the payloads for missing and unexpected embeddings in the vector store. An embedding
        is missing if its annotation is in the documents but the embedding is not in the vector
        store. An embedding is unexpected if it is in the vector store but the annotation is not in
        the documents.

        Returns:
            A dictionary with the missing and unexpected payloads.
        """
        expected_payloads = []
        for document in self.documents.values():
            for annotation in document[self.span_layer_name].predictions:
                annotation_id = labeled_span_to_id(annotation)
                payload = self.construct_embedding_payload(document, annotation_id)
                expected_payloads.append(payload)
        vector_sore_payloads = self.vector_store.as_indices_vectors_payloads()[2]
        # construct mappings from ids to payloads to compare the expected and actual payloads
        expected_mapping = {
            json.dumps(payload, sort_keys=True): payload for payload in expected_payloads
        }
        vector_store_mapping = {
            json.dumps(payload, sort_keys=True): payload for payload in vector_sore_payloads
        }
        missing = set(expected_mapping) - set(vector_store_mapping)
        unexpected = set(vector_store_mapping) - set(expected_mapping)

        # return the missing and unexpected payloads
        return {
            "missing": {payload: expected_mapping[payload] for payload in missing},
            "unexpected": {payload: vector_store_mapping[payload] for payload in unexpected},
        }

    def add_documents_from_zip(self, file_path: str) -> None:
        temp_dir = os.path.join(tempfile.gettempdir(), "document_store")
        # remove the temporary directory if it already exists
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
        with zipfile.ZipFile(file_path, "r") as zipf:
            # extract all files to the temporary directory
            zipf.extractall(temp_dir)
        json_file_path = os.path.join(temp_dir, self.JSON_FILE_NAME)
        self.add_documents_from_json(json_file_path)
        # load the vector store from the temporary directory
        self.vector_store.load_from_directory(temp_dir)
        # delete the temporary directory
        shutil.rmtree(temp_dir)

    def add_documents_from_file(self, file_path: str) -> None:
        if file_path.endswith(".json"):
            self.add_documents_from_json(file_path)
        elif file_path.endswith(".zip"):
            self.add_documents_from_zip(file_path)
        else:
            raise gr.Error(f"Unsupported file format: {file_path}")

    def save_to_json(self, file_path: str, include_embeddings: bool = True, **kwargs) -> None:
        with open(file_path, "w", encoding="utf-8") as f:
            json.dump(self.as_dict(include_embeddings=include_embeddings), f, **kwargs)

    def save_to_zip(self, file_path: str, **kwargs) -> None:
        # first create a new temporary directory and save the documents as json file in it
        temp_dir = os.path.join(tempfile.gettempdir(), "document_store")
        # remove the temporary directory if it already exists
        if os.path.exists(temp_dir):
            shutil.rmtree(temp_dir)
        os.makedirs(temp_dir)
        temp_file_path = os.path.join(temp_dir, self.JSON_FILE_NAME)
        self.save_to_json(temp_file_path, include_embeddings=False, **kwargs)
        self.vector_store.save_to_directory(temp_dir)
        # then zip all files in the temporary directory and write them to the target file
        with zipfile.ZipFile(file_path, "w") as zipf:
            for root, _, files in os.walk(temp_dir):
                for file in files:
                    zipf.write(
                        os.path.join(root, file),
                        os.path.relpath(os.path.join(root, file), temp_dir),
                    )
        # delete the temporary directory
        shutil.rmtree(temp_dir)

    def save_to_file(self, file_path: str, **kwargs) -> None:
        if file_path.endswith(".json"):
            self.save_to_json(file_path, **kwargs)
        elif file_path.endswith(".zip"):
            self.save_to_zip(file_path, **kwargs)
        else:
            raise gr.Error(f"Unsupported file format: {file_path}")

    def get_document(self, doc_id: str, with_embeddings: bool = False) -> TextBasedDocument:
        document = self.documents[doc_id]
        if not with_embeddings:
            return document

        # TODO: is this really required?
        # copy because we add the embeddings to the metadata
        document = document.copy()
        # get the embeddings from the vector store
        embeddings = {}
        for annotation in document[self.span_layer_name].predictions:
            annotation_id = labeled_span_to_id(annotation)
            payload = self.construct_embedding_payload(document, annotation_id)
            embedding = self.vector_store.get(payload=payload)
            if embedding is not None:
                embeddings[annotation_id] = embedding
        document.metadata["embeddings"] = embeddings

        return document

    def overview(
        self, with_max_cross_doc_sims: bool = False, min_similarity: float = 0.9
    ) -> pd.DataFrame:
        rows = []
        for doc_id, document in self.documents.items():
            layers = {
                caption: document[layer_name]
                for layer_name, caption in self.layer_captions.items()
            }
            if self.use_predictions:
                layers = {caption: layer.predictions for caption, layer in layers.items()}
            layer_sizes = {f"num_{caption}s": len(layer) for caption, layer in layers.items()}
            rows.append({"doc_id": doc_id, **layer_sizes})
        df = pd.DataFrame(rows)

        # add highest cross-document similarity score for each document
        if with_max_cross_doc_sims and len(self.documents) > 1:
            all2all_adu_similarities = self.get_all2all_adu_similarities(
                min_similarity=min_similarity, columns=["doc_id", "other_doc_id", "sim_score"]
            )
            max_doc2doc_similarities = all2all_adu_similarities.pivot_table(
                values="sim_score", index="doc_id", columns="other_doc_id", aggfunc="max"
            )

            max_doc2doc_similarities.sort_index(axis="index", inplace=True)
            max_doc2doc_similarities.sort_index(axis="columns", inplace=True)
            # check that the index and columns are the same
            if (max_doc2doc_similarities.index != max_doc2doc_similarities.columns).any():
                raise gr.Error("Index and columns of max_doc2doc_similarities are not the same.")
            # set diagonal entries to minus infinity to exclude them from the maximum
            np.fill_diagonal(max_doc2doc_similarities.values, -np.inf)

            max_doc_ids = max_doc2doc_similarities.idxmax(axis="columns")
            max_similarities = max_doc2doc_similarities.max(axis="columns")

            # set the index to the doc_id to correctly join the series
            df.set_index("doc_id", inplace=True)
            df["max_cross_doc_sim_doc_id"] = max_doc_ids
            df["max_cross_doc_sim_score"] = max_similarities
            df.reset_index(inplace=True)

        return df

    def as_dict(self, include_embeddings: bool = True) -> dict:
        result = {}
        for doc_id, document in self.documents.items():
            doc_dict = document.asdict()
            if not include_embeddings and "embeddings" in (doc_dict.get("metadata") or {}):
                doc_dict["metadata"] = {
                    k: v for k, v in doc_dict["metadata"].items() if k != "embeddings"
                }
            result[doc_id] = doc_dict
        return result

    def get_all2all_adu_similarities(
        self,
        min_similarity: Optional[float] = 0.5,
        columns: Optional[List[str]] = None,
    ) -> pd.DataFrame:
        """Get the similarities between all ADUs in the store.

        Args:
            min_similarity: The minimum similarity score to consider. If None, all similarities are included.
            columns: The columns to include in the result DataFrame. If None, all columns are included.

        Returns:
            A DataFrame with the columns: doc_id, text, other_doc_id, other_text, sim_score.
        """

        # shape of all_embeddings: (num_embeddings, embedding_dim)
        (
            all_embed_ids,
            all_embeddings,
            all_payloads,
        ) = self.vector_store.as_indices_vectors_payloads()

        doc_id_and_annotation_id2annotation_text = {}
        for doc in self.documents.values():
            for annotation in doc[self.span_layer_name]:
                doc_id_and_annotation_id2annotation_text[
                    (doc.id, labeled_span_to_id(annotation))
                ] = str(annotation)
            for annotation in doc[self.span_layer_name].predictions:
                doc_id_and_annotation_id2annotation_text[
                    (doc.id, labeled_span_to_id(annotation))
                ] = str(annotation)

        # calculate cosine similarities between all embeddings
        dot_prod = np.dot(all_embeddings, all_embeddings.T)
        norm = np.linalg.norm(all_embeddings, axis=1)
        norm_prod = np.outer(norm, norm)
        similarities = dot_prod / norm_prod

        gr.Info(f"Similarities shape: {similarities.shape}")

        if min_similarity is not None:
            gr.Info(f"Filtering similarities below {min_similarity}.")
            # set similarities below min_similarity to 0
            similarities[similarities < min_similarity] = 0.0

        # set triangular part to 0
        similarities = np.triu(similarities, k=1)
        # create a sparse matrix
        sparse_matrix = csr_matrix(similarities)
        sparse_matrix.eliminate_zeros()
        # Get the non-zero values and their indices
        non_zero_idx = sparse_matrix.nonzero()
        scores = sparse_matrix.data

        gr.Info(f"Number of similarities above {min_similarity}: {len(scores)}")

        # construct the DataFrame
        records = []
        for idx1, idx2 in zip(non_zero_idx[0], non_zero_idx[1]):
            if idx1 < idx2:
                doc_id1 = all_payloads[idx1]["doc_id"]
                doc_id2 = all_payloads[idx2]["doc_id"]
                annotation_id1 = all_payloads[idx1]["annotation_id"]
                annotation_id2 = all_payloads[idx2]["annotation_id"]
                annotation_text1 = doc_id_and_annotation_id2annotation_text[
                    (doc_id1, annotation_id1)
                ]
                annotation_text2 = doc_id_and_annotation_id2annotation_text[
                    (doc_id2, annotation_id2)
                ]
                records.append(
                    {
                        "sim_score": scores[idx1],
                        "doc_id": doc_id1,
                        "other_doc_id": doc_id2,
                        "adu_id": annotation_id1,
                        "other_adu_id": annotation_id2,
                        "text": annotation_text1,
                        "other_text": annotation_text2,
                    }
                )
        result_df = pd.DataFrame(records)
        gr.Info(f"DataFrame shape: {result_df.shape}")

        if columns is not None:
            result_df = result_df[columns]
        return result_df