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,91 +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(strings):
-    embeddings = [get_embedding(text) for text in strings 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 origin point
-    fig.add_trace(go.Scatter3d(x=[0], y=[0], z=[0], mode='markers', name='Origin',
-                               marker=dict(color='black', size=8)))
-
-    # Add lines and points for each string
-    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', name=f'Text {i+1}',
-                                   line=dict(color=color), marker=dict(color=color)))
-
-    fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
+    emb1_3d = reduce_to_3d(emb1)
+    emb2_3d = reduce_to_3d(emb2)
     
-    return fig
-
-def add_string(string_list):
-    string_list.append("")
-    return string_list
-
-def remove_string(string_list):
-    if len(string_list) > 1:
-        string_list.pop()
-    return string_list
-
-with gr.Blocks(title="3D Embedding Comparison") as iface:
-    gr.Markdown("# 3D Embedding Comparison")
-    gr.Markdown("Compare the embeddings of multiple strings visualized in 3D space using Mistral 7B.")
-    
-    with gr.Row():
-        with gr.Column(scale=1):
-            string_inputs = gr.State([""])
-            with gr.Column(variant="panel") as textbox_container:
-                textboxes = gr.Textbox(label="Text 1")
-            add_button = gr.Button("Add String")
-            remove_button = gr.Button("Remove String")
-        
-        with gr.Column(scale=2):
-            plot_output = gr.Plot()
+    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')
+    ])
     
-    submit_button = gr.Button("Compare Embeddings")
-
-    def update_textboxes(string_list):
-        components = []
-        for i, s in enumerate(string_list):
-            components.append(gr.Textbox(value=s, label=f"Text {i+1}"))
-        return gr.Column.update(visible=True, children=components)
-
-    add_button.click(
-        add_string,
-        inputs=[string_inputs],
-        outputs=[string_inputs]
-    ).then(
-        update_textboxes,
-        inputs=[string_inputs],
-        outputs=[textbox_container]
-    )
-
-    remove_button.click(
-        remove_string,
-        inputs=[string_inputs],
-        outputs=[string_inputs]
-    ).then(
-        update_textboxes,
-        inputs=[string_inputs],
-        outputs=[textbox_container]
-    )
+    fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
     
-    def process_and_compare(strings):
-        processed_strings = [t.strip() for t in strings if t.strip()]
-        return compare_embeddings(processed_strings)
+    return fig
 
-    submit_button.click(
-        process_and_compare,
-        inputs=[string_inputs],
-        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()