Update app.py

app.py

@@ -1,10 +1,13 @@
 import gradio as gr
-from transformers import pipeline, AutoTokenizer
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from typing import List, Dict, Any, Tuple
 import torch
+import warnings
+
+# Подавляем ненужные предупреждения
+warnings.filterwarnings("ignore", message=".*low_cpu_mem_usage.*")
 
 # CPU-модели (только одна маленькая модель для экономии памяти)
-# Исправлено: используем словарь вместо множества
 MODELS = {
     "Qwen2.5-0.5B": "Qwen/Qwen2.5-0.5B-Instruct",
     "Qwen2.5-1.5B": "Qwen/Qwen2.5-1.5B-Instruct",
@@ -13,23 +16,39 @@ MODELS = {
 def load_model(model_key: str):
     model_id = MODELS[model_key]
     print(f"🚀 Загрузка {model_id}...")
+    
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
     
+    # Сначала загружаем модель отдельно с оптимизацией памяти
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.float32,
+            device_map=None,  # Используем CPU
+            low_cpu_mem_usage=True,
+            trust_remote_code=True
+        )
+    except Exception as e:
+        print(f"⚠️ Не удалось загрузить с low_cpu_mem_usage: {e}")
+        print("🔄 Пробуем без low_cpu_mem_usage...")
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.float32,
+            trust_remote_code=True
+        )
+    
+    # Затем создаем pipeline
     pipe = pipeline(
         "text-generation",
-        model=model_id,
+        model=model,
         tokenizer=tokenizer,
-        torch_dtype=torch.float32,  # Use float32 for CPU
-        device_map=None,  # Explicitly set to CPU
+        device=-1,  # Явно указываем CPU (-1 означает CPU)
         max_new_tokens=128,  # Ещё меньше токенов для экономии памяти
         do_sample=True,
         temperature=0.7,
-        pad_token_id=tokenizer.eos_token_id,
-        # Memory optimization parameters
-        low_cpu_mem_usage=True,
-        trust_remote_code=True
+        pad_token_id=tokenizer.eos_token_id
     )
     print(f"✅ {model_id} загружена!")
     return pipe
@@ -57,50 +76,112 @@ def respond(message: str,
         messages.append({"role": "user", "content": message})
         
         tokenizer = pipe.tokenizer
-        prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         
-        outputs = pipe(prompt, max_new_tokens=512, do_sample=True, temperature=0.7)
-        bot_reply = outputs[0]["generated_text"][len(prompt):].strip()
+        # Используем чат-шаблон для Qwen моделей
+        try:
+            prompt = tokenizer.apply_chat_template(
+                messages, 
+                tokenize=False, 
+                add_generation_prompt=True
+            )
+        except Exception as e:
+            print(f"⚠️ Ошибка применения чат-шаблона: {e}")
+            # Альтернативный способ форматирования
+            prompt = ""
+            for msg in messages:
+                if msg["role"] == "system":
+                    prompt += f"System: {msg['content']}\n\n"
+                elif msg["role"] == "user":
+                    prompt += f"User: {msg['content']}\n\n"
+                elif msg["role"] == "assistant":
+                    prompt += f"Assistant: {msg['content']}\n\n"
+            prompt += "Assistant:"
+        
+        outputs = pipe(
+            prompt, 
+            max_new_tokens=256,  # Уменьшил для экономии памяти
+            do_sample=True, 
+            temperature=0.7,
+            repetition_penalty=1.1
+        )
+        
+        # Извлекаем ответ
+        generated_text = outputs[0]["generated_text"]
+        if generated_text.startswith(prompt):
+            bot_reply = generated_text[len(prompt):].strip()
+        else:
+            bot_reply = generated_text.strip()
         
         print(f"✅ Ответ: {bot_reply[:50]}...")
         
-        new_history = history + [{"role": "user", "content": message}, {"role": "assistant", "content": bot_reply}]
+        new_history = history + [
+            {"role": "user", "content": message}, 
+            {"role": "assistant", "content": bot_reply}
+        ]
         return new_history, "", gr.update(value="")
     
     except Exception as e:
         error_msg = f"❌ {model_key}: {str(e)}"
         print(f"💥 {error_msg}")
-        new_history = history + [{"role": "user", "content": message}, {"role": "assistant", "content": error_msg}]
+        new_history = history + [
+            {"role": "user", "content": message}, 
+            {"role": "assistant", "content": error_msg}
+        ]
         return new_history, error_msg, gr.update(value="")
 
-with gr.Blocks(title="🚀 Локальный HF Чат (на слабом CPU!)") as demo:
-    gr.Markdown("# Локальный Inference (без API!)\n**Маленькие модели** — 1-3 сек CPU. Большие думают ооочень долго. Нет limits/token. В качестве примера.")
+with gr.Blocks(title="🚀 Локальный HF Чат (на слабом CPU!)", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("""
+    # 🚀 Локальный Inference (без API!)
+    **Маленькие модели** — 1-3 сек CPU. Большие думают ооочень долго. Нет limits/token. В качестве примера.
+    
+    ⚠️ **Внимание**: Модели загружаются при первшем выборе и могут занять несколько минут!
+    """)
     
     with gr.Row():
-        # Исправлено: первое значение должно быть ключом из словаря MODELS
-        model_dropdown = gr.Dropdown(choices=list(MODELS.keys()), value="Qwen2.5-0.5B", label="🧠 Модель")
-        system_prompt = gr.Textbox(label="📝 System", placeholder="Ты весёлый ИИ.", lines=2)
+        model_dropdown = gr.Dropdown(
+            choices=list(MODELS.keys()), 
+            value="Qwen2.5-0.5B", 
+            label="🧠 Модель",
+            info="Выберите модель (загрузка при первшем использовании)"
+        )
+        system_prompt = gr.Textbox(
+            label="📝 System Prompt", 
+            placeholder="Ты весёлый и полезный ИИ-ассистент.", 
+            lines=2,
+            value="Ты весёлый и полезный ИИ-ассистент."
+        )
     
-    chatbot = gr.Chatbot(height=500, label="Чат")  # ← без type
+    chatbot = gr.Chatbot(
+        height=400, 
+        label="Чат",
+        avatar_images=(None, "🤖")
+    )
     
     with gr.Row():
-        msg_input = gr.Textbox(placeholder="Привет! (Enter)", show_label=False, lines=1)
-        send_btn = gr.Button("📤 Отправить", variant="primary")
+        msg_input = gr.Textbox(
+            placeholder="Напишите сообщение... (Enter для отправки)", 
+            show_label=False, 
+            lines=2,
+            scale=4
+        )
+        send_btn = gr.Button("📤 Отправить", variant="primary", scale=1)
     
     with gr.Row():
-        clear_btn = gr.Button("🗑️ Очистить")
-        retry_btn = gr.Button("🔄 Повторить")
+        clear_btn = gr.Button("🗑️ Очистить историю", variant="secondary")
+        retry_btn = gr.Button("🔄 Повторить последнее", variant="secondary")
     
-    status = gr.Textbox(label="Логи", interactive=False, lines=4)
-    
-    send_btn.click(fn=respond, inputs=[msg_input, chatbot, model_dropdown, system_prompt], outputs=[chatbot, status, msg_input])
-    msg_input.submit(fn=respond, inputs=[msg_input, chatbot, model_dropdown, system_prompt], outputs=[chatbot, status, msg_input])
+    status = gr.Textbox(
+        label="Статус", 
+        interactive=False, 
+        lines=3,
+        placeholder="Здесь будут отображаться логи работы..."
+    )
     
-    def clear():
+    # Обработчики событий
+    def clear_chat():
         return [], "", gr.update(value="")
-    clear_btn.click(clear, outputs=[chatbot, status, msg_input])
     
-    def retry(history: List[Dict[str, str]]):
+    def retry_last(history: List[Dict[str, str]]):
         if history:
             last_user_msg = None
             for msg in reversed(history):
@@ -109,7 +190,48 @@ with gr.Blocks(title="🚀 Локальный HF Чат (на слабом CPU!)
                     break
             return last_user_msg if last_user_msg else ""
         return ""
-    retry_btn.click(retry, inputs=[chatbot], outputs=[msg_input])
+    
+    # Привязка событий
+    send_btn.click(
+        fn=respond, 
+        inputs=[msg_input, chatbot, model_dropdown, system_prompt], 
+        outputs=[chatbot, status, msg_input]
+    )
+    
+    msg_input.submit(
+        fn=respond, 
+        inputs=[msg_input, chatbot, model_dropdown, system_prompt], 
+        outputs=[chatbot, status, msg_input]
+    )
+    
+    clear_btn.click(
+        fn=clear_chat, 
+        outputs=[chatbot, status, msg_input]
+    )
+    
+    retry_btn.click(
+        fn=retry_last, 
+        inputs=[chatbot], 
+        outputs=[msg_input]
+    )
+    
+    # Информация о состоянии
+    gr.Markdown("""
+    ### 💡 Советы:
+    1. Первая загрузка модели может занять 1-5 минут
+    2. Ответы генерируются на CPU, будьте терпеливы
+    3. Для более быстрых ответов используйте Qwen2.5-0.5B
+    4. Очищайте историю, если чат становится медленным
+    """)
 
 if __name__ == "__main__":
-    demo.queue(max_size=10).launch(debug=True, ssr_mode=False)
+    print("=" * 50)
+    print("🚀 Запуск локального чат-бота на CPU")
+    print("=" * 50)
+    
+    demo.queue(max_size=5).launch(
+        debug=False, 
+        show_error=True,
+        server_name="0.0.0.0",
+        server_port=7860
+    )