Update app.py

app.py

@@ -21,6 +21,63 @@ model = AutoModelForCausalLM.from_pretrained(
 )
 model.eval()
 
+# ─────────────────────────────────────────────
+# 1b. THINKING ENFORCEMENT
+#
+#  This model isn't guaranteed to emit <think>...</think> on its own
+#  (especially as an abliterated finetune). To make it 100% reliable:
+#
+#    1. Try passing enable_thinking=True to apply_chat_template
+#       (native support on Qwen3-style templates, if present).
+#    2. ALWAYS force-append "<think>\n" onto the templated prompt
+#       so generation is physically forced to begin inside a think
+#       block, regardless of whether enable_thinking worked.
+#    3. Append a short mandatory instruction onto whatever system
+#       prompt is passed in, telling the model to close the tag.
+#    4. Manually re-prepend "<think>\n" onto the decoded output,
+#       since it was part of the forced prompt and gets stripped
+#       out by skip_prompt / prompt-slicing.
+# ─────────────────────────────────────────────
+THINK_INSTRUCTION = (
+    "\n\nAlways reason through the problem step by step inside <think> "
+    "and </think> tags first. After the closing </think> tag, give your "
+    "final answer. Never skip the opening or closing think tags."
+)
+
+
+def build_prompt(prompt: str, system_prompt: str) -> str:
+    full_system = (system_prompt or "You are a helpful assistant.") + THINK_INSTRUCTION
+
+    messages = [
+        {"role": "system", "content": full_system},
+        {"role": "user", "content": prompt},
+    ]
+
+    try:
+        # Native Qwen3-style thinking toggle, if the tokenizer supports it
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=True,
+        )
+    except TypeError:
+        # Tokenizer/template doesn't accept enable_thinking — fall back
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+        )
+
+    # Force the assistant turn to start inside a <think> block,
+    # guaranteeing the tag appears no matter what the model would
+    # have done on its own.
+    if not text.rstrip().endswith("<think>"):
+        text = text + "<think>\n"
+
+    return text
+
+
 # ─────────────────────────────────────────────
 # 2. INFERENCE FUNCTION (ZeroGPU decorated)
 # ─────────────────────────────────────────────
@@ -34,32 +91,10 @@ def generate(
     do_sample: bool = True,
 ) -> str:
     """
-    Generate a text response from Mira-1-Large.
-
-    Args:
-        prompt: The user message / prompt to send to the model.
-        system_prompt: System-level instruction for the model.
-        max_new_tokens: Maximum number of tokens to generate.
-        temperature: Sampling temperature (higher = more creative).
-        top_p: Nucleus sampling probability mass.
-        do_sample: Whether to use sampling (True) or greedy decoding (False).
-
-    Returns:
-        The model's text response as a string.
+    Generate a text response from the model.
+    Response is guaranteed to start with a <think> block.
     """
-    # Build chat-style messages (Qwen uses apply_chat_template)
-    messages = [
-        {"role": "system", "content": system_prompt},
-        {"role": "user",   "content": prompt},
-    ]
-
-    # Qwen / Mira chat template
-    text = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True,
-    )
-
+    text = build_prompt(prompt, system_prompt)
     inputs = tokenizer(text, return_tensors="pt").to(model.device)
 
     with torch.no_grad():
@@ -72,9 +107,12 @@ def generate(
             pad_token_id=tokenizer.eos_token_id,
         )
 
-    # Decode only the newly generated tokens
     new_tokens = output_ids[0][inputs["input_ids"].shape[1]:]
     response = tokenizer.decode(new_tokens, skip_special_tokens=True)
+
+    # Re-add the <think> tag we force-injected into the prompt —
+    # it was stripped out because it's technically part of the input.
+    response = "<think>\n" + response
     return response
 
 
@@ -88,30 +126,15 @@ def generate_stream(
     max_new_tokens: int = 512,
     temperature: float = 0.7,
     top_p: float = 0.9,
-) -> str:
+):
     """
-    Stream a text response token-by-token from Mira-1-Large via SSE.
-
-    Args:
-        prompt: The user message / prompt.
-        system_prompt: System-level instruction for the model.
-        max_new_tokens: Maximum number of tokens to generate.
-        temperature: Sampling temperature.
-        top_p: Nucleus sampling probability mass.
-
-    Yields:
-        Partial response strings, growing with each new token.
+    Stream a text response token-by-token via SSE.
+    Guaranteed to start with a <think> block.
     """
     from transformers import TextIteratorStreamer
     from threading import Thread
 
-    messages = [
-        {"role": "system", "content": system_prompt},
-        {"role": "user",   "content": prompt},
-    ]
-    text = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
+    text = build_prompt(prompt, system_prompt)
     inputs = tokenizer(text, return_tensors="pt").to(model.device)
 
     streamer = TextIteratorStreamer(
@@ -131,7 +154,12 @@ def generate_stream(
     thread = Thread(target=model.generate, kwargs=gen_kwargs)
     thread.start()
 
-    partial = ""
+    # Re-inject the forced <think> prefix as the very first chunk,
+    # since it was part of the prompt and streamer.skip_prompt=True
+    # will not emit it on its own.
+    partial = "<think>\n"
+    yield partial
+
     for new_text in streamer:
         partial += new_text
         yield partial
@@ -146,15 +174,15 @@ app = Server(
     version="1.0.0",
 )
 
-# Register as Gradio API endpoints (queued, SSE-streaming capable)
-app.api(generate,        name="generate")        # POST /gradio_api/call/generate
-app.api(generate_stream, name="generate_stream") # POST /gradio_api/call/generate_stream
+app.api(generate,        name="generate")
+app.api(generate_stream, name="generate_stream")
+
 
-# Optional: plain FastAPI GET health-check route
 @app.get("/health")
 def health():
     return {"status": "ok", "model": MODEL_ID}
 
+
 # ─────────────────────────────────────────────
 # 5. LAUNCH
 # ─────────────────────────────────────────────