Create app.py

app.py

@@ -0,0 +1,186 @@
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import time
+import re
+from typing import Tuple, Dict
+
+# ------------------------------------------------------------
+# Конфигурация
+# ------------------------------------------------------------
+MODEL_NAMES = [
+    "tinkoff-ai/ruDialoGPT-small",
+    "tinkoff-ai/ruDialoGPT-medium"
+]
+DEFAULT_MODEL = MODEL_NAMES[0]
+
+# Лимиты на длину ввода (в символах)
+MAX_DOCUMENT_CHARS = 2000
+MAX_QUESTION_CHARS = 1000
+MAX_TOTAL_CHARS = MAX_DOCUMENT_CHARS + MAX_QUESTION_CHARS
+
+# Кэш для моделей и токенизаторов
+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
+) -> Tuple[str, 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}\nВопрос: {question}\nОтвет:"
+
+    # Загрузка модели
+    try:
+        tokenizer, model = load_model(model_name)
+    except Exception as e:
+        return f"Ошибка загрузки модели: {type(e).__name__}: {e}", 0.0
+
+    # Токенизация с учётом максимальной длины модели
+    try:
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            truncation=True,
+            max_length=tokenizer.model_max_length
+        )
+    except Exception as e:
+        return f"Ошибка токенизации: {type(e).__name__}: {e}", 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,
+                do_sample=True,
+                top_p=0.95,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        latency = time.time() - start_time
+    except Exception as e:
+        return f"Ошибка генерации: {type(e).__name__}: {e}", time.time() - start_time
+
+    # Декодирование ответа
+    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
+    if not response.strip():
+        response = "[модель не дала ответа]"
+
+    return response.strip(), latency
+
+# ------------------------------------------------------------
+# Интерфейс 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.1, 2.0, value=0.7, 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
+            )
+
+    # Примеры (заполняют документ и вопрос, остальные параметры остаются текущими)
+    gr.Examples(
+        examples=[
+            [
+                "Кофе эспрессо готовится путём пропускания горячей воды под давлением через молотые зёрна. Температура воды 90-96°C, давление 9 бар. Выход напитка 25-35 мл.",
+                "Как приготовить эспрессо?"
+            ],
+            [
+                "Солнечная система состоит из Солнца и планет: Меркурий, Венера, Земля, Марс, Юпитер, Сатурн, Уран, Нептун. Земля — третья планета от Солнца, единственная известная планета с жизнью.",
+                "Какая планета третья от Солнца?"
+            ],
+            [
+                "Для сборки стола необходимо: столешница, 4 ножки, 8 шурупов, отвёртка. Сначала прикрутить ножки к столешнице, затянув шурупы крест-накрест.",
+                "Какие инструменты нужны для сборки стола?"
+            ]
+        ],
+        inputs=[document_input, question_input],
+        label="Примеры запросов"
+    )
+
+    # Функция обработки
+    def process(document, question, model_name, max_tokens, temperature):
+        answer, latency = generate_response(
+            document, question, model_name,
+            max_tokens, temperature
+        )
+        return answer, f"{latency:.3f}"
+
+    submit_btn.click(
+        fn=process,
+        inputs=[document_input, question_input, model_selector, max_tokens_slider, temperature_slider],
+        outputs=[answer_output, latency_output]
+    )
+
+demo.launch()