import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset

# 1. Загрузка датасета (используем реальный существующий датасет)
try:
    dataset = load_dataset("blinoff/ru_customer_support", split="train[:50]")
    examples = [d["question"] for d in dataset]
except Exception as e:
    print(f"Ошибка загрузки датасета: {e}")
    examples = [
        "Мой заказ #12345 не пришел",
        "Как оформить возврат товара?",
        "Не приходит SMS-код подтверждения",
        "Ошибка при оплате картой"
    ]

# 2. Загрузка модели (используем локальное выполнение)
try:
    model_name = "ai-forever/rugpt3small_based_on_gpt2"  # Рабочая альтернатива
    
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    
    generator = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        device="cpu"
    )
except Exception as e:
    raise RuntimeError(f"Ошибка загрузки модели: {str(e)}")

# 3. Функция генерации ответа
def generate_response(message):
    prompt = f"""Ты оператор поддержки. Ответь клиенту вежливо на русском.
    
Клиент: {message}
Оператор:"""
    
    try:
        response = generator(
            prompt,
            max_new_tokens=150,
            temperature=0.4,
            do_sample=True,
            top_p=0.9
        )
        return response[0]["generated_text"].split("Оператор:")[-1].strip()
    except Exception as e:
        return f"Извините, произошла ошибка. ({str(e)})"

# 4. Интерфейс Gradio
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""<h1><center>📞 Поддержка клиентов</center></h1>""")
    
    with gr.Row():
        with gr.Column():
            chatbot = gr.Chatbot(height=350)
            msg = gr.Textbox(label="Опишите проблему")
            btn = gr.Button("Отправить", variant="primary")
        
        with gr.Column():
            gr.Examples(examples, inputs=msg, label="Примеры обращений")
            gr.Markdown("**Совет:** Укажите номер заказа для быстрого решения")

    btn.click(lambda m, c: (m, generate_response(m)), [msg, chatbot], [msg, chatbot])

demo.launch()