File size: 2,330 Bytes
7103ccc 5a0b505 748826b 7103ccc 5a0b505 748826b 5a0b505 4250d36 1ba32de 4250d36 88b0945 4250d36 1ba32de 5e78e4f 4250d36 9006e63 1ba32de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import gradio as gr
import spaces
import torch
from transformers import AutoTokenizer, AutoModel
import plotly.graph_objects as go
model_name = "mistralai/Mistral-7B-v0.1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = None
# Set pad token to eos token if not defined
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
@spaces.GPU
def get_embedding(text):
global model
if model is None:
model = AutoModel.from_pretrained(model_name).cuda()
model.resize_token_embeddings(len(tokenizer))
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512).to('cuda')
with torch.no_grad():
outputs = model(**inputs)
return outputs.last_hidden_state.mean(dim=1).squeeze().cpu().numpy()
def reduce_to_3d(embedding):
return embedding[:3]
@spaces.GPU
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()
for i, emb in enumerate(embeddings_3d):
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}'))
fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
return fig
def generate_text_boxes(n):
return [gr.Textbox(label=f"Text {i+1}") for i in range(n)]
def update_interface(n):
new_inputs = generate_text_boxes(n)
return new_inputs, gr.Plot()
with gr.Blocks() 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():
num_texts = gr.Slider(minimum=2, maximum=10, step=1, value=2, label="Number of texts to compare")
update_button = gr.Button("Update")
with gr.Column() as input_column:
inputs = generate_text_boxes(2)
output = gr.Plot()
compare_button = gr.Button("Compare Embeddings")
update_button.click(
update_interface,
inputs=[num_texts],
outputs=[input_column, output]
)
compare_button.click(
compare_embeddings,
inputs=inputs,
outputs=output
)
iface.launch() |