return embeddings from storage retrieval

document_qa/document_qa_engine.py

@@ -1,23 +1,43 @@
 import copy
 import os
 from pathlib import Path
-from typing import Union, Any
+from typing import Union, Any, Optional, List, Dict, Tuple, ClassVar, Collection
 
 import tiktoken
-from grobid_client.grobid_client import GrobidClient
 from langchain.chains import create_extraction_chain
 from langchain.chains.question_answering import load_qa_chain, stuff_prompt, refine_prompts, map_reduce_prompt, \
     map_rerank_prompt
 from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
 from langchain.retrievers import MultiQueryRetriever
 from langchain.schema import Document
-from langchain.vectorstores import Chroma
+from langchain_community.vectorstores.chroma import Chroma, DEFAULT_K
+from langchain_community.vectorstores.faiss import FAISS
+from langchain_core.callbacks import CallbackManagerForRetrieverRun
+from langchain_core.utils import xor_args
+from langchain_core.vectorstores import VectorStore, VectorStoreRetriever
 from tqdm import tqdm
 
 from document_qa.grobid_processors import GrobidProcessor
 
 
+def _results_to_docs_scores_and_embeddings(results: Any) -> List[Tuple[Document, float, List[float]]]:
+    return [
+        (Document(page_content=result[0], metadata=result[1] or {}), result[2], result[3])
+        for result in zip(
+            results["documents"][0],
+            results["metadatas"][0],
+            results["distances"][0],
+            results["embeddings"][0],
+        )
+    ]
+
+
 class TextMerger:
+    """
+    This class tries to replicate the RecursiveTextSplitter from LangChain, to preserve and merge the
+    coordinate information from the PDF document.
+    """
+
     def __init__(self, model_name=None, encoding_name="gpt2"):
         if model_name is not None:
             self.enc = tiktoken.encoding_for_model(model_name)
@@ -85,52 +105,187 @@ class TextMerger:
 
         return new_passages_struct
 
-class DataStorage:
 
+class BaseRetrieval:
 
+    def __init__(
+            self,
+            persist_directory: Path,
+            embedding_function
+    ):
+        self.embedding_function = embedding_function
+        self.persist_directory = persist_directory
+
+
+class AdvancedVectorStoreRetriever(VectorStoreRetriever):
+    allowed_search_types: ClassVar[Collection[str]] = (
+        "similarity",
+        "similarity_score_threshold",
+        "mmr",
+        "similarity_with_embeddings"
+    )
+
+    def _get_relevant_documents(
+            self, query: str, *, run_manager: CallbackManagerForRetrieverRun
+    ) -> List[Document]:
+        if self.search_type == "similarity":
+            docs = self.vectorstore.similarity_search(query, **self.search_kwargs)
+        elif self.search_type == "similarity_score_threshold":
+            docs_and_similarities = (
+                self.vectorstore.similarity_search_with_relevance_scores(
+                    query, **self.search_kwargs
+                )
+            )
+            for doc, similarity in docs_and_similarities:
+                if '__similarity' not in doc.metadata.keys():
+                    doc.metadata['__similarity'] = similarity
+
+            docs = [doc for doc, _ in docs_and_similarities]
+        elif self.search_type == "mmr":
+            docs = self.vectorstore.max_marginal_relevance_search(
+                query, **self.search_kwargs
+            )
+        elif self.search_type == "similarity_with_embeddings":
+            docs_scores_and_embeddings = (
+                self.vectorstore.advanced_similarity_search(
+                    query, **self.search_kwargs
+                )
+            )
 
-class DocumentQAEngine:
-    llm = None
-    qa_chain_type = None
-    embedding_function = None
+            for doc, score, embeddings in docs_scores_and_embeddings:
+                if '__embeddings' not in doc.metadata.keys():
+                    doc.metadata['__embeddings'] = embeddings
+                if '__similarity' not in doc.metadata.keys():
+                    doc.metadata['__similarity'] = score
+
+            docs = [doc for doc, _, _ in docs_scores_and_embeddings]
+        else:
+            raise ValueError(f"search_type of {self.search_type} not allowed.")
+        return docs
+
+
+class AdvancedVectorStore(VectorStore):
+    def as_retriever(self, **kwargs: Any) -> AdvancedVectorStoreRetriever:
+        tags = kwargs.pop("tags", None) or []
+        tags.extend(self._get_retriever_tags())
+        return AdvancedVectorStoreRetriever(vectorstore=self, **kwargs, tags=tags)
+
+
+class ChromaAdvancedRetrieval(Chroma, AdvancedVectorStore):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    @xor_args(("query_texts", "query_embeddings"))
+    def __query_collection(
+            self,
+            query_texts: Optional[List[str]] = None,
+            query_embeddings: Optional[List[List[float]]] = None,
+            n_results: int = 4,
+            where: Optional[Dict[str, str]] = None,
+            where_document: Optional[Dict[str, str]] = None,
+            **kwargs: Any,
+    ) -> List[Document]:
+        """Query the chroma collection."""
+        try:
+            import chromadb  # noqa: F401
+        except ImportError:
+            raise ValueError(
+                "Could not import chromadb python package. "
+                "Please install it with `pip install chromadb`."
+            )
+        return self._collection.query(
+            query_texts=query_texts,
+            query_embeddings=query_embeddings,
+            n_results=n_results,
+            where=where,
+            where_document=where_document,
+            **kwargs,
+        )
+
+    def advanced_similarity_search(
+            self,
+            query: str,
+            k: int = DEFAULT_K,
+            filter: Optional[Dict[str, str]] = None,
+            **kwargs: Any,
+    ) -> [List[Document], float, List[float]]:
+        docs_scores_and_embeddings = self.similarity_search_with_scores_and_embeddings(query, k, filter=filter)
+        return docs_scores_and_embeddings
+
+    def similarity_search_with_scores_and_embeddings(
+            self,
+            query: str,
+            k: int = DEFAULT_K,
+            filter: Optional[Dict[str, str]] = None,
+            where_document: Optional[Dict[str, str]] = None,
+            **kwargs: Any,
+    ) -> List[Tuple[Document, float, List[float]]]:
+
+        if self._embedding_function is None:
+            results = self.__query_collection(
+                query_texts=[query],
+                n_results=k,
+                where=filter,
+                where_document=where_document,
+                include=['metadatas', 'documents', 'embeddings', 'distances']
+            )
+        else:
+            query_embedding = self._embedding_function.embed_query(query)
+            results = self.__query_collection(
+                query_embeddings=[query_embedding],
+                n_results=k,
+                where=filter,
+                where_document=where_document,
+                include=['metadatas', 'documents', 'embeddings', 'distances']
+            )
+
+        return _results_to_docs_scores_and_embeddings(results)
+
+
+class FAISSAdvancedRetrieval(FAISS):
+    pass
+
+
+class NER_Retrival(VectorStore):
+    """
+    This class implement a retrieval based on NER models.
+    This is an alternative retrieval to embeddings that relies on extracted entities.
+    """
+    pass
+
+
+engines = {
+    'chroma': ChromaAdvancedRetrieval,
+    'faiss': FAISSAdvancedRetrieval,
+    'ner': NER_Retrival
+}
+
+
+class DataStorage:
     embeddings_dict = {}
     embeddings_map_from_md5 = {}
     embeddings_map_to_md5 = {}
 
-    default_prompts = {
-        'stuff': stuff_prompt,
-        'refine': refine_prompts,
-        "map_reduce": map_reduce_prompt,
-        "map_rerank": map_rerank_prompt
-    }
-
-    def __init__(self,
-                 llm,
-                 embedding_function,
-                 qa_chain_type="stuff",
-                 embeddings_root_path=None,
-                 grobid_url=None,
-                 memory=None
-                 ):
+    def __init__(
+            self,
+            embedding_function,
+            root_path: Path = None,
+            engine=ChromaAdvancedRetrieval,
+    ) -> None:
+        self.root_path = root_path
+        self.engine = engine
         self.embedding_function = embedding_function
-        self.llm = llm
-        self.memory = memory
-        self.chain = load_qa_chain(llm, chain_type=qa_chain_type)
-        self.text_merger = TextMerger()
 
-        if embeddings_root_path is not None:
-            self.embeddings_root_path = embeddings_root_path
-            if not os.path.exists(embeddings_root_path):
-                os.makedirs(embeddings_root_path)
+        if root_path is not None:
+            self.embeddings_root_path = root_path
+            if not os.path.exists(root_path):
+                os.makedirs(root_path)
             else:
                 self.load_embeddings(self.embeddings_root_path)
 
-        if grobid_url:
-            self.grobid_processor = GrobidProcessor(grobid_url)
-
     def load_embeddings(self, embeddings_root_path: Union[str, Path]) -> None:
         """
-        Load the embeddings assuming they are all persisted and stored in a single directory.
+        Load the vector storage assuming they are all persisted and stored in a single directory.
         The root path of the embeddings containing one data store for each document in each subdirectory
         """
 
@@ -141,8 +296,10 @@ class DocumentQAEngine:
             return
 
         for embedding_document_dir in embeddings_directories:
-            self.embeddings_dict[embedding_document_dir.name] = Chroma(persist_directory=embedding_document_dir.path,
-                                                                       embedding_function=self.embedding_function)
+            self.embeddings_dict[embedding_document_dir.name] = self.engine(
+                persist_directory=embedding_document_dir.path,
+                embedding_function=self.embedding_function
+            )
 
             filename_list = list(Path(embedding_document_dir).glob('*.storage_filename'))
             if filename_list:
@@ -161,9 +318,60 @@ class DocumentQAEngine:
     def get_filename_from_md5(self, md5):
         return self.embeddings_map_from_md5[md5]
 
-    def query_document(self, query: str, doc_id, output_parser=None, context_size=4, extraction_schema=None,
-                       verbose=False) -> (
-            Any, str):
+    def embed_document(self, doc_id, texts, metadatas):
+        if doc_id not in self.embeddings_dict.keys():
+            self.embeddings_dict[doc_id] = self.engine.from_texts(texts,
+                                                                  embedding=self.embedding_function,
+                                                                  metadatas=metadatas,
+                                                                  collection_name=doc_id)
+        else:
+            # Workaround Chroma (?) breaking change
+            self.embeddings_dict[doc_id].delete_collection()
+            self.embeddings_dict[doc_id] = self.engine.from_texts(texts,
+                                                                  embedding=self.embedding_function,
+                                                                  metadatas=metadatas,
+                                                                  collection_name=doc_id)
+
+        self.embeddings_root_path = None
+
+
+class DocumentQAEngine:
+    llm = None
+    qa_chain_type = None
+
+    default_prompts = {
+        'stuff': stuff_prompt,
+        'refine': refine_prompts,
+        "map_reduce": map_reduce_prompt,
+        "map_rerank": map_rerank_prompt
+    }
+
+    def __init__(self,
+                 llm,
+                 data_storage: DataStorage,
+                 qa_chain_type="stuff",
+                 grobid_url=None,
+                 memory=None
+                 ):
+
+        self.llm = llm
+        self.memory = memory
+        self.chain = load_qa_chain(llm, chain_type=qa_chain_type)
+        self.text_merger = TextMerger()
+        self.data_storage = data_storage
+
+        if grobid_url:
+            self.grobid_processor = GrobidProcessor(grobid_url)
+
+    def query_document(
+            self,
+            query: str,
+            doc_id,
+            output_parser=None,
+            context_size=4,
+            extraction_schema=None,
+            verbose=False
+    ) -> (Any, str):
         # self.load_embeddings(self.embeddings_root_path)
 
         if verbose:
@@ -192,16 +400,22 @@ class DocumentQAEngine:
         else:
             return None, response, coordinates
 
-    def query_storage(self, query: str, doc_id, context_size=4):
-        documents = self._get_context(doc_id, query, context_size)
+    def query_storage(self, query: str, doc_id, context_size=4) -> (List[Document], list):
+        """
+        Returns the context related to a given query
+        """
+        documents, coordinates = self._get_context(doc_id, query, context_size)
 
         context_as_text = [doc.page_content for doc in documents]
-        return context_as_text
+        return context_as_text, coordinates
 
     def query_storage_and_embeddings(self, query: str, doc_id, context_size=4):
-        db = self.embeddings_dict[doc_id]
-        retriever = db.as_retriever(search_kwargs={"k": context_size})
-        relevant_documents = retriever.get_relevant_documents(query, include=["embeddings"])
+        """
+        Returns both the context and the embedding information from a given query
+        """
+        db = self.data_storage.embeddings_dict[doc_id]
+        retriever = db.as_retriever(search_kwargs={"k": context_size}, search_type="similarity_with_embeddings")
+        relevant_documents = retriever.get_relevant_documents(query)
 
         context_as_text = [doc.page_content for doc in relevant_documents]
         return context_as_text
@@ -229,11 +443,11 @@ class DocumentQAEngine:
 
         return parsed_output
 
-    def _run_query(self, doc_id, query, context_size=4):
+    def _run_query(self, doc_id, query, context_size=4) -> (List[Document], list):
         relevant_documents = self._get_context(doc_id, query, context_size)
         relevant_document_coordinates = [doc.metadata['coordinates'].split(";") if 'coordinates' in doc.metadata else []
                                          for doc in
-                                         relevant_documents]  # filter(lambda d: d['type'] == "sentence", relevant_documents)]
+                                         relevant_documents]
         response = self.chain.run(input_documents=relevant_documents,
                                   question=query)
 
@@ -241,33 +455,40 @@ class DocumentQAEngine:
             self.memory.save_context({"input": query}, {"output": response})
         return response, relevant_document_coordinates
 
-    def _get_context(self, doc_id, query, context_size=4):
-        db = self.embeddings_dict[doc_id]
+    def _get_context(self, doc_id, query, context_size=4) -> (List[Document], list):
+        db = self.data_storage.embeddings_dict[doc_id]
         retriever = db.as_retriever(search_kwargs={"k": context_size})
         relevant_documents = retriever.get_relevant_documents(query)
+        relevant_document_coordinates = [doc.metadata['coordinates'].split(";") if 'coordinates' in doc.metadata else []
+                                         for doc in
+                                         relevant_documents]
         if self.memory and len(self.memory.buffer_as_messages) > 0:
             relevant_documents.append(
                 Document(
                     page_content="""Following, the previous question and answers. Use these information only when in the question there are unspecified references:\n{}\n\n""".format(
                         self.memory.buffer_as_str))
             )
-        return relevant_documents
+        return relevant_documents, relevant_document_coordinates
 
-    def get_all_context_by_document(self, doc_id):
-        """Return the full context from the document"""
-        db = self.embeddings_dict[doc_id]
+    def get_full_context_by_document(self, doc_id):
+        """
+        Return the full context from the document
+        """
+        db = self.data_storage.embeddings_dict[doc_id]
         docs = db.get()
         return docs['documents']
 
     def _get_context_multiquery(self, doc_id, query, context_size=4):
-        db = self.embeddings_dict[doc_id].as_retriever(search_kwargs={"k": context_size})
+        db = self.data_storage.embeddings_dict[doc_id].as_retriever(search_kwargs={"k": context_size})
         multi_query_retriever = MultiQueryRetriever.from_llm(retriever=db, llm=self.llm)
         relevant_documents = multi_query_retriever.get_relevant_documents(query)
         return relevant_documents
 
     def get_text_from_document(self, pdf_file_path, chunk_size=-1, perc_overlap=0.1, verbose=False):
         """
-        Extract text from documents using Grobid, if chunk_size is < 0 it keeps each paragraph separately
+        Extract text from documents using Grobid.
+        - if chunk_size is < 0, keeps each paragraph separately
+        - if chunk_size > 0, aggregate all paragraphs and split them again using an approximate chunk size
         """
         if verbose:
             print("File", pdf_file_path)
@@ -307,7 +528,13 @@ class DocumentQAEngine:
 
         return texts, metadatas, ids
 
-    def create_memory_embeddings(self, pdf_path, doc_id=None, chunk_size=500, perc_overlap=0.1):
+    def create_memory_embeddings(
+            self,
+            pdf_path,
+            doc_id=None,
+            chunk_size=500,
+            perc_overlap=0.1
+    ):
         texts, metadata, ids = self.get_text_from_document(
             pdf_path,
             chunk_size=chunk_size,
@@ -317,25 +544,17 @@ class DocumentQAEngine:
         else:
             hash = metadata[0]['hash']
 
-        if hash not in self.embeddings_dict.keys():
-            self.embeddings_dict[hash] = Chroma.from_texts(texts,
-                                                           embedding=self.embedding_function,
-                                                           metadatas=metadata,
-                                                           collection_name=hash)
-        else:
-            # if 'documents' in self.embeddings_dict[hash].get() and len(self.embeddings_dict[hash].get()['documents']) == 0:
-            #     self.embeddings_dict[hash].delete(ids=self.embeddings_dict[hash].get()['ids'])
-            self.embeddings_dict[hash].delete_collection()
-            self.embeddings_dict[hash] = Chroma.from_texts(texts,
-                                                           embedding=self.embedding_function,
-                                                           metadatas=metadata,
-                                                           collection_name=hash)
-
-        self.embeddings_root_path = None
+        self.data_storage.embed_document(hash, texts, metadata)
 
         return hash
 
-    def create_embeddings(self, pdfs_dir_path: Path, chunk_size=500, perc_overlap=0.1, include_biblio=False):
+    def create_embeddings(
+            self,
+            pdfs_dir_path: Path,
+            chunk_size=500,
+            perc_overlap=0.1,
+            include_biblio=False
+    ):
         input_files = []
         for root, dirs, files in os.walk(pdfs_dir_path, followlinks=False):
             for file_ in files:
@@ -347,17 +566,16 @@ class DocumentQAEngine:
                                desc="Grobid + embeddings processing"):
 
             md5 = self.calculate_md5(input_file)
-            data_path = os.path.join(self.embeddings_root_path, md5)
+            data_path = os.path.join(self.data_storage.embeddings_root_path, md5)
 
             if os.path.exists(data_path):
                 print(data_path, "exists. Skipping it ")
                 continue
-            include = ["biblio"] if include_biblio else []
+            # include = ["biblio"] if include_biblio else []
             texts, metadata, ids = self.get_text_from_document(
                 input_file,
                 chunk_size=chunk_size,
-                perc_overlap=perc_overlap,
-                include=include)
+                perc_overlap=perc_overlap)
             filename = metadata[0]['filename']
 
             vector_db_document = Chroma.from_texts(texts,

requirements.txt

@@ -4,10 +4,10 @@ grobid-client-python==0.0.7
 grobid_tei_xml==0.1.3
 
 # Utils
-tqdm==4.66.1
+tqdm==4.66.2
 pyyaml==6.0.1
-pytest==7.4.3
-streamlit==1.29.0
+pytest==8.1.1
+streamlit==1.33.0
 lxml
 Beautifulsoup4
 python-dotenv
@@ -15,13 +15,13 @@ watchdog
 dateparser
 
 # LLM
-chromadb==0.4.19
-tiktoken==0.4.0
-openai==0.27.7
-langchain==0.0.350
-langchain-core==0.1.0
+chromadb==0.4.24
+tiktoken==0.6.0
+openai==1.16.2
+langchain==0.1.14
+langchain-core==0.1.40
 typing-inspect==0.9.0
-typing_extensions==4.8.0
-pydantic==2.4.2
-sentence_transformers==2.2.2
+typing_extensions==4.11.0
+pydantic==2.6.4
+sentence_transformers==2.6.1
 streamlit-pdf-viewer

streamlit_app.py

@@ -9,15 +9,16 @@ from langchain.llms.huggingface_hub import HuggingFaceHub
 from langchain.memory import ConversationBufferWindowMemory
 from streamlit_pdf_viewer import pdf_viewer
 
+from document_qa.ner_client_generic import NERClientGeneric
+
 dotenv.load_dotenv(override=True)
 
 import streamlit as st
 from langchain.chat_models import ChatOpenAI
 from langchain.embeddings import OpenAIEmbeddings, HuggingFaceEmbeddings
 
-from document_qa.document_qa_engine import DocumentQAEngine
+from document_qa.document_qa_engine import DocumentQAEngine, DataStorage
 from document_qa.grobid_processors import GrobidAggregationProcessor, decorate_text_with_annotations
-from grobid_client_generic import GrobidClientGeneric
 
 OPENAI_MODELS = ['gpt-3.5-turbo',
                  "gpt-4",
@@ -168,14 +169,15 @@ def init_qa(model, api_key=None):
         st.stop()
         return
 
-    return DocumentQAEngine(chat, embeddings, grobid_url=os.environ['GROBID_URL'], memory=st.session_state['memory'])
+    storage = DataStorage(embeddings)
+    return DocumentQAEngine(chat, storage, grobid_url=os.environ['GROBID_URL'], memory=st.session_state['memory'])
 
 
 @st.cache_resource
 def init_ner():
     quantities_client = QuantitiesAPI(os.environ['GROBID_QUANTITIES_URL'], check_server=True)
 
-    materials_client = GrobidClientGeneric(ping=True)
+    materials_client = NERClientGeneric(ping=True)
     config_materials = {
         'grobid': {
             "server": os.environ['GROBID_MATERIALS_URL'],
@@ -190,10 +192,8 @@ def init_ner():
 
     materials_client.set_config(config_materials)
 
-    gqa = GrobidAggregationProcessor(None,
-                                     grobid_quantities_client=quantities_client,
-                                     grobid_superconductors_client=materials_client
-                                     )
+    gqa = GrobidAggregationProcessor(grobid_quantities_client=quantities_client,
+                                     grobid_superconductors_client=materials_client)
     return gqa
 
 
@@ -340,9 +340,12 @@ with st.sidebar:
 
     st.session_state['pdf_rendering'] = st.radio(
         "PDF rendering mode",
-        {"PDF.JS", "Native browser engine"},
-        index=1,
+        ("unwrap", "legacy_embed"),
+        index=0,
         disabled=not uploaded_file,
+        help="PDF rendering engine."
+             "Note: The Legacy PDF viewer does not support annotations and might not work on Chrome.",
+        format_func=lambda q: "Legacy PDF Viewer" if q == "legacy_embed" else "Streamlit PDF Viewer (Pdf.js)"
     )
 
     st.divider()
@@ -441,7 +444,8 @@ with right_column:
         text_response = None
         if mode == "Embeddings":
             with st.spinner("Generating LLM response..."):
-                text_response = st.session_state['rqa'][model].query_storage_and_embeddings(question, st.session_state.doc_id,
+                text_response, coordinates = st.session_state['rqa'][model].query_storage(question,
+                                                                             st.session_state.doc_id,
                                                                              context_size=context_size)
         elif mode == "LLM":
             with st.spinner("Generating response..."):
@@ -449,14 +453,14 @@ with right_column:
                                                                                               st.session_state.doc_id,
                                                                                               context_size=context_size)
 
-                annotations = [[GrobidAggregationProcessor.box_to_dict([cs for cs in c.split(",")]) for c in coord_doc]
-                               for coord_doc in coordinates]
-                gradients = generate_color_gradient(len(annotations))
-                for i, color in enumerate(gradients):
-                    for annotation in annotations[i]:
-                        annotation['color'] = color
-                st.session_state['annotations'] = [annotation for annotation_doc in annotations for annotation in
-                                                   annotation_doc]
+        annotations = [[GrobidAggregationProcessor.box_to_dict([cs for cs in c.split(",")]) for c in coord_doc]
+                       for coord_doc in coordinates]
+        gradients = generate_color_gradient(len(annotations))
+        for i, color in enumerate(gradients):
+            for annotation in annotations[i]:
+                annotation['color'] = color
+        st.session_state['annotations'] = [annotation for annotation_doc in annotations for annotation in
+                                           annotation_doc]
 
         if not text_response:
             st.error("Something went wrong. Contact Luca Foppiano (Foppiano.Luca@nims.co.jp) to report the issue.")
@@ -486,5 +490,5 @@ with left_column:
             height=800,
             annotation_outline_size=1,
             annotations=st.session_state['annotations'],
-            rendering='unwrap' if st.session_state['pdf_rendering'] == 'PDF.JS' else 'legacy_embed'
+            rendering=st.session_state['pdf_rendering']
         )