Update app.py

app.py

@@ -3,7 +3,6 @@ import spaces
 import torch
 from transformers import AutoTokenizer, AutoModel
 import plotly.graph_objects as go
-import random
 
 model_name = "mistralai/Mistral-7B-v0.1"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -28,77 +27,32 @@ def get_embedding(text):
 def reduce_to_3d(embedding):
     return embedding[:3]
 
-def generate_random_color():
-    return f'rgb({random.randint(0, 255)}, {random.randint(0, 255)}, {random.randint(0, 255)})'
-
 @spaces.GPU
-def compare_embeddings(*texts):
-    embeddings = [get_embedding(text) for text in texts if text.strip()]
-    embeddings_3d = [reduce_to_3d(emb) for emb in embeddings]
-    
-    fig = go.Figure()
+def compare_embeddings(text1, text2):
+    emb1 = get_embedding(text1)
+    emb2 = get_embedding(text2)
     
-    # Add origin point
-    fig.add_trace(go.Scatter3d(x=[0], y=[0], z=[0], mode='markers', marker=dict(size=5, color='black'), name='Origin'))
+    emb1_3d = reduce_to_3d(emb1)
+    emb2_3d = reduce_to_3d(emb2)
     
-    # Add lines and points for each text
-    for i, emb in enumerate(embeddings_3d):
-        color = generate_random_color()
-        fig.add_trace(go.Scatter3d(x=[0, emb[0]], y=[0, emb[1]], z=[0, emb[2]], mode='lines+markers',
-                                   line=dict(color=color), marker=dict(color=color), name=f'Text {i+1}'))
+    fig = go.Figure(data=[
+        go.Scatter3d(x=[0, emb1_3d[0]], y=[0, emb1_3d[1]], z=[0, emb1_3d[2]], mode='lines+markers', name='Text 1'),
+        go.Scatter3d(x=[0, emb2_3d[0]], y=[0, emb2_3d[1]], z=[0, emb2_3d[2]], mode='lines+markers', name='Text 2')
+    ])
     
     fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
     
     return fig
 
-class DynamicInterface:
-    def __init__(self):
-        self.num_strings = 2
-        self.interface = self.build_interface()
-
-    def build_interface(self):
-        with gr.Blocks() as interface:
-            gr.Markdown("# 3D Embedding Comparison")
-            gr.Markdown("Compare the embeddings of strings visualized in 3D space using Mistral 7B.")
-            
-            with gr.Column() as self.string_column:
-                self.string_boxes = [gr.Textbox(label=f"Text {i+1}") for i in range(self.num_strings)]
-            
-            with gr.Row():
-                add_btn = gr.Button("Add String")
-                remove_btn = gr.Button("Remove String")
-            
-            with gr.Row():
-                clear_btn = gr.Button("Clear")
-                submit_btn = gr.Button("Submit")
-            
-            output = gr.Plot()
-            
-            add_btn.click(self.add_string, outputs=[self.string_column])
-            remove_btn.click(self.remove_string, outputs=[self.string_column])
-            clear_btn.click(self.clear_inputs, outputs=self.string_boxes)
-            submit_btn.click(compare_embeddings, inputs=self.string_boxes, outputs=output)
-        
-        return interface
-
-    def add_string(self):
-        self.num_strings += 1
-        new_textbox = gr.Textbox(label=f"Text {self.num_strings}")
-        self.string_boxes.append(new_textbox)
-        return gr.Column.update(visible=True, children=self.string_boxes)
-
-    def remove_string(self):
-        if self.num_strings > 2:
-            self.num_strings -= 1
-            self.string_boxes.pop()
-            return gr.Column.update(visible=True, children=self.string_boxes)
-        return gr.Column.update()
-
-    def clear_inputs(self):
-        return ["" for _ in self.string_boxes]
-
-    def launch(self):
-        self.interface.launch()
-
-dynamic_interface = DynamicInterface()
-dynamic_interface.launch()
+iface = gr.Interface(
+    fn=compare_embeddings,
+    inputs=[
+        gr.Textbox(label="Text 1"),
+        gr.Textbox(label="Text 2")
+    ],
+    outputs=gr.Plot(),
+    title="3D Embedding Comparison",
+    description="Compare the embeddings of two strings visualized in 3D space using Mistral 7B."
+)
+
+iface.launch()