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", "google/mt5-small")
+DEFAULT_GEN_MODEL = os.getenv("GEN_MODEL", "cointegrated/rut5-base")
 
 # 3) QA over checklist
 DEFAULT_QA_MODEL = os.getenv("QA_MODEL", "deepset/xlm-roberta-base-squad2")
@@ -161,8 +161,40 @@ def generate_checklist(user_goal: str, category: str, style: str, constraints: s
         prompt,
         max_new_tokens=450,
         do_sample=False,
+        num_beams=4,
     )
-    text = out[0]["generated_text"].strip()
+    text = (out[0].get("generated_text") or "").strip()
+
+    # Если модель вернула пусто/слишком коротко — повторим с более "толкающими" параметрами
+    if len(text) < 80:
+        short_prompt = (
+            "Сделай практичный чек-лист на русском.\n"
+            "Формат строго:\n"
+            "- [ ] пункт\n"
+            "Подпункты: '  - ...'\n"
+            "В конце: 'Проверка готовности' (3–5 вопросов) и 'Риски и как снизить'.\n\n"
+            f"Цель: {user_goal}\n"
+            f"Контекст: {constraints}\n"
+            f"Тема: {chosen_theme}\n\n"
+            "Чек-лист:\n"
+        )
+        out2 = gen_pipe(
+            short_prompt,
+            max_new_tokens=500,
+            do_sample=True,
+            temperature=0.9,
+            top_p=0.9,
+        )
+        text2 = (out2[0].get("generated_text") or "").strip()
+        if len(text2) > len(text):
+            text = text2
+
+    # Если всё равно пусто — покажем явное сообщение (а не “как будто ничего нет”)
+    if len(text) < 20:
+        text = (
+            "- [ ] Не удалось сгенерировать чек-лист этой моделью.\n"
+            "- [ ] Попробуйте сменить GEN_MODEL (см. ниже) или сократить описание цели.\n"
+        )
 
     # Store in state: checklist text + theme + original goal
     meta = {
@@ -271,7 +303,7 @@ with gr.Blocks(title="Умный чек-лист (3 Transformers)") as demo:
                 gen_btn = gr.Button("Сгенерировать чек-лист", variant="primary")
 
             with gr.Column(scale=3):
-                checklist_out = gr.Markdown(label="Чек-лист")
+                checklist_out = gr.Code(label="Чек-лист", language="markdown")
                 intent_debug = gr.Textbox(label="Диагностика распознавания намерения", lines=10)
 
         theme_out = gr.Textbox(label="Выбранная/распознанная тема (если определилась)", interactive=False)