import gradio as gr
from transformers import pipeline, AutoTokenizer

# 1. Используем стабильную публичную модель
MODEL_NAME = "sberbank-ai/rugpt3small"  # Русскоязычная модель от Сбера

# 2. Загрузка модели с обработкой ошибок
try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = pipeline(
        "text-generation",
        model=MODEL_NAME,
        tokenizer=tokenizer,
        device="cpu",
        torch_dtype="auto"
    )
except Exception as e:
    raise RuntimeError(f"Ошибка загрузки модели: {str(e)}")

# 3. Примеры обращений (без датасета)
examples = [
    "Где мой заказ #12345?",
    "Как вернуть товар?",
    "Не приходит SMS-код подтверждения",
    "Ошибка при оплате картой",
    "Как изменить адрес доставки?"
]

# 4. Генерация ответа
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)}"

# 5. Интерфейс Gradio
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="Опишите проблему")
            output_text = gr.Textbox(label="Ответ", lines=5)
            btn = gr.Button("Отправить")
        
        with gr.Column():
            gr.Examples(examples, inputs=input_text, label="Примеры обращений")
            gr.Markdown("**Совет:** Указывайте номер заказа для быстрого решения")

    btn.click(generate_response, input_text, output_text)

demo.launch()