Sergidev commited on
Commit
1c026a2
·
verified ·
1 Parent(s): 9006e63

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +21 -14
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(text1, text2):
32
- emb1 = get_embedding(text1)
33
- emb2 = get_embedding(text2)
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="Text 1"),
51
- gr.Textbox(label="Text 2")
52
  ],
53
  outputs=gr.Plot(),
54
  title="3D Embedding Comparison",
55
- description="Compare the embeddings of two strings visualized in 3D space using Mistral 7B."
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()