Update app.py
Browse files
app.py
CHANGED
@@ -3,6 +3,7 @@ import spaces
|
|
3 |
import torch
|
4 |
from transformers import AutoTokenizer, AutoModel
|
5 |
import plotly.graph_objects as go
|
|
|
6 |
|
7 |
model_name = "mistralai/Mistral-7B-v0.1"
|
8 |
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
@@ -28,18 +29,25 @@ def reduce_to_3d(embedding):
|
|
28 |
return embedding[:3]
|
29 |
|
30 |
@spaces.GPU
|
31 |
-
def compare_embeddings(
|
32 |
-
|
33 |
-
|
34 |
-
|
35 |
-
emb1_3d = reduce_to_3d(emb1)
|
36 |
-
emb2_3d = reduce_to_3d(emb2)
|
37 |
-
|
38 |
-
fig = go.Figure(data=[
|
39 |
-
go.Scatter3d(x=[0, emb1_3d[0]], y=[0, emb1_3d[1]], z=[0, emb1_3d[2]], mode='lines+markers', name='Text 1'),
|
40 |
-
go.Scatter3d(x=[0, emb2_3d[0]], y=[0, emb2_3d[1]], z=[0, emb2_3d[2]], mode='lines+markers', name='Text 2')
|
41 |
-
])
|
42 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
43 |
fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
|
44 |
|
45 |
return fig
|
@@ -47,12 +55,11 @@ def compare_embeddings(text1, text2):
|
|
47 |
iface = gr.Interface(
|
48 |
fn=compare_embeddings,
|
49 |
inputs=[
|
50 |
-
gr.Textbox(label="
|
51 |
-
gr.Textbox(label="Text 2")
|
52 |
],
|
53 |
outputs=gr.Plot(),
|
54 |
title="3D Embedding Comparison",
|
55 |
-
description="Compare the embeddings of
|
56 |
)
|
57 |
|
58 |
iface.launch()
|
|
|
3 |
import torch
|
4 |
from transformers import AutoTokenizer, AutoModel
|
5 |
import plotly.graph_objects as go
|
6 |
+
import numpy as np
|
7 |
|
8 |
model_name = "mistralai/Mistral-7B-v0.1"
|
9 |
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
|
29 |
return embedding[:3]
|
30 |
|
31 |
@spaces.GPU
|
32 |
+
def compare_embeddings(text_input):
|
33 |
+
texts = text_input.split('\n')
|
34 |
+
embeddings = [get_embedding(text) for text in texts]
|
35 |
+
embeddings_3d = [reduce_to_3d(emb) for emb in embeddings]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
36 |
|
37 |
+
fig = go.Figure()
|
38 |
+
|
39 |
+
# Add origin point (black)
|
40 |
+
fig.add_trace(go.Scatter3d(x=[0], y=[0], z=[0], mode='markers', name='Origin',
|
41 |
+
marker=dict(size=5, color='black')))
|
42 |
+
|
43 |
+
# Add lines and points for each text embedding
|
44 |
+
colors = ['red', 'blue', 'green', 'purple', 'orange', 'cyan', 'magenta', 'yellow']
|
45 |
+
for i, emb in enumerate(embeddings_3d):
|
46 |
+
color = colors[i % len(colors)]
|
47 |
+
fig.add_trace(go.Scatter3d(x=[0, emb[0]], y=[0, emb[1]], z=[0, emb[2]],
|
48 |
+
mode='lines+markers', name=f'Text {i+1}',
|
49 |
+
line=dict(color=color), marker=dict(color=color)))
|
50 |
+
|
51 |
fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
|
52 |
|
53 |
return fig
|
|
|
55 |
iface = gr.Interface(
|
56 |
fn=compare_embeddings,
|
57 |
inputs=[
|
58 |
+
gr.Textbox(label="Input Texts", lines=5, placeholder="Enter multiple texts, each on a new line")
|
|
|
59 |
],
|
60 |
outputs=gr.Plot(),
|
61 |
title="3D Embedding Comparison",
|
62 |
+
description="Compare the embeddings of multiple strings visualized in 3D space using Mistral 7B."
|
63 |
)
|
64 |
|
65 |
iface.launch()
|