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

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

# 2. Оптимизированная загрузка модели
try:
    tokenizer = AutoTokenizer.from_pretrained("IlyaGusev/saiga_mistral_7b-lora")
    model = pipeline(
        "text-generation",
        model="IlyaGusev/saiga_mistral_7b-lora",
        tokenizer=tokenizer,
        device="cpu",
        torch_dtype="auto",
        model_kwargs={"load_in_4bit": True}  # Критически важно для экономии памяти
    )
except Exception as e:
    raise RuntimeError(f"Ошибка загрузки модели: {str(e)}")

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

# 4. Интерфейс с примерами
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""<h1><center>🎯 Поддержка клиентов</center></h1>""")
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="Опишите проблему", placeholder="Мой заказ не прибыл...")
            output_text = gr.Textbox(label="Ответ поддержки", lines=5)
            btn = gr.Button("Отправить", variant="primary")
        
        with gr.Column():
            gr.Examples(
                examples=examples,
                inputs=input_text,
                label="Примеры обращений"
            )
            gr.Markdown("""
            **Советы для быстрого решения:**
            1. Указывайте номер заказа
            2. Прикладывайте скриншоты ошибок
            3. Проверяйте спам-папку
            """)

    btn.click(fn=generate_response, inputs=input_text, outputs=output_text)

demo.launch()