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, TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions, 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 in ["labeled_spans", "labeled_multi_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 not isinstance(document, self.document_type): raise gr.Error( f"The document to add must be of type {self.document_type}, but is of type " f"{type(document)}." ) 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") # entries where max_similarities is -inf are documents with no entries in other documents # with similarity > min_similarity mask = max_similarities == -np.inf max_doc_ids[mask] = np.nan max_similarities[mask] = np.nan # 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 because we only want the upper triangular part which # contains entries with idx1 < idx2 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 sparse_idx, (idx1, idx2) in enumerate(zip(non_zero_idx[0], non_zero_idx[1])): payload1 = all_payloads[idx1] payload2 = all_payloads[idx2] doc_id1 = payload1["doc_id"] doc_id2 = payload2["doc_id"] annotation_id1 = payload1["annotation_id"] annotation_id2 = payload2["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[sparse_idx], "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