Update rag_pipeline.py

rag_pipeline.py

@@ -2,25 +2,39 @@ 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_community.document_loaders import PyMuPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
 
+def load_documents(pdf_dir):
+    docs = []
+    for pdf_file in Path(pdf_dir).glob("*.pdf"):
+        loader = PyMuPDFLoader(str(pdf_file))
+        docs.extend(loader.load())
+    return docs
+
 def load_rag_chain():
+    # Ensure directory exists
     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)
+    pdf_dir.mkdir(parents=True, exist_ok=True)
+
+    # Load and split PDFs
+    raw_docs = load_documents(pdf_dir)
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    pages = splitter.split_documents(raw_docs)
 
+    # Embedding model
     embeddings = HuggingFaceEmbeddings(
         model_name="sentence-transformers/all-MiniLM-L6-v2",
         model_kwargs={"device": "cpu"},
     )
 
+    # Vector database
     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
     hf_pipeline = pipeline(
         "text-generation",
         model="mistralai/Mistral-7B-Instruct-v0.2",
@@ -28,9 +42,10 @@ def load_rag_chain():
         max_new_tokens=512,
         temperature=0.3,
         return_full_text=True,
-        device=-1,
+        device=-1  # CPU
     )
     llm = HuggingFacePipeline(pipeline=hf_pipeline)
 
+    # QA Chain
     qa_chain = RetrievalQA.from_llm(llm=llm, retriever=retriever)
-    return qa_chain
+    return qa_chain