File size: 1,785 Bytes
7103ccc 748826b 7103ccc c5fcdf3 580a16d 748826b 416fea8 748826b d590a55 748826b d590a55 748826b 416fea8 748826b 7103ccc 748826b 416fea8 748826b 7103ccc 748826b afc5a87 d590a55 |
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 |
import gradio as gr
import spaces
import torch
from transformers import AutoTokenizer, AutoModel
import plotly.graph_objects as go
from huggingface_hub import HfApi
from huggingface_hub import hf_hub_download
import os
import sys
HF_TOKEN = os.getenv("HF_TOKEN")
# Update the model name to Llama 3 8B
model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = None
@spaces.GPU
def get_embedding(text):
global model
if model is None:
model = AutoModel.from_pretrained(model_name).cuda()
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):
# Instead of PCA, we'll just take the first 3 dimensions
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')
])
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."
)
iface.launch() |