Update app.py

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from datasets import load_dataset
+import torch
 
-# 1. Загрузка датасета (используем реальный существующий датасет)
+# 1. Загрузка датасета
 try:
     dataset = load_dataset("blinoff/ru_customer_support", split="train[:50]")
     examples = [d["question"] for d in dataset]
@@ -15,40 +16,55 @@ except Exception as e:
         "Ошибка при оплате картой"
     ]
 
-# 2. Загрузка модели (используем локальное выполнение)
+# 2. Загрузка модели с обработкой ошибок
 try:
-    model_name = "ai-forever/rugpt3small_based_on_gpt2"  # Рабочая альтернатива
+    model_name = "ai-forever/rugpt3small_based_on_gpt2"
     
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name)
+    # Явно указываем доверенный источник
+    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="cpu"
+        device="cuda" if torch.cuda.is_available() else "cpu"
     )
+    
 except Exception as e:
     raise RuntimeError(f"Ошибка загрузки модели: {str(e)}")
 
 # 3. Функция генерации ответа
-def generate_response(message):
+def generate_response(message, history):
     prompt = f"""Ты оператор поддержки. Ответь клиенту вежливо на русском.
     
+История диалога:
+{history}
 Клиент: {message}
 Оператор:"""
     
     try:
         response = generator(
             prompt,
-            max_new_tokens=150,
-            temperature=0.4,
+            max_new_tokens=200,
+            temperature=0.7,
             do_sample=True,
-            top_p=0.9
+            top_p=0.9,
+            repetition_penalty=1.1
         )
         return response[0]["generated_text"].split("Оператор:")[-1].strip()
     except Exception as e:
-        return f"Извините, произошла ошибка. ({str(e)})"
+        return f"Ошибка генерации ответа: {str(e)}"
 
 # 4. Интерфейс Gradio
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
@@ -56,14 +72,15 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     
     with gr.Row():
         with gr.Column():
-            chatbot = gr.Chatbot(height=350)
-            msg = gr.Textbox(label="Опишите проблему")
+            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("**Совет:** Укажите номер заказа для быстрого решения")
+            gr.Markdown("**Рекомендации:**\n1. Укажите номер заказа\n2. Опишите проблему подробно")
 
-    btn.click(lambda m, c: (m, generate_response(m)), [msg, chatbot], [msg, chatbot])
+    btn.click(generate_response, [msg, chatbot], [chatbot])
+    msg.submit(generate_response, [msg, chatbot], [chatbot])
 
 demo.launch()