Update app.py

app.py

@@ -1,57 +1,18 @@
-import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
-import json
-from pathlib import Path
-from datetime import datetime
-import requests
-
-# === ЗАГРУЗКА МОДЕЛИ ===
+# === ЗАГРУЗКА МОДЕЛИ (ИСПРАВЛЕННАЯ ЧАСТЬ) ===
 print("🚀 Загрузка модели...")
 model_id = "OpenRussianAI/OpenAirAI-X"
 tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+# ВАЖНО: Убеждаемся, что у токенизатора есть pad_token
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
 model = AutoModelForCausalLM.from_pretrained(model_id)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model = model.to(device).eval()
 print(f"✅ Модель загружена на {device}")
 
-# === ХРАНИЛИЩЕ ИСТОРИИ ===
-HISTORY_DIR = Path("chat_history")
-HISTORY_DIR.mkdir(exist_ok=True)
-
-def get_history_path(username):
-    safe_name = "".join(c for c in username if c.isalnum() or c in ('-', '_')).lower()
-    return HISTORY_DIR / f"{safe_name}.json"
-
-def load_history(username):
-    path = get_history_path(username)
-    if path.exists():
-        with open(path, "r", encoding="utf-8") as f:
-            return json.load(f)
-    return {}
-
-def save_history(username, history_data):
-    path = get_history_path(username)
-    with open(path, "w", encoding="utf-8") as f:
-        json.dump(history_data, f, ensure_ascii=False, indent=2)
-
-def check_if_pro(token):
-    if not token or not token.startswith("hf_"):
-        return False
-    try:
-        response = requests.get(
-            "https://huggingface.co/api/whoami-v2",
-            headers={"Authorization": f"Bearer {token}"},
-            timeout=5
-        )
-        if response.status_code == 200:
-            data = response.json()
-            return data.get("isPro", False) or data.get("plan", {}).get("name") == "PRO"
-    except Exception as e:
-        print(f"Ошибка проверки PRO: {e}")
-    return False
-
-# === ГЕНЕРАЦИЯ ОТВЕТА ===
+# === ГЕНЕРАЦИЯ ОТВЕТА (ПОЛНОСТЬЮ ПЕРЕПИСАННАЯ) ===
 def generate_response(message, history, username, current_chat_id):
     if not username:
         gr.Warning("Сначала введите имя пользователя!")
@@ -60,8 +21,10 @@ def generate_response(message, history, username, current_chat_id):
     if not message.strip():
         return history, gr.update(), current_chat_id, gr.update()
     
+    # Добавляем сообщение пользователя в историю
     history = history + [{"role": "user", "content": message}]
     
+    # Формируем промпт
     prompt = ""
     for msg in history:
         if msg["role"] == "user":
@@ -71,19 +34,37 @@ def generate_response(message, history, username, current_chat_id):
     prompt += "AI:"
     
     inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
-            max_new_tokens=256,
+            max_new_tokens=128, # Уменьшим, чтобы модель не уходила в бред
             temperature=0.7,
             top_p=0.9,
             do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
+            pad_token_id=tokenizer.eos_token_id,
+            eos_token_id=tokenizer.eos_token_id, # Важно!
+            repetition_penalty=1.2 # Штраф за повторения, чтобы убрать цикл "Я OpenAI"
         )
     
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    ai_response = response.split("AI:")[-1].strip()
+    # Декодируем только НОВУЮ часть текста
+    generated_ids = outputs[0][inputs.input_ids.shape[-1]:]
+    response_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
+    
+    # Очистка от мусора и обрезка
+    # Обычно модель может начать ответ с пробела или новой строки
+    ai_response = response_text.strip()
     
+    # Если модель все же написала "Пользователь:" или "AI:" в конце, обрезаем
+    stop_words = ["Пользователь:", "User:", "\n\n"]
+    for stop_word in stop_words:
+        if stop_word in ai_response:
+            ai_response = ai_response.split(stop_word)[0].strip()
+            
+    # Если ответ пустой после очистки
+    if not ai_response:
+        ai_response = "..."
+
     history = history + [{"role": "assistant", "content": ai_response}]
     
     if username and current_chat_id:
@@ -98,237 +79,4 @@ def generate_response(message, history, username, current_chat_id):
         save_history(username, history_data)
     
     chat_list_choices = get_chat_list(username)
-    # Возвращаем обновленный список через gr.update
-    return history, gr.update(value=""), current_chat_id, gr.update(choices=chat_list_choices)
-
-# === УПРАВЛЕНИЕ ЧАТАМИ ===
-def new_chat(username):
-    if not username:
-        return [], None, gr.update(choices=[])
-    chat_id = datetime.now().strftime("%Y%m%d_%H%M%S")
-    history_data = load_history(username)
-    history_data[chat_id] = {
-        "title": "Новый чат",
-        "created": datetime.now().isoformat(),
-        "messages": []
-    }
-    save_history(username, history_data)
-    chat_list_choices = get_chat_list(username)
-    return [], chat_id, gr.update(choices=chat_list_choices, value=None)
-
-def load_chat(chat_title, username, current_chat_id):
-    if not username or not chat_title:
-        return [], current_chat_id
-    history_data = load_history(username)
-    for cid, data in history_data.items():
-        if data["title"] == chat_title:
-            return data["messages"], cid
-    return [], current_chat_id
-
-def get_chat_list(username):
-    if not username:
-        return []
-    history_data = load_history(username)
-    sorted_chats = sorted(
-        history_data.items(),
-        key=lambda x: x[1].get("created", ""),
-        reverse=True
-    )
-    return [data["title"] for _, data in sorted_chats]
-
-def delete_chat(chat_title, username):
-    if not username or not chat_title:
-        return [], {}, None, gr.update(choices=[])
-    history_data = load_history(username)
-    for cid, data in list(history_data.items()):
-        if data["title"] == chat_title:
-            del history_data[cid]
-            break
-    save_history(username, history_data)
-    chat_list_choices = get_chat_list(username)
-    return [], history_data, None, gr.update(choices=chat_list_choices, value=None)
-
-# === CSS ===
-CUSTOM_CSS = """
-.main-header {
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-    color: white;
-    padding: 20px;
-    border-radius: 12px;
-    margin-bottom: 20px;
-    text-align: center;
-}
-.main-header h1 { margin: 0; font-size: 2em; }
-.pro-badge {
-    background: linear-gradient(135deg, #ffd700 0%, #ffed4e 100%);
-    color: #333;
-    padding: 4px 12px;
-    border-radius: 20px;
-    font-weight: bold;
-    display: inline-block;
-    margin-left: 10px;
-    box-shadow: 0 2px 8px rgba(255, 215, 0, 0.4);
-}
-.user-info {
-    padding: 10px;
-    background: white;
-    border-radius: 8px;
-    margin-bottom: 10px;
-    text-align: center;
-    font-weight: 600;
-}
-.privacy-banner {
-    background: linear-gradient(135deg, #10b981 0%, #059669 100%);
-    color: white;
-    padding: 15px 20px;
-    border-radius: 10px;
-    margin: 15px 0;
-    display: flex;
-    align-items: center;
-    gap: 12px;
-    font-size: 0.95em;
-    box-shadow: 0 2px 8px rgba(16, 185, 129, 0.3);
-}
-.privacy-banner .icon { font-size: 1.8em; }
-.privacy-banner strong { display: block; font-size: 1.1em; margin-bottom: 4px; }
-"""
-
-PRIVACY_BANNER = """
-<div class="privacy-banner">
-    <div class="icon">🔒</div>
-    <div>
-        <strong>Ваши данные в безопасности</strong>
-        Все сообщения обрабатываются локально. 
-        <b>Данные НЕ передаются в OpenRussianAI</b>. 
-        История хранится только внутри контейнера.
-    </div>
-</div>
-"""
-
-# === ИНТЕРФЕЙС ===
-with gr.Blocks(title="OpenAirAI-X Chat") as demo:
-    username_state = gr.State("")
-    history_state = gr.State({})
-    current_chat_id_state = gr.State(None)
-    
-    with gr.Column(visible=False) as main_interface:
-        gr.HTML("<div class='main-header'><h1>🤖 OpenAirAI-X Chat</h1><p>Русскоязычный AI-ассистент</p></div>")
-        gr.HTML(PRIVACY_BANNER)
-        
-        with gr.Row():
-            with gr.Column(scale=1, min_width=250):
-                user_info = gr.Textbox(label="Пользователь", value="Не авторизован", interactive=False, elem_classes="user-info")
-                pro_badge = gr.HTML('<div class="pro-badge">👑 PRO</div>', visible=False)
-                new_chat_btn = gr.Button("➕ Новый чат", variant="primary")
-                gr.Markdown("### 📚 История чатов")
-                # Используем Dropdown с allow_custom_value=False для стабильности
-                chat_list = gr.Dropdown(choices=[], label="Ваши чаты", interactive=True, allow_custom_value=False, value=None)
-                delete_chat_btn = gr.Button("🗑️ Удалить выбранный чат", variant="stop", size="sm")
-                logout_btn = gr.Button("🚪 Выйти", size="sm")
-            
-            with gr.Column(scale=3):
-                chatbot = gr.Chatbot(label="Диалог", height=500)
-                with gr.Row():
-                    msg_input = gr.Textbox(placeholder="Напишите сообщение...", lines=2, scale=5, show_label=False)
-                    send_btn = gr.Button("📤 Отправить", variant="primary", scale=1)
-    
-    with gr.Column(visible=True) as login_screen:
-        gr.HTML("<div class='main-header'><h1>🤖 OpenAirAI-X Chat</h1><p>Введите данные для входа</p></div>")
-        gr.HTML(PRIVACY_BANNER)
-        gr.Markdown("### 🔐 Вход в систему\nВведите ваше имя пользователя. История чатов будет привязана к этому имени.")
-        username_input = gr.Textbox(label="Имя пользователя (HF Username)", placeholder="Например: RootLinux21")
-        token_input = gr.Textbox(label="HF Token (необязательно, для PRO)", placeholder="hf_...")
-        login_btn = gr.Button("🔑 Войти", variant="primary")
-    
-    # === ОБРАБОТЧИКИ ===
-    
-    def handle_login(username, token):
-        if not username.strip():
-            gr.Warning("Введите имя пользователя!")
-            return gr.update(), gr.update(), "Не авторизован", "", {}, None, gr.update(choices=[], value=None), gr.update(visible=False)
-        
-        is_pro = check_if_pro(token)
-        history_data = load_history(username)
-        chat_list_choices = get_chat_list(username)
-        
-        current_chat_id = None
-        if not history_data:
-            chat_id = datetime.now().strftime("%Y%m%d_%H%M%S")
-            history_data[chat_id] = {
-                "title": "Новый чат",
-                "created": datetime.now().isoformat(),
-                "messages": []
-            }
-            save_history(username, history_data)
-            current_chat_id = chat_id
-            chat_list_choices = get_chat_list(username)
-        
-        pro_badge_html = '<div class="pro-badge">👑 PRO</div>'
-        display_name = f"👤 {username}"
-        
-        return (
-            gr.update(visible=True),
-            gr.update(visible=False),
-            display_name,
-            username,
-            history_data,
-            current_chat_id,
-            gr.update(choices=chat_list_choices, value=None),
-            gr.update(visible=is_pro, value=pro_badge_html)
-        )
-    
-    def handle_logout():
-        return (
-            gr.update(visible=False),
-            gr.update(visible=True),
-            "Не авторизован",
-            "",
-            {},
-            None,
-            gr.update(choices=[], value=None),
-            gr.update(visible=False)
-        )
-    
-    login_btn.click(
-        fn=handle_login,
-        inputs=[username_input, token_input],
-        outputs=[main_interface, login_screen, user_info, username_state, history_state, current_chat_id_state, chat_list, pro_badge]
-    )
-    
-    send_btn.click(
-        fn=generate_response,
-        inputs=[msg_input, chatbot, username_state, current_chat_id_state],
-        outputs=[chatbot, msg_input, current_chat_id_state, chat_list]
-    )
-    
-    msg_input.submit(
-        fn=generate_response,
-        inputs=[msg_input, chatbot, username_state, current_chat_id_state],
-        outputs=[chatbot, msg_input, current_chat_id_state, chat_list]
-    )
-    
-    new_chat_btn.click(
-        fn=new_chat,
-        inputs=[username_state],
-        outputs=[chatbot, current_chat_id_state, chat_list]
-    )
-    
-    chat_list.change(
-        fn=load_chat,
-        inputs=[chat_list, username_state, current_chat_id_state],
-        outputs=[chatbot, current_chat_id_state]
-    )
-    
-    delete_chat_btn.click(
-        fn=delete_chat,
-        inputs=[chat_list, username_state],
-        outputs=[chatbot, history_state, current_chat_id_state, chat_list]
-    )
-    
-    logout_btn.click(
-        fn=handle_logout,
-        inputs=None,
-        outputs=[main_interface, login_screen, user_info, username_state, history_state, current_chat_id_state, chat_list, pro_badge]
-    )
-
-demo.launch(server_name="0.0.0.0", server_port=7860, css=CUSTOM_CSS)
+    return history, gr.update(value=""), current_chat_id, gr.update(choices=chat_list_choices)