Clean code and add readme

LISA_mini.ipynb

@@ -1,8 +1,16 @@
 {
  "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "9267529d",
+   "metadata": {},
+   "source": [
+    "A mini version of LISA in a Jupyter notebook for easier testing and playing around."
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "id": "adcfdba2",
    "metadata": {},
    "outputs": [],
@@ -18,14 +26,13 @@
     "from langchain.chains import ConversationalRetrievalChain\n",
     "from langchain.llms import HuggingFaceTextGenInference\n",
     "from langchain.chains.conversation.memory import (\n",
-    "    ConversationBufferMemory,\n",
     "    ConversationBufferWindowMemory,\n",
     ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 3,
    "id": "2d85c6d9",
    "metadata": {},
    "outputs": [],
@@ -68,7 +75,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 5,
    "id": "2d5bacd5",
    "metadata": {},
    "outputs": [],
@@ -107,7 +114,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 6,
    "id": "8cd31248",
    "metadata": {},
    "outputs": [],
@@ -140,21 +147,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
-   "id": "73d560de",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Create retrievers"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 8,
    "id": "e5796990",
    "metadata": {},
    "outputs": [],
    "source": [
+    "# Create retrievers\n",
     "# Some advanced RAG, with parent document retriever, hybrid-search and rerank\n",
     "\n",
     "# 1. ParentDocumentRetriever. Note: this will take a long time (~several minutes)\n",
@@ -178,7 +176,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 9,
    "id": "bc299740",
    "metadata": {},
    "outputs": [],
@@ -191,7 +189,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 10,
    "id": "2eb8bc8f",
    "metadata": {},
    "outputs": [],
@@ -214,7 +212,7 @@
     "\n",
     "from sentence_transformers import CrossEncoder\n",
     "\n",
-    "model_name = \"BAAI/bge-reranker-large\"  #\n",
+    "model_name = \"BAAI/bge-reranker-large\"\n",
     "\n",
     "class BgeRerank(BaseDocumentCompressor):\n",
     "    model_name:str = model_name\n",
@@ -273,7 +271,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 11,
    "id": "af780912",
    "metadata": {},
    "outputs": [],
@@ -283,7 +281,7 @@
     "# Ensemble all above\n",
     "ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever, parent_doc_retriver], weights=[0.5, 0.5])\n",
     "\n",
-    "# Re-rank\n",
+    "# Rerank\n",
     "compressor = BgeRerank()\n",
     "rerank_retriever = ContextualCompressionRetriever(\n",
     "    base_compressor=compressor, base_retriever=ensemble_retriever\n",
@@ -292,7 +290,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 12,
    "id": "beb9ab21",
    "metadata": {},
    "outputs": [],
@@ -307,7 +305,7 @@
     "        self.return_messages = return_messages\n",
     "\n",
     "    def create(self, retriver, llm):\n",
-    "        memory = ConversationBufferWindowMemory(  # ConversationBufferMemory(\n",
+    "        memory = ConversationBufferWindowMemory(\n",
     "            memory_key=self.memory_key,\n",
     "            return_messages=self.return_messages,\n",
     "            output_key=self.output_key,\n",
@@ -634,7 +632,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "lisa",
    "language": "python",
    "name": "python3"
   },
@@ -648,7 +646,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.11.10"
   }
  },
  "nbformat": 4,

README.md

@@ -11,3 +11,34 @@ startup_duration_timeout: 2h
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+
+LISA (Lithium Ion Solid-state Assistant) is a question-and-answer (Q&A) research assistant designed for efficient knowledge management with a primary focus on battery science, yet versatile enough to support broader scientific domains. Built on a Retrieval-Augmented Generation (RAG) architecture, LISA uses advanced Large Language Models (LLMs) to provide reliable, detailed answers to research questions.
+
+DEMO: https://huggingface.co/spaces/Kadi-IAM/LISA
+
+### Installation
+1. Clone the Repository:
+    ```bash
+    git clone "link of this repo"
+    cd LISA
+    ```
+
+2. Install Dependencies:
+    ```bash
+    pip install -r requirements.txt
+    ```
+
+3. Set Up the Knowledge Base
+Populate the knowledge base with relevant documents or research papers. Ensure that documents are in a format (pdf or xml) compatible with the RAG pipeline. By default documents should be located at `data/documents`. After running the following comand, some caches files are saved into `data/db`. ATTENTION: pickle is used to save these caches, be careful with potential security risks.
+    ```bash
+    python preprocess_documents.py
+    ```
+
+4. Running LISA
+Once setup is complete, run the following command to launch LISA:
+    ```bash
+    python app.py
+    ```
+
+### About
+For more information on our work in intelligent research data management systems, please visit [KadiAI](https://kadi.iam.kit.edu/kadi-ai).

app.py

@@ -1,12 +1,15 @@
+"""
+Main app for LISA RAG chatbot based on langchain.
+"""
+
 import os
 import time
 import re
-from pathlib import Path
-from dotenv import load_dotenv
+import gradio as gr
 import pickle
 
-
-import gradio as gr
+from pathlib import Path
+from dotenv import load_dotenv
 
 from huggingface_hub import login
 from langchain.vectorstores import FAISS
@@ -15,24 +18,21 @@ from llms import get_groq_chat
 from documents import load_pdf_as_docs, load_xml_as_docs
 from vectorestores import get_faiss_vectorestore
 
-
 # For debug
 # from langchain.globals import set_debug
 # set_debug(True)
 
-
 # Load and set env variables
 load_dotenv()
 
+# Set API keys
 HUGGINGFACEHUB_API_TOKEN = os.environ["HUGGINGFACEHUB_API_TOKEN"]
 login(HUGGINGFACEHUB_API_TOKEN)
 TAVILY_API_KEY = os.environ["TAVILY_API_KEY"]  # Search engine
 
-# Other settings
-os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
 
 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
-
+# Set database path
 database_root = "./data/db"
 document_path = "./data/documents"
 
@@ -80,12 +80,13 @@ from langchain.retrievers import BM25Retriever, EnsembleRetriever
 
 bm25_retriever = BM25Retriever.from_documents(
     document_chunks, k=5
-)  # 1/2 of dense retriever, experimental value
+)  # k = 1/2 of dense retriever, experimental value
 
-# Ensemble all above
+# Ensemble all above retrievers
 ensemble_retriever = EnsembleRetriever(
     retrievers=[bm25_retriever, parent_doc_retriver], weights=[0.5, 0.5]
 )
+
 # Reranker
 from rerank import BgeRerank
 
@@ -98,7 +99,7 @@ print("rerank loaded")
 llm = get_groq_chat(model_name="llama-3.1-70b-versatile")
 
 
-# # # Create conversation qa chain (Note: conversation is not supported yet)
+# Create conversation qa chain (Note: conversation is not supported yet)
 from ragchain import RAGChain
 
 rag_chain = RAGChain()
@@ -108,13 +109,11 @@ lisa_qa_conversation = rag_chain.create(rerank_retriever, llm, add_citation=True
 from langchain_community.retrievers import TavilySearchAPIRetriever
 from langchain.chains import RetrievalQAWithSourcesChain
 
-web_search_retriever = TavilySearchAPIRetriever(
-    k=4
-)  # , include_raw_content=True)#, include_raw_content=True)
+web_search_retriever = TavilySearchAPIRetriever(k=4)  # , include_raw_content=True)
 web_qa_chain = RetrievalQAWithSourcesChain.from_chain_type(
     llm, retriever=web_search_retriever, return_source_documents=True
 )
-print("chain loaded")
+print("chains loaded")
 
 
 # Gradio utils
@@ -136,7 +135,7 @@ def add_text(history, text):
 
 
 def postprocess_remove_cite_misinfo(text, allowed_max_cite_num=6):
-    """Exp.-based removal of misinfo. of citations."""
+    """Heuristic removal of misinfo. of citations."""
 
     # Remove trailing references at end of text
     if "References:\n[" in text:
@@ -480,7 +479,7 @@ def main():
                         # flag_web_search = gr.Checkbox(label="Search web", info="Search information from Internet")
                         gr.Markdown("More in DEV...")
 
-            # Manage functions
+            # Action functions
             user_txt.submit(check_input_text, user_txt, None).success(
                 add_text, [chatbot, user_txt], [chatbot, user_txt]
             ).then(bot_lisa, [chatbot, flag_web_search], [chatbot, doc_citation])
@@ -575,6 +574,7 @@ def main():
                     with gr.Tab("Setting"):
                         gr.Markdown("More in DEV...")
 
+            # Actions
             load_document.click(
                 document_changes,
                 inputs=[uploaded_doc],  # , repo_id],
@@ -606,8 +606,9 @@ def main():
             )
 
         ##########################
-        # Preview tab
+        # Preview tabs
         with gr.Tab("Preview feature 🔬"):
+            # VLM model
             with gr.Tab("Vision LM 🖼"):
                 vision_tmp_link = (
                     "https://kadi-iam-lisa-vlm.hf.space/"  # vision model link
@@ -620,6 +621,7 @@ def main():
                     )
                     # gr.Markdown("placeholder")
 
+            # OAuth2 linkage to Kadi-demo
             with gr.Tab("KadiChat 💬"):
                 kadichat_tmp_link = (
                     "https://kadi-iam-kadichat.hf.space/"  # vision model link
@@ -631,9 +633,12 @@ def main():
                         )
                     )
 
+            # Knowledge graph-enhanced RAG
             with gr.Tab("RAG enhanced with Knowledge Graph (dev) 🔎"):
                 kg_tmp_link = "https://kadi-iam-kadikgraph.static.hf.space/index.html"
-                gr.Markdown("[If rendering fails, look at the graph here](https://kadi-iam-kadikgraph.static.hf.space)")
+                gr.Markdown(
+                    "[If rendering fails, look at the graph here](https://kadi-iam-kadikgraph.static.hf.space)"
+                )
                 with gr.Blocks(css="""footer {visibility: hidden};""") as preview_tab:
                     gr.HTML(
                         """<iframe

documents.py

@@ -1,25 +1,30 @@
+"""
+Parse documents, currently pdf and xml are supported.
+"""
+
 import os
-import shutil
 
 from langchain.document_loaders import (
     PyMuPDFLoader,
 )
 from langchain.docstore.document import Document
-
-from langchain.vectorstores import Chroma
-
 from langchain.text_splitter import (
-    RecursiveCharacterTextSplitter,
+    # RecursiveCharacterTextSplitter,
     SpacyTextSplitter,
 )
 
+
 def load_pdf_as_docs(pdf_path, loader_module=None, load_kwargs=None):
     """Load and parse pdf file(s)."""
-    
-    if pdf_path.endswith('.pdf'):  # single file
+
+    if pdf_path.endswith(".pdf"):  # single file
         pdf_docs = [pdf_path]
     else:  # a directory
-        pdf_docs = [os.path.join(pdf_path, f) for f in os.listdir(pdf_path) if f.endswith('.pdf')]
+        pdf_docs = [
+            os.path.join(pdf_path, f)
+            for f in os.listdir(pdf_path)
+            if f.endswith(".pdf")
+        ]
 
     if load_kwargs is None:
         load_kwargs = {}
@@ -31,180 +36,96 @@ def load_pdf_as_docs(pdf_path, loader_module=None, load_kwargs=None):
         loader = loader_module(pdf, **load_kwargs)
         doc = loader.load()
         docs.extend(doc)
-        
+
     return docs
 
+
 def load_xml_as_docs(xml_path, loader_module=None, load_kwargs=None):
     """Load and parse xml file(s)."""
-    
+
     from bs4 import BeautifulSoup
     from unstructured.cleaners.core import group_broken_paragraphs
-    
-    if xml_path.endswith('.xml'):  # single file
+
+    if xml_path.endswith(".xml"):  # single file
         xml_docs = [xml_path]
     else:  # a directory
-        xml_docs = [os.path.join(xml_path, f) for f in os.listdir(xml_path) if f.endswith('.xml')]
-    
+        xml_docs = [
+            os.path.join(xml_path, f)
+            for f in os.listdir(xml_path)
+            if f.endswith(".xml")
+        ]
+
     if load_kwargs is None:
         load_kwargs = {}
 
     docs = []
     for xml_file in xml_docs:
-        # print("now reading file...")
         with open(xml_file) as fp:
-            soup = BeautifulSoup(fp, features="xml")    # txt is simply the a string with your XML file
+            soup = BeautifulSoup(
+                fp, features="xml"
+            )  # txt is simply the a string with your XML file
             pageText = soup.findAll(string=True)
-            parsed_text = '\n'.join(pageText)  # or " ".join, seems similar
-            # # Clean text
+            parsed_text = "\n".join(pageText)  # or " ".join, seems similar
+            # Clean text
             parsed_text_grouped = group_broken_paragraphs(parsed_text)
-            
+
             # get metadata
             try:
                 from lxml import etree as ET
+
                 tree = ET.parse(xml_file)
 
                 # Define namespace
                 ns = {"tei": "http://www.tei-c.org/ns/1.0"}
                 # Read Author personal names as an example
-                pers_name_elements = tree.xpath("tei:teiHeader/tei:fileDesc/tei:titleStmt/tei:author/tei:persName", namespaces=ns)
+                pers_name_elements = tree.xpath(
+                    "tei:teiHeader/tei:fileDesc/tei:titleStmt/tei:author/tei:persName",
+                    namespaces=ns,
+                )
                 first_per = pers_name_elements[0].text
                 author_info = first_per + " et al"
 
-                title_elements = tree.xpath("tei:teiHeader/tei:fileDesc/tei:titleStmt/tei:title", namespaces=ns)
+                title_elements = tree.xpath(
+                    "tei:teiHeader/tei:fileDesc/tei:titleStmt/tei:title", namespaces=ns
+                )
                 title = title_elements[0].text
 
                 # Combine source info
                 source_info = "_".join([author_info, title])
             except:
                 source_info = "unknown"
-                
-            # maybe even better TODO: discuss with Jens
+
+            # maybe even better parsing method. TODO: discuss with TUD
             # first_author = soup.find("author")
             # publication_year = soup.find("date", attrs={'type': 'published'})
             # title = soup.find("title")
             # source_info = [first_author, publication_year, title]
             # source_info_str = "_".join([info.text.strip() if info is not None else "unknown" for info in source_info])
-    
-            doc =  [Document(page_content=parsed_text_grouped, metadata={"source": source_info})]#, metadata={"source": "local"})
+
+            doc = [
+                Document(
+                    page_content=parsed_text_grouped, metadata={"source": source_info}
+                )
+            ]
 
             docs.extend(doc)
-            
+
     return docs
 
 
 def get_doc_chunks(docs, splitter=None):
     """Split docs into chunks."""
-    
+
     if splitter is None:
-        # splitter = RecursiveCharacterTextSplitter(
+        # splitter = RecursiveCharacterTextSplitter(  # original default
         #    # separators=["\n\n", "\n"], chunk_size=1024, chunk_overlap=256
         #    separators=["\n\n", "\n"], chunk_size=256, chunk_overlap=128
         # )
+        # Spacy seems better
         splitter = SpacyTextSplitter.from_tiktoken_encoder(
             chunk_size=512,
             chunk_overlap=128,
         )
     chunks = splitter.split_documents(docs)
-    
-    return chunks
-
 
-def persist_vectorstore(document_chunks, embeddings, persist_directory="db", overwrite=False):
-    # embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-    # vectorstore = FAISS.from_texts(texts=text_chunks, embedding=embeddings)
-    if overwrite:
-        shutil.rmtree(persist_directory)  # Empty and reset db
-    db = Chroma.from_documents(documents=document_chunks, embedding=embeddings, persist_directory=persist_directory)
-    # db.delete_collection()
-    db.persist()
-    # db = None
-    # db = Chroma(persist_directory="db", embedding_function = embeddings, client_settings=CHROMA_SETTINGS)
-    # vectorstore = FAISS.from_documents(documents=document_chunks, embedding=embeddings)
-    return db
-    
-    
-class VectorstoreManager:
-    
-    def __init__(self):
-        self.vectorstore_class = Chroma
-        
-    def create_db(self, embeddings):
-        db = self.vectorstore_class(embedding_function=embeddings)
-        
-        self.db = db
-        return db
-    
-    
-    def load_db(self, persist_directory, embeddings):
-        """Load local vectorestore."""
-        
-        db = self.vectorstore_class(persist_directory=persist_directory, embedding_function=embeddings)
-        self.db = db
-        
-        return db
-    
-    def create_db_from_documents(self, document_chunks, embeddings, persist_directory="db", overwrite=False):
-        """Create db from documents."""
-        
-        if overwrite:
-            shutil.rmtree(persist_directory)  # Empty and reset db
-        db = self.vectorstore_class.from_documents(documents=document_chunks, embedding=embeddings, persist_directory=persist_directory)
-        self.db = db
-        
-        return db
-    
-    def persist_db(self, persist_directory="db"):
-        """Persist db."""
-        
-        assert self.db
-        self.db.persist()  # Chroma
-        
-class RetrieverManager:
-    # some other retrievers Using Advanced Retrievers in LangChain https://www.comet.com/site/blog/using-advanced-retrievers-in-langchain/
-    
-    def __init__(self, vectorstore, k=10):
-        
-        self.vectorstore = vectorstore
-        self.retriever = vectorstore.as_retriever(search_kwargs={"k": k})  #search_kwargs={"k": 8}),
-        
-    def get_rerank_retriver(self, base_retriever=None):
-        
-        if base_retriever is None:
-            base_retriever = self.retriever
-        # with rerank
-        from rerank import BgeRerank
-        from langchain.retrievers import ContextualCompressionRetriever
-        
-        compressor = BgeRerank()
-        compression_retriever = ContextualCompressionRetriever(
-            base_compressor=compressor, base_retriever=base_retriever
-        )
-        
-        return compression_retriever
-        
-    def get_parent_doc_retriver(self, documents, store_file="./store_location"):
-        # TODO need better design
-        # Ref: explain how it works: https://clusteredbytes.pages.dev/posts/2023/langchain-parent-document-retriever/
-        from langchain.storage.file_system import LocalFileStore
-        from langchain.storage import InMemoryStore
-        from langchain.storage._lc_store import create_kv_docstore
-        from langchain.retrievers import ParentDocumentRetriever
-        # Ref: https://stackoverflow.com/questions/77385587/persist-parentdocumentretriever-of-langchain
-        # fs = LocalFileStore("./store_location")
-        # store = create_kv_docstore(fs)
-        docstore = InMemoryStore()
-        
-        # TODO: how to better set this?
-        parent_splitter = RecursiveCharacterTextSplitter(separators=["\n\n", "\n"], chunk_size=1024, chunk_overlap=256)
-        child_splitter = RecursiveCharacterTextSplitter(separators=["\n\n", "\n"], chunk_size=256, chunk_overlap=128)
-        
-        retriever = ParentDocumentRetriever(
-            vectorstore=self.vectorstore,
-            docstore=docstore,
-            child_splitter=child_splitter,
-            parent_splitter=parent_splitter,
-            search_kwargs={"k":10}  # Better settings?
-        )
-        retriever.add_documents(documents)#, ids=None)
-        
-        return retriever
+    return chunks

embeddings.py

@@ -1,39 +1,50 @@
+"""
+Load embedding models from huggingface.
+"""
 
 import torch
 from langchain.embeddings import HuggingFaceEmbeddings
 
 
 def get_hf_embeddings(model_name=None):
-    """Get huggingface embedding."""
-    
+    """Get huggingface embedding by name."""
+
     if model_name is None:
-        # Some candiates: 
+        # Some candiates
         # "BAAI/bge-m3" (good, though large and slow)
-        # "BAAI/bge-base-en-v1.5" -> seems not that good with current settings
-        # "sentence-transformers/all-mpnet-base-v2", "maidalun1020/bce-embedding-base_v1", "intfloat/multilingual-e5-large"
-        # Ref: https://huggingface.co/spaces/mteb/leaderboard, https://huggingface.co/maidalun1020/bce-embedding-base_v1
-        model_name = "BAAI/bge-large-en-v1.5"  # or ""
-    
+        # "BAAI/bge-base-en-v1.5" -> also good
+        # "sentence-transformers/all-mpnet-base-v2"
+        # "maidalun1020/bce-embedding-base_v1"
+        # "intfloat/multilingual-e5-large"
+        # Ref: https://huggingface.co/spaces/mteb/leaderboard
+        # https://huggingface.co/maidalun1020/bce-embedding-base_v1
+        model_name = "BAAI/bge-large-en-v1.5"
+
     embeddings = HuggingFaceEmbeddings(model_name=model_name)
-    
+
     return embeddings
 
 
-def get_jinaai_embeddings(model_name="jinaai/jina-embeddings-v2-base-en", device="auto"):
+def get_jinaai_embeddings(
+    model_name="jinaai/jina-embeddings-v2-base-en", device="auto"
+):
     """Get jinaai embedding."""
-    
+
     # device: cpu or cuda
     if device == "auto":
         device = "cuda" if torch.cuda.is_available() else "cpu"
     # For jinaai. Ref: https://github.com/langchain-ai/langchain/issues/6080
     from transformers import AutoModel
-    model = AutoModel.from_pretrained(model_name, trust_remote_code=True)  # -> will yield error, need bug fixing
+
+    model = AutoModel.from_pretrained(
+        model_name, trust_remote_code=True
+    )  # -> will yield error, need bug fixing
 
     model_name = model_name
-    model_kwargs = {'device': device, "trust_remote_code": True}
+    model_kwargs = {"device": device, "trust_remote_code": True}
     embeddings = HuggingFaceEmbeddings(
         model_name=model_name,
         model_kwargs=model_kwargs,
     )
-    
-    return embeddings
+
+    return embeddings

llms.py

@@ -1,22 +1,22 @@
-# from langchain import HuggingFaceHub, LLMChain
-from langchain.llms import HuggingFacePipeline
+"""
+Load LLMs from huggingface, Groq, etc.
+"""
+
 from transformers import (
-    AutoModelForCausalLM,
+    # AutoModelForCausalLM,
     AutoTokenizer,
     pipeline,
 )
-from transformers import LlamaForCausalLM, AutoModelForCausalLM, LlamaTokenizer
-from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
+from langchain.llms import HuggingFacePipeline
 from langchain_groq import ChatGroq
-
-
-from langchain.chat_models import ChatOpenAI
 from langchain.llms import HuggingFaceTextGenInference
 
+# from langchain.chat_models import ChatOpenAI  # oai model
+
 
 def get_llm_hf_online(inference_api_url=""):
     """Get LLM using huggingface inference."""
-    
+
     if not inference_api_url:  # default api url
         inference_api_url = (
             "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta"
@@ -35,20 +35,16 @@ def get_llm_hf_online(inference_api_url=""):
 
 
 def get_llm_hf_local(model_path):
-    """Get local LLM."""
-    
-    model = LlamaForCausalLM.from_pretrained(
-        model_path, device_map="auto"
-    )
+    """Get local LLM from huggingface."""
+
+    model = LlamaForCausalLM.from_pretrained(model_path, device_map="auto")
     tokenizer = AutoTokenizer.from_pretrained(model_path)
 
-    # print('making a pipeline...')
-    # max_length has typically been deprecated for max_new_tokens
     pipe = pipeline(
         "text-generation",
         model=model,
         tokenizer=tokenizer,
-        max_new_tokens=1024,  # better setting?
+        max_new_tokens=2048,  # better setting?
         model_kwargs={"temperature": 0.1},  # better setting?
     )
     llm = HuggingFacePipeline(pipeline=pipe)
@@ -56,22 +52,8 @@ def get_llm_hf_local(model_path):
     return llm
 
 
-
-def get_llm_openai_chat(model_name, inference_server_url):
-    """Get openai-like LLM."""
-    
-    llm = ChatOpenAI(
-        model=model_name,
-        openai_api_key="EMPTY",
-        openai_api_base=inference_server_url,
-        max_tokens=1024,  # better setting?
-        temperature=0,
-    )
-
-    return llm
-
-
-def  get_groq_chat(model_name="llama-3.1-70b-versatile"):
+def get_groq_chat(model_name="llama-3.1-70b-versatile"):
+    """Get LLM from Groq."""
 
     llm = ChatGroq(temperature=0, model_name=model_name)
-    return llm
+    return llm

preprocess_documents.py

@@ -1,15 +1,17 @@
 """
-Load and parse files (pdf) in the data/documents and save cached pkl files.
+Load and parse files (pdf) in the "data/documents" and save cached pkl files.
+It will load and parse files and save 4 caches:
+1. "docs.pkl" for loaded text documents
+2. "docs_chunks.pkl" for chunked text
+3. "docstore.pkl" for small-to-big retriever
+4. faiss_index for FAISS vectore store
 """
 
 import os
 import pickle
 
 from dotenv import load_dotenv
-
-
 from huggingface_hub import login
-
 from documents import load_pdf_as_docs, get_doc_chunks
 from embeddings import get_jinaai_embeddings
 
@@ -23,11 +25,14 @@ login(HUGGINGFACEHUB_API_TOKEN)
 
 
 def save_to_pickle(obj, filename):
+    """Save obj to disk using pickle."""
+
     with open(filename, "wb") as file:
         pickle.dump(obj, file, pickle.HIGHEST_PROTOCOL)
 
 
 # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
+# Set database path, should be same as defined in "app.py"
 database_root = "./data/db"
 document_path = "./data/documents"
 

ragchain.py

@@ -1,3 +1,7 @@
+"""
+Main RAG chain based on langchain.
+"""
+
 from langchain.chains import LLMChain
 
 from langchain.prompts import (
@@ -11,17 +15,17 @@ from langchain.chains import ConversationalRetrievalChain
 from langchain.chains.conversation.memory import (
     ConversationBufferWindowMemory,
 )
-
-
 from langchain.chains import StuffDocumentsChain
 
 
 def get_cite_combine_docs_chain(llm):
+    """Get doc chain which adds metadata to text chunks."""
 
     # Ref: https://github.com/langchain-ai/langchain/issues/7239
     # Function to format each document with an index, source, and content.
     def format_document(doc, index, prompt):
         """Format a document into a string based on a prompt template."""
+
         # Create a dictionary with document content and metadata.
         base_info = {
             "page_content": doc.page_content,
@@ -40,7 +44,11 @@ def get_cite_combine_docs_chain(llm):
 
     # Custom chain class to handle document combination with source indices.
     class StuffDocumentsWithIndexChain(StuffDocumentsChain):
+        """Custom chain class to handle document combination with source indices."""
+
         def _get_inputs(self, docs, **kwargs):
+            """Overwrite _get_inputs to add metadata for text chunks."""
+
             # Format each document and combine them.
             doc_strings = [
                 format_document(doc, i, self.document_prompt)
@@ -58,6 +66,7 @@ def get_cite_combine_docs_chain(llm):
             )
             return inputs
 
+    # Main prompt for RAG chain with citation
     # Ref: https://huggingface.co/spaces/Ekimetrics/climate-question-answering/blob/main/climateqa/engine/prompts.py
     # Define a chat prompt with instructions for citing documents.
     combine_doc_prompt = PromptTemplate(
@@ -103,6 +112,8 @@ def get_cite_combine_docs_chain(llm):
 
 
 class RAGChain:
+    """Main RAG chain."""
+
     def __init__(
         self, memory_key="chat_history", output_key="answer", return_messages=True
     ):
@@ -111,14 +122,17 @@ class RAGChain:
         self.return_messages = return_messages
 
     def create(self, retriever, llm, add_citation=False):
-        memory = ConversationBufferWindowMemory(  # ConversationBufferMemory(
+        """Create a rag chain instance."""
+
+        # Memory is kept for later support of conversational chat
+        memory = ConversationBufferWindowMemory(  # Or ConversationBufferMemory
             k=2,
             memory_key=self.memory_key,
             return_messages=self.return_messages,
             output_key=self.output_key,
         )
 
-        # https://github.com/langchain-ai/langchain/issues/4608
+        # Ref: https://github.com/langchain-ai/langchain/issues/4608
         conversation_chain = ConversationalRetrievalChain.from_llm(
             llm=llm,
             retriever=retriever,
@@ -127,7 +141,6 @@ class RAGChain:
             rephrase_question=False,  # disable rephrase, for test purpose
             get_chat_history=lambda x: x,
             # return_generated_question=True,  # for debug
-            # verbose=True,
             # combine_docs_chain_kwargs={"prompt": PROMPT},  # additional prompt control
             # condense_question_prompt=CONDENSE_QUESTION_PROMPT,  # additional prompt control
         )

requirements.txt

@@ -5,7 +5,7 @@ langchain-community==0.2.4
 text-generation
 pypdf
 pymupdf
-gradio
+gradio==4.44.1
 faiss-cpu
 chromadb
 rank-bm25

rerank.py

@@ -1,5 +1,6 @@
 """
-Retrank with cross encoder.
+Rerank with cross encoder.
+Ref:
 https://medium.aiplanet.com/advanced-rag-cohere-re-ranker-99acc941601c
 https://github.com/langchain-ai/langchain/issues/13076
 """
@@ -7,7 +8,7 @@ https://github.com/langchain-ai/langchain/issues/13076
 from __future__ import annotations
 from typing import Optional, Sequence
 from langchain.schema import Document
-from langchain.pydantic_v1 import Extra, root_validator
+from langchain.pydantic_v1 import Extra
 
 from langchain.callbacks.manager import Callbacks
 from langchain.retrievers.document_compressors.base import BaseDocumentCompressor

retrievers.py

@@ -1,7 +1,10 @@
+"""
+Retrievers for text chunks.
+"""
+
 import os
 
 from langchain.text_splitter import (
-    CharacterTextSplitter,
     RecursiveCharacterTextSplitter,
     SpacyTextSplitter,
 )
@@ -9,6 +12,7 @@ from langchain.text_splitter import (
 from rerank import BgeRerank
 from langchain.retrievers import ContextualCompressionRetriever
 
+
 def get_parent_doc_retriever(
     documents,
     vectorstore,
@@ -40,12 +44,14 @@ def get_parent_doc_retriever(
         from langchain_rag.storage import SQLStore
 
         # Instantiate the SQLStore with the root path
-        docstore = SQLStore(namespace="test", db_url="sqlite:///parent_retrieval_db.db")  # TODO: WIP
+        docstore = SQLStore(
+            namespace="test", db_url="sqlite:///parent_retrieval_db.db"
+        )  # TODO: WIP
     else:
         docstore = docstore  # TODO: add check
-        # raise  # TODO implement
+        # raise  # TODO implement other docstores
 
-    # TODO: how to better set this?
+    # TODO: how to better set these values?
     # parent_splitter = RecursiveCharacterTextSplitter(separators=["\n\n", "\n"], chunk_size=1024, chunk_overlap=256)
     # child_splitter = RecursiveCharacterTextSplitter(separators=["\n\n", "\n"], chunk_size=256, chunk_overlap=64)
     parent_splitter = SpacyTextSplitter.from_tiktoken_encoder(
@@ -62,11 +68,11 @@ def get_parent_doc_retriever(
         docstore=docstore,
         child_splitter=child_splitter,
         parent_splitter=parent_splitter,
-        search_kwargs={"k": k},  # Better settings?
+        search_kwargs={"k": k},
     )
 
     if add_documents:
-        retriever.add_documents(documents)  # , ids=None)
+        retriever.add_documents(documents)
 
     if save_vectorstore:
         vectorstore.save_local(os.path.join(save_path_root, "faiss_index"))
@@ -80,7 +86,6 @@ def get_parent_doc_retriever(
 
         save_to_pickle(docstore, os.path.join(save_path_root, "docstore.pkl"))
 
-
     return retriever
 
 

vectorestores.py

@@ -1,8 +1,13 @@
-from langchain.vectorstores import Chroma, FAISS
+"""
+Vector stores.
+"""
+
+from langchain.vectorstores import FAISS
+
 
 def get_faiss_vectorestore(embeddings):
     # Add extra text to init
     texts = ["LISA - Lithium Ion Solid-state Assistant"]
     vectorstore = FAISS.from_texts(texts, embeddings)
-    
-    return vectorstore
+
+    return vectorstore