app.py

import gradio as gr
import spaces
import torch
from transformers import AutoTokenizer, AutoModel
import plotly.graph_objects as go
import random

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]

def random_color():
    return f'rgb({random.randint(0,255)}, {random.randint(0,255)}, {random.randint(0,255)})'

@spaces.GPU
def compare_embeddings(*texts):
    embeddings = [get_embedding(text) for text in texts if text.strip()]
    embeddings_3d = [reduce_to_3d(emb) for emb in embeddings]
    
    fig = go.Figure()

    # Add black origin point
    fig.add_trace(go.Scatter3d(x=[0], y=[0], z=[0], mode='markers', marker=dict(size=5, color='black'), name='Origin'))

    # Add lines and points for each text
    for i, emb in enumerate(embeddings_3d):
        color = random_color()
        fig.add_trace(go.Scatter3d(x=[0, emb[0]], y=[0, emb[1]], z=[0, emb[2]], mode='lines+markers', 
                                   line=dict(color=color), marker=dict(color=color), name=f'Text {i+1}'))
    
    fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
    
    return fig

def add_textbox(num_textboxes):
    return gr.Textbox.update(visible=True, label=f"Text {num_textboxes + 1}"), num_textboxes + 1

def remove_textbox(num_textboxes):
    return gr.Textbox.update(visible=False), max(2, num_textboxes - 1)

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.")
    
    num_textboxes = gr.State(value=2)
    
    with gr.Column() as textbox_container:
        textboxes = [gr.Textbox(label=f"Text {i+1}") for i in range(10)]  # Create 10 textboxes
        for i in range(2, 10):  # Hide textboxes 3-10 initially
            textboxes[i].visible = False
    
    with gr.Row():
        add_btn = gr.Button("Add String")
        remove_btn = gr.Button("Remove String")
    
    plot_output = gr.Plot()
    
    submit_btn = gr.Button("Submit")
    clear_btn = gr.ClearButton(components=textboxes + [plot_output], value="Clear")
    
    add_btn.click(add_textbox, inputs=[num_textboxes], outputs=[textboxes[-1], num_textboxes])
    remove_btn.click(remove_textbox, inputs=[num_textboxes], outputs=[textboxes[-1], num_textboxes])
    submit_btn.click(compare_embeddings, inputs=textboxes, outputs=[plot_output])

iface.launch()