app.py

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

# 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,
        trust_remote_code=True
    )
    
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        torch_dtype=torch.float16,
        device_map="auto"
    )
    
    # Создаем pipeline с правильными параметрами
    generator = pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        device="cuda" if torch.cuda.is_available() else "cpu"
    )
    
except Exception as e:
    raise RuntimeError(f"Ошибка загрузки модели: {str(e)}")

# 3. Функция генерации ответа
def generate_response(message, history):
    prompt = f"""Ты оператор поддержки. Ответь клиенту вежливо на русском.
    
История диалога:
{history}
Клиент: {message}
Оператор:"""
    
    try:
        response = generator(
            prompt,
            max_new_tokens=200,
            temperature=0.7,
            do_sample=True,
            top_p=0.9,
            repetition_penalty=1.1
        )
        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, label="Диалог")
            msg = gr.Textbox(label="Опишите проблему", placeholder="Введите ваше сообщение...")
            btn = gr.Button("Отправить", variant="primary")
        
        with gr.Column():
            gr.Examples(examples, inputs=msg, label="Примеры обращений")
            gr.Markdown("**Рекомендации:**\n1. Укажите номер заказа\n2. Опишите проблему подробно")

    btn.click(generate_response, [msg, chatbot], [chatbot])
    msg.submit(generate_response, [msg, chatbot], [chatbot])

demo.launch()