Upload 3 files

app.py

@@ -0,0 +1,16 @@
+import gradio as gr
+from rag_pipeline import load_rag_chain
+
+rag_chain = load_rag_chain()
+
+def ask_question(query):
+    result = rag_chain.invoke(query)
+    return result['result']
+
+iface = gr.Interface(fn=ask_question,
+                     inputs=gr.Textbox(lines=3, label="Ask a Question"),
+                     outputs="text",
+                     title="Custom PDF RAG Chatbot")
+
+if __name__ == "__main__":
+    iface.launch()

rag_pipeline.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+from langchain.chains import RetrievalQA
+from transformers import pipeline, AutoTokenizer
+from langchain_community.vectorstores import Chroma
+from langchain_community.document_loaders import DirectoryLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
+
+def load_rag_chain():
+    pdf_dir = Path("data")
+    loader = DirectoryLoader(str(pdf_dir), glob="*.pdf")
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    pages = loader.load_and_split(text_splitter=text_splitter)
+
+    embeddings = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-MiniLM-L6-v2",
+        model_kwargs={"device": "cpu"},
+    )
+
+    vectordb_dir = "chroma_db"
+    vectordb = Chroma.from_documents(pages, embeddings, persist_directory=vectordb_dir)
+    retriever = vectordb.as_retriever(search_type="mmr", search_kwargs={"k": 5})
+
+    hf_pipeline = pipeline(
+        "text-generation",
+        model="mistralai/Mistral-7B-Instruct-v0.2",
+        tokenizer=AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2"),
+        max_new_tokens=512,
+        temperature=0.3,
+        return_full_text=True,
+        device=-1,
+    )
+    llm = HuggingFacePipeline(pipeline=hf_pipeline)
+
+    qa_chain = RetrievalQA.from_llm(llm=llm, retriever=retriever)
+    return qa_chain

requirements.txt

@@ -0,0 +1,7 @@
+gradio
+transformers
+sentence-transformers
+langchain
+langchain-community
+langchain-huggingface
+chromadb