v2p3

app.py

@@ -3,6 +3,7 @@ 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)
@@ -27,32 +28,78 @@ 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(text1, text2):
-    emb1 = get_embedding(text1)
-    emb2 = get_embedding(text2)
+def compare_embeddings(*texts):
+    embeddings = [get_embedding(text) for text in texts]
+    embeddings_3d = [reduce_to_3d(emb) for emb in embeddings]
+    
+    fig = go.Figure()
     
-    emb1_3d = reduce_to_3d(emb1)
-    emb2_3d = reduce_to_3d(emb2)
+    # Add origin point
+    fig.add_trace(go.Scatter3d(x=[0], y=[0], z=[0], mode='markers', marker=dict(size=5, color='black'), name='Origin'))
     
-    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')
-    ])
+    # 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.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
     
     return fig
 
-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()
+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.")
+            
+            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, inputs=[], outputs=[interface])
+            remove_btn.click(self.remove_string, inputs=[], outputs=[interface])
+            clear_btn.click(self.clear_inputs, inputs=self.string_boxes, 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
+        self.string_boxes.append(gr.Textbox(label=f"Text {self.num_strings}"))
+        return self.rebuild_interface()
+
+    def remove_string(self):
+        if self.num_strings > 2:
+            self.num_strings -= 1
+            self.string_boxes.pop()
+            return self.rebuild_interface()
+        return None
+
+    def clear_inputs(self, *inputs):
+        return ["" for _ in inputs]
+
+    def rebuild_interface(self):
+        return gr.update(visible=False), self.build_interface()
+
+    def launch(self):
+        self.interface.launch()
+
+dynamic_interface = DynamicInterface()
+dynamic_interface.launch()