Update rag_pipeline.py

rag_pipeline.py

@@ -1,31 +1,53 @@
 import time
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-import chromadb
-from chromadb.config import Settings
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import numpy as np
 
 class RAGPipeline:
     def __init__(self):
-        self.tokenizer = AutoTokenizer.from_pretrained("aubmindlab/aragpt2-mega")
-        self.model = AutoModelForCausalLM.from_pretrained("aubmindlab/aragpt2-mega")
-        self.generator = pipeline("text-generation", model=self.model, tokenizer=self.tokenizer, device=0 if torch.cuda.is_available() else -1)
-        self.client = chromadb.Client(Settings(allow_reset=True))
-        self.collection = None
-
-    def build_index(self, passages):
-        self.client.reset()
-        self.collection = self.client.create_collection(name="documents")
-        documents = [p["text"] for p in passages]
-        metadatas = [{"source": p["source"]} for p in passages]
-        ids = [str(i) for i in range(len(documents))]
-        self.collection.add(documents=documents, metadatas=metadatas, ids=ids)
-
-    def generate_answer(self, question):
-        if not self.collection:
-            return "⚠️ لا يوجد فهرس معرف.", []
-        results = self.collection.query(query_texts=[question], n_results=3)
-        retrieved = [d for d in results['documents'][0]]
-        metadatas = results['metadatas'][0]
-        prompt = question + "\n" + "\n".join(retrieved)
-        result = self.generator(prompt, max_new_tokens=150, do_sample=True)[0]['generated_text']
-        return result, [f"{md['source']}: {text[:60]}..." for md, text in zip(metadatas, retrieved)]
+        print("[RAG] Initializing tokenizer and model...")
+        self.tokenizer = AutoTokenizer.from_pretrained("aubmindlab/aragpt2-mega", trust_remote_code=True)
+        self.generator = AutoModelForCausalLM.from_pretrained("aubmindlab/aragpt2-mega", trust_remote_code=True)
+        self.chunk_embeddings = []
+        self.index = []
+        print("[RAG] Initialization done.")
+
+    def build_index(self, chunks):
+        start_time = time.time()
+        print(f"[RAG] Building index for {len(chunks)} chunks...")
+        self.chunk_embeddings = []
+        self.index = []
+
+        for i, chunk in enumerate(chunks):
+            # هنا يمكن إضافة embedding حقيقي عبر نموذج أو طريقة أخرى
+            embedding = self._dummy_embedding(chunk)
+            self.chunk_embeddings.append(embedding)
+            self.index.append(chunk)
+            if i % 10 == 0:
+                print(f"[RAG] Processed {i+1}/{len(chunks)} chunks.")
+
+        self.chunk_embeddings = np.array(self.chunk_embeddings)
+        dim = self.chunk_embeddings.shape[1]
+        print(f"[RAG] Index built with dimension {dim}. Took {time.time()-start_time:.2f} seconds.")
+        return "Index built successfully."
+
+    def _dummy_embedding(self, text):
+        # مؤقتاً فقط: تمثيل نصي عشوائي كـ embedding
+        return np.random.rand(768)
+
+    def generate_answer(self, question, passages):
+        start_time = time.time()
+        print(f"[RAG] Generating answer for question: {question}")
+
+        # دمج المقتطفات مع السؤال لصياغة سؤال موجه للنموذج
+        prompt = question + "\n\nمراجع:\n" + "\n".join(passages)
+
+        output = self.generator.generate(
+            self.tokenizer(prompt, return_tensors="pt").input_ids,
+            max_new_tokens=150,
+            do_sample=True,
+        )
+        response = self.tokenizer.decode(output[0], skip_special_tokens=True)
+
+        elapsed = time.time() - start_time
+        print(f"[RAG] Answer generated in {elapsed:.2f} seconds.")
+        return response, passages