Refactor predict function to streamline prompt building and enhance output formatting; update UI components for better user experience and status reporting

app.py

@@ -147,38 +147,26 @@ def predict(instruction: str,
             mode: str):
     """Generate text using the cached pipeline and return output or error message."""
     if not instruction or not instruction.strip():
-        return "⚠️ مهرباني وکړئ یوه لارښوونه ولیکئ."  # please provide an instruction
+        return "⚠️ مهرباني وکړئ یوه لارښوونه ولیکئ."
 
-    # Fast path: echo/useless mode avoids loading large models during testing.
+    def build_prompt() -> str:
+        base = instruction.strip()
+        if input_text and input_text.strip():
+            return base + "\n" + input_text.strip()
+        return base
+
+    prompt = build_prompt()
     active_mode = (mode or "").strip().lower() or ECHO_MODE
     if active_mode in ("echo", "useless"):
-        prompt = instruction.strip()
-        if input_text and input_text.strip():
-            prompt += "\n" + input_text.strip()
         if active_mode == "echo":
-            return prompt
-        else:
-            return "This is a useless placeholder response."
-
-    # Build a simple prompt: instruction (+ input if provided)
-    prompt = instruction.strip()
-    if input_text and input_text.strip():
-        prompt += "\n" + input_text.strip()
-
-    def _filter_generation_kwargs(kwargs: dict) -> dict:
-        allowed = {
-            "max_new_tokens",
-            "num_beams",
-            "do_sample",
-            "temperature",
-            "top_p",
-            "num_return_sequences",
-        }
-        return {k: v for k, v in kwargs.items() if k in allowed}
+            return f"### Prompt\n\n````\n{prompt}\n````\n\n### Output\n\n````\n{prompt}\n````"
+        return f"### Prompt\n\n````\n{prompt}\n````\n\n### Output\n\nThis is a useless placeholder response."
+
+    allowed_keys = {"max_new_tokens", "num_beams", "do_sample", "temperature", "top_p", "num_return_sequences"}
 
     try:
         gen = get_generator()
-        gen_kwargs = {
+        raw_kwargs = {
             "max_new_tokens": int(max_new_tokens),
             "num_beams": int(num_beams) if not do_sample else 1,
             "do_sample": bool(do_sample),
@@ -186,19 +174,23 @@ def predict(instruction: str,
             "top_p": float(top_p),
             "num_return_sequences": max(1, int(num_return_sequences)),
         }
-
-        gen_kwargs = _filter_generation_kwargs(gen_kwargs)
+        gen_kwargs = {k: v for k, v in raw_kwargs.items() if k in allowed_keys}
         outputs = gen(prompt, **gen_kwargs)
 
+        if not isinstance(outputs, list):
+            outputs = [outputs]
         texts = []
-        for out in outputs if isinstance(outputs, list) else [outputs]:
-            text = out.get("generated_text", "").strip()
-            texts.append(text)
-
+        for out in outputs:
+            if isinstance(out, dict):
+                text = out.get("generated_text", "").strip()
+            else:
+                text = str(out).strip()
+            if text:
+                texts.append(text)
         if not texts:
-            return "⚠️ No response generated."
-        return "\n\n---\n\n".join(texts)
-
+            return f"### Prompt\n\n````\n{prompt}\n````\n\n### Output\n\n⚠️ No response generated."
+        joined = "\n\n---\n\n".join(texts)
+        return f"### Prompt\n\n````\n{prompt}\n````\n\n### Output\n\n{joined}"
     except Exception as e:
         logger.exception("Generation failed: %s", e)
         return f"⚠️ Generation failed: {e}"
@@ -215,7 +207,6 @@ def build_ui():
             که د موډ بدلول غواړئ لاندې د Mode selector څخه استفاده وکړئ.
             """
         )
-
         with gr.Row():
             with gr.Column(scale=2):
                 instruction_dropdown = gr.Dropdown(
@@ -230,7 +221,7 @@ def build_ui():
                     label="لارښوونه",
                 )
                 input_text = gr.Textbox(lines=2, placeholder="اختیاري متن...", label="متن")
-                output = gr.Textbox(label="ځواب", interactive=False, lines=8)
+                output = gr.Markdown(label="ځواب")
                 generate_btn = gr.Button("جوړول", variant="primary")
                 mode_selector = gr.Dropdown(
                     choices=["off", "echo", "useless"],
@@ -238,6 +229,8 @@ def build_ui():
                     label="Mode (off=real, echo=return prompt, useless=fixed)",
                     interactive=True,
                 )
+                status_box = gr.Markdown(value="Loading status pending...", label="Status")
+                refresh_status = gr.Button("Refresh Status")
 
             with gr.Column(scale=1):
                 gr.Markdown("### د تولید تنظیمات")
@@ -250,12 +243,22 @@ def build_ui():
 
         instruction_dropdown.change(lambda x: x, inputs=instruction_dropdown, outputs=instruction_textbox)
 
+        def refresh():
+            return f"**Device:** {'GPU' if _detect_device() != -1 else 'CPU'} | **Offline:** {os.getenv('HF_HUB_OFFLINE','0')} | **Env Mode:** {ECHO_MODE}"
+
+        refresh_status.click(fn=refresh, inputs=None, outputs=status_box)
+
         generate_btn.click(
             fn=predict,
             inputs=[instruction_textbox, input_text, max_new_tokens, num_beams, do_sample, temperature, top_p, num_return_sequences, mode_selector],
             outputs=output,
         )
 
+        # Model load banner shown after interface loads (async)
+        def _post_load():
+            return "✅ Model interface ready. If this is the first run and model wasn't cached, initial generation may still warm up."
+        demo.load(_post_load, inputs=None, outputs=status_box, every=None)
+
     return demo