Update rag_pipeline.py

rag_pipeline.py

@@ -6,80 +6,34 @@ from typing import List, Dict
 
 class ArabicRAGSystem:
     def __init__(self):
-        """Initialize with guaranteed-accessible Arabic models"""
-        # Verified embedding models (publicly available)
-        self.embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
-        
-        # Main Arabic LLM with local fallback
-        try:
-            self.tokenizer = AutoTokenizer.from_pretrained("inception-mbzuai/jais-13b-chat")
-            self.llm = AutoModelForCausalLM.from_pretrained("inception-mbzuai/jais-13b-chat")
-        except:
-            try:
-                self.tokenizer = AutoTokenizer.from_pretrained("aubmindlab/aragpt2-base")
-                self.llm = AutoModelForCausalLM.from_pretrained("aubmindlab/aragpt2-base")
-            except:
-                raise Exception("Failed to load any Arabic language model. Please check internet connection and try again.")
-        
-        self.index = faiss.IndexFlatL2(384)  # Multilingual MiniLM uses 384-dim
+        """Initialize with dependency-safe models"""
+        # Verified working embedding model
+        self.embedding_model = SentenceTransformer(
+            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
+            device="cpu"
+        )
+        
+        # Load Arabic LLM with safe tokenizer settings
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            "aubmindlab/aragpt2-base",
+            use_safetensors=True
+        )
+        self.llm = AutoModelForCausalLM.from_pretrained(
+            "aubmindlab/aragpt2-base",
+            use_safetensors=True,
+            device_map="auto",
+            torch_dtype="auto"
+        )
+        
+        self.index = faiss.IndexFlatL2(384)  # Matching embedding dim
 
     def generate_answer(self, question: str, documents: List[Dict], 
                       top_k: int = 3, temperature: float = 0.7) -> tuple:
-        """Robust generation with guaranteed fallbacks"""
-        # Index documents
+        """Optimized generation with memory safety"""
+        # Convert documents to embeddings
         texts = [doc["text"] for doc in documents]
-        self.index.add(np.array(self.embedding_model.encode(texts)))
-        
-        # Simple semantic search (no cross-encoder dependency)
-        query_embedding = self.embedding_model.encode([question])
-        distances, indices = self.index.search(query_embedding, top_k)
-        
-        # Prepare context with metadata
-        context = "\n\n".join([
-            f"المرجع: {documents[idx]['source']}\n"
-            f"الصفحة: {documents[idx].get('page', 'N/A')}\n"
-            f"النص: {documents[idx]['text']}\n"
-            for idx in indices[0]
-        ])
-        
-        # Generation with bulletproof prompt
-        prompt = f"""
-        أنت مساعد ذكي متخصص في النصوص الدينية العربية. أجب على السؤال بناءً على السياق التالي فقط:
-        
-        السياق:
-        {context}
-        
-        السؤال: {question}
-        
-        التعليمات:
-        1. استخدم المعلومات من السياق فقط
-        2. أجب باللغة العربية الفصحى
-        3. أشر إلى المصادر بهذا الشكل: [المرجع: اسم الملف، الصفحة]
-        4. إذا لم تجد إجابة واضحة قل "لا توجد معلومات كافية في النصوص المقدمة"
-        
-        الإجابة:
-        """.strip()
-        
-        try:
-            inputs = self.tokenizer(prompt, return_tensors="pt")
-            outputs = self.llm.generate(
-                inputs.input_ids,
-                max_new_tokens=512,
-                temperature=temperature,
-                do_sample=True,
-                pad_token_id=self.tokenizer.eos_token_id
-            )
-            answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            answer = answer.split("الإجابة:")[-1].strip()
-        except Exception as e:
-            answer = f"عذراً، حدث خطأ في معالجة السؤال. التفاصيل: {str(e)}"
-        
-        # Prepare sources
-        sources = [{
-            "text": documents[idx]["text"],
-            "source": documents[idx]["source"],
-            "page": documents[idx].get("page", "N/A"),
-            "score": float(1 - distances[0][i])
-        } for i, idx in enumerate(indices[0])]
+        embeddings = self.embedding_model.encode(texts, convert_to_numpy=True)
+        self.index.add(embeddings)
         
-        return answer, sources
+        # Semantic search
+        query_embedding = self.embedding_model