Update app.py

app.py

@@ -1,9 +1,7 @@
 import gradio as gr
-import spaces
 import torch
 from transformers import AutoTokenizer, AutoModel
 import plotly.graph_objects as go
-import numpy as np
 
 model_name = "mistralai/Mistral-7B-v0.1"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -13,22 +11,20 @@ model = None
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
 
-@spaces.GPU
 def get_embedding(text):
     global model
     if model is None:
-        model = AutoModel.from_pretrained(model_name).cuda()
+        model = AutoModel.from_pretrained(model_name)
         model.resize_token_embeddings(len(tokenizer))
     
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to('cuda')
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
     with torch.no_grad():
         outputs = model(**inputs)
-    return outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()
+    return outputs.last_hidden_state.mean(dim=1).squeeze().detach().numpy()
 
 def reduce_to_3d(embedding):
     return embedding[:3]
 
-@spaces.GPU
 def compare_embeddings(*texts):
     embeddings = [get_embedding(text) for text in texts if text.strip()]  # Only process non-empty texts
     embeddings_3d = [reduce_to_3d(emb) for emb in embeddings]
@@ -51,27 +47,23 @@ def compare_embeddings(*texts):
     
     return fig
 
-def create_interface(num_inputs):
-    with gr.Blocks() as new_interface:
-        text_inputs = [gr.Textbox(label=f"Text {i+1}") for i in range(num_inputs)]
-        output = gr.Plot()
-        submit_btn = gr.Button("Compare Embeddings")
-        submit_btn.click(fn=compare_embeddings, inputs=text_inputs, outputs=output)
-    return new_interface
-
 with gr.Blocks() as iface:
     gr.Markdown("# 3D Embedding Comparison")
     gr.Markdown("Compare the embeddings of multiple strings visualized in 3D space using Mistral 7B.")
     
-    num_inputs = gr.Slider(minimum=2, maximum=10, step=1, value=2, label="Number of texts to compare")
-    interface_container = gr.HTML()
-    
-    def update_interface(num):
-        return create_interface(num)
+    with gr.Row():
+        num_inputs = gr.Slider(minimum=2, maximum=10, step=1, value=2, label="Number of texts to compare")
     
-    num_inputs.change(fn=update_interface, inputs=[num_inputs], outputs=[interface_container])
+    with gr.Row() as text_container:
+        text_inputs = [gr.Textbox(label=f"Text {i+1}") for i in range(2)]
     
-    # Initialize the interface with 2 text boxes
-    interface_container.update(create_interface(2))
+    output = gr.Plot()
+    submit_btn = gr.Button("Compare Embeddings")
+
+    def update_text_inputs(num):
+        return {text_container: gr.Row.update(children=[gr.Textbox(label=f"Text {i+1}") for i in range(num)])}
+
+    num_inputs.change(fn=update_text_inputs, inputs=[num_inputs], outputs=[text_container])
+    submit_btn.click(fn=compare_embeddings, inputs=text_container.children, outputs=output)
 
 iface.launch()