import gradio as gr
import time
from transformers import pipeline

# Настройка моделей
models = {
    "ChatGPT-like (FRED-T5)": pipeline("text2text-generation", model="ai-forever/FRED-T5-1.7B", tokenizer="ai-forever/FRED-T5-1.7B", device=-1),
    "DeepSeek-like (Qwen7B)": pipeline("text-generation", model="lightblue/DeepSeek-R1-Distill-Qwen-7B-Multilingual", tokenizer="lightblue/DeepSeek-R1-Distill-Qwen-7B-Multilingual", device=-1),
    "GigaChat-like (GigaChat-20B)": pipeline("text-generation", model="ai-sage/GigaChat-20B-A3B-instruct", tokenizer="ai-sage/GigaChat-20B-A3B-instruct", device=-1)
}

# Промпты
cot_instruction = """Ты — банковский специалист. Клиент задал вопрос.
Проанализируй обращение пошагово:
1. В чём проблема?
2. Возможные причины?
3. Как решить?
Вывод: категория обращения."""

simple_instruction = "Ты — банковский специалист. Определи категорию обращения клиента кратко."

# Генерация промптов

def build_prompt(instruction, user_input):
    return f"{instruction}\n\nКлиент: {user_input}"

# Обработка ответа

def get_output(pipe, prompt, max_tokens=300):
    try:
        output = pipe(prompt, max_new_tokens=max_tokens, truncation=True, do_sample=True, temperature=0.7)[0]
        return output.get("generated_text") or output.get("output_text") or "(нет ответа)"
    except Exception as e:
        return f"Ошибка: {e}"

# Основная функция

def generate_comparison(user_input):
    result = {}
    prompt_cot = build_prompt(cot_instruction, user_input)
    prompt_simple = build_prompt(simple_instruction, user_input)

    for name, pipe in models.items():
        # CoT ответ
        start_cot = time.time()
        answer_cot = get_output(pipe, prompt_cot)
        time_cot = round(time.time() - start_cot, 2)

        # Обычный ответ
        start_simple = time.time()
        answer_simple = get_output(pipe, prompt_simple)
        time_simple = round(time.time() - start_simple, 2)

        result[name] = {
            "cot_answer": answer_cot.strip(),
            "cot_time": f"{time_cot} сек",
            "simple_answer": answer_simple.strip(),
            "simple_time": f"{time_simple} сек"
        }

    return (
        result["ChatGPT-like (FRED-T5)"]["cot_answer"], result["ChatGPT-like (FRED-T5)"]["cot_time"],
        result["ChatGPT-like (FRED-T5)"]["simple_answer"], result["ChatGPT-like (FRED-T5)"]["simple_time"],
        result["DeepSeek-like (Qwen7B)"]["cot_answer"], result["DeepSeek-like (Qwen7B)"]["cot_time"],
        result["DeepSeek-like (Qwen7B)"]["simple_answer"], result["DeepSeek-like (Qwen7B)"]["simple_time"],
        result["GigaChat-like (GigaChat-20B)"]["cot_answer"], result["GigaChat-like (GigaChat-20B)"]["cot_time"],
        result["GigaChat-like (GigaChat-20B)"]["simple_answer"], result["GigaChat-like (GigaChat-20B)"]["simple_time"]
    )

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

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

    # ChatGPT-like
    gr.Markdown("### ChatGPT-like (FRED-T5)")
    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 (Qwen7B)")
    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 (GigaChat-20B)")
    cot3 = gr.Textbox(label="CoT ответ")
    cot3_time = gr.Textbox(label="Время CoT")
    simple3 = gr.Textbox(label="Обычный ответ")
    simple3_time = gr.Textbox(label="Время обычного")

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

if __name__ == '__main__':
    demo.launch()