Update app.py

app.py

@@ -11,14 +11,18 @@ import sys
 
 model_name = "sentence-transformers/all-MiniLM-L6-v2"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModel.from_pretrained(model_name)
+model = None  # We'll load the model inside the GPU-enabled function
 
 @spaces.GPU
 def get_embedding(text):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+    global model
+    if model is None:
+        model = AutoModel.from_pretrained(model_name).cuda()
+    
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to('cuda')
     with torch.no_grad():
         outputs = model(**inputs)
-    return outputs.last_hidden_state.mean(dim=1).squeeze().numpy()
+    return outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()
 
 def compress_to_3d(embedding):
     pca = PCA(n_components=3)
@@ -52,5 +56,4 @@ iface = gr.Interface(
     description="Compare the embeddings of two strings visualized in 3D space."
 )
 
-iface.launch()
-demo.launch()
+iface.launch()