app.py

import gradio as gr
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",
}

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,
        tokenizer=tokenizer,
        device=-1,  # Явно указываем CPU (-1 означает CPU)
        max_new_tokens=128,  # Ещё меньше токенов для экономии памяти
        do_sample=True,
        temperature=0.7,
        pad_token_id=tokenizer.eos_token_id
    )
    print(f"✅ {model_id} загружена!")
    return pipe

model_cache = {}

def respond(message: str, 
            history: List[Dict[str, str]], 
            model_key: str, 
            system_prompt: str) -> Tuple[List[Dict[str, str]], str, Dict[str, Any]]:
    try:
        if model_key not in model_cache:
            model_cache[model_key] = load_model(model_key)
        pipe = model_cache[model_key]
        
        print(f"🚀 Генерация: {model_key}, Msg='{message[:30]}...'")
        
        messages = []
        if system_prompt.strip():
            messages.append({"role": "system", "content": system_prompt})
        
        for msg in history:
            messages.append({"role": msg["role"], "content": msg["content"]})
        
        messages.append({"role": "user", "content": message})
        
        tokenizer = pipe.tokenizer
        
        # Используем чат-шаблон для 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}
        ]
        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}
        ]
        return new_history, error_msg, gr.update(value="")

# Исправлено: убран параметр theme для совместимости со старой версией Gradio
with gr.Blocks(title="🚀 Локальный HF Чат (на слабом CPU!)") as demo:
    gr.Markdown("""
    # 🚀 Локальный Inference (без API!)
   
    ⚠️ **Внимание**: Модели загружаются при первом выборе и могут занять несколько минут! Работают на слабом CPU - запаситесь терпением.
    """)
    
    with gr.Row():
        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=400, 
        label="Чат"
    )
    
    with gr.Row():
        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("🗑️ Очистить историю", variant="secondary")
        retry_btn = gr.Button("🔄 Повторить последнее", variant="secondary")
    
    status = gr.Textbox(
        label="Статус", 
        interactive=False, 
        lines=3,
        placeholder="Здесь будут отображаться логи работы..."
    )
    
    # Обработчики событий
    def clear_chat():
        return [], "", gr.update(value="")
    
    def retry_last(history: List[Dict[str, str]]):
        if history:
            last_user_msg = None
            for msg in reversed(history):
                if msg["role"] == "user":
                    last_user_msg = msg["content"]
                    break
            return last_user_msg if last_user_msg else ""
        return ""
    
    # Привязка событий
    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__":
    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,
        share=False  # Отключаем share для локального использования
    )