Update rag_pipeline.py

rag_pipeline.py

@@ -1,31 +1,47 @@
+# rag_pipeline.py
+import time
+import logging
+import numpy as np
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel
 from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from langchain.vectorstores import Chroma
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from utils import extract_text_from_files
+import chromadb
+from chromadb.config import Settings
+
+logger = logging.getLogger("RAG")
 
 class RAGPipeline:
     def __init__(self):
-        print("[RAG] جاري تحميل النموذج والمحول...")
-        self.embedding_model = HuggingFaceEmbeddings(model_name="intfloat/multilingual-e5-base")
-        self.generator = pipeline("text-generation", model="tiiuae/falcon-7b-instruct", trust_remote_code=True, device_map="auto")
-        self.db = None
-        print("[RAG] تم التحميل بنجاح.")
+        logger.info("[RAG] جاري تحميل النموذج والمحول...")
+        self.embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+        self.chunk_embeddings = []
+        self.chunks = []
+        self.client = chromadb.Client(Settings(chroma_db_impl="memory", persist_directory=None))
+        self.collection = self.client.create_collection(name="rag_collection")
+        self.tokenizer = AutoTokenizer.from_pretrained("aubmindlab/aragpt2-mega", trust_remote_code=True)
+        self.lm = AutoModelForCausalLM.from_pretrained("aubmindlab/aragpt2-mega", trust_remote_code=True)
+        logger.info("[RAG] تم التحميل بنجاح.")
+
+    def build_index(self, chunks):
+        start_time = time.time()
+        self.chunks = chunks
+        self.chunk_embeddings = self.embedding_model.encode(chunks, show_progress_bar=True)
+        logger.info(f"[RAG] تم بناء الفهرس بأبعاد {self.chunk_embeddings.shape[1]} في {time.time() - start_time:.2f} ثانية.")
+        for i, chunk in enumerate(chunks):
+            self.collection.add(documents=[chunk], ids=[str(i)], embeddings=[self.chunk_embeddings[i].tolist()])
 
-    def load_and_index(self, files):
-        text = extract_text_from_files(files)
-        splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-        chunks = splitter.split_text(text)
-        self.db = Chroma.from_texts(chunks, embedding=self.embedding_model)
-        return f"[RAG] تم بناء الفهرس لـ {len(chunks)} مقاطع."
+    def retrieve(self, query, k=5):
+        logger.info("[RAG] استرجاع المقاطع الأكثر صلة بالسؤال...")
+        query_embedding = self.embedding_model.encode([query])[0].tolist()
+        results = self.collection.query(query_embeddings=[query_embedding], n_results=k)
+        return results["documents"][0], results["ids"][0]
 
-    def answer_question(self, question):
-        if self.db is None:
-            return "⚠️ لم يتم تحميل مستندات.", []
-        docs = self.db.similarity_search(question, k=3)
-        context = "\n".join([doc.page_content for doc in docs])
-        prompt = f"أجب عن السؤال التالي بناءً على المراجع التالية فقط:\n{context}\n\nالسؤال: {question}\nالإجابة:"
-        result = self.generator(prompt, max_new_tokens=200)[0]["generated_text"]
-        answer = result.split("الإجابة:")[-1].strip()
-        return answer, [doc.page_content for doc in docs]
+    def generate_answer(self, query):
+        docs, ids = self.retrieve(query)
+        context = "\n\n".join(docs)
+        prompt = f"السياق:\n{context}\n\nالسؤال: {query}\nالإجابة:"
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        with torch.no_grad():
+            outputs = self.lm.generate(**inputs, max_new_tokens=200)
+        answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return answer, context