Update app.py

app.py

@@ -1,69 +1,76 @@
 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]
-except Exception as e:
-    print(f"Ошибка загрузки датасета: {e}")
-    examples = [
-        "Мой заказ #12345 не пришел",
-        "Как оформить возврат товара?",
-        "Не приходит SMS-код подтверждения",
-        "Ошибка при оплате картой"
-    ]
-
-# 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)
-    
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16,
+        device_map="auto",
+        load_in_8bit=True
+    )
     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):
-    prompt = f"""Ты оператор поддержки. Ответь клиенту вежливо на русском.
-    
-Клиент: {message}
-Оператор:"""
+# 2. Примеры обращений
+examples = [
+    "Мой заказ #12345 не пришел",
+    "Как оформить возврат товара?",
+    "Не приходит SMS-подтверждение",
+    "Ошибка при оплате картой"
+]
+
+# 3. Функция генерации ответа с правильным форматом сообщений
+def generate_response(message, chat_history):
+    # Формируем промпт с историей диалога
+    prompt = "Ты оператор поддержки. Вежливо отвечай клиенту на русском.\n\n"
+    for user_msg, bot_msg in chat_history:
+        prompt += f"Клиент: {user_msg}\nОператор: {bot_msg}\n"
+    prompt += f"Клиент: {message}\nОператор:"
     
     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
         )
-        return response[0]["generated_text"].split("Оператор:")[-1].strip()
+        bot_message = response[0]["generated_text"].split("Оператор:")[-1].strip()
+        
+        # Возвращаем обновленную историю диалога в правильном формате
+        return chat_history + [(message, bot_message)]
     except Exception as e:
-        return f"Извините, произошла ошибка. ({str(e)})"
+        print(f"Ошибка генерации: {str(e)}")
+        return chat_history + [(message, f"Извините, произошла ошибка. {str(e)}")]
 
-# 4. Интерфейс Gradio
+# 4. Создание интерфейса с правильным форматом Chatbot
 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, type="messages")
-            msg = gr.Textbox(label="Опишите проблему")
+            chatbot = gr.Chatbot(height=350)
+            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()
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)