import gradio as gr
from sentence_transformers import SentenceTransformer, util

# Загружаем русскоязычную модель
model = SentenceTransformer("ai-forever/sbert_large_nlu_ru")

def embed_text(text):
    embedding = model.encode([text])[0]
    return str(embedding.tolist()[:10]) + " ...", len(embedding)

def compare_texts(text1, text2):
    emb1 = model.encode([text1])[0]
    emb2 = model.encode([text2])[0]
    similarity = util.cos_sim(emb1, emb2).item()
    return similarity

with gr.Blocks() as demo:
    gr.Markdown("## 🔎 Russian SBERT Embeddings Demo (ai-forever/sbert_large_nlu_ru)")

    with gr.Tab("Эмбеддинг одного текста"):
        with gr.Row():
            with gr.Column():
                input_text = gr.Textbox(label="Введите текст", placeholder="Например: Я люблю машинное обучение")
                run_btn = gr.Button("Создать эмбеддинг")
            with gr.Column():
                output_embed = gr.Textbox(label="Вектор (первые 10 значений)")
                output_dim = gr.Number(label="Размерность эмбеддинга")

        run_btn.click(embed_text, inputs=[input_text], outputs=[output_embed, output_dim])

    with gr.Tab("Сравнение двух текстов"):
        with gr.Row():
            with gr.Column():
                text1 = gr.Textbox(label="Текст 1", placeholder="Например: Я люблю машинное обучение")
                text2 = gr.Textbox(label="Текст 2", placeholder="Например: Машинное обучение — моя страсть")
                compare_btn = gr.Button("Сравнить тексты")
            with gr.Column():
                similarity_score = gr.Number(label="Косинусное сходство")

        compare_btn.click(compare_texts, inputs=[text1, text2], outputs=[similarity_score])

demo.launch()