Update app.py

app.py

@@ -13,6 +13,7 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import PyPDF2
 import io
+import spaces  # Added for @spaces.GPU decorator
 
 # Set up logging
 logging.basicConfig(
@@ -42,13 +43,11 @@ class Vision2030Assistant:
         logger.info("Assistant initialized successfully")
 
     def load_embedding_models(self):
-        """Load Arabic and English embedding models with fallback mechanism."""
+        """Load Arabic and English embedding models on CPU."""
         try:
             self.arabic_embedder = SentenceTransformer('CAMeL-Lab/bert-base-arabic-camelbert-ca')
             self.english_embedder = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
-            if has_gpu:
-                self.arabic_embedder = self.arabic_embedder.to('cuda')
-                self.english_embedder = self.english_embedder.to('cuda')
+            # Models remain on CPU; GPU usage handled in decorated functions
             logger.info("Embedding models loaded successfully")
         except Exception as e:
             logger.error(f"Failed to load embedding models: {e}")
@@ -58,7 +57,7 @@ class Vision2030Assistant:
         """Fallback method for embedding models using a simple random vector approach."""
         logger.warning("Using fallback embedding method")
         class SimpleEmbedder:
-            def encode(self, text):
+            def encode(self, text, device=None):  # Added device parameter for compatibility
                 import hashlib
                 hash_obj = hashlib.md5(text.encode())
                 np.random.seed(int(hash_obj.hexdigest(), 16) % 2**32)
@@ -67,17 +66,15 @@ class Vision2030Assistant:
         self.english_embedder = SimpleEmbedder()
 
     def load_language_model(self):
-        """Load the DistilGPT-2 language model for response generation."""
+        """Load the DistilGPT-2 language model on CPU."""
         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
+                'text-generation',
+                model=self.model,
+                tokenizer=self.tokenizer,
+                device=-1  # CPU
             )
             logger.info("Language model loaded successfully")
         except Exception as e:
@@ -100,7 +97,7 @@ class Vision2030Assistant:
         self.pdf_arabic_texts = []
 
     def _create_indices(self):
-        """Create FAISS indices for the initial knowledge base."""
+        """Create FAISS indices for the initial knowledge base on CPU."""
         try:
             # English index
             english_vectors = [self.english_embedder.encode(text) for text in self.english_texts]
@@ -123,21 +120,21 @@ class Vision2030Assistant:
     def _create_sample_eval_data(self):
         """Create sample evaluation data for testing factual accuracy."""
         self.eval_data = [
-            {"question": "What are the key pillars of Vision 2030?", 
-             "lang": "en", 
+            {"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", 
+            {"question": "ما هي الركائز الرئيسية لرؤية 2030؟",
+             "lang": "ar",
              "reference": "الركائز الرئيسية لرؤية 2030 هي مجتمع حيوي، واقتصاد مزدهر، ووطن طموح."}
         ]
 
-    def retrieve_context(self, query, lang, session_id):
-        """Retrieve relevant context based on the query and session history."""
+    def retrieve_context(self, query, lang, session_id, device='cpu'):
+        """Retrieve relevant context using the specified device for encoding."""
         try:
             history = self.session_history.get(session_id, [])
             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)
+            query_vec = embedder.encode(query, device=device)
 
             if lang == "ar":
                 if self.has_pdf_content and self.pdf_arabic_texts:
@@ -161,25 +158,30 @@ class Vision2030Assistant:
             logger.error(f"Retrieval error: {e}")
             return "Error retrieving context."
 
+    @spaces.GPU
     def generate_response(self, query, session_id):
-        """Generate a response to the user's query using context and session history."""
+        """Generate a response using GPU resources when available."""
         if not query.strip():
             return "Please enter a valid question."
-        
+
         start_time = time.time()
         try:
             lang = "ar" if any('\u0600' <= c <= '\u06FF' for c in query) else "en"
-            context = self.retrieve_context(query, lang, session_id)
-            
+            context = self.retrieve_context(query, lang, session_id, device='cuda')
+
             if "Error" in context or "No relevant" in context:
                 reply = context
             elif self.generator:
+                # Move the language model to GPU
+                self.generator.model.to('cuda')
                 prompt = f"Context: {context}\nQuestion: {query}\nAnswer:"
                 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()
+                # Move the language model back to CPU
+                self.generator.model.to('cpu')
             else:
                 reply = context
-            
+
             self.session_history.setdefault(session_id, []).append((query, reply))
             self.metrics["response_times"].append(time.time() - start_time)
             return reply
@@ -199,25 +201,26 @@ class Vision2030Assistant:
             logger.error(f"Evaluation error: {e}")
             return 0.0
 
+    @spaces.GPU
     def process_pdf(self, file):
-        """Process an uploaded PDF file and update the knowledge base."""
+        """Process a PDF file and update the knowledge base using GPU for encoding."""
         if not file:
             return "Please upload a PDF file."
-        
+
         try:
             pdf_reader = PyPDF2.PdfReader(io.BytesIO(file))
             text = "".join([page.extract_text() or "" for page in pdf_reader.pages])
             if not text.strip():
                 return "No extractable text found in PDF."
-            
+
             # 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)]
 
-            # Create indices for PDF content
+            # Create indices for PDF content using GPU
             if self.pdf_english_texts:
-                english_vectors = [self.english_embedder.encode(text) for text in self.pdf_english_texts]
+                english_vectors = [self.english_embedder.encode(text, device='cuda') for text in self.pdf_english_texts]
                 dim = len(english_vectors[0])
                 nlist = max(1, len(english_vectors) // 10)
                 quantizer = faiss.IndexFlatL2(dim)
@@ -226,7 +229,7 @@ class Vision2030Assistant:
                 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]
+                arabic_vectors = [self.arabic_embedder.encode(text, device='cuda') for text in self.pdf_arabic_texts]
                 dim = len(arabic_vectors[0])
                 nlist = max(1, len(arabic_vectors) // 10)
                 quantizer = faiss.IndexFlatL2(dim)