Update rag_pipeline.py

rag_pipeline.py

@@ -16,40 +16,48 @@ class RAGPipeline:
 
         self.index = None
         self.chunks = []
+        self.summaries = []
         self.chunk_embeddings = []
 
         print("[RAG] تم تحميل النماذج بنجاح.")
 
+    def summarize_text(self, text):
+        prompt = f"لخص النص التالي باللغة العربية:\n\n{text}"
+        try:
+            inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
+            summary_ids = self.model.generate(inputs["input_ids"], max_length=256)
+            return self.tokenizer.decode(summary_ids[0], skip_special_tokens=True).strip()
+        except Exception as e:
+            print(f"[RAG] خطأ في التلخيص: {e}")
+            return ""
+
     def build_index(self, chunks, logs=None):
         self.chunks = chunks
+        self.summaries = [self.summarize_text(chunk) for chunk in chunks]
         self.chunk_embeddings = self.embedder.encode(chunks, convert_to_numpy=True)
-        if logs is not None:
-            logs.append(f"[RAG] تم بناء الفهرس بـ {self.chunk_embeddings.shape[0]} مقطع.")
         self.index = np.array(self.chunk_embeddings)
+        if logs is not None:
+            logs.append(f"[RAG] تم بناء الفهرس بـ {len(self.chunk_embeddings)} مقطع.")
 
     def answer(self, question):
-        # Step 1: استرجاع المقاطع الأكثر صلة
         question_embedding = self.embedder.encode([question], convert_to_numpy=True)
         similarities = np.dot(self.index, question_embedding.T).squeeze()
         top_idx = similarities.argsort()[-5:][::-1]
+
+        relevant_summaries = [self.summaries[i] for i in top_idx if self.summaries[i].strip()]
         sources = [self.chunks[i] for i in top_idx]
-        context = " ".join(sources)
 
-        # Step 2: تلخيص النص المسترجع
+        combined_summary = " ".join(relevant_summaries).strip()
+        if not combined_summary:
+            combined_summary = " ".join(sources)
+
+        qa_prompt = f"أجب عن السؤال التالي بناء على النص:\n\n{combined_summary}\n\nالسؤال: {question}\nالإجابة:"
         try:
-            summary_prompt = f"لخص النص التالي باللغة العربية:\n\n{context}"
-            inputs = self.tokenizer(summary_prompt, return_tensors="pt", truncation=True, max_length=512)
-            summary_ids = self.model.generate(inputs["input_ids"], max_length=256)
-            summary = self.tokenizer.decode(summary_ids[0], skip_special_tokens=True).strip()
+            inputs = self.tokenizer(qa_prompt, return_tensors="pt", truncation=True, max_length=512)
+            output_ids = self.model.generate(inputs["input_ids"], max_length=200)
+            answer = self.tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
         except Exception as e:
-            summary = ""
-            print(f"[RAG] خطأ في التلخيص: {e}")
-
-        # Step 3: توليد الإجابة من الملخص أو من النص الأصلي
-        qa_context = summary if summary else context
-        qa_prompt = f"أجب عن السؤال التالي بناء على النص:\n\n{qa_context}\n\nالسؤال: {question}\nالإجابة:"
-        qa_inputs = self.tokenizer(qa_prompt, return_tensors="pt", truncation=True, max_length=512)
-        answer_ids = self.model.generate(qa_inputs["input_ids"], max_length=200)
-        answer = self.tokenizer.decode(answer_ids[0], skip_special_tokens=True).strip()
+            print(f"[RAG] خطأ في توليد الإجابة: {e}")
+            answer = ""
 
-        return answer, sources, summary
+        return answer, sources, combined_summary