app.py

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()