app.py

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import time
import re
from typing import Tuple, Dict

# ------------------------------------------------------------
# Конфигурация
# ------------------------------------------------------------
MODEL_NAMES = [
    "Dilana/Llama-3.2-1B-Adaptive-RAG-v3",
    "LiquidAI/LFM2-1.2B-RAG",
    "HuggingFaceTB/SmolLM3-3B",
    "thelamapi/next2.5",
    "Qwen/Qwen3-4B-Instruct-2507",
    "utter-project/EuroLLM-1.7B-Instruct",
    "ai-sage/GigaChat3.1-10B-A1.8B"
]
DEFAULT_MODEL = MODEL_NAMES[0]

# Лимиты на длину ввода (в символах)
MAX_DOCUMENT_CHARS = 2000
MAX_QUESTION_CHARS = 1000
MAX_TOTAL_CHARS = MAX_DOCUMENT_CHARS + MAX_QUESTION_CHARS
MAX_PROMPT_TOKENS = 512

# Кэш для моделей и токенизаторов
model_cache: Dict[str, Tuple] = {}  # имя -> (tokenizer, model)

def load_model(model_name: str):
    """Загружает токенизатор и модель, если ещё не загружены."""
    if model_name not in model_cache:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForCausalLM.from_pretrained(model_name)
        model_cache[model_name] = (tokenizer, model)
    return model_cache[model_name]

def truncate_text(text: str, max_chars: int) -> str:
    """Обрезает текст до указанного количества символов (грубо, по символам)."""
    if len(text) > max_chars:
        return text[:max_chars] + "..."
    return text

def generate_response(
    document: str,
    question: str,
    model_name: str,
    max_new_tokens: int,
    temperature: float,
    repetition_penalty: float
) -> Tuple[str, float, float, float]:
    """
    Генерирует ответ модели на основе документа и вопроса.
    Возвращает (ответ, время_генерации_сек).
    """
    # Проверка на пустые входные данные
    if not document.strip():
        return "Ошибка: документ не может быть пустым.", 0.0
    if not question.strip():
        return "Ошибка: вопрос не может быть пустым.", 0.0

    # Обрезка по длине
    document = truncate_text(document, MAX_DOCUMENT_CHARS)
    question = truncate_text(question, MAX_QUESTION_CHARS)

    # Формирование промпта (простая инструкция)
    prompt = f"Ты бот, который должен чётко и точно ответить на вопрос пользователя или сообщить, что требуемой информации не обнаружено, по документу:<document>{document}</document>\nВопрос пользователя:<question>{question}</question>\nОтвет на вопрос:<answer>\n"
    
    messages = [
        {"role": "system", "content": f"Ты бот, который дает короткий и чёткий ответ пользователю на русском языке строго по данным из документа не выдумывая ничего лишнего. Задача выполняется на 1000\10. Данные из документа: <document>{document}</document>"},
        {"role": "user", "content": f"Вопрос по документу: {question}"}
    ]
    
    # Загрузка модели
    try:
        tokenizer, model = load_model(model_name)
    except Exception as e:
        return f"Ошибка загрузки модели: {type(e).__name__}: {e}", 0.0

    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    start_tok = time.time()
    # Токенизация с учётом максимальной длины модели
    try:
        inputs = tokenizer(
            prompt,
            return_tensors="pt",
            truncation=True,
            max_length=MAX_PROMPT_TOKENS
        )
        tok_time = time.time() - start_tok
    except Exception as e:
        return f"Ошибка токенизации: {type(e).__name__}: {e}", 0.0

    char_count = len(prompt)
    tok_time_per_char = tok_time / char_count if char_count > 0 else 0.0
    
    # Генерация
    start_time = time.time()
    try:
        with torch.no_grad():
            outputs = model.generate(
                inputs.input_ids,
                max_new_tokens=max_new_tokens,
                temperature=temperature if temperature > 0 else None,
                do_sample=False,
                top_p=0.95,
                repetition_penalty=repetition_penalty,      # штраф за повторяющиеся токены
                early_stopping=True,          # остановка при достижении eos_token
                pad_token_id=tokenizer.eos_token_id
            )
        gen_time = time.time() - start_time
    except Exception as e:
        return f"Ошибка генерации: {type(e).__name__}: {e}", time.time() - start_time

    # Количество сгенерированных токенов
    generated_tokens = outputs[0].shape[0] - inputs.input_ids.shape[1]
    gen_time_per_token = gen_time / generated_tokens if generated_tokens > 0 else 0.0
    
    # Декодирование ответа
    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
    if not response.strip():
        response = "[модель не дала ответа]"

    return response.strip(), gen_time, tok_time_per_char, gen_time_per_token

# ------------------------------------------------------------
# Интерфейс Gradio
# ------------------------------------------------------------
with gr.Blocks(title="Мини-чат по документу (русский язык)") as demo:
    gr.Markdown("""
    ## Чат с моделью на основе одного документа
    Задайте вопрос по предоставленному тексту. Модель ответит, используя только информацию из документа.
    """)

    with gr.Row():
        with gr.Column(scale=2):
            document_input = gr.Textbox(
                label="Документ (контекст)",
                lines=6,
                placeholder="Вставьте текст документа здесь..."
            )
            question_input = gr.Textbox(
                label="Ваш вопрос",
                lines=2,
                placeholder="Например: О чём говорится в документе?"
            )
            with gr.Row():
                model_selector = gr.Dropdown(
                    choices=MODEL_NAMES,
                    value=DEFAULT_MODEL,
                    label="Модель"
                )
                max_tokens_slider = gr.Slider(
                    10, 200, value=50, step=5,
                    label="Макс. новых токенов"
                )
                temperature_slider = gr.Slider(
                    0.0, 2.0, value=0.7, step=0.1,
                    label="Температура"
                )
                repetition_penalty_slider = gr.Slider(
                    0.1, 2.0, value=1.0, step=0.1,
                    label="Штраф за повторение"
                )
            submit_btn = gr.Button("Спросить", variant="primary")

        with gr.Column(scale=1):
            answer_output = gr.Textbox(
                label="Ответ модели",
                lines=6,
                interactive=False
            )
            latency_output = gr.Textbox(
                label="Время генерации (сек)",
                lines=1,
                interactive=False
            )
            tok_time_output = gr.Textbox(
                label="Ср. время токенизации на символ (сек)",
                lines=1,
                interactive=False
            )
            gen_time_output = gr.Textbox(
                label="Ср. время генерации на токен (сек)",
                lines=1,
                interactive=False
            )

    # Примеры (заполняют документ и вопрос, остальные параметры остаются текущими)
    gr.Examples(
        examples=[
            [
                "Кофе эспрессо готовится путём пропускания горячей воды под давлением через молотые зёрна. Температура воды 90-96°C, давление 9 бар. Выход напитка 25-35 мл.",
                "Как приготовить эспрессо?"
            ],
            [
                "Солнечная система состоит из Солнца и планет: Меркурий, Венера, Земля, Марс, Юпитер, Сатурн, Уран, Нептун. Земля — третья планета от Солнца, единственная известная планета с жизнью.",
                "Какая планета третья от Солнца?"
            ],
            [
                "Для сборки стола необходимо: столешница, 4 ножки, 8 шурупов, отвёртка. Сначала прикрутить ножки к столешнице, затянув шурупы крест-накрест.",
                "Какие инструменты нужны для сборки стола?"
            ],
            [
                "Исследования последних лет показывают, что регулярное употребление зелёного чая может снижать риск сердечно-сосудистых заболеваний. В состав зелёного чая входят катехины – антиоксиданты, которые нейтрализуют свободные радикалы и уменьшают окислительный стресс. Кроме того, зелёный чай содержит L-теанин – аминокислоту, способствующую расслаблению и улучшению когнитивных функций. Однако важно помнить, что чрезмерное потребление (более 5 чашек в день) может привести к негативным эффектам из-за кофеина. Врачи рекомендуют употреблять 2–3 чашки качественного зелёного чая в день для поддержания здоровья.",
                "Какие полезные вещества содержатся в зелёном чае и как они влияют на организм?"
            ],
            [
                "Для установки программы «Калькулятор v2.0» скачайте установочный файл с официального сайта. Запустите скачанный файл и следуйте инструкциям мастера установки. На первом этапе выберите язык интерфейса (русский или английский). Затем укажите папку для установки (по умолчанию C:\\Program Files\\Calculator). После завершения установки на рабочем столе появится ярлык. Для запуска программы дважды кликните по ярлыку. В главном окне доступны базовые арифметические операции: сложение, вычитание, умножение, деление. Для продвинутых вычислений откройте меню «Вид» и выберите «Инженерный режим».",
                "Как переключить программу в инженерный режим?"
            ],
            [
                "Старик сажал яблони. Ему говорили: «Зачем тебе эти яблони? Долго ждать с них плода, и ты не съешь с них яблочка». Старик ответил: «Я не съем – другие съедят, мне спасибо скажут». Эта притча отражает идею бескорыстной заботы о будущем. Многие философы считали, что смысл жизни заключается не только в личном благополучии, но и в том, чтобы оставить след, помочь следующим поколениям. В современном мире этот принцип можно увидеть в экологических движениях, посадке деревьев, создании общественных парков.",
                "Какой смысл вложен в притчу о старике, сажающем яблони?"
            ]
        ],
        inputs=[document_input, question_input],
        label="Примеры запросов"
    )

    # Функция обработки
    def process(document, question, model_name, max_tokens, temperature, repetition_penalty):
        answer, gen_time, tok_time_char, gen_time_token = generate_response(
            document, question, model_name, max_tokens, temperature, repetition_penalty
        )
        return answer, f"{gen_time:.3f}", f"{tok_time_char:.6f}", f"{gen_time_token:.6f}"

    submit_btn.click(
        fn=process,
        inputs=[document_input, question_input, model_selector, max_tokens_slider, temperature_slider, repetition_penalty_slider],
        outputs=[answer_output, latency_output, tok_time_output, gen_time_output]
    )

demo.launch()