Update app.py (#11)

- Update app.py (6c0e639f1c6795c66f42b01b8b6e103c71578d0e)

app.py

@@ -7,8 +7,8 @@ from transformers import pipeline
 import spaces
 
 # === Config (override via Space secrets/env vars) ===
-MODEL_ID = os.environ.get("MODEL_ID", "openai/gpt-oss-20b")
-STATIC_PROMPT = """"""
+MODEL_ID = os.environ.get("MODEL_ID", "tlhv/osb-minier")
+STATIC_PROMPT = ""  # optional extra system message
 DEFAULT_MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", 512))
 DEFAULT_TEMPERATURE = float(os.environ.get("TEMPERATURE", 0.7))
 DEFAULT_TOP_P = float(os.environ.get("TOP_P", 0.95))
@@ -16,30 +16,72 @@ DEFAULT_REPETITION_PENALTY = float(os.environ.get("REPETITION_PENALTY", 1.0))
 ZGPU_DURATION = int(os.environ.get("ZGPU_DURATION", 120))  # seconds
 
 _pipe = None  # cached pipeline
-_tok = None  # tokenizer for parsing Harmony format
 
 
+# ----------------------------
+# Helpers (simple & explicit)
+# ----------------------------
+
 def _to_messages(policy: str, user_prompt: str) -> List[Dict[str, str]]:
-    messages = []
+    msgs: List[Dict[str, str]] = []
     if policy.strip():
-        messages.append({"role": "system", "content": policy.strip()})
-    # if STATIC_PROMPT:
-    #     messages.append({"role": "system", "content": STATIC_PROMPT})
-    messages.append({"role": "user", "content": user_prompt})
-    return messages
-
-
-def _parse_harmony_output(last, tokenizer):
-    analysis, content = None, None
-    if isinstance(last, dict) and ("content" in last or "thinking" in last):
-        analysis = last.get("thinking")
-        content = last.get("content")
-    else:
-        parsed = tokenizer.parse_response(last)
-        analysis = parsed.get("thinking")
-        content = parsed.get("content")
-    return analysis, content
-
+        msgs.append({"role": "system", "content": policy.strip()})
+    if STATIC_PROMPT:
+        msgs.append({"role": "system", "content": STATIC_PROMPT})
+    msgs.append({"role": "user", "content": user_prompt})
+    return msgs
+
+
+def _extract_assistant_content(outputs) -> str:
+    """Extract the assistant's content from the known shape:
+    outputs = [
+      {
+        'generated_text': [
+          {'role': 'system', 'content': ...},
+          {'role': 'user', 'content': ...},
+          {'role': 'assistant', 'content': 'analysis...assistantfinal...'}
+        ]
+      }
+    ]
+    Keep this forgiving and minimal.
+    """
+    try:
+        msgs = outputs[0]["generated_text"]
+        for m in reversed(msgs):
+            if isinstance(m, dict) and m.get("role") == "assistant":
+                return m.get("content", "")
+        last = msgs[-1]
+        return last.get("content", "") if isinstance(last, dict) else str(last)
+    except Exception:
+        return str(outputs)
+
+
+def _parse_harmony_output_from_string(s: str) -> Tuple[str, str]:
+    """Split a Harmony-style concatenated string into (analysis, final).
+    Expects markers 'analysis' ... 'assistantfinal'.
+    No heavy parsing — just string finds.
+    """
+    if not isinstance(s, str):
+        s = str(s)
+    final_key = "assistantfinal"
+    j = s.find(final_key)
+    if j != -1:
+        final_text = s[j + len(final_key):].strip()
+        i = s.find("analysis")
+        if i != -1 and i < j:
+            analysis_text = s[i + len("analysis"): j].strip()
+        else:
+            analysis_text = s[:j].strip()
+        return analysis_text, final_text
+    # no explicit final marker
+    if s.startswith("analysis"):
+        return s[len("analysis"):].strip(), ""
+    return "", s.strip()
+
+
+# ----------------------------
+# Inference
+# ----------------------------
 
 @spaces.GPU(duration=ZGPU_DURATION)
 def generate_long_prompt(
@@ -50,7 +92,7 @@ def generate_long_prompt(
     top_p: float,
     repetition_penalty: float,
 ) -> Tuple[str, str, str]:
-    global _pipe, _tok
+    global _pipe
     start = time.time()
 
     if _pipe is None:
@@ -60,7 +102,6 @@ def generate_long_prompt(
             torch_dtype="auto",
             device_map="auto",
         )
-        _tok = _pipe.tokenizer
 
     messages = _to_messages(policy, prompt)
 
@@ -72,42 +113,54 @@ def generate_long_prompt(
         top_p=top_p,
         repetition_penalty=repetition_penalty,
     )
-    print(outputs)
-    res = outputs[0]
-    print(res)
-    last = res.get("generated_text", [])
-    print(last)
-    if isinstance(last, list) and last:
-        last = last[-1]
 
-    analysis, content = _parse_harmony_output(last, _tok)
+    assistant_str = _extract_assistant_content(outputs)
+    analysis_text, final_text = _parse_harmony_output_from_string(assistant_str)
 
     elapsed = time.time() - start
     meta = f"Model: {MODEL_ID} | Time: {elapsed:.1f}s | max_new_tokens={max_new_tokens}"
-    return analysis or "(No analysis)", content or "(No answer)", meta
+    return analysis_text or "(No analysis)", final_text or "(No answer)", meta
+
 
+# ----------------------------
+# UI
+# ----------------------------
 
-CUSTOM_CSS = "/** Simple styling **/\n.gradio-container {font-family: ui-sans-serif, system-ui, Inter, Roboto;}\ntextarea {font-family: ui-monospace, monospace;}\nfooter {display:none;}"
+CUSTOM_CSS = "/** Pretty but simple **/\n:root { --radius: 14px; }\n.gradio-container { font-family: ui-sans-serif, system-ui, Inter, Roboto, Arial; }\n#hdr h1 { font-weight: 700; letter-spacing: -0.02em; }\ntextarea { font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, 'Liberation Mono', 'Courier New', monospace; }\nfooter { display:none; }\n"
 
 with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
-    gr.Markdown("""# GPT‑OSS Harmony Demo\nProvide a **Policy**, a **Prompt**, and see both **Analysis** and **Answer** separately.""")
+    with gr.Column(elem_id="hdr"):
+        gr.Markdown("""
+        # GPT‑OSS (ZeroGPU) — Harmony View
+        Provide a **Policy**, a **Prompt**, and see **Analysis** and **Answer** separately.
+        """)
 
     with gr.Row():
-        with gr.Column(scale=1):
-            policy = gr.Textbox(label="Policy (system)", lines=20, placeholder="Enter the guiding rules and tone…")
-            prompt = gr.Textbox(label="Prompt (user)", lines=10, placeholder="Enter your main prompt…")
+        with gr.Column(scale=1, min_width=380):
+            policy = gr.Textbox(
+                label="Policy (system)",
+                lines=20,  # bigger than prompt
+                placeholder="Rules, tone, and constraints…",
+            )
+            prompt = gr.Textbox(
+                label="Prompt (user)",
+                lines=10,
+                placeholder="Your request…",
+            )
             with gr.Accordion("Advanced settings", open=False):
                 max_new_tokens = gr.Slider(16, 4096, value=DEFAULT_MAX_NEW_TOKENS, step=8, label="max_new_tokens")
                 temperature = gr.Slider(0.0, 1.5, value=DEFAULT_TEMPERATURE, step=0.05, label="temperature")
                 top_p = gr.Slider(0.0, 1.0, value=DEFAULT_TOP_P, step=0.01, label="top_p")
                 repetition_penalty = gr.Slider(0.8, 2.0, value=DEFAULT_REPETITION_PENALTY, step=0.05, label="repetition_penalty")
-            generate = gr.Button("Generate", variant="primary")
-        with gr.Column(scale=1):
-            analysis = gr.Textbox(label="Analysis (Harmony thinking)", lines=10)
-            answer = gr.Textbox(label="Answer", lines=10)
+            with gr.Row():
+                btn = gr.Button("Generate", variant="primary")
+                clr = gr.Button("Clear", variant="secondary")
+        with gr.Column(scale=1, min_width=380):
+            analysis = gr.Textbox(label="Analysis", lines=12)
+            answer = gr.Textbox(label="Answer", lines=12)
             meta = gr.Markdown()
 
-    generate.click(
+    btn.click(
         fn=generate_long_prompt,
         inputs=[policy, prompt, max_new_tokens, temperature, top_p, repetition_penalty],
         outputs=[analysis, answer, meta],
@@ -115,5 +168,21 @@ with gr.Blocks(css=CUSTOM_CSS, theme=gr.themes.Soft()) as demo:
         api_name="generate",
     )
 
+    def _clear():
+        return "", "", "", ""
+
+    clr.click(_clear, outputs=[policy, prompt, analysis, answer])
+
+    gr.Examples(
+        examples=[
+            ["Be concise and refuse unsafe requests.", "Explain transformers in 2 lines."],
+            [
+                "Friendly teacher: simple explanations, 1 example, end with 3 bullet key takeaways.",
+                "What is attention, briefly?",
+            ],
+        ],
+        inputs=[policy, prompt],
+    )
+
 if __name__ == "__main__":
     demo.queue(max_size=32).launch()