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

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 if text.strip()]  # Only process non-empty texts
    embeddings_3d = [reduce_to_3d(emb) for emb in embeddings]
    
    fig = go.Figure()
    
    colors = ['red', 'blue', 'green', 'purple', 'orange', 'cyan', 'magenta', 'yellow']
    
    for i, emb in enumerate(embeddings_3d):
        color = colors[i % len(colors)]
        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}',
            line=dict(color=color),
            marker=dict(color=color)
        ))
    
    fig.update_layout(scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'))
    
    return fig

def create_interface(num_inputs):
    with gr.Blocks() as new_interface:
        text_inputs = [gr.Textbox(label=f"Text {i+1}") for i in range(num_inputs)]
        output = gr.Plot()
        submit_btn = gr.Button("Compare Embeddings")
        submit_btn.click(fn=compare_embeddings, inputs=text_inputs, outputs=output)
    return new_interface

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_inputs = gr.Slider(minimum=2, maximum=10, step=1, value=2, label="Number of texts to compare")
    interface_container = gr.HTML()
    
    def update_interface(num):
        return create_interface(num)
    
    num_inputs.change(fn=update_interface, inputs=[num_inputs], outputs=[interface_container])
    
    # Initialize the interface with 2 text boxes
    interface_container.update(create_interface(2))

iface.launch()