Update rag_pipeline.py

rag_pipeline.py

@@ -1,53 +1,56 @@
 import time
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import numpy as np
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 class RAGPipeline:
-    def __init__(self):
-        print("[RAG] Initializing tokenizer and model...")
+    def __init__(self, logger):
+        self.logger = logger
+        self.logger("[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.")
+        self.logger("[RAG] Initialization done.")
 
     def build_index(self, chunks):
-        start_time = time.time()
-        print(f"[RAG] Building index for {len(chunks)} chunks...")
+        start = time.time()
+        self.logger(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.")
+            if (i+1) % 10 == 0 or (i+1) == len(chunks):
+                self.logger(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.")
+        dim = self.chunk_embeddings.shape[1] if len(self.chunk_embeddings) > 0 else 0
+        elapsed = time.time() - start
+        self.logger(f"[RAG] Index built with dimension {dim} in {elapsed:.2f}s.")
         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}")
+    def generate_answer(self, question, top_k=3):
+        start = time.time()
+        self.logger(f"[RAG] Generating answer for question:\n{question}")
+
+        if len(self.index) == 0:
+            self.logger("[RAG] Warning: index is empty, please build index first.")
+            return "لم يتم بناء الفهرس بعد.", []
+
+        # بحث مبسط لأقرب النصوص (dummy - عشوائي)
+        passages = self.index[:top_k]
 
-        # دمج المقتطفات مع السؤال لصياغة سؤال موجه للنموذج
-        prompt = question + "\n\nمراجع:\n" + "\n".join(passages)
+        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,
-        )
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        output = self.generator.generate(inputs.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.")
+        elapsed = time.time() - start
+        self.logger(f"[RAG] Answer generated in {elapsed:.2f}s.")
         return response, passages