Upload app.py with huggingface_hub

app.py

@@ -19,8 +19,8 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 
 MODEL_ID = "XHToken/Spark-X2.5-4B"
 MAX_CONTEXT_TOKENS = 32_768
-MIN_NEW_TOKENS = 128
-MAX_NEW_TOKENS = 3072
+MIN_NEW_TOKENS = 256
+MAX_NEW_TOKENS = 4096
 
 # The model ships a custom `spark2_5` architecture via `trust_remote_code`.
 # Its custom attention path currently requires the eager implementation.
@@ -50,18 +50,19 @@ TASK_PRESETS = {
 }
 
 
-def split_reasoning(text: str) -> tuple[str, str]:
+def split_reasoning(text: str, enable_thinking: bool) -> tuple[str, str]:
     """Split a partial/complete generation into (reasoning, answer).
 
-    Handles three states: still thinking (no closing tag yet), finished
-    thinking, and thinking disabled (template emits `</think>` immediately).
+    With thinking enabled the chat template appends `<think>` to the prompt, so
+    the generated text is pure reasoning until it emits `</think>` and then the
+    answer. With thinking disabled the model answers directly.
     """
-    if "<think>" in text and "</think>" not in text:
-        return text.split("<think>", 1)[1].strip(), ""
+    if not enable_thinking:
+        return "", text.strip()
     if "</think>" in text:
         reasoning, answer = text.split("</think>", 1)
-        return reasoning.replace("<think>", "").strip(), answer.strip()
-    return "", text.replace("<think>", "").strip()
+        return reasoning.strip(), answer.strip()
+    return text.strip(), ""
 
 
 def history_to_messages(history: list[object] | None) -> list[dict[str, str]]:
@@ -91,7 +92,7 @@ def _estimate_duration(
     history=None,
     system_prompt=None,
     enable_thinking=None,
-    max_new_tokens=1024,
+    max_new_tokens=2048,
     temperature=None,
     top_p=None,
     *args,
@@ -173,8 +174,8 @@ def respond(
     last_emit = 0.0
     for new_text in streamer:
         completion += new_text
-        reasoning, answer = split_reasoning(completion)
-        display = answer if answer else ("_Thinking…_" if reasoning else "")
+        reasoning, answer = split_reasoning(completion, bool(enable_thinking))
+        display = answer or ("_Thinking…_" if reasoning else "")
         now = time.perf_counter()
         if now - last_emit >= 0.1:
             last_emit = now
@@ -183,9 +184,20 @@ def respond(
             )
     thread.join()
 
-    reasoning, answer = split_reasoning(completion)
-    final = base + [{"role": "assistant", "content": answer or reasoning}]
-    yield final, reasoning if enable_thinking else ""
+    reasoning, answer = split_reasoning(completion, bool(enable_thinking))
+    if answer:
+        final_text = answer
+    elif reasoning:
+        final_text = (
+            reasoning
+            + "\n\n> ⚠️ The model stopped before finishing its reasoning — "
+            "increase **Max new tokens** and try again."
+        )
+    else:
+        final_text = ""
+    yield base + [{"role": "assistant", "content": final_text}], (
+        reasoning if enable_thinking else ""
+    )
 
 
 with gr.Blocks(title="Spark-X2.5-4B Code Assistant") as demo:
@@ -224,7 +236,7 @@ with gr.Blocks(title="Spark-X2.5-4B Code Assistant") as demo:
         )
         with gr.Row():
             max_new_tokens = gr.Slider(
-                MIN_NEW_TOKENS, MAX_NEW_TOKENS, value=1024, step=128, label="Max new tokens"
+                MIN_NEW_TOKENS, MAX_NEW_TOKENS, value=2048, step=256, label="Max new tokens"
             )
             temperature = gr.Slider(0, 1.5, value=1.0, step=0.05, label="Temperature")
             top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")