Item Search

app.py

@@ -1,16 +1,17 @@
+# app.py
 import os
 import re
 import gradio as gr
 import torch
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
 
-# ----------------------------
-# Model config (3 Transformers)
-# ----------------------------
+# ============================
+# 3 Transformers (pipelines)
+# ============================
 # 1) Intent / zero-shot
 DEFAULT_INTENT_MODEL = os.getenv("INTENT_MODEL", "joeddav/xlm-roberta-large-xnli")
 
-# 2) Checklist generator
+# 2) Checklist generator (instruct)
 DEFAULT_GEN_MODEL = os.getenv("GEN_MODEL", "Qwen/Qwen2.5-0.5B-Instruct")
 
 # 3) QA over checklist
@@ -20,28 +21,21 @@ DEVICE = 0 if torch.cuda.is_available() else -1
 
 
 def safe_make_pipeline(task: str, model_name: str, **kwargs):
-    """
-    Tries to load a pipeline; if fails, uses a smaller/safer fallback.
-    This keeps the Space alive even if the preferred model name is unavailable.
-    """
+    """Load pipeline with small fallbacks if needed."""
     try:
         return pipeline(task, model=model_name, device=DEVICE, **kwargs), model_name
-    except Exception as e:
-        # Fallbacks (kept simple)
+    except Exception:
         if task == "zero-shot-classification":
             fallback = "facebook/bart-large-mnli"
-        elif task == "text2text-generation":
-            fallback = "google/flan-t5-base"
+        elif task == "text-generation":
+            fallback = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
         elif task == "question-answering":
             fallback = "distilbert-base-cased-distilled-squad"
         else:
-            raise e
+            raise
+        return pipeline(task, model=fallback, device=DEVICE, **kwargs), fallback
 
-        pipe = pipeline(task, model=fallback, device=DEVICE, **kwargs)
-        return pipe, fallback
 
-
-# Create 3 pipelines (3 transformers)
 intent_pipe, intent_model_used = safe_make_pipeline(
     "zero-shot-classification",
     DEFAULT_INTENT_MODEL,
@@ -55,19 +49,24 @@ qa_pipe, qa_model_used = safe_make_pipeline(
     DEFAULT_QA_MODEL,
 )
 
+# For Qwen chat formatting (works for many instruct LMs too)
+try:
+    gen_tokenizer = AutoTokenizer.from_pretrained(gen_model_used, use_fast=True)
+except Exception:
+    gen_tokenizer = None
 
-# ----------------------------
-# App logic
-# ----------------------------
+# ============================
+# Labels / UI choices
+# ============================
 DEFAULT_LABELS = [
     "обучение",
     "переезд",
+    "релокация/иммиграция",
     "путешествие",
     "карьера/поиск работы",
     "финансы/покупка",
     "здоровье/фитнес",
     "ремонт/быт",
-    "личный проект",
 ]
 
 CATEGORY_CHOICES = ["Авто (определить по тексту)"] + DEFAULT_LABELS
@@ -80,15 +79,10 @@ def normalize_text(s: str) -> str:
 
 
 def infer_intent(user_goal: str, labels: list[str]):
-    """
-    Returns (top_label, score, all_labels_scores_as_text).
-    """
     if not user_goal:
         return "не задано", 0.0, "Нет входного текста."
 
-    # zero-shot expects candidate_labels
     result = intent_pipe(user_goal, candidate_labels=labels, multi_label=False)
-    # result: {'sequence': ..., 'labels': [...], 'scores': [...]}
     top_label = result["labels"][0]
     top_score = float(result["scores"][0])
 
@@ -98,23 +92,36 @@ def infer_intent(user_goal: str, labels: list[str]):
     return top_label, top_score, "\n".join(lines)
 
 
-def build_checklist_prompt(user_goal: str, theme: str | None, style: str, constraints: str):
+def build_checklist_prompt(user_goal: str, theme: str | None, style: str, constraints: str) -> str:
     theme_part = f"Тема: {theme}\n" if theme else ""
     constraints_part = f"Контекст: {constraints}\n" if constraints else ""
 
+    style_hint = {
+        "кратко": "короткие пункты без лишних слов",
+        "подробно": "чуть более подробные пункты + подпункты где уместно",
+        "с акцентом на риски": "сильный акцент на предотвращение ошибок и рисков",
+        "с акцентом на сроки": "добавляй дедлайны/временные окна там, где уместно",
+    }.get(style, "короткие пункты")
+
     return (
-        "Ты — помощник, который делает практичные чек-листы.\n"
+        "Составь практичный чек-лист на русском языке.\n"
         "Верни ТОЛЬКО чек-лист без вступлений.\n"
-        "Формат:\n"
+        "Запрещено:\n"
+        "- плейсхолдеры и шаблоны: никаких '[секунды]', '[какой-то текст]', '{...}', '<...>'\n"
+        "- таблицы, 'Расчёт', 'Банк', 'Сообщение', поля для заполнения\n"
+        "Формат ��трого:\n"
         "- [ ] пункт\n"
         "  - подпункт (если нужно)\n"
         "Требования:\n"
-        "- 10–15 пунктов\n"
-        "- конкретно и измеримо\n"
-        "- в конце добавь разделы:\n"
-        "Проверка готовности: (3–5 вопросов)\n"
-        "Риски и как снизить: (3–6 пунктов)\n"
-        "Пиши по-русски.\n\n"
+        f"- стиль: {style_hint}\n"
+        "- 12–18 пунктов\n"
+        "- без нумерации (никаких '1.')\n"
+        "- пункты конкретные и выполнимые\n"
+        "- в конце 2 блока:\n"
+        "Проверка готовности:\n"
+        "- ... (3–5 вопросов)\n"
+        "Риски и как снизить:\n"
+        "- ... (3–6 пунктов)\n\n"
         f"{theme_part}"
         f"{constraints_part}"
         f"Цель: {user_goal}\n\n"
@@ -122,90 +129,168 @@ def build_checklist_prompt(user_goal: str, theme: str | None, style: str, constr
     )
 
 
+def apply_chat_template_if_available(user_prompt: str) -> str:
+    """Wrap prompt into chat template if tokenizer supports it (Qwen works best with this)."""
+    if gen_tokenizer is None or not hasattr(gen_tokenizer, "apply_chat_template"):
+        return user_prompt
+    messages = [
+        {"role": "system", "content": "Ты пишешь грамотно по-русски и строго соблюдаешь формат чек-листа."},
+        {"role": "user", "content": user_prompt},
+    ]
+    return gen_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+
+
+def clean_checklist(text: str) -> str:
+    text = (text or "").strip()
+    # Remove T5-style tokens (just in case)
+    text = re.sub(r"<extra_id_\d+>", "", text).strip()
+
+    cleaned = []
+    for ln in text.splitlines():
+        s = ln.rstrip()
+
+        if not s.strip():
+            continue
+
+        # Cut obvious "form-like" garbage blocks
+        low = s.strip().lower()
+        if low.startswith(("расч", "банк", "сообщение", "время ", "план поездки", "расход")):
+            continue
+
+        # Remove square-bracket placeholders but keep "[ ]" checkbox
+        s2 = re.sub(r"\[(?!\s*\])([^\]]+)\]", "", s).strip()
+
+        # Normalize bullets
+        s2 = s2.replace("•", "-").replace("–", "-")
+
+        # Keep only checklist or the two required blocks
+        if s2.startswith("- [ ]") or s2.startswith("  -") or s2.startswith("Проверка готовности") or s2.startswith("Риски"):
+            cleaned.append(s2)
+
+    out = "\n".join(cleaned).strip()
+
+    # If model didn't follow checkbox format, force it
+    if "- [ ]" not in out:
+        raw = [l.strip() for l in text.splitlines() if l.strip()]
+        raw = [re.sub(r"^\d+[\).\s]+", "", r).strip(" -•\t") for r in raw]
+        raw = [r for r in raw if r]
+        raw = [f"- [ ] {r}" for r in raw[:18]]
+        out = "\n".join(raw).strip()
+
+    return out
+
+
+def polish_russian_same_model(checklist_text: str) -> str:
+    """Second pass with SAME generator (still transformer #2) to polish grammar."""
+    if not checklist_text or "- [ ]" not in checklist_text:
+        return checklist_text
+
+    polish_prompt = (
+        "Исправь орфографию, грамматику и стилистику русского языка.\n"
+        "НИЧЕГО НЕ ДОБАВЛЯЙ и не меняй смысл.\n"
+        "Сохрани формат чек-листа:\n"
+        "- [ ] пункт\n"
+        "  - подпункт\n"
+        "А также блоки 'Проверка готовности' и 'Риски и как снизить'.\n\n"
+        f"{checklist_text}\n"
+    )
+    polish_prompt = apply_chat_template_if_available(polish_prompt)
+
+    outp = gen_pipe(
+        polish_prompt,
+        max_new_tokens=260,
+        do_sample=False,
+        return_full_text=False,
+    )
+    text2 = (outp[0].get("generated_text") or "").strip()
+    text2 = clean_checklist(text2)
+
+    # keep only if not worse
+    if text2.count("- [ ]") >= checklist_text.count("- [ ]"):
+        return text2
+    return checklist_text
+
+
 def generate_checklist(user_goal: str, category: str, style: str, constraints: str):
     user_goal = normalize_text(user_goal)
     constraints = normalize_text(constraints)
 
     if not user_goal:
         return (
-            "Введите описание цели (например: 'Хочу переехать в другой город за 2 месяца').",
             "",
+            "Введите описание цели.",
             "",
+            f"intent={intent_model_used}\ngen={gen_model_used}\nqa={qa_model_used}",
             None,
             None,
         )
 
-    # Decide labels for intent detection (we keep it to 8)
-    labels = DEFAULT_LABELS
-
-    inferred_label, inferred_score, intent_debug = infer_intent(user_goal, labels)
+    inferred_label, inferred_score, intent_debug = infer_intent(user_goal, DEFAULT_LABELS)
 
-    chosen_theme = None
     if category and category != "Авто (определить по тексту)":
         chosen_theme = category
     else:
-        # Use inferred label only if confidence is decent; otherwise keep theme=None
-        chosen_theme = inferred_label if inferred_score >= 0.35 else None
+        chosen_theme = inferred_label if inferred_score >= 0.30 else None
 
-    prompt = build_checklist_prompt(
+    base_prompt = build_checklist_prompt(
         user_goal=user_goal,
         theme=chosen_theme,
         style=style,
         constraints=constraints,
     )
+    prompt = apply_chat_template_if_available(base_prompt)
 
-    # Generation parameters: conservative to avoid rambling
+    # Generation (primary)
     out = gen_pipe(
         prompt,
-        max_new_tokens=450,
+        max_new_tokens=520,
         do_sample=True,
-        temperature=0.7,
+        temperature=0.6,
         top_p=0.9,
-        repetition_penalty=1.05,
-        return_full_text=False,  # критично, чтобы не возвращался промпт целиком
+        repetition_penalty=1.08,
+        return_full_text=False,
     )
     text = (out[0].get("generated_text") or "").strip()
-
-    text = re.sub(r"<extra_id_\d+>", "", text).strip()
-
-    # Если модель вдруг не вывела чекбоксы — принудительно форматируем строки
-    if "- [ ]" not in text:
-        lines = [ln.strip("-• ").strip() for ln in text.splitlines() if ln.strip()]
-        lines = [f"- [ ] {ln}" for ln in lines[:15]]
-        text = "\n".join(lines)
-
-    # Если модель вернула пусто/слишком коротко — повторим с более "толкающими" параметрами
-    if len(text) < 80:
-        short_prompt = (
-            "Сделай практичный чек-лист на русском.\n"
-            "Формат строго:\n"
-            "- [ ] пункт\n"
-            "Подпункты: '  - ...'\n"
-            "В конце: 'Проверка готовности' (3–5 вопросов) и 'Риски и как снизить'.\n\n"
+    text = clean_checklist(text)
+
+    # Retry if too short/poor format
+    if text.count("- [ ]") < 8:
+        retry_prompt = (
+            "Сделай чек-лист строго в формате '- [ ] ...' (12–18 пунктов) на русском.\n"
+            "Без плейсхолдеров в квадратных скобках. Без таблиц. Без 'Расчёт/Банк/Сообщение'.\n"
+            "В конце добавь:\n"
+            "Проверка готовности: (3–5 вопросов)\n"
+            "Риски и как снизить: (3–6 пунктов)\n\n"
             f"Цель: {user_goal}\n"
             f"Контекст: {constraints}\n"
             f"Тема: {chosen_theme}\n\n"
             "Чек-лист:\n"
         )
+        retry_prompt = apply_chat_template_if_available(retry_prompt)
         out2 = gen_pipe(
-            short_prompt,
-            max_new_tokens=500,
+            retry_prompt,
+            max_new_tokens=620,
             do_sample=True,
-            temperature=0.9,
-            top_p=0.9,
+            temperature=0.75,
+            top_p=0.92,
+            repetition_penalty=1.10,
+            return_full_text=False,
         )
         text2 = (out2[0].get("generated_text") or "").strip()
-        if len(text2) > len(text):
+        text2 = clean_checklist(text2)
+        if text2.count("- [ ]") > text.count("- [ ]"):
             text = text2
 
-    # Если всё равно пусто — покажем явное сообщение (а не “как будто ничего нет”)
-    if len(text) < 20:
+    # Polish Russian (still same generator = transformer #2)
+    text = polish_russian_same_model(text)
+
+    # Final guard
+    if text.count("- [ ]") < 6:
         text = (
-            "- [ ] Не удалось сгенерировать чек-лист этой моделью.\n"
-            "- [ ] Попробуйте сменить GEN_MODEL (см. ниже) или сократить описание цели.\n"
+            "- [ ] Не удалось корректно сформировать чек-лист.\n"
+            "- [ ] Попробуйте уточнить цель (страна/срок/статус визы/бюджет) и нажать ещё раз.\n"
         )
 
-    # Store in state: checklist text + theme + original goal
     meta = {
         "goal": user_goal,
         "theme": chosen_theme,
@@ -216,16 +301,10 @@ def generate_checklist(user_goal: str, category: str, style: str, constraints: s
         "qa_model": qa_model_used,
     }
 
-    # A small header for UX
-    header = []
-    header.append(f"**Цель:** {user_goal}")
-    if chosen_theme:
-        header.append(f"**Тема:** {chosen_theme}")
-    header.append(f"**Модели:** intent=`{intent_model_used}`, gen=`{gen_model_used}`, qa=`{qa_model_used}`")
-    header.append("")
-    checklist_md = "\n".join(header) + text
+    models_info = f"intent={intent_model_used}\ngen={gen_model_used}\nqa={qa_model_used}"
+    theme_info = f"{chosen_theme or '(не определилась)'}"
 
-    return checklist_md, intent_debug, chosen_theme or "", checklist_md, meta
+    return text, intent_debug, theme_info, models_info, text, meta
 
 
 def answer_question(question: str, checklist_state: str, meta_state: dict | None):
@@ -234,61 +313,66 @@ def answer_question(question: str, checklist_state: str, meta_state: dict | None
         return "Сначала сгенерируйте чек-лист на первой вкладке.", ""
 
     if not question:
-        return "Введите вопрос (например: 'Какие документы подготовить?').", ""
+        return "Введите вопрос.", ""
 
-    # Use extractive QA first
-    context = checklist_state
-    qa_res = qa_pipe(question=question, context=context)
+    # Extractive QA first
+    qa_res = qa_pipe(question=question, context=checklist_state)
     answer = (qa_res.get("answer") or "").strip()
     score = float(qa_res.get("score") or 0.0)
 
     evidence = f"QA score: {score:.3f}\n"
     if answer:
-        evidence += f"Extracted span: {answer}\n"
+        evidence += f"Span: {answer}\n"
 
-    # If QA is weak or empty -> fallback to generator (still transformer #2, already loaded)
+    # If QA weak -> use generator (still transformer #2)
     if (not answer) or score < 0.20 or len(answer) < 3:
         goal = (meta_state or {}).get("goal", "")
         theme = (meta_state or {}).get("theme", "")
 
-        prompt = (
-            "Ты — помощник по уточняющим вопросам к чек-листу.\n"
-            "Ответь кратко и практично. Ссылайся на пункты чек-листа (если можно).\n"
-            "Если в чек-листе этого нет — предложи, какими 2–5 пунктами его дополнить.\n"
-            "Пиши по-русски.\n\n"
+        user_prompt = (
+            "Ответь на вопрос по чек-листу. Пиши по-русски, кратко и практично.\n"
+            "Если ответа нет в чек-листе — предложи 3–6 дополнительных пунктов, которые стоит добавить.\n"
+            "Не используй плейсхолдеры в квадратных скобках.\n\n"
             f"Цель: {goal}\n"
             f"Тема: {theme}\n\n"
             f"Чек-лист:\n{checklist_state}\n\n"
             f"Вопрос: {question}\n"
             "Ответ:\n"
         )
-        gen_out = gen_pipe(prompt, max_new_tokens=220, do_sample=False)[0]["generated_text"].strip()
+        prompt = apply_chat_template_if_available(user_prompt)
+        gen_out = gen_pipe(
+            prompt,
+            max_new_tokens=260,
+            do_sample=False,
+            return_full_text=False,
+        )[0]["generated_text"].strip()
+        gen_out = re.sub(r"<extra_id_\d+>", "", gen_out).strip()
         return gen_out, evidence + "Fallback: generator used (QA confidence low)."
 
-    # Otherwise return extracted answer with a bit of framing
-    final = f"{answer}\n\n_(Найдено в чек-листе; уверенность: {score:.2f})_"
-    return final, evidence
+    return f"{answer}\n\n_(Найдено в чек-листе; уверенность: {score:.2f})_", evidence
 
 
-# ----------------------------
-# Gradio UI
-# ----------------------------
+# ============================
+# UI: Checklist big at bottom,
+# helper info on the right
+# ============================
 with gr.Blocks(title="Умный чек-лист (3 Transformers)") as demo:
     gr.Markdown(
         "# ✅ Умный чек-лист (3 Transformers)\n"
-        "1) Распознаём намерение (zero-shot) → 2) Генерируем чек-лист → 3) Отвечаем на вопросы по чек-листу\n"
+        "1) распознаём намерение → 2) генерируем чек-лист → 3) отвечаем на вопросы по чек-листу\n"
     )
 
-    checklist_state = gr.State(value=None)  # stores checklist markdown
-    meta_state = gr.State(value=None)       # stores dict
+    checklist_state = gr.State(value=None)
+    meta_state = gr.State(value=None)
 
     with gr.Tab("1) Создать чек-лист"):
+        # Top row: inputs left, helper mini-blocks right
         with gr.Row():
-            with gr.Column(scale=2):
+            with gr.Column(scale=3):
                 user_goal = gr.Textbox(
                     label="Опишите, что вы хотите сделать",
-                    placeholder="Например: 'Хочу переехать в другой город за 2 месяца и не забыть важное'",
-                    lines=3,
+                    placeholder="Напр.: Хочу переехать в Китай на 2 года. Бюджет 200 000 ₽. Есть загранпаспорт.",
+                    lines=4,
                 )
 
                 category = gr.Dropdown(
@@ -305,30 +389,51 @@ with gr.Blocks(title="Умный чек-лист (3 Transformers)") as demo:
 
                 constraints = gr.Textbox(
                     label="Контекст/ограничения (необязательно)",
-                    placeholder="Напр.: бюджет, срок, страна/город, семейное положение, уровень опыта...",
-                    lines=2,
+                    placeholder="Напр.: бюджет, срок, город/страна, семья/дети, удалённая работа, уровень языка...",
+                    lines=3,
                 )
 
                 gen_btn = gr.Button("Сгенерировать чек-лист", variant="primary")
 
-            with gr.Column(scale=3):
-                checklist_out = gr.Code(label="Чек-лист", language="markdown")
-                intent_debug = gr.Textbox(label="Диагностика распознавания намерения", lines=10)
+            with gr.Column(scale=2):
+                theme_box = gr.Textbox(
+                    label="Выбранная/распознанная тема",
+                    lines=2,
+                    interactive=False,
+                )
+                models_box = gr.Textbox(
+                    label="Модели",
+                    lines=3,
+                    interactive=False,
+                )
+                intent_debug = gr.Textbox(
+                    label="Диагностика намерения",
+                    lines=10,
+                    interactive=False,
+                )
 
-        theme_out = gr.Textbox(label="Выбранная/распознанная тема (если определилась)", interactive=False)
+        # Big checklist output BELOW (full width)
+        checklist_out = gr.Code(label="Чек-лист", language="markdown")
 
         gen_btn.click(
             fn=generate_checklist,
             inputs=[user_goal, category, style, constraints],
-            outputs=[checklist_out, intent_debug, theme_out, checklist_state, meta_state],
+            outputs=[checklist_out, intent_debug, theme_box, models_box, checklist_state, meta_state],
         )
 
     with gr.Tab("2) Уточняющие вопросы по чек-листу"):
-        gr.Markdown("Задайте вопрос по уже сгенерированному чек-листу (например: *'Какие документы подготовить?'*).")
-        question = gr.Textbox(label="Ваш вопрос", placeholder="Введите вопрос...", lines=2)
-        ask_btn = gr.Button("Ответить", variant="primary")
-        answer_out = gr.Markdown(label="Ответ")
-        evidence_out = gr.Textbox(label="Тех. детали (score и режим ответа)", lines=6)
+        with gr.Row():
+            with gr.Column(scale=3):
+                gr.Markdown("Задайте вопрос по уже сгенерированному чек-листу.")
+                question = gr.Textbox(
+                    label="Ваш вопрос",
+                    placeholder="Напр.: Какие документы подготовить? Как снизить расходы? Что сделать первым делом?",
+                    lines=2,
+                )
+                ask_btn = gr.Button("Ответить", variant="primary")
+                answer_out = gr.Markdown(label="Ответ")
+            with gr.Column(scale=2):
+                evidence_out = gr.Textbox(label="Тех. детали", lines=10)
 
         ask_btn.click(
             fn=answer_question,
@@ -336,11 +441,5 @@ with gr.Blocks(title="Умный чек-лист (3 Transformers)") as demo:
             outputs=[answer_out, evidence_out],
         )
 
-    gr.Markdown(
-        "### Примечания\n"
-        "- Режим **QA** сначала пытается извлечь ответ прямо из чек-листа.\n"
-        "- Если уверенность низкая, включается генератор и предлагает уточнение/дополнение чек-листа.\n"
-    )
-
 if __name__ == "__main__":
     demo.launch()