Revert

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,57 +27,32 @@ def get_embedding(text):
 def reduce_to_3d(embedding):
     return embedding[:3]
 
-def 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]
+def compare_embeddings(text1, text2):
+    emb1 = get_embedding(text1)
+    emb2 = get_embedding(text2)
     
-    fig = go.Figure()
-
-    # Add black origin point
-    fig.add_trace(go.Scatter3d(x=[0], y=[0], z=[0], mode='markers', marker=dict(size=5, color='black'), name='Origin'))
-
-    # Add lines and points for each text
-    for i, emb in enumerate(embeddings_3d):
-        color = 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}'))
+    emb1_3d = reduce_to_3d(emb1)
+    emb2_3d = reduce_to_3d(emb2)
+    
+    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
 
-def add_textbox(num_textboxes):
-    return gr.Textbox.update(visible=True, label=f"Text {num_textboxes + 1}"), num_textboxes + 1
-
-def remove_textbox(num_textboxes):
-    return gr.Textbox.update(visible=False), max(2, num_textboxes - 1)
-
-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_textboxes = gr.State(value=2)
-    
-    with gr.Column() as textbox_container:
-        textboxes = [gr.Textbox(label=f"Text {i+1}") for i in range(10)]  # Create 10 textboxes
-        for i in range(2, 10):  # Hide textboxes 3-10 initially
-            textboxes[i].visible = False
-    
-    with gr.Row():
-        add_btn = gr.Button("Add String")
-        remove_btn = gr.Button("Remove String")
-    
-    plot_output = gr.Plot()
-    
-    submit_btn = gr.Button("Submit")
-    clear_btn = gr.ClearButton(components=textboxes + [plot_output], value="Clear")
-    
-    add_btn.click(add_textbox, inputs=[num_textboxes], outputs=[textboxes[-1], num_textboxes])
-    remove_btn.click(remove_textbox, inputs=[num_textboxes], outputs=[textboxes[-1], num_textboxes])
-    submit_btn.click(compare_embeddings, inputs=textboxes, outputs=[plot_output])
+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()