import gradio as gr
import time
from transformers import pipeline

# Инициализация моделей
models = {
    "ChatGPT-like": pipeline(
        "text-generation",
        model="OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
        tokenizer="OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5",
        device=-1
    ),
    "DeepSeek-like": pipeline(
        "text-generation",
        model="deepseek-ai/DeepSeek-Coder-1.3B-instruct",
        tokenizer="deepseek-ai/DeepSeek-Coder-1.3B-instruct",
        device=-1
    ),
    "GigaChat-like": pipeline(
        "text-generation",
        model="tinkoff-ai/ruDialoGPT-medium",
        tokenizer="tinkoff-ai/ruDialoGPT-medium",
        device=-1
    )
}

# Промпты
def build_simple_prompt(user_input):
    return f"Клиент: {user_input}\nКатегория обращения:"

def build_cot_prompt(user_input):
    return (
        f"Клиент: {user_input}\n"
        "Проанализируй обращение пошагово:\n"
        "1. В чём суть проблемы?\n"
        "2. Возможные причины?\n"
        "3. Решение?\n"
        "Вывод: категория обращения:"
    )

# Генерация ответов
def generate_classification(user_input):
    prompt_simple = build_simple_prompt(user_input)
    prompt_cot = build_cot_prompt(user_input)

    results = {}
    for name, pipe in models.items():
        # CoT
        start_cot = time.time()
        cot_out = pipe(prompt_cot, max_new_tokens=150, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
        time_cot = round(time.time() - start_cot, 2)

        # Simple
        start_simple = time.time()
        simple_out = pipe(prompt_simple, max_new_tokens=80, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
        time_simple = round(time.time() - start_simple, 2)

        results[name] = {
            "cot": cot_out.strip(),
            "cot_time": time_cot,
            "simple": simple_out.strip(),
            "simple_time": time_simple
        }

    return (
        results["ChatGPT-like"]["cot"], f"{results['ChatGPT-like']['cot_time']} сек",
        results["ChatGPT-like"]["simple"], f"{results['ChatGPT-like']['simple_time']} сек",
        results["DeepSeek-like"]["cot"], f"{results['DeepSeek-like']['cot_time']} сек",
        results["DeepSeek-like"]["simple"], f"{results['DeepSeek-like']['simple_time']} сек",
        results["GigaChat-like"]["cot"], f"{results['GigaChat-like']['cot_time']} сек",
        results["GigaChat-like"]["simple"], f"{results['GigaChat-like']['simple_time']} сек"
    )

# Интерфейс
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 Сравнение моделей: ChatGPT, DeepSeek, GigaChat — классификация банковских обращений")

    inp = gr.Textbox(label="Вопрос клиента", placeholder="Например: Я не могу перевести деньги", lines=2)
    btn = gr.Button("Сгенерировать")

    # ChatGPT
    gr.Markdown("### ChatGPT-like (OpenAssistant)")
    cot1 = gr.Textbox(label="CoT ответ")
    cot1_time = gr.Textbox(label="Время CoT")
    simple1 = gr.Textbox(label="Обычный ответ")
    simple1_time = gr.Textbox(label="Время обычного")

    # DeepSeek
    gr.Markdown("### DeepSeek-like")
    cot2 = gr.Textbox(label="CoT ответ")
    cot2_time = gr.Textbox(label="Время CoT")
    simple2 = gr.Textbox(label="Обычный ответ")
    simple2_time = gr.Textbox(label="Время обычного")

    # GigaChat
    gr.Markdown("### GigaChat-like (ruDialoGPT)")
    cot3 = gr.Textbox(label="CoT ответ")
    cot3_time = gr.Textbox(label="Время CoT")
    simple3 = gr.Textbox(label="Обычный ответ")
    simple3_time = gr.Textbox(label="Время обычного")

    btn.click(generate_classification, inputs=[inp], outputs=[
        cot1, cot1_time, simple1, simple1_time,
        cot2, cot2_time, simple2, simple2_time,
        cot3, cot3_time, simple3, simple3_time
    ])

demo.launch()