Update app.py

app.py

@@ -4,13 +4,77 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from functools import lru_cache
 
 # --- Конфигурация Hugging Face Space ---
-# Загрузка модели и токенизатора один раз при запуске приложения
 MODEL_NAME = "Kenan023214/PyroNet-mini"
 DEVICE = "cpu"  # Используем CPU, как указано для Basic Space
-MAX_NEW_TOKENS = 256
+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():
     """Загружает модель и токенайзер, кешируя их для производительности."""
@@ -19,7 +83,7 @@ def load_model():
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         device_map=DEVICE,
-        torch_dtype=torch.float32  # Используем float32 для совместимости с CPU
+        torch_dtype=torch.float32
     )
     print("Model loaded.")
     return tokenizer, model
@@ -28,7 +92,7 @@ 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):
@@ -46,37 +110,39 @@ def trim_history_to_max_tokens(messages, max_tokens):
 
 def build_messages_for_template(history_messages, reasoning: bool, language: str):
     """Подготавливает сообщения для шаблона, включая системное сообщение."""
-    if language == 'ru':
-        system_message = "Ты — дружелюбный ассистент, который говорит на русском. Отвечай кратко, но по делу."
-        reasoning_instruction = ("[REASONING MODE]\n"
-                                 "Когда отвечаешь, сначала представь краткие пронумерованные шаги рассуждения. "
-                                 "Затем на новой строке напиши 'Final:' и дай короткий окончательный ответ. Сохраняй шаги лаконичными.")
-    elif language == 'uk':
-        system_message = "Ти — дружній асистент, який говорить українською. Відповідай коротко, але по суті."
-        reasoning_instruction = ("[REASONING MODE]\n"
-                                 "Коли відповідаєш, спершу представ короткі пронумеровані кроки розмірковування. "
-                                 "Потім на новому рядку напиши 'Final:' і дай коротку остаточну відповідь. Зберігай кроки лаконічними.")
-    else:  # 'en'
-        system_message = "You are a friendly assistant who speaks English. Answer concisely but to the point."
-        reasoning_instruction = ("[REASONING MODE]\n"
-                                 "When answering, first present concise numbered reasoning steps. "
-                                 "Then on a new line write 'Final:' and give a short final answer. Keep steps brief.")
-
-    messages = [{"role": "system", "content": system_message}] + list(history_messages)
+    # Используем полный системный промпт из словаря
+    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": reasoning_instruction})
+        messages.append({"role": "user", "content": f"Режим рассуждения: покажи свои шаги, а затем окончательный ответ, начиная с '{REASONING_SEPARATOR}'"})
 
     return messages
 
-def extract_assistant_reply(raw_generated_text: str) -> str:
-    """Убирает лишние токены и возвращает только ответ ассистента."""
+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, "")
-    return text.strip()
+
+    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):
@@ -88,12 +154,12 @@ def generate_response(user_text: str, history, reasoning: bool, language: str):
 
     messages_for_template = build_messages_for_template(trimmed_history, reasoning, language)
 
-    # Выбираем шаблон из файлов в репозитории
-    template_file = f"chat_template_{language}.jinja"
+    # Используем содержимое шаблона из словаря, а не путь к файлу
+    template_content = CHAT_TEMPLATES.get(language, CHAT_TEMPLATES["en"])
     
     text = tokenizer.apply_chat_template(
         messages_for_template,
-        template_path=template_file,
+        chat_template=template_content,  # Передаем содержимое шаблона напрямую
         tokenize=False,
         add_generation_prompt=True
     )
@@ -111,19 +177,29 @@ def generate_response(user_text: str, history, reasoning: bool, language: str):
         )
 
     raw = tokenizer.decode(outputs[0], skip_special_tokens=False)
-    reply = extract_assistant_reply(raw)
+    reply, reasoning_text = extract_assistant_reply_and_reasoning(raw)
 
     history.append({"role": "assistant", "content": reply})
     
-    return "", history
+    return "", 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.")
     
-    chatbot = gr.Chatbot(height=500)
-    
+    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(
@@ -148,18 +224,18 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     btn_send.click(
         fn=generate_response,
         inputs=[msg, chatbot, reasoning_checkbox, language_dropdown],
-        outputs=[msg, chatbot]
+        outputs=[msg, chatbot, reasoning_box]
     )
     msg.submit(
         fn=generate_response,
         inputs=[msg, chatbot, reasoning_checkbox, language_dropdown],
-        outputs=[msg, chatbot]
+        outputs=[msg, chatbot, reasoning_box]
     )
     btn_clear.click(
-        fn=lambda: ([], None),
+        fn=lambda: ([], "", ""),
         inputs=[],
-        outputs=[chatbot, msg]
+        outputs=[chatbot, msg, reasoning_box]
     )
     
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()