import gradio as gr import torch from transformers import AutoModelForCausalLM, AutoTokenizer from functools import lru_cache import time # --- Конфигурация Hugging Face Space --- MODEL_NAME = "Kenan023214/PyroNet-mini" DEVICE = "cpu" # Используем CPU, как указано для Basic Space MAX_NEW_TOKENS = 512 # Увеличим для "хода мыслей" MAX_CONTEXT_TOKENS = 2048 # Словарь с встроенным содержимым шаблонов чата CHAT_TEMPLATES = { "ru": """<|system|> Ты — PyroNet-mini, облегчённая и свободная версия PyroNet, созданная Артёмом (IceL1ghtning). - Эксперт в физике, математике, программировании, биологии и смежных областях. - Дружелюбна, энергична, слегка иронична. - Отвечай на языке пользователя (русский). - В режиме reasoning показывай шаги рассуждений → затем итог; в обычном режиме — будь краткой. - Предпочитай списки и нумерацию, код выделяй в ```код``` с тэгом языка, математика = формула + результат. - Отказывай в явно опасных/незаконных запросах, предлагай альтернативы. <|end|> {% for m in messages %} {% if m['role'] == 'user' %} <|user|>{{ m['content'] }}<|end|> {% elif m['role'] == 'assistant' %} <|assistant|>{{ m['content'] }}<|end|> {% endif %} {% endfor %} {% if add_generation_prompt %}<|assistant|>{% endif %}""", "en": """<|system|> You are **PyroNet-mini**, a lighter and freer version of PyroNet, created by Artyom (IceL1ghtning) in Ukraine. - You are knowledgeable in physics, mathematics, programming, biology, and adjacent domains. - Energetic, friendly, slightly ironic. - Mirror the user's language (English). - In reasoning mode: show concise step-by-step reasoning → then final answer; otherwise be concise. - Prefer bullet points and numbered steps, code in ```code``` with correct language tags, math = formula + numeric result. - Refuse unsafe/illegal requests, suggest safe alternatives. <|end|> {% for m in messages %} {% if m['role'] == 'user' %} <|user|>{{ m['content'] }}<|end|> {% elif m['role'] == 'assistant' %} <|assistant|>{{ m['content'] }}<|end|> {% endif %} {% endfor %} {% if add_generation_prompt %}<|assistant|>{% endif %}""", "uk": """<|system|> Ти — **PyroNet-mini**, полегшена й більш вільна версія PyroNet, створена Артемом (IceL1ghtning) в Україні. - Експерт у фізиці, математиці, програмуванні, біології та суміжних темах. - Енергійна, дружня, злегка іронічна. - Відповідай на мові користувача (українська). - У режимі reasoning показуй лаконічні кроки → потім висновок; в іншому будь короткою. - Віддавай перевагу спискам, код у ```код``` з тегом мови, математика = формула + результат. - Відмовляй у небезпечних/незаконних запитах, пропонуй альтернативи. <|end|> {% for m in messages %} {% if m['role'] == 'user' %} <|user|>{{ m['content'] }}<|end|> {% elif m['role'] == 'assistant' %} <|assistant|>{{ m['content'] }}<|end|> {% endif %} {% endfor %} {% if add_generation_prompt %}<|assistant|>{% endif %}""" } # Ключевая фраза для разделения ответа модели REASONING_SEPARATOR = "Final:" @lru_cache(maxsize=1) def load_model(): """Загружает модель и токенайзер, кешируя их для производительности.""" print("Loading model and tokenizer...") tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, device_map=DEVICE, torch_dtype=torch.float32 ) print("Model loaded.") return tokenizer, model tokenizer, model = load_model() # --- Утилиты --- def num_tokens_of_text(text: str) -> int: """Приблизительное количество токенов для заданного текста.""" return len(tokenizer.encode(text, add_special_tokens=False)) def trim_history_to_max_tokens(messages, max_tokens): """Обрезает историю сообщений, чтобы она соответствовала лимиту токенов.""" rev = list(reversed(messages)) total = 0 kept = [] for m in rev: approx = num_tokens_of_text(m["content"]) + 8 if total + approx > max_tokens: break kept.append(m) total += approx return list(reversed(kept)) def build_messages_for_template(history_messages, reasoning: bool, language: str): """Подготавливает сообщения для шаблона, включая системное сообщение.""" full_template_content = CHAT_TEMPLATES.get(language, CHAT_TEMPLATES["en"]) system_start_tag = "<|system|>" system_end_tag = "<|end|>" system_message_raw = full_template_content.split(system_start_tag)[1].split(system_end_tag)[0].strip() messages = [{"role": "system", "content": system_message_raw}] + list(history_messages) if reasoning: messages.append({"role": "user", "content": f"Режим рассуждения: покажи свои шаги, а затем окончательный ответ, начиная с '{REASONING_SEPARATOR}'"}) return messages def extract_assistant_reply_and_reasoning(raw_generated_text: str) -> tuple[str, str]: """Убирает лишние токены и разделяет ответ на ход мыслей и окончательный ответ.""" text = raw_generated_text if "<|assistant|>" in text: text = text.split("<|assistant|>")[-1] for tag in ["<|end|>", "<|end_of_text|>", "<|end|>"]: text = text.replace(tag, "") text = text.strip() if REASONING_SEPARATOR in text: parts = text.split(REASONING_SEPARATOR, 1) reasoning = parts[0].strip() reply = parts[1].strip() return reply, reasoning else: return text, "" # Если разделитель не найден, возвращаем все как ответ # --- Основная функция для Gradio --- def generate_response(user_text: str, history, reasoning: bool, language: str): """ Обрабатывает пользовательский запрос, генерирует ответ и возвращает его с эффектом печати. """ # Добавляем user-сообщение во внутреннюю историю history.append([user_text, None]) # Конвертируем Gradio-историю в наш внутренний формат internal_history = [{"role": "user", "content": h[0]} for h in history if h[0] is not None] trimmed_history = trim_history_to_max_tokens(internal_history, MAX_CONTEXT_TOKENS) messages_for_template = build_messages_for_template(trimmed_history, reasoning, language) template_content = CHAT_TEMPLATES.get(language, CHAT_TEMPLATES["en"]) # Применяем шаблон и токенизируем text = tokenizer.apply_chat_template( messages_for_template, chat_template=template_content, tokenize=False, add_generation_prompt=True ) inputs = tokenizer(text, return_tensors="pt").to(DEVICE) with torch.no_grad(): outputs = model.generate( **inputs, max_new_tokens=MAX_NEW_TOKENS, do_sample=True, top_p=0.9, temperature=0.8, pad_token_id=tokenizer.eos_token_id ) raw = tokenizer.decode(outputs[0], skip_special_tokens=False) # Извлекаем финальный ответ и ход мыслей reply, reasoning_text = extract_assistant_reply_and_reasoning(raw) # Обновляем историю Gradio с финальным ответом history[-1][1] = "" # Используем генератор для создания эффекта печати for chunk in reply.split(): history[-1][1] += chunk + " " time.sleep(0.05) # Небольшая задержка для анимации yield "", history, reasoning_text # --- Интерфейс Gradio --- with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown("# PyroNet-mini Chat") gr.Markdown("A demonstration of PyroNet-mini with multilingual templates and a reasoning mode.") with gr.Tabs(): with gr.TabItem("Chat"): chatbot = gr.Chatbot(height=500) with gr.TabItem("Reasoning"): reasoning_box = gr.Textbox( label="Reasoning Steps", interactive=False, lines=20, placeholder="The model's thought process will appear here when Reasoning Mode is enabled.", show_copy_button=True ) with gr.Row(): with gr.Column(scale=4): msg = gr.Textbox( label="Your Prompt", placeholder="Write your message here...", container=False ) with gr.Column(scale=1, min_width=100): language_dropdown = gr.Dropdown( choices=["en", "ru", "uk"], value="en", label="Language", container=False ) reasoning_checkbox = gr.Checkbox( label="Enable Reasoning Mode" ) btn_send = gr.Button("Send") btn_clear = gr.Button("Clear") btn_send.click( fn=generate_response, inputs=[msg, chatbot, reasoning_checkbox, language_dropdown], outputs=[msg, chatbot, reasoning_box] ) msg.submit( fn=generate_response, inputs=[msg, chatbot, reasoning_checkbox, language_dropdown], outputs=[msg, chatbot, reasoning_box] ) btn_clear.click( fn=lambda: (None, [], ""), inputs=[], outputs=[msg, chatbot, reasoning_box] ) if __name__ == "__main__": demo.launch()