v1

.gitignore

@@ -0,0 +1,3 @@
+/__pycache__
+/temp
+/models

app.py

@@ -0,0 +1,64 @@
+import streamlit as st
+from llm import load_llm, response_generator
+from vector_store import load_vector_store, process_pdf
+from uuid import uuid4
+
+# repo_id = "Qwen/Qwen2.5-0.5B-Instruct-GGUF"
+# filename = "qwen2.5-0.5b-instruct-q8_0.gguf"
+repo_id = "MaziyarPanahi/Qwen2.5-7B-Instruct-GGUF"
+filename = "Qwen2.5-7B-Instruct.Q4_K_S.gguf"
+
+llm = load_llm(repo_id, filename)
+
+st.title("PDF QA")
+# Initialize chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+# Display chat messages from history on app rerun
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        if message["role"] == "user":
+            st.markdown(message["content"])
+        else:
+            st.code(message["content"])
+
+# Accept user input
+if prompt := st.chat_input("What is up?"):
+    # Add user message to chat history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    # Display user message in chat message container
+    with st.chat_message("user"):
+        st.markdown(prompt)
+
+    # Display assistant response in chat message container
+    with st.chat_message("assistant"):
+        vector_store = load_vector_store()
+        retriever = vector_store.as_retriever()
+        docs = retriever.get_relevant_documents(prompt)
+
+        response = response_generator(llm, st.session_state.messages, prompt, retriever)
+
+        st.markdown(response["answer"])
+
+    # Add assistant response to chat history
+    st.session_state.messages.append(
+        {"role": "assistant", "content": response["answer"]}
+    )
+
+with st.sidebar:
+    st.title("PDFs")
+    st.write("Upload your pdfs here")
+    uploaded_files = st.file_uploader(
+        "Choose a PDF file", accept_multiple_files=True, type="pdf"
+    )
+    if uploaded_files is not None:
+        vector_store = load_vector_store()
+        for uploaded_file in uploaded_files:
+            temp_file = f"./temp/{uploaded_file.name}-{uuid4()}.pdf"
+            with open(temp_file, "wb") as file:
+                file.write(uploaded_file.getvalue())
+
+            st.write("filename:", uploaded_file.name)
+            process_pdf(temp_file, vector_store)
+            st.success("PDFs uploaded successfully. ✅")

llm.py

@@ -0,0 +1,61 @@
+import streamlit as st
+import pathlib
+
+from huggingface_hub import hf_hub_download
+from langchain_community.llms import LlamaCpp
+from langchain.chains import create_retrieval_chain
+from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain_core.prompts import ChatPromptTemplate
+
+
+@st.cache_resource()
+def load_llm(repo_id, filename):
+    # Create a directory for models if it doesn't exist
+    models_folder = pathlib.Path("models")
+    models_folder.mkdir(exist_ok=True)
+
+    # Download the model
+    model_path = hf_hub_download(
+        repo_id=repo_id, filename=filename, local_dir=models_folder
+    )
+
+    llm = LlamaCpp(
+        model_path=model_path,
+        repo_id=repo_id,
+        filename=filename,
+        verbose=False,
+        use_mmap=True,
+        use_mlock=True,
+        n_threads=4,
+        n_threads_batch=4,
+        n_ctx=8000,
+    )
+    print(f"{repo_id} loaded successfully. ✅")
+    return llm
+
+
+# Streamed response emulator
+def response_generator(llm, messages, question, retriever):
+    system_prompt = (
+        "You are an assistant for question-answering tasks. "
+        "Use the following pieces of retrieved context to answer "
+        "the question. If you don't know the answer, say that you "
+        "don't know. Use three sentences maximum and keep the "
+        "answer concise."
+        "\n\n"
+        "{context}"
+    )
+
+    prompt = ChatPromptTemplate.from_messages(
+        [
+            ("system", system_prompt),
+            ("user", "{input}"),
+        ]
+    )
+
+    question_answer_chain = create_stuff_documents_chain(llm, prompt)
+    rag_chain = create_retrieval_chain(retriever, question_answer_chain)
+
+    results = rag_chain.invoke({"input": question})
+
+    return results

requirements.txt

@@ -0,0 +1,7 @@
+huggingface_hub
+llama-cpp-python
+pypdf
+langchain_community
+chromadb
+langchain-huggingface
+langchain-chroma

vector_store.py

@@ -0,0 +1,40 @@
+import streamlit as st
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_chroma import Chroma
+from langchain_community.document_loaders import PyPDFLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+
+@st.cache_resource()
+def load_embedding_model(model):
+    """
+    sentence-transformers/all-mpnet-base-v2
+    sentence-transformers/all-MiniLM-L6-v2
+    """
+    model = HuggingFaceEmbeddings(model_name=model)
+    return model
+
+
+def load_vector_store():
+    """
+    Loads a simple vector store
+    I didn't use @st.cache because I want to
+    load vector store on every page load
+    """
+    model = load_embedding_model("sentence-transformers/all-MiniLM-L6-v2")
+    vector_store = Chroma(
+        collection_name="main_store",
+        embedding_function=model,
+    )
+    return vector_store
+
+
+def process_pdf(pdf, vector_store):
+    """
+    Loads a pdf and splits it into chunks
+    """
+    loader = PyPDFLoader(pdf)
+    docs = loader.load()
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    splits = text_splitter.split_documents(docs)
+    vector_store.add_documents(splits)