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 для локального использования )