Update rag_pipeline.py

rag_pipeline.py

@@ -3,37 +3,33 @@ from langchain_community.document_loaders import PyMuPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain_huggingface import HuggingFaceEmbeddings
-from transformers import pipeline, AutoModelForQuestionAnswering, AutoTokenizer
+from transformers import pipeline
 
-# Load and index Arabic documents
-def load_and_index():
-    pdf_dir = Path("data")
-    pdf_dir.mkdir(exist_ok=True)
+# Retriever for top-5 relevant document chunks
+def init_retriever():
+    Path("data").mkdir(exist_ok=True)
     docs = []
-    for pdf_file in pdf_dir.glob("*.pdf"):
-        loader = PyMuPDFLoader(str(pdf_file))
+    for pdf in Path("data").glob("*.pdf"):
+        loader = PyMuPDFLoader(str(pdf))
         docs.extend(loader.load())
     splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    split_docs = splitter.split_documents(docs)
+    chunks = splitter.split_documents(docs)
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/LaBSE", model_kwargs={"device": "cpu"})
+    vectordb = Chroma.from_documents(chunks, embeddings, persist_directory="chroma_db")
+    return vectordb.as_retriever(search_kwargs={"k": 5})
 
-    embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/LaBSE", model_kwargs={"device": "cpu"})
-    db = Chroma.from_documents(split_docs, embedding, persist_directory="chroma_db")
-    return db.as_retriever(search_kwargs={"k": 5})
+retriever = init_retriever()
 
-# Load Arabic QA model
+# Arabic QA pipeline (extractive)
 qa_pipeline = pipeline(
     "question-answering",
-    model=AutoModelForQuestionAnswering.from_pretrained("alyaa82/aravec-bert-base-qa"),
-    tokenizer=AutoTokenizer.from_pretrained("alyaa82/aravec-bert-base-qa"),
+    model="ZeyadAhmed/AraElectra-Arabic-SQuADv2-QA",
+    tokenizer="ZeyadAhmed/AraElectra-Arabic-SQuADv2-QA",
     device=-1
 )
 
-# Get retriever once
-retriever = load_and_index()
-
-# Perform retrieval + QA
-def answer_question(question: str) -> str:
+def answer(question: str) -> str:
     docs = retriever.get_relevant_documents(question)
-    context = "\n\n".join(doc.page_content for doc in docs)
-    result = qa_pipeline(question=question, context=context)
-    return result['answer']
+    context = "\n\n".join(d.page_content for d in docs)
+    out = qa_pipeline(question=question, context=context)
+    return out.get("answer", "عفواً، لم أجد إجابة واضحة.")