commit name

.gitattributes

@@ -1,35 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +0,0 @@
----
-title: Test
-emoji: 🦀
-colorFrom: red
-colorTo: indigo
-sdk: streamlit
-sdk_version: 1.44.0
-app_file: app.py
-pinned: false
-license: apache-2.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,198 @@
+import streamlit as st
+from langchain_community.callbacks import StreamlitCallbackHandler
+from langchain_core.runnables import RunnableConfig
+
+from src import CFG
+from src.query_expansion import build_multiple_queries_expansion_chain
+from src.retrieval_qa import (
+    build_retrieval_qa,
+    build_base_retriever,
+    build_rerank_retriever,
+    build_compression_retriever,
+)
+from src.vectordb import build_vectordb, load_faiss, load_chroma
+from streamlit_app.pdf_display import get_doc_highlighted, display_pdf
+from streamlit_app.utils import load_base_embeddings, load_llm, load_reranker
+
+st.set_page_config(page_title="Retrieval QA", layout="wide")
+
+LLM = load_llm()
+BASE_EMBEDDINGS = load_base_embeddings()
+RERANKER = load_reranker()
+
+
+@st.cache_resource
+def load_vectordb():
+    if CFG.VECTORDB_TYPE == "faiss":
+        return load_faiss(BASE_EMBEDDINGS)
+    if CFG.VECTORDB_TYPE == "chroma":
+        return load_chroma(BASE_EMBEDDINGS)
+    raise NotImplementedError
+
+
+@st.cache_resource
+def load_retriever(_vectordb, retrieval_mode):
+    if retrieval_mode == "Base":
+        return build_base_retriever(_vectordb)
+    if retrieval_mode == "Rerank":
+        return build_rerank_retriever(_vectordb, RERANKER)
+    if retrieval_mode == "Contextual compression":
+        return build_compression_retriever(_vectordb, BASE_EMBEDDINGS)
+    raise NotImplementedError
+
+
+def init_sess_state():
+    if "uploaded_filename" not in st.session_state:
+        st.session_state["uploaded_filename"] = ""
+
+    if "last_form" not in st.session_state:
+        st.session_state["last_form"] = list()
+
+    if "last_query" not in st.session_state:
+        st.session_state["last_query"] = ""
+
+    if "last_response" not in st.session_state:
+        st.session_state["last_response"] = dict()
+
+    if "last_related" not in st.session_state:
+        st.session_state["last_related"] = list()
+
+
+def doc_qa():
+    init_sess_state()
+
+    with st.sidebar:
+        st.header("RAG with quantized LLM")
+        st.info(f"LLM: `{CFG.LLM_PATH}`")
+        st.info(f"Embeddings: `{CFG.EMBEDDINGS_PATH}`")
+        st.info(f"Reranker: `{CFG.RERANKER_PATH}`")
+
+        uploaded_file = st.file_uploader(
+            "Upload a PDF and build VectorDB", type=["pdf"]
+        )
+        if st.button("Build VectorDB"):
+            if uploaded_file is None:
+                st.error("No PDF uploaded")
+            else:
+                uploaded_filename = f"./data/{uploaded_file.name}"
+                with open(uploaded_filename, "wb") as f:
+                    f.write(uploaded_file.getvalue())
+                with st.spinner("Building VectorDB..."):
+                    build_vectordb(uploaded_filename)
+                st.session_state.uploaded_filename = uploaded_filename
+
+        if st.session_state.uploaded_filename != "":
+            st.info(f"Current document: {st.session_state.uploaded_filename}")
+
+        try:
+            with st.status("Load VectorDB", expanded=False) as status:
+                st.write("Loading VectorDB ...")
+                vectordb = load_vectordb()
+                status.update(
+                    label="Loading complete!", state="complete", expanded=False
+                )
+
+            st.success("Reading from existing VectorDB")
+        except Exception as e:
+            st.error(f"No existing VectorDB found: {e}")
+
+    c0, c1 = st.columns(2)
+
+    with c0.form("qa_form"):
+        user_query = st.text_area("Your query")
+        with st.expander("Settings"):
+            mode = st.radio(
+                "Mode",
+                ["Retrieval only", "Retrieval QA"],
+                index=1,
+                help="""Retrieval only will output extracts related to your query immediately, \
+                while Retrieval QA will output an answer to your query and will take a while on CPU.""",
+            )
+            retrieval_mode = st.radio(
+                "Retrieval method",
+                ["Base", "Rerank", "Contextual compression"],
+                index=1,
+            )
+
+        submitted = st.form_submit_button("Query")
+        if submitted:
+            if user_query == "":
+                st.error("Please enter a query.")
+
+    if user_query != "" and (
+        st.session_state.last_query != user_query
+    ):
+        st.session_state.last_query = user_query
+
+        if mode == "Retrieval only":
+            retriever = load_retriever(vectordb, retrieval_mode)
+            with c0:
+                with st.spinner("Retrieving ..."):
+                    relevant_docs = retriever.get_relevant_documents(user_query)
+
+            st.session_state.last_response = {
+                "query": user_query,
+                "source_documents": relevant_docs,
+            }
+
+            chain = build_multiple_queries_expansion_chain(LLM)
+            res = chain.invoke(user_query)
+            st.session_state.last_related = [
+                x.strip() for x in res.split("\n") if x.strip()
+            ]
+        else:
+            retriever = load_retriever(db, retrieval_mode)
+            retrieval_qa = build_retrieval_qa(LLM, retriever)
+
+            st_callback = StreamlitCallbackHandler(
+                parent_container=c0.container(),
+                expand_new_thoughts=True,
+                collapse_completed_thoughts=True,
+            )
+            st.session_state.last_response = retrieval_qa.invoke(
+                user_query, config=RunnableConfig(callbacks=[st_callback])
+            )
+            st_callback._complete_current_thought()
+
+    if st.session_state.last_response:
+        with c0:
+            st.warning(f"##### {st.session_state.last_query}")
+            if st.session_state.last_response.get("result") is not None:
+                st.success(st.session_state.last_response["result"])
+
+            if st.session_state.last_related:
+                st.write("#### Related")
+                for r in st.session_state.last_related:
+                    st.write(f"```\n{r}\n```")
+
+        with c1:
+            st.write("#### Sources")
+            for row in st.session_state.last_response["source_documents"]:
+                st.write("**Page {}**".format(row.metadata["page"] + 1))
+                st.info(row.page_content.replace("$", r"\$"))
+
+            # Display PDF
+            st.write("---")
+            _display_pdf_from_docs(st.session_state.last_response["source_documents"])
+
+
+def _display_pdf_from_docs(source_documents):
+    n = len(source_documents)
+    i = st.radio(
+        "View in PDF", list(range(n)), format_func=lambda x: f"Extract {x + 1}"
+    )
+    row = source_documents[i]
+    try:
+        extracted_doc, page_nums = get_doc_highlighted(
+            row.metadata["source"], row.page_content
+        )
+        if extracted_doc is None:
+            st.error("No page found")
+        else:
+            display_pdf(extracted_doc, page_nums[0] + 1)
+    except Exception as e:
+        st.error(e)
+
+
+if __name__ == "__main__":
+    doc_qa()

condaenv.khtzjjyc.requirements.txt

@@ -0,0 +1,12 @@
+chromadb==0.4.22
+ctransformers==0.2.27
+faiss-cpu==1.7.4
+fastapi==0.104.1
+langchain==0.1.3
+PyMuPDF==1.23.8
+python-box==7.1.1
+rank-bm25==0.2.2
+sentence-transformers==2.2.2
+simsimd==3.3.0
+streamlit==1.30.0
+umap-learn==0.5.5

requirements.txt

@@ -0,0 +1,12 @@
+chromadb==0.4.22
+ctransformers==0.2.27
+faiss-cpu==1.7.4
+fastapi==0.104.1
+langchain==0.1.3
+PyMuPDF==1.23.8
+python-box==7.1.1
+rank-bm25==0.2.2
+sentence-transformers==2.2.2
+simsimd==3.3.0
+streamlit==1.30.0
+umap-learn==0.5.5