update libraries and code to support Modal as OpenAI-based server

document_qa/document_qa_engine.py

@@ -5,7 +5,8 @@ from typing import Union, Any, List
 
 import tiktoken
 from langchain.chains import create_extraction_chain
-from langchain.chains.question_answering import load_qa_chain, stuff_prompt, refine_prompts, map_reduce_prompt, \
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain.chains.question_answering import stuff_prompt, refine_prompts, map_reduce_prompt, \
     map_rerank_prompt
 from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTemplate, ChatPromptTemplate
 from langchain.retrievers import MultiQueryRetriever
@@ -14,7 +15,6 @@ from langchain_community.vectorstores.chroma import Chroma
 from langchain_core.vectorstores import VectorStore
 from tqdm import tqdm
 
-# from document_qa.embedding_visualiser import QueryVisualiser
 from document_qa.grobid_processors import GrobidProcessor
 from document_qa.langchain import ChromaAdvancedRetrieval
 
@@ -177,17 +177,19 @@ class DataStorage:
 
     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)
+            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_dict[doc_id] = self.engine.from_texts(
+                texts,
+                embedding=self.embedding_function,
+                metadatas=metadatas,
+                collection_name=doc_id)
 
         self.embeddings_root_path = None
 
@@ -206,14 +208,13 @@ class DocumentQAEngine:
     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.chain = create_stuff_documents_chain(llm, self.default_prompts['stuff'].PROMPT)
         self.text_merger = TextMerger()
         self.data_storage = data_storage
 
@@ -271,7 +272,10 @@ class DocumentQAEngine:
         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")
+        retriever = db.as_retriever(
+            search_kwargs={"k": context_size},
+            search_type="similarity_with_embeddings"
+        )
         relevant_documents = retriever.invoke(query)
 
         return relevant_documents
@@ -327,20 +331,18 @@ class DocumentQAEngine:
 
     def _run_query(self, doc_id, query, context_size=4) -> (List[Document], list):
         relevant_documents, relevant_document_coordinates = self._get_context(doc_id, query, context_size)
-        response = self.chain.run(input_documents=relevant_documents,
-                                  question=query)
-
-        if self.memory:
-            self.memory.save_context({"input": query}, {"output": response})
+        response = self.chain.invoke({"context": relevant_documents, "question": query})
         return response, relevant_document_coordinates
 
     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.invoke(query)
-        relevant_document_coordinates = [doc.metadata['coordinates'].split(";") if 'coordinates' in doc.metadata else []
-                                         for doc in
-                                         relevant_documents]
+        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(

document_qa/langchain.py

@@ -1,4 +1,3 @@
-from pathlib import Path
 from typing import Any, Optional, List, Dict, Tuple, ClassVar, Collection
 
 from langchain.schema import Document

requirements.txt

@@ -1,32 +1,33 @@
 # Grobid
 grobid-quantities-client==0.4.0
 grobid-client-python==0.0.9
-grobid_tei_xml==0.1.3
+grobid-tei-xml==0.1.3
 
 # Utils
-tqdm==4.66.2
+tqdm==4.66.3
 pyyaml==6.0.1
 pytest==8.1.1
-streamlit==1.37.1
-lxml
-Beautifulsoup4
-python-dotenv
-watchdog
-dateparser
+streamlit==1.45.1
+lxml==5.2.1
+beautifulsoup4==4.12.3
+python-dotenv==1.0.1
+watchdog==4.0.0
+dateparser==1.2.0
+requests>=2.31.0
 
 # LLM
 chromadb==0.4.24
-tiktoken==0.7.0
-openai==1.42.0
-langchain==0.2.14
-langchain-core==0.2.34
-langchain-openai==0.1.22
-langchain-huggingface==0.0.3
-langchain-community==0.2.12
+tiktoken==0.9.0
+openai==1.82.0
+langchain==0.3.25
+langchain-core==0.3.61
+langchain-openai==0.3.18
+langchain-huggingface==0.2.0
+langchain-community==0.3.21
 typing-inspect==0.9.0
-typing_extensions==4.11.0
-pydantic==2.6.4
-sentence_transformers==2.6.1
-streamlit-pdf-viewer==0.0.22
-umap-learn
-plotly
+typing_extensions==4.12.2
+pydantic==2.10.6
+sentence-transformers==2.6.1
+streamlit-pdf-viewer==0.0.22rc0
+umap-learn==0.5.6
+plotly==5.20.0

streamlit_app.py

@@ -5,11 +5,9 @@ from tempfile import NamedTemporaryFile
 
 import dotenv
 from grobid_quantities.quantities import QuantitiesAPI
-from langchain.memory import ConversationBufferWindowMemory
-from langchain_community.chat_models import ChatOpenAI
-from langchain_community.llms.huggingface_endpoint import HuggingFaceEndpoint
+from langchain.memory import ConversationBufferMemory
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_openai import OpenAIEmbeddings
+from langchain_openai import ChatOpenAI
 from streamlit_pdf_viewer import pdf_viewer
 
 from document_qa.ner_client_generic import NERClientGeneric
@@ -20,30 +18,14 @@ import streamlit as st
 from document_qa.document_qa_engine import DocumentQAEngine, DataStorage
 from document_qa.grobid_processors import GrobidAggregationProcessor, decorate_text_with_annotations
 
-OPENAI_MODELS = ['gpt-3.5-turbo',
-                 "gpt-4",
-                 "gpt-4-1106-preview"]
-
-OPENAI_EMBEDDINGS = [
-    'text-embedding-ada-002',
-    'text-embedding-3-large',
-    'openai-text-embedding-3-small'
-]
-
-OPEN_MODELS = {
-    'Mistral-Nemo-Instruct-2407': 'mistralai/Mistral-Nemo-Instruct-2407',
-    'mistral-7b-instruct-v0.3': 'mistralai/Mistral-7B-Instruct-v0.3',
-    'Phi-3-mini-4k-instruct': "microsoft/Phi-3-mini-4k-instruct"
+API_MODELS = {
+    "microsoft/Phi-4-mini-instruct": os.environ["MODAL_1_URL"]
 }
 
-DEFAULT_OPEN_EMBEDDING_NAME = 'Default (all-MiniLM-L6-v2)'
-OPEN_EMBEDDINGS = {
-    DEFAULT_OPEN_EMBEDDING_NAME: 'all-MiniLM-L6-v2',
-    'SFR-Embedding-Mistral': 'Salesforce/SFR-Embedding-Mistral',
-    'SFR-Embedding-2_R': 'Salesforce/SFR-Embedding-2_R',
-    'NV-Embed': 'nvidia/NV-Embed-v1',
-    'e5-mistral-7b-instruct': 'intfloat/e5-mistral-7b-instruct',
-    'gte-large-en-v1.5': 'Alibaba-NLP/gte-large-en-v1.5'
+API_EMBEDDINGS = {
+    'intfloat/e5-large-v2': 'intfloat/e5-large-v2',
+    'intfloat/multilingual-e5-large-instruct': 'intfloat/multilingual-e5-large-instruct:',
+    'Salesforce/SFR-Embedding-2_R': 'Salesforce/SFR-Embedding-2_R'
 }
 
 if 'rqa' not in st.session_state:
@@ -141,48 +123,20 @@ def clear_memory():
 
 
 # @st.cache_resource
-def init_qa(model, embeddings_name=None, api_key=None):
-    ## For debug add: callbacks=[PromptLayerCallbackHandler(pl_tags=["langchain", "chatgpt", "document-qa"])])
-    if model in OPENAI_MODELS:
-        if embeddings_name is None:
-            embeddings_name = 'text-embedding-ada-002'
-
-        st.session_state['memory'] = ConversationBufferWindowMemory(k=4)
-        if api_key:
-            chat = ChatOpenAI(model_name=model,
-                              temperature=0,
-                              openai_api_key=api_key,
-                              frequency_penalty=0.1)
-            if embeddings_name not in OPENAI_EMBEDDINGS:
-                st.error(f"The embeddings provided {embeddings_name} are not supported by this model {model}.")
-                st.stop()
-                return
-            embeddings = OpenAIEmbeddings(model=embeddings_name, openai_api_key=api_key)
+def init_qa(model_name, embeddings_name):
+    st.session_state['memory'] = ConversationBufferMemory(
+        memory_key="chat_history",
+        return_messages=True
+    )
+    chat = ChatOpenAI(
+        model=model_name,
+        temperature=0.0,
+        base_url=API_MODELS[model_name],
+        api_key=os.environ.get('API_KEY')
+    )
 
-        else:
-            chat = ChatOpenAI(model_name=model,
-                              temperature=0,
-                              frequency_penalty=0.1)
-            embeddings = OpenAIEmbeddings(model=embeddings_name)
-
-    elif model in OPEN_MODELS:
-        if embeddings_name is None:
-            embeddings_name = DEFAULT_OPEN_EMBEDDING_NAME
-
-        chat = HuggingFaceEndpoint(
-            repo_id=OPEN_MODELS[model],
-            temperature=0.01,
-            max_new_tokens=4092,
-            model_kwargs={"max_length": 8192},
-            # callbacks=[PromptLayerCallbackHandler(pl_tags=[model, "document-qa"])]
-        )
-        embeddings = HuggingFaceEmbeddings(
-            model_name=OPEN_EMBEDDINGS[embeddings_name])
-        # st.session_state['memory'] = ConversationBufferWindowMemory(k=4) if model not in DISABLE_MEMORY else None
-    else:
-        st.error("The model was not loaded properly. Try reloading. ")
-        st.stop()
-        return
+    embeddings = HuggingFaceEmbeddings(
+        model_name=API_EMBEDDINGS[embeddings_name])
 
     storage = DataStorage(embeddings)
     return DocumentQAEngine(chat, storage, grobid_url=os.environ['GROBID_URL'], memory=st.session_state['memory'])
@@ -246,65 +200,31 @@ with st.sidebar:
     st.divider()
     st.session_state['model'] = model = st.selectbox(
         "Model:",
-        options=OPENAI_MODELS + list(OPEN_MODELS.keys()),
-        index=(OPENAI_MODELS + list(OPEN_MODELS.keys())).index(
+        options=API_MODELS.keys(),
+        index=(list(API_MODELS.keys())).index(
             os.environ["DEFAULT_MODEL"]) if "DEFAULT_MODEL" in os.environ and os.environ["DEFAULT_MODEL"] else 0,
         placeholder="Select model",
         help="Select the LLM model:",
         disabled=st.session_state['doc_id'] is not None or st.session_state['uploaded']
     )
-    embedding_choices = OPENAI_EMBEDDINGS if model in OPENAI_MODELS else OPEN_EMBEDDINGS
 
     st.session_state['embeddings'] = embedding_name = st.selectbox(
         "Embeddings:",
-        options=embedding_choices,
-        index=0,
+        options=API_EMBEDDINGS.keys(),
+        index=(list(API_EMBEDDINGS.keys())).index(
+            os.environ["DEFAULT_EMBEDDING"]) if "DEFAULT_EMBEDDING" in os.environ and os.environ[
+            "DEFAULT_EMBEDDING"] else 0,
         placeholder="Select embedding",
         help="Select the Embedding function:",
         disabled=st.session_state['doc_id'] is not None or st.session_state['uploaded']
     )
 
-    if (model in OPEN_MODELS) and model not in st.session_state['api_keys']:
-        if 'HUGGINGFACEHUB_API_TOKEN' not in os.environ:
-            api_key = st.text_input('Huggingface API Key', type="password")
+    api_key = os.environ['API_KEY']
 
-            st.markdown("Get it [here](https://huggingface.co/docs/hub/security-tokens)")
-        else:
-            api_key = os.environ['HUGGINGFACEHUB_API_TOKEN']
-
-        if api_key:
-            # st.session_state['api_key'] = is_api_key_provided = True
-            if model not in st.session_state['rqa'] or model not in st.session_state['api_keys']:
-                with st.spinner("Preparing environment"):
-                    st.session_state['api_keys'][model] = api_key
-                    # if 'HUGGINGFACEHUB_API_TOKEN' not in os.environ:
-                    #     os.environ["HUGGINGFACEHUB_API_TOKEN"] = api_key
-                    st.session_state['rqa'][model] = init_qa(model, embedding_name)
-
-    elif model in OPENAI_MODELS and model not in st.session_state['api_keys']:
-        if 'OPENAI_API_KEY' not in os.environ:
-            api_key = st.text_input('OpenAI API Key', type="password")
-            st.markdown("Get it [here](https://platform.openai.com/account/api-keys)")
-        else:
-            api_key = os.environ['OPENAI_API_KEY']
-
-        if api_key:
-            if model not in st.session_state['rqa'] or model not in st.session_state['api_keys']:
-                with st.spinner("Preparing environment"):
-                    st.session_state['api_keys'][model] = api_key
-                    if 'OPENAI_API_KEY' not in os.environ:
-                        st.session_state['rqa'][model] = init_qa(model, st.session_state['embeddings'], api_key)
-                    else:
-                        st.session_state['rqa'][model] = init_qa(model, st.session_state['embeddings'])
-    # else:
-    #     is_api_key_provided = st.session_state['api_key']
-
-    # st.button(
-    #     'Reset chat memory.',
-    #     key="reset-memory-button",
-    #     on_click=clear_memory,
-    #     help="Clear the conversational memory. Currently implemented to retrain the 4 most recent messages.",
-    #     disabled=model in st.session_state['rqa'] and st.session_state['rqa'][model].memory is None)
+    if model not in st.session_state['rqa'] or model not in st.session_state['api_keys']:
+        with st.spinner("Preparing environment"):
+            st.session_state['rqa'][model] = init_qa(model, st.session_state['embeddings'])
+            st.session_state['api_keys'][model] = api_key
 
 left_column, right_column = st.columns([5, 4])
 right_column = right_column.container(border=True)
@@ -390,15 +310,16 @@ if uploaded_file and not st.session_state.loaded_embeddings:
         st.stop()
 
     with left_column:
-        with st.spinner('Reading file, calling Grobid, and creating memory embeddings...'):
+        with st.spinner('Reading file, calling Grobid, and creating in-memory embeddings...'):
             binary = uploaded_file.getvalue()
             tmp_file = NamedTemporaryFile()
             tmp_file.write(bytearray(binary))
             st.session_state['binary'] = binary
 
-            st.session_state['doc_id'] = hash = st.session_state['rqa'][model].create_memory_embeddings(tmp_file.name,
-                                                                                                        chunk_size=chunk_size,
-                                                                                                        perc_overlap=0.1)
+            st.session_state['doc_id'] = hash = st.session_state['rqa'][model].create_memory_embeddings(
+                tmp_file.name,
+                chunk_size=chunk_size,
+                perc_overlap=0.1)
             st.session_state['loaded_embeddings'] = True
             st.session_state.messages = []
 
@@ -477,7 +398,7 @@ with right_column:
                                            annotation_doc]
 
         if not text_response:
-            st.error("Something went wrong. Contact Luca Foppiano (Foppiano.Luca@nims.co.jp) to report the issue.")
+            st.error("Something went wrong. Contact info AT sciencialab.com to report the issue through GitHub.")
 
         if mode == "llm":
             if st.session_state['ner_processing']:
@@ -503,5 +424,6 @@ with left_column:
                 annotation_outline_size=2,
                 annotations=st.session_state['annotations'] if st.session_state['annotations'] else [],
                 render_text=True,
-                scroll_to_annotation=1 if (st.session_state['annotations'] and st.session_state['scroll_to_first_annotation']) else None
+                scroll_to_annotation=1 if (st.session_state['annotations'] and st.session_state[
+                    'scroll_to_first_annotation']) else None
             )