Update rag_pipeline.py

rag_pipeline.py

@@ -14,10 +14,8 @@ def load_documents(pdf_dir):
     return docs
 
 def load_rag_chain():
-    pdf_dir = Path("data")
-    pdf_dir.mkdir(parents=True, exist_ok=True)
-
-    raw_docs = load_documents(pdf_dir)
+    Path("data").mkdir(exist_ok=True)
+    raw_docs = load_documents("data")
     splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
     pages = splitter.split_documents(raw_docs)
 
@@ -26,19 +24,17 @@ def load_rag_chain():
         model_kwargs={"device": "cpu"},
     )
 
-    vectordb_dir = "chroma_db"
-    vectordb = Chroma.from_documents(pages, embeddings, persist_directory=vectordb_dir)
+    vectordb = Chroma.from_documents(pages, embeddings, persist_directory="chroma_db")
     retriever = vectordb.as_retriever(search_type="mmr", search_kwargs={"k": 5})
 
     hf_pipeline = pipeline(
         "text2text-generation",
-        model="csebuetnlp/mT5_small_finetuned_squad",
-        tokenizer=AutoTokenizer.from_pretrained("csebuetnlp/mT5_small_finetuned_squad"),
+        model="ArabicNLP/mT5-base_ar",
+        tokenizer=AutoTokenizer.from_pretrained("ArabicNLP/mT5-base_ar"),
         max_new_tokens=512,
         temperature=0.3,
-        device=-1
+        device=-1,
     )
     llm = HuggingFacePipeline(pipeline=hf_pipeline)
 
-    qa_chain = RetrievalQA.from_llm(llm=llm, retriever=retriever)
-    return qa_chain
+    return RetrievalQA.from_llm(llm=llm, retriever=retriever)