Update rag_pipeline.py

rag_pipeline.py

@@ -6,6 +6,7 @@ from langchain_community.document_loaders import PyMuPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
 
+# Load all PDFs from the data folder
 def load_documents(pdf_dir):
     docs = []
     for pdf_file in Path(pdf_dir).glob("*.pdf"):
@@ -14,7 +15,7 @@ def load_documents(pdf_dir):
     return docs
 
 def load_rag_chain():
-    # Ensure directory exists
+    # Make sure the data directory exists
     pdf_dir = Path("data")
     pdf_dir.mkdir(parents=True, exist_ok=True)
 
@@ -23,29 +24,28 @@ def load_rag_chain():
     splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
     pages = splitter.split_documents(raw_docs)
 
-    # Embedding model
+    # Load sentence transformer for embeddings
     embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2",
         model_kwargs={"device": "cpu"},
     )
 
-    # Vector database
+    # Vector store
     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})
 
-    # LLM pipeline using free model
+    # Load a completely free and CPU-compatible model
     hf_pipeline = pipeline(
-        "text-generation",
-        model="mistralai/Mistral-7B-Instruct-v0.2",
-        tokenizer=AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2"),
+        "text2text-generation",
+        model="google/flan-t5-base",
+        tokenizer=AutoTokenizer.from_pretrained("google/flan-t5-base"),
         max_new_tokens=512,
         temperature=0.3,
-        return_full_text=True,
-        device=-1  # CPU
+        device=-1  # -1 means CPU
     )
     llm = HuggingFacePipeline(pipeline=hf_pipeline)
 
-    # QA Chain
+    # Build RetrievalQA chain
     qa_chain = RetrievalQA.from_llm(llm=llm, retriever=retriever)
     return qa_chain