import gradio as gr
import time
from transformers import pipeline

# Инициализация моделей
models = {
    "ChatGPT-like": pipeline(
        "text-generation",
        model="Gre8o8nta/ChatGPT",
        tokenizer="Gre8o8nta/ChatGPT",
        device=-1
    ),
    "DeepSeek-like": pipeline(
        "text-generation",
        model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
        tokenizer="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
        device=-1
    ),
    "GigaChat-like": pipeline(
        "text-generation",
        model="IlyaGusev/saiga_mistral_7b_text",
        tokenizer="IlyaGusev/saiga_mistral_7b_text",
        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_output = pipe(prompt_cot, max_new_tokens=150, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
        end_cot = round(time.time() - start_cot, 2)

        # Обычный ответ
        start_simple = time.time()
        simple_output = pipe(prompt_simple, max_new_tokens=80, do_sample=True, top_p=0.9, temperature=0.7)[0]["generated_text"]
        end_simple = round(time.time() - start_simple, 2)

        results[name] = {
            "cot_answer": cot_output.strip(),
            "cot_time": end_cot,
            "simple_answer": simple_output.strip(),
            "simple_time": end_simple
        }

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

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("🏦 **Сравнение моделей ChatGPT, DeepSeek и GigaChat по классификации банковских обращений (CoT и обычный)**")

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

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

    # DeepSeek-like
    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-like
    gr.Markdown("### GigaChat-like (заменена на Saiga)")
    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()