import gradio as gr
from transformers import AutoTokenizer, AutoModel
import torch

model_name = "Qwen/Qwen-Embedding-0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)

def get_embedding(text):
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        embeddings = model(**inputs).last_hidden_state.mean(dim=1).squeeze().tolist()
    return embeddings

demo = gr.Interface(fn=get_embedding,
                    inputs=gr.Textbox(lines=2, placeholder="Enter text..."),
                    outputs="json",
                    title="Qwen Embedding Generator")
demo.launch()