|
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}": 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} |
|
|
