Update app.py

app.py

@@ -1,5 +1,5 @@
 # Vision 2030 Virtual Assistant with RAG and Evaluation Framework
-# Modified for Hugging Face Spaces compatibility
+# Modified for Hugging Face Spaces compatibility with GPU support
 
 import gradio as gr
 import time
@@ -16,6 +16,8 @@ import json
 from langdetect import detect
 from sentence_transformers import SentenceTransformer
 import faiss
+import torch
+import spaces
 
 # Configure logging
 logging.basicConfig(
@@ -27,6 +29,10 @@ logging.basicConfig(
 )
 logger = logging.getLogger('vision2030_assistant')
 
+# Check for GPU availability
+has_gpu = torch.cuda.is_available()
+logger.info(f"GPU available: {has_gpu}")
+
 class Vision2030Assistant:
     def __init__(self, pdf_path=None, eval_data_path=None):
         """
@@ -64,18 +70,53 @@ class Vision2030Assistant:
         self.response_history = []
         logger.info("Vision 2030 Assistant initialized successfully")
     
+    @spaces.GPU
     def load_embedding_models(self):
-        """Load embedding models for retrieval"""
-        logger.info("Loading embedding models...")
+        """Load embedding models for retrieval with GPU support"""
+        logger.info("Loading embedding models with GPU support...")
         
         try:
             # Load embedding models
             self.arabic_embedder = SentenceTransformer('CAMeL-Lab/bert-base-arabic-camelbert-ca')
             self.english_embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+            
+            # Move to GPU if available
+            if has_gpu:
+                self.arabic_embedder = self.arabic_embedder.to('cuda')
+                self.english_embedder = self.english_embedder.to('cuda')
+                logger.info("Models moved to GPU")
+            
             logger.info("Embedding models loaded successfully")
         except Exception as e:
             logger.error(f"Error loading embedding models: {str(e)}")
-            raise
+            # Create simple placeholder models if loading fails
+            self._create_fallback_embedders()
+
+    def _create_fallback_embedders(self):
+        """Create fallback embedding methods if model loading fails"""
+        logger.warning("Using fallback embedding methods")
+        
+        # Simple fallback using character-level encoding (not a real embedding, just for demo)
+        def simple_encode(text, dim=384):
+            import hashlib
+            # Create a hash of the text
+            hash_object = hashlib.md5(text.encode())
+            # Use the hash to seed a random number generator
+            import numpy as np
+            np.random.seed(int(hash_object.hexdigest(), 16) % 2**32)
+            # Generate a random vector
+            return np.random.randn(dim).astype(np.float32)
+        
+        # Create embedding function objects
+        class SimpleEmbedder:
+            def __init__(self, dim=384):
+                self.dim = dim
+            
+            def encode(self, text):
+                return simple_encode(text, self.dim)
+        
+        self.arabic_embedder = SimpleEmbedder()
+        self.english_embedder = SimpleEmbedder()
 
     def load_and_process_documents(self, pdf_path):
         """Load and process the Vision 2030 document from PDF"""
@@ -152,16 +193,28 @@ class Vision2030Assistant:
             "تتضمن رؤية 2030 خططًا لتطوير البنية التحتية الرقمية والدعم للشركات الناشئة التكنولوجية في المملكة العربية السعودية."
         ]
 
+    @spaces.GPU
     def _create_indices(self):
-        """Create FAISS indices for fast text retrieval"""
+        """Create FAISS indices for fast text retrieval with GPU support"""
         logger.info("Creating FAISS indices for text retrieval")
         
         try:
             # Process and embed English texts
             self.english_vectors = []
             for text in self.english_texts:
-                vec = self.english_embedder.encode(text)
-                self.english_vectors.append(vec)
+                try:
+                    if has_gpu and hasattr(self.english_embedder, 'to') and callable(getattr(self.english_embedder, 'to')):
+                        # If it's a real model on GPU
+                        with torch.no_grad():
+                            vec = self.english_embedder.encode(text)
+                    else:
+                        # If it's our fallback
+                        vec = self.english_embedder.encode(text)
+                    self.english_vectors.append(vec)
+                except Exception as e:
+                    logger.error(f"Error encoding English text: {str(e)}")
+                    # Use a random vector as fallback
+                    self.english_vectors.append(np.random.randn(384).astype(np.float32))
             
             # Create English index
             if self.english_vectors:
@@ -174,8 +227,19 @@ class Vision2030Assistant:
             # Process and embed Arabic texts
             self.arabic_vectors = []
             for text in self.arabic_texts:
-                vec = self.arabic_embedder.encode(text)
-                self.arabic_vectors.append(vec)
+                try:
+                    if has_gpu and hasattr(self.arabic_embedder, 'to') and callable(getattr(self.arabic_embedder, 'to')):
+                        # If it's a real model on GPU
+                        with torch.no_grad():
+                            vec = self.arabic_embedder.encode(text)
+                    else:
+                        # If it's our fallback
+                        vec = self.arabic_embedder.encode(text)
+                    self.arabic_vectors.append(vec)
+                except Exception as e:
+                    logger.error(f"Error encoding Arabic text: {str(e)}")
+                    # Use a random vector as fallback
+                    self.arabic_vectors.append(np.random.randn(384).astype(np.float32))
             
             # Create Arabic index
             if self.arabic_vectors:
@@ -225,17 +289,28 @@ class Vision2030Assistant:
         ]
         logger.info(f"Created {len(self.eval_data)} sample evaluation examples")
 
+    @spaces.GPU
     def retrieve_context(self, query, lang):
-        """Retrieve relevant context for a query based on language"""
+        """Retrieve relevant context for a query based on language with GPU support"""
         start_time = time.time()
         
         try:
             if lang == "ar":
-                query_vec = self.arabic_embedder.encode(query)
+                if has_gpu and hasattr(self.arabic_embedder, 'to') and callable(getattr(self.arabic_embedder, 'to')):
+                    with torch.no_grad():
+                        query_vec = self.arabic_embedder.encode(query)
+                else:
+                    query_vec = self.arabic_embedder.encode(query)
+                
                 D, I = self.arabic_index.search(np.array([query_vec]), k=2)  # Get top 2 most relevant chunks
                 context = "\n".join([self.arabic_texts[i] for i in I[0] if i < len(self.arabic_texts) and i >= 0])
             else:
-                query_vec = self.english_embedder.encode(query)
+                if has_gpu and hasattr(self.english_embedder, 'to') and callable(getattr(self.english_embedder, 'to')):
+                    with torch.no_grad():
+                        query_vec = self.english_embedder.encode(query)
+                else:
+                    query_vec = self.english_embedder.encode(query)
+                
                 D, I = self.english_index.search(np.array([query_vec]), k=2)  # Get top 2 most relevant chunks
                 context = "\n".join([self.english_texts[i] for i in I[0] if i < len(self.english_texts) and i >= 0])
             
@@ -345,8 +420,9 @@ class Vision2030Assistant:
             
         return accuracy
 
+    @spaces.GPU
     def evaluate_on_test_set(self):
-        """Evaluate the assistant on the test set"""
+        """Evaluate the assistant on the test set with GPU support"""
         logger.info("Running evaluation on test set")
         
         eval_results = []
@@ -442,102 +518,113 @@ class Vision2030Assistant:
 
 # Create the Gradio interface
 def create_gradio_interface():
-    # Initialize the assistant
-    assistant = Vision2030Assistant()
-    
-    def chat(message, history):
-        if not message.strip():
-            return history, ""
+    try:
+        # Initialize the assistant
+        assistant = Vision2030Assistant()
         
-        # Generate response
-        reply = assistant.generate_response(message)
-        
-        # Update history
-        history.append((message, reply))
-        
-        return history, ""
-    
-    def provide_feedback(history, rating, feedback_text):
-        # Record feedback for the last conversation
-        if history and len(history) > 0:
-            last_interaction = history[-1]
-            assistant.record_user_feedback(last_interaction[0], last_interaction[1], rating, feedback_text)
-            return f"Thank you for your feedback! (Rating: {rating}/5)"
-        return "No conversation found to rate."
-    
-    def run_evaluation():
-        results = assistant.evaluate_on_test_set()
-        
-        # Create summary text
-        summary = f"""
-        Evaluation Results:
-        ------------------
-        Total questions evaluated: {len(results['detailed_results'])}
-        Overall factual accuracy: {results['average_factual_accuracy']:.2f}
-        Average response time: {results['average_response_time']:.4f} seconds
-        
-        Detailed Results:
-        """
-        
-        for i, result in enumerate(results['detailed_results']):
-            summary += f"\nQ{i+1}: {result['question']}\n"
-            summary += f"Reference: {result['reference']}\n"
-            summary += f"Response: {result['response']}\n"
-            summary += f"Accuracy: {result['factual_accuracy']:.2f}\n"
-            summary += "-" * 40 + "\n"
-        
-        # Return both the results summary and visualization
-        fig = assistant.visualize_evaluation_results(results)
+        def chat(message, history):
+            if not message.strip():
+                return history, ""
+            
+            # Generate response
+            reply = assistant.generate_response(message)
+            
+            # Update history
+            history.append((message, reply))
+            
+            return history, ""
         
-        return summary, fig
+        def provide_feedback(history, rating, feedback_text):
+            # Record feedback for the last conversation
+            if history and len(history) > 0:
+                last_interaction = history[-1]
+                assistant.record_user_feedback(last_interaction[0], last_interaction[1], rating, feedback_text)
+                return f"Thank you for your feedback! (Rating: {rating}/5)"
+            return "No conversation found to rate."
+        
+        @spaces.GPU
+        def run_evaluation():
+            results = assistant.evaluate_on_test_set()
+            
+            # Create summary text
+            summary = f"""
+            Evaluation Results:
+            ------------------
+            Total questions evaluated: {len(results['detailed_results'])}
+            Overall factual accuracy: {results['average_factual_accuracy']:.2f}
+            Average response time: {results['average_response_time']:.4f} seconds
+            
+            Detailed Results:
+            """
+            
+            for i, result in enumerate(results['detailed_results']):
+                summary += f"\nQ{i+1}: {result['question']}\n"
+                summary += f"Reference: {result['reference']}\n"
+                summary += f"Response: {result['response']}\n"
+                summary += f"Accuracy: {result['factual_accuracy']:.2f}\n"
+                summary += "-" * 40 + "\n"
+            
+            # Return both the results summary and visualization
+            fig = assistant.visualize_evaluation_results(results)
+            
+            return summary, fig
 
-    def process_uploaded_file(file):
-        if file is not None:
-            # Create a new assistant with the uploaded PDF
-            global assistant
-            assistant = Vision2030Assistant(pdf_path=file.name)
-            return f"Successfully processed {file.name}. The assistant is ready to use."
-        return "No file uploaded. Using sample data."
-    
-    # Create the Gradio interface
-    with gr.Blocks() as demo:
-        gr.Markdown("# Vision 2030 Virtual Assistant 🌟")
-        gr.Markdown("Ask questions about Saudi Arabia's Vision 2030 in both Arabic and English")
-        
-        with gr.Tab("Chat"):
-            chatbot = gr.Chatbot(height=400)
-            msg = gr.Textbox(label="Your Question", placeholder="Ask about Vision 2030...")
-            with gr.Row():
-                submit_btn = gr.Button("Submit")
-                clear_btn = gr.Button("Clear Chat")
-            
-            gr.Markdown("### Provide Feedback")
-            with gr.Row():
-                rating = gr.Slider(minimum=1, maximum=5, step=1, value=3, label="Rate the Response (1-5)")
-                feedback_text = gr.Textbox(label="Additional Comments (Optional)")
-            feedback_btn = gr.Button("Submit Feedback")
-            feedback_result = gr.Textbox(label="Feedback Status")
-        
-        with gr.Tab("Evaluation"):
-            evaluate_btn = gr.Button("Run Evaluation on Test Set")
-            eval_output = gr.Textbox(label="Evaluation Results", lines=20)
-            eval_chart = gr.Plot(label="Evaluation Metrics")
-        
-        with gr.Tab("Upload PDF"):
-            file_input = gr.File(label="Upload Vision 2030 PDF")
-            upload_result = gr.Textbox(label="Upload Status")
-            upload_btn = gr.Button("Process PDF")
-        
-        # Set up event handlers
-        msg.submit(chat, [msg, chatbot], [chatbot, msg])
-        submit_btn.click(chat, [msg, chatbot], [chatbot, msg])
-        clear_btn.click(lambda: [], None, chatbot)
-        feedback_btn.click(provide_feedback, [chatbot, rating, feedback_text], feedback_result)
-        evaluate_btn.click(run_evaluation, None, [eval_output, eval_chart])
-        upload_btn.click(process_uploaded_file, [file_input], upload_result)
-    
-    return demo
+        @spaces.GPU
+        def process_uploaded_file(file):
+            if file is not None:
+                # Create a new assistant with the uploaded PDF
+                global assistant
+                assistant = Vision2030Assistant(pdf_path=file.name)
+                return f"Successfully processed {file.name}. The assistant is ready to use."
+            return "No file uploaded. Using sample data."
+        
+        # Create the Gradio interface
+        with gr.Blocks() as demo:
+            gr.Markdown("# Vision 2030 Virtual Assistant 🌟")
+            gr.Markdown("Ask questions about Saudi Arabia's Vision 2030 in both Arabic and English")
+            
+            with gr.Tab("Chat"):
+                chatbot = gr.Chatbot(height=400)
+                msg = gr.Textbox(label="Your Question", placeholder="Ask about Vision 2030...")
+                with gr.Row():
+                    submit_btn = gr.Button("Submit")
+                    clear_btn = gr.Button("Clear Chat")
+                
+                gr.Markdown("### Provide Feedback")
+                with gr.Row():
+                    rating = gr.Slider(minimum=1, maximum=5, step=1, value=3, label="Rate the Response (1-5)")
+                    feedback_text = gr.Textbox(label="Additional Comments (Optional)")
+                feedback_btn = gr.Button("Submit Feedback")
+                feedback_result = gr.Textbox(label="Feedback Status")
+            
+            with gr.Tab("Evaluation"):
+                evaluate_btn = gr.Button("Run Evaluation on Test Set")
+                eval_output = gr.Textbox(label="Evaluation Results", lines=20)
+                eval_chart = gr.Plot(label="Evaluation Metrics")
+            
+            with gr.Tab("Upload PDF"):
+                file_input = gr.File(label="Upload Vision 2030 PDF")
+                upload_result = gr.Textbox(label="Upload Status")
+                upload_btn = gr.Button("Process PDF")
+            
+            # Set up event handlers
+            msg.submit(chat, [msg, chatbot], [chatbot, msg])
+            submit_btn.click(chat, [msg, chatbot], [chatbot, msg])
+            clear_btn.click(lambda: [], None, chatbot)
+            feedback_btn.click(provide_feedback, [chatbot, rating, feedback_text], feedback_result)
+            evaluate_btn.click(run_evaluation, None, [eval_output, eval_chart])
+            upload_btn.click(process_uploaded_file, [file_input], upload_result)
+        
+        return demo
+    except Exception as e:
+        logger.error(f"Error creating Gradio interface: {str(e)}")
+        # Create a simple demo for fallback
+        with gr.Blocks() as demo:
+            gr.Markdown("# Vision 2030 Virtual Assistant")
+            gr.Markdown("There was an error initializing the assistant. Please check the logs.")
+            gr.Markdown(f"Error: {str(e)}")
+        return demo
 
-# Launch the app
+# Launch the app with proper GPU initialization
 demo = create_gradio_interface()
 demo.launch()