Item Search

app.py

@@ -11,7 +11,7 @@ from transformers import pipeline
 DEFAULT_INTENT_MODEL = os.getenv("INTENT_MODEL", "joeddav/xlm-roberta-large-xnli")
 
 # 2) Checklist generator
-DEFAULT_GEN_MODEL = os.getenv("GEN_MODEL", "cointegrated/rut5-base")
+DEFAULT_GEN_MODEL = os.getenv("GEN_MODEL", "Qwen/Qwen2.5-0.5B-Instruct")
 
 # 3) QA over checklist
 DEFAULT_QA_MODEL = os.getenv("QA_MODEL", "deepset/xlm-roberta-base-squad2")
@@ -47,7 +47,7 @@ intent_pipe, intent_model_used = safe_make_pipeline(
     DEFAULT_INTENT_MODEL,
 )
 gen_pipe, gen_model_used = safe_make_pipeline(
-    "text2text-generation",
+    "text-generation",
     DEFAULT_GEN_MODEL,
 )
 qa_pipe, qa_model_used = safe_make_pipeline(
@@ -99,27 +99,25 @@ def infer_intent(user_goal: str, labels: list[str]):
 
 
 def build_checklist_prompt(user_goal: str, theme: str | None, style: str, constraints: str):
-    """
-    Prompt for generator model.
-    """
-    theme_part = f"Тема (если помогает): {theme}\n" if theme else ""
-    constraints_part = f"Ограничения/контекст: {constraints}\n" if constraints else ""
+    theme_part = f"Тема: {theme}\n" if theme else ""
+    constraints_part = f"Контекст: {constraints}\n" if constraints else ""
 
-    # Works for mt5/flan-t5 style models; they respond better to clear structure.
     return (
         "Ты — помощник, который делает практичные чек-листы.\n"
-        "Сформируй чек-лист так, чтобы обычный пользователь мог выполнить задачу.\n"
+        "Верни ТОЛЬКО чек-лист без вступлений.\n"
+        "Формат:\n"
+        "- [ ] пункт\n"
+        "  - подпункт (если нужно)\n"
         "Требования:\n"
-        "- Выведи 8–15 пунктов максимум.\n"
-        "- Каждый пункт в формате: '- [ ] ...'\n"
-        "- Где уместно, добавляй краткие подпункты (через '  - ...').\n"
-        "- Делай пункты измеримыми и конкретными.\n"
-        "- В конце добавь блок 'Проверка готовности' (3–5 вопросов) и блок 'Риски и как снизить'.\n"
-        "- Пиши по-русски.\n\n"
-        f"Стиль: {style}\n"
+        "- 10–15 пунктов\n"
+        "- конкретно и измеримо\n"
+        "- в конце добавь разделы:\n"
+        "Проверка готовности: (3–5 вопросов)\n"
+        "Риски и как снизить: (3–6 пунктов)\n"
+        "Пиши по-русски.\n\n"
         f"{theme_part}"
         f"{constraints_part}"
-        f"Задача пользователя: {user_goal}\n\n"
+        f"Цель: {user_goal}\n\n"
         "Чек-лист:\n"
     )
 
@@ -160,11 +158,22 @@ def generate_checklist(user_goal: str, category: str, style: str, constraints: s
     out = gen_pipe(
         prompt,
         max_new_tokens=450,
-        do_sample=False,
-        num_beams=4,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+        repetition_penalty=1.05,
+        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 = (