Update app.py

app.py

@@ -3,6 +3,7 @@ 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)
@@ -28,18 +29,25 @@ def reduce_to_3d(embedding):
     return embedding[:3]
 
 @spaces.GPU
-def compare_embeddings(text1, text2):
-    emb1 = get_embedding(text1)
-    emb2 = get_embedding(text2)
-    
-    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')
-    ])
+def compare_embeddings(text_input):
+    texts = text_input.split('\n')
+    embeddings = [get_embedding(text) for text in texts]
+    embeddings_3d = [reduce_to_3d(emb) for emb in embeddings]
     
+    fig = go.Figure()
+
+    # Add origin point (black)
+    fig.add_trace(go.Scatter3d(x=[0], y=[0], z=[0], mode='markers', name='Origin',
+                               marker=dict(size=5, color='black')))
+
+    # Add lines and points for each text embedding
+    colors = ['red', 'blue', 'green', 'purple', 'orange', 'cyan', 'magenta', 'yellow']
+    for i, emb in enumerate(embeddings_3d):
+        color = colors[i % len(colors)]
+        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'))
     
     return fig
@@ -47,12 +55,11 @@ def compare_embeddings(text1, text2):
 iface = gr.Interface(
     fn=compare_embeddings,
     inputs=[
-        gr.Textbox(label="Text 1"),
-        gr.Textbox(label="Text 2")
+        gr.Textbox(label="Input Texts", lines=5, placeholder="Enter multiple texts, each on a new line")
     ],
     outputs=gr.Plot(),
     title="3D Embedding Comparison",
-    description="Compare the embeddings of two strings visualized in 3D space using Mistral 7B."
+    description="Compare the embeddings of multiple strings visualized in 3D space using Mistral 7B."
 )
 
 iface.launch()