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

# Загрузка модели
model_name = "Alibaba-NLP/gte-Qwen2-7B-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

def embed(text: str):
    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
    with torch.no_grad():
        outputs = model(**inputs)
    embedding = outputs.last_hidden_state[:, 0, :].squeeze().tolist()
    return {"embedding": embedding}

iface = gr.Interface(
    fn=embed,
    inputs=gr.Textbox(lines=2, placeholder="Введите текст..."),
    outputs="json",
    title="GTE Qwen2 Embedder"
)

iface.launch()