Update app.py

app.py

@@ -10,6 +10,7 @@ import plotly.graph_objects as go
 from collections import defaultdict
 import json
 import traceback
+import spaces # Import the spaces library
 
 class MultiClientThemeClassifier:
     def __init__(self):
@@ -18,7 +19,7 @@ class MultiClientThemeClassifier:
         self.model_loaded = False
         
     def load_model(self, model_name: str = 'Qwen/Qwen3-Embedding-0.6B'):
-        """Load the embedding model"""
+        """Load the embedding model onto the GPU"""
         try:
             if self.model_loaded:
                 # If switching models, reset everything
@@ -26,10 +27,11 @@ class MultiClientThemeClassifier:
                 self.client_themes = {}
                 self.model_loaded = False
             
-            print(f"Loading model: {model_name}")
-            self.model = SentenceTransformer(model_name,trust_remote_code=True)
+            print(f"Loading model: {model_name} onto CUDA device")
+            # Load the model directly onto the GPU
+            self.model = SentenceTransformer(model_name, device='cuda', trust_remote_code=True)
             self.model_loaded = True
-            return f"✅ Model '{model_name}' loaded successfully!"
+            return f"✅ Model '{model_name}' loaded successfully onto GPU!"
         except Exception as e:
             self.model_loaded = False
             error_details = traceback.format_exc()
@@ -208,11 +210,13 @@ class MultiClientThemeClassifier:
 # Initialize the classifier
 classifier = MultiClientThemeClassifier()
 
+@spaces.GPU
 def load_model_interface(model_name: str):
     if not model_name.strip():
         model_name = 'Qwen/Qwen3-Embedding-0.6B'  # Default
     return classifier.load_model(model_name.strip())
 
+@spaces.GPU
 def add_themes_interface(client_id: str, themes_text: str):
     if not themes_text.strip():
         return "❌ Please enter themes!"
@@ -220,6 +224,7 @@ def add_themes_interface(client_id: str, themes_text: str):
     themes = [theme.strip() for theme in themes_text.split('\n') if theme.strip()]
     return classifier.add_client_themes(client_id, themes)
 
+@spaces.GPU
 def classify_interface(text: str, client_id: str, confidence_threshold: float):
     if not text.strip():
         return "Please enter text to classify!", ""
@@ -241,6 +246,7 @@ def classify_interface(text: str, client_id: str, confidence_threshold: float):
     
     return result, ""
 
+@spaces.GPU(duration=300) # Request GPU for 300 seconds for longer benchmark jobs
 def benchmark_interface(csv_file, client_id: str):
     if csv_file is None:
         return "Please upload a CSV file!", "", ""