Update app.py

app.py

@@ -1,3 +1,4 @@
+# Import necessary libraries
 import gradio as gr
 import time
 import logging
@@ -6,16 +7,14 @@ import re
 from datetime import datetime
 import numpy as np
 import pandas as pd
-import matplotlib.pyplot as plt
 from sentence_transformers import SentenceTransformer, util
 import faiss
 import torch
-import spaces
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import PyPDF2
 import io
 
-# Configure logging for debugging and monitoring
+# Set up logging
 logging.basicConfig(
     level=logging.INFO,
     format='%(asctime)s - %(levelname)s - %(message)s',
@@ -27,36 +26,23 @@ logger = logging.getLogger('Vision2030Assistant')
 has_gpu = torch.cuda.is_available()
 logger.info(f"GPU available: {has_gpu}")
 
+# Define the Vision2030Assistant class
 class Vision2030Assistant:
     def __init__(self):
-        """Initialize the assistant with enhanced features"""
+        """Initialize the Vision 2030 Assistant with models, knowledge base, and indices."""
         logger.info("Initializing Vision 2030 Assistant...")
-        
-        # Load models with error handling
         self.load_embedding_models()
         self.load_language_model()
-        
-        # Initialize knowledge base and indices
         self._create_knowledge_base()
         self._create_indices()
-        
-        # Sample evaluation data
         self._create_sample_eval_data()
-        
-        # Metrics storage
         self.metrics = {"response_times": [], "user_ratings": [], "factual_accuracy": []}
-        
-        # Session management
-        self.session_history = {}
-        
-        # PDF content flag
-        self.has_pdf_content = False
-        
+        self.session_history = {}  # Dictionary to store session history
+        self.has_pdf_content = False  # Flag to indicate if PDF content is available
         logger.info("Assistant initialized successfully")
 
-    @spaces.GPU
     def load_embedding_models(self):
-        """Load embedding models with fallback"""
+        """Load Arabic and English embedding models with fallback mechanism."""
         try:
             self.arabic_embedder = SentenceTransformer('CAMeL-Lab/bert-base-arabic-camelbert-ca')
             self.english_embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
@@ -69,37 +55,37 @@ class Vision2030Assistant:
             self._fallback_embedding()
 
     def _fallback_embedding(self):
-        """Fallback to simple embedding if model loading fails"""
+        """Fallback method for embedding models using a simple random vector approach."""
         logger.warning("Using fallback embedding method")
-        def simple_embed(text):
-            import hashlib
-            hash_obj = hashlib.md5(text.encode())
-            np.random.seed(int(hash_obj.hexdigest(), 16) % 2**32)
-            return np.random.randn(384).astype(np.float32)
-        
         class SimpleEmbedder:
             def encode(self, text):
-                return simple_embed(text)
-        
+                import hashlib
+                hash_obj = hashlib.md5(text.encode())
+                np.random.seed(int(hash_obj.hexdigest(), 16) % 2**32)
+                return np.random.randn(384).astype(np.float32)
         self.arabic_embedder = SimpleEmbedder()
         self.english_embedder = SimpleEmbedder()
 
-    @spaces.GPU
     def load_language_model(self):
-        """Load language model for advanced response generation"""
+        """Load the DistilGPT-2 language model for response generation."""
         try:
             self.tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
             self.model = AutoModelForCausalLM.from_pretrained("distilgpt2")
             if has_gpu:
                 self.model = self.model.to('cuda')
-            self.generator = pipeline('text-generation', model=self.model, tokenizer=self.tokenizer, device=0 if has_gpu else -1)
+            self.generator = pipeline(
+                'text-generation', 
+                model=self.model, 
+                tokenizer=self.tokenizer, 
+                device=0 if has_gpu else -1
+            )
             logger.info("Language model loaded successfully")
         except Exception as e:
             logger.error(f"Failed to load language model: {e}")
             self.generator = None
 
     def _create_knowledge_base(self):
-        """Create initial knowledge base"""
+        """Initialize the knowledge base with basic Vision 2030 information."""
         self.english_texts = [
             "Vision 2030 is Saudi Arabia's strategic framework to reduce dependence on oil, diversify the economy, and develop public sectors.",
             "The key pillars of Vision 2030 are a vibrant society, a thriving economy, and an ambitious nation.",
@@ -113,11 +99,10 @@ class Vision2030Assistant:
         self.pdf_english_texts = []
         self.pdf_arabic_texts = []
 
-    @spaces.GPU
     def _create_indices(self):
-        """Create scalable FAISS indices"""
+        """Create FAISS indices for the initial knowledge base."""
         try:
-            # English index with IVF for scalability
+            # English index
             english_vectors = [self.english_embedder.encode(text) for text in self.english_texts]
             dim = len(english_vectors[0])
             nlist = max(1, len(english_vectors) // 10)
@@ -125,44 +110,50 @@ class Vision2030Assistant:
             self.english_index = faiss.IndexIVFFlat(quantizer, dim, nlist)
             self.english_index.train(np.array(english_vectors))
             self.english_index.add(np.array(english_vectors))
-            
+
             # Arabic index
             arabic_vectors = [self.arabic_embedder.encode(text) for text in self.arabic_texts]
             self.arabic_index = faiss.IndexIVFFlat(quantizer, dim, nlist)
             self.arabic_index.train(np.array(arabic_vectors))
             self.arabic_index.add(np.array(arabic_vectors))
-            
             logger.info("FAISS indices created successfully")
         except Exception as e:
             logger.error(f"Error creating indices: {e}")
 
     def _create_sample_eval_data(self):
-        """Sample evaluation data"""
+        """Create sample evaluation data for testing factual accuracy."""
         self.eval_data = [
-            {"question": "What are the key pillars of Vision 2030?", "lang": "en", "reference": "The key pillars of Vision 2030 are a vibrant society, a thriving economy, and an ambitious nation."},
-            {"question": "ما هي الركائز الرئيسية لرؤية 2030؟", "lang": "ar", "reference": "الركائز الرئيسية لرؤية 2030 هي مجتمع حيوي، واقتصاد مزدهر، ووطن طموح."}
+            {"question": "What are the key pillars of Vision 2030?", 
+             "lang": "en", 
+             "reference": "The key pillars of Vision 2030 are a vibrant society, a thriving economy, and an ambitious nation."},
+            {"question": "ما هي الركائز الرئيسية لرؤية 2030؟", 
+             "lang": "ar", 
+             "reference": "الركائز الرئيسية لرؤية 2030 هي مجتمع حيوي، واقتصاد مزدهر، ووطن طموح."}
         ]
 
-    @spaces.GPU
     def retrieve_context(self, query, lang, session_id):
-        """Retrieve context with session history integration"""
+        """Retrieve relevant context based on the query and session history."""
         try:
-            # Incorporate session history
             history = self.session_history.get(session_id, [])
-            history_context = " ".join([f"Q: {q} A: {a}" for q, a in history[-2:]])  # Last 2 interactions
-            
-            # Embed query
+            history_context = " ".join([f"Q: {q} A: {a}" for q, a in history[-2:]])
             embedder = self.arabic_embedder if lang == "ar" else self.english_embedder
             query_vec = embedder.encode(query)
-            
-            # Search appropriate index
-            index = self.pdf_arabic_index if (lang == "ar" and self.has_pdf_content) else \
-                    self.pdf_english_index if (lang == "en" and self.has_pdf_content) else \
-                    self.arabic_index if lang == "ar" else self.english_index
-            texts = self.pdf_arabic_texts if (lang == "ar" and self.has_pdf_content) else \
-                    self.pdf_english_texts if (lang == "en" and self.has_pdf_content) else \
-                    self.arabic_texts if lang == "ar" else self.english_texts
-            
+
+            if lang == "ar":
+                if self.has_pdf_content and self.pdf_arabic_texts:
+                    index = self.pdf_arabic_index
+                    texts = self.pdf_arabic_texts
+                else:
+                    index = self.arabic_index
+                    texts = self.arabic_texts
+            else:
+                if self.has_pdf_content and self.pdf_english_texts:
+                    index = self.pdf_english_index
+                    texts = self.pdf_english_texts
+                else:
+                    index = self.english_index
+                    texts = self.english_texts
+
             D, I = index.search(np.array([query_vec]), k=2)
             context = "\n".join([texts[i] for i in I[0] if i >= 0]) + f"\nHistory: {history_context}"
             return context if context.strip() else "No relevant information found."
@@ -170,9 +161,8 @@ class Vision2030Assistant:
             logger.error(f"Retrieval error: {e}")
             return "Error retrieving context."
 
-    @spaces.GPU
     def generate_response(self, query, session_id):
-        """Generate advanced responses with error handling"""
+        """Generate a response to the user's query using context and session history."""
         if not query.strip():
             return "Please enter a valid question."
         
@@ -188,9 +178,8 @@ class Vision2030Assistant:
                 response = self.generator(prompt, max_length=150, num_return_sequences=1, do_sample=True, temperature=0.7)
                 reply = response[0]['generated_text'].split("Answer:")[-1].strip()
             else:
-                reply = context  # Fallback
+                reply = context
             
-            # Update session history
             self.session_history.setdefault(session_id, []).append((query, reply))
             self.metrics["response_times"].append(time.time() - start_time)
             return reply
@@ -199,9 +188,9 @@ class Vision2030Assistant:
             return "Sorry, an error occurred. Please try again."
 
     def evaluate_factual_accuracy(self, response, reference):
-        """Evaluate using semantic similarity"""
+        """Evaluate the factual accuracy of a response using semantic similarity."""
         try:
-            embedder = self.english_embedder  # Assuming reference is in English; extend for Arabic if needed
+            embedder = self.english_embedder  # Assuming reference is in English for simplicity
             response_vec = embedder.encode(response)
             reference_vec = embedder.encode(reference)
             similarity = util.cos_sim(response_vec, reference_vec).item()
@@ -210,9 +199,8 @@ class Vision2030Assistant:
             logger.error(f"Evaluation error: {e}")
             return 0.0
 
-    @spaces.GPU
     def process_pdf(self, file):
-        """Process PDF with scalability and error handling"""
+        """Process an uploaded PDF file and update the knowledge base."""
         if not file:
             return "Please upload a PDF file."
         
@@ -222,56 +210,61 @@ class Vision2030Assistant:
             if not text.strip():
                 return "No extractable text found in PDF."
             
-            # Chunk text for scalability
+            # Split text into chunks
             chunks = [text[i:i+300] for i in range(0, len(text), 300)]
             self.pdf_english_texts = [c for c in chunks if not any('\u0600' <= char <= '\u06FF' for char in c)]
             self.pdf_arabic_texts = [c for c in chunks if any('\u0600' <= char <= '\u06FF' for char in c)]
-            
-            # Batch process embeddings
-            batch_size = 32
-            for lang, texts, embedder in [("en", self.pdf_english_texts, self.english_embedder), 
-                                         ("ar", self.pdf_arabic_texts, self.arabic_embedder)]:
-                if texts:
-                    vectors = []
-                    for i in range(0, len(texts), batch_size):
-                        batch = texts[i:i+batch_size]
-                        vectors.extend(embedder.encode(batch))
-                    dim = len(vectors[0])
-                    nlist = max(1, len(vectors) // 10)
-                    quantizer = faiss.IndexFlatL2(dim)
-                    index = faiss.IndexIVFFlat(quantizer, dim, nlist)
-                    index.train(np.array(vectors))
-                    index.add(np.array(vectors))
-                    setattr(self, f"pdf_{lang}_index", index)
-            
+
+            # Create indices for PDF content
+            if self.pdf_english_texts:
+                english_vectors = [self.english_embedder.encode(text) for text in self.pdf_english_texts]
+                dim = len(english_vectors[0])
+                nlist = max(1, len(english_vectors) // 10)
+                quantizer = faiss.IndexFlatL2(dim)
+                self.pdf_english_index = faiss.IndexIVFFlat(quantizer, dim, nlist)
+                self.pdf_english_index.train(np.array(english_vectors))
+                self.pdf_english_index.add(np.array(english_vectors))
+
+            if self.pdf_arabic_texts:
+                arabic_vectors = [self.arabic_embedder.encode(text) for text in self.pdf_arabic_texts]
+                dim = len(arabic_vectors[0])
+                nlist = max(1, len(arabic_vectors) // 10)
+                quantizer = faiss.IndexFlatL2(dim)
+                self.pdf_arabic_index = faiss.IndexIVFFlat(quantizer, dim, nlist)
+                self.pdf_arabic_index.train(np.array(arabic_vectors))
+                self.pdf_arabic_index.add(np.array(arabic_vectors))
+
             self.has_pdf_content = True
             return f"PDF processed: {len(self.pdf_english_texts)} English, {len(self.pdf_arabic_texts)} Arabic chunks."
         except Exception as e:
             logger.error(f"PDF processing error: {e}")
             return f"Error processing PDF: {e}"
 
-# Gradio Interface
+# Create the Gradio interface
 def create_interface():
+    """Set up the Gradio interface for chatting and PDF uploading."""
     assistant = Vision2030Assistant()
-    
+
     def chat(query, history, session_id):
         reply = assistant.generate_response(query, session_id)
         history.append((query, reply))
         return history, ""
-    
+
     with gr.Blocks() as demo:
         gr.Markdown("# Vision 2030 Virtual Assistant")
-        session_id = gr.State(value="user1")  # Simple session ID; enhance with authentication
+        session_id = gr.State(value="user1")  # Fixed session ID for simplicity
         chatbot = gr.Chatbot()
         msg = gr.Textbox(label="Ask a question")
         submit = gr.Button("Submit")
         pdf_upload = gr.File(label="Upload PDF", type="binary")
         upload_status = gr.Textbox(label="Upload Status")
-        
+
         submit.click(chat, [msg, chatbot, session_id], [chatbot, msg])
         pdf_upload.upload(assistant.process_pdf, pdf_upload, upload_status)
-    
+
     return demo
 
-demo = create_interface()
-demo.launch()
+# Launch the interface
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()