Upload folder using huggingface_hub

README.md

@@ -1,20 +1,23 @@
 ---
-title: Arithmetic-SLM Playground
+title: Arithmetic-SLM
 emoji: 🧮
 colorFrom: green
 colorTo: pink
 sdk: gradio
 sdk_version: 6.19.0
 app_file: app.py
-short_description: Tiny 31.7M-param arithmetic model playground
+short_description: Tiny 31.7M SLM that solves arithmetic expressions
 python_version: "3.12"
 startup_duration_timeout: 30m
 ---
 
-# Arithmetic-SLM Playground
+# Arithmetic-SLM
 
-A Gradio demo for [WhirlwindAI/Arithmetic-SLM](https://huggingface.co/WhirlwindAI/Arithmetic-SLM),
-a tiny 31.7M-parameter language model specialized for arithmetic continuation.
+An interactive demo for [WhirlwindAI/Arithmetic-SLM](https://huggingface.co/WhirlwindAI/Arithmetic-SLM),
+a tiny (31.7M parameter) specialized language model that completes arithmetic
+expressions token by token. It handles operator precedence, parentheses, and
+decimals.
 
-Enter an arithmetic expression ending with `=` (e.g. `(10 + 28) * 3 =`) and the model
-completes it with the answer. Supports `+`, `-`, `*`, `/`, parentheses, and decimals.
+Inference is ported 1:1 from the model repo's `inference.py` reference
+implementation (custom manual sampling loop with the pure-torch attention
+backend), running on ZeroGPU.

app.py

@@ -1,58 +1,44 @@
-import spaces  # MUST come before torch / any CUDA-touching import
-import os
-import sys
 import random
-from typing import List, Optional, Set, Dict
-
-# Add custom_model to path before importing torch/transformers
-sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from typing import Dict, List, Set
 
+import spaces
 import torch
-import torch.nn as nn
 import torch.nn.functional as F
 import gradio as gr
-from transformers import AutoTokenizer
-from custom_model.modeling_tiny_gpt import TinyGPTForCausalLM
-from custom_model.configuration_tiny_gpt import TinyGPTConfig
-from huggingface_hub import hf_hub_download
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 MODEL_ID = "WhirlwindAI/Arithmetic-SLM"
 
 IM_START = "[IM_START]"
 IM_END = "[IM_END]"
 NO_THINK = "/no think"
+CTX_LEN = 2048
 
-# ============================================================
-# Model + tokenizer loaded at module scope
-# ============================================================
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+STOP_STRINGS = [IM_END, IM_START]
 
-# Load config from the hub, then instantiate our patched custom model
-# We use trust_remote_code=True for the config (to download the custom
-# config class) but load weights manually to avoid the remote custom
-# modeling code shadowing our patched local version.
-config = TinyGPTConfig.from_pretrained(MODEL_ID, trust_remote_code=True)
-
-# Download the safetensors weights
-from safetensors.torch import load_file
-weights_path = hf_hub_download(MODEL_ID, "model.safetensors")
-state_dict = load_file(weights_path)
-
-# Instantiate the model from config and load weights
-model = TinyGPTForCausalLM(config)
-model.load_state_dict(state_dict, strict=False)
-model = model.to(dtype=torch.bfloat16).to("cuda")
+# ---------------------------------------------------------------------------
+# Load model + tokenizer once, at module scope, moved eagerly to CUDA so
+# ZeroGPU can pack the weights and stream them into VRAM on the first call.
+# The model uses custom code (TinyGPTForCausalLM) with the pure-torch attention
+# backend (config: attention_backend="torch", torch_fallback=True) so no flash
+# kernels are needed at runtime.
+# ---------------------------------------------------------------------------
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    dtype=torch.bfloat16,
+    trust_remote_code=True,
+).to("cuda")
 model.eval()
 
-# ============================================================
-# Sampling utilities (ported from the model's inference.py)
-# ============================================================
 
-def apply_repetition_penalty(logits: torch.Tensor, generated_ids: List[int], penalty: float) -> torch.Tensor:
+# ---------------------------------------------------------------------------
+# Sampling helpers — ported 1:1 from the model repo's inference.py so output
+# matches the authors' reference path exactly.
+# ---------------------------------------------------------------------------
+def apply_repetition_penalty(logits, generated_ids, penalty):
     if penalty is None or penalty == 1.0:
         return logits
-    if penalty <= 0:
-        raise ValueError("repetition_penalty must be > 0")
     for tid in set(generated_ids):
         if tid < 0 or tid >= logits.numel():
             continue
@@ -63,12 +49,7 @@ def apply_repetition_penalty(logits: torch.Tensor, generated_ids: List[int], pen
     return logits
 
 
-def apply_frequency_presence_penalty(
-    logits: torch.Tensor,
-    generated_ids: List[int],
-    frequency_penalty: float,
-    presence_penalty: float,
-) -> torch.Tensor:
+def apply_frequency_presence_penalty(logits, generated_ids, frequency_penalty, presence_penalty):
     if not generated_ids:
         return logits
     if frequency_penalty == 0.0 and presence_penalty == 0.0:
@@ -86,7 +67,7 @@ def apply_frequency_presence_penalty(
     return logits
 
 
-def get_banned_ngram_tokens(generated_ids: List[int], no_repeat_ngram_size: int) -> Set[int]:
+def get_banned_ngram_tokens(generated_ids, no_repeat_ngram_size) -> Set[int]:
     n = no_repeat_ngram_size
     banned: Set[int] = set()
     if n <= 0:
@@ -97,16 +78,14 @@ def get_banned_ngram_tokens(generated_ids: List[int], no_repeat_ngram_size: int)
     current_prefix = tuple(generated_ids[-prefix_len:])
     ngram_map: Dict[tuple, Set[int]] = {}
     for i in range(len(generated_ids) - n + 1):
-        prefix = tuple(generated_ids[i : i + prefix_len])
+        prefix = tuple(generated_ids[i:i + prefix_len])
         next_token = generated_ids[i + prefix_len]
-        if prefix not in ngram_map:
-            ngram_map[prefix] = set()
-        ngram_map[prefix].add(next_token)
+        ngram_map.setdefault(prefix, set()).add(next_token)
     banned.update(ngram_map.get(current_prefix, set()))
     return banned
 
 
-def apply_no_repeat_ngram(logits: torch.Tensor, generated_ids: List[int], no_repeat_ngram_size: int) -> torch.Tensor:
+def apply_no_repeat_ngram(logits, generated_ids, no_repeat_ngram_size):
     if no_repeat_ngram_size <= 0:
         return logits
     banned = get_banned_ngram_tokens(generated_ids, no_repeat_ngram_size)
@@ -116,7 +95,7 @@ def apply_no_repeat_ngram(logits: torch.Tensor, generated_ids: List[int], no_rep
     return logits
 
 
-def apply_top_k(logits: torch.Tensor, top_k: int) -> torch.Tensor:
+def apply_top_k(logits, top_k):
     if top_k is None or top_k <= 0:
         return logits
     top_k = min(top_k, logits.size(-1))
@@ -126,11 +105,11 @@ def apply_top_k(logits: torch.Tensor, top_k: int) -> torch.Tensor:
     return logits
 
 
-def apply_top_p(logits: torch.Tensor, top_p: float) -> torch.Tensor:
+def apply_top_p(logits, top_p):
     if top_p is None or top_p >= 1.0:
         return logits
     if top_p <= 0:
-        raise ValueError("top_p must be > 0")
+        return logits
     sorted_logits, sorted_indices = torch.sort(logits, descending=True)
     sorted_probs = F.softmax(sorted_logits, dim=-1)
     cumulative = torch.cumsum(sorted_probs, dim=-1)
@@ -142,131 +121,54 @@ def apply_top_p(logits: torch.Tensor, top_p: float) -> torch.Tensor:
     return logits
 
 
-def build_stop_sequences(stop_strings: List[str]) -> List[List[int]]:
-    out: List[List[int]] = []
-    for s in stop_strings:
-        ids = tokenizer.encode(s, add_special_tokens=False)
-        if ids:
-            out.append(ids)
-    return out
-
-
-def endswith_sequence(ids: List[int], suffix: List[int]) -> bool:
-    if not suffix:
-        return False
-    if len(ids) < len(suffix):
-        return False
-    return ids[-len(suffix) :] == suffix
-
-
-def sample_next_token(
-    logits: torch.Tensor,
-    generated_ids: List[int],
-    temperature: float,
-    top_k: int,
-    top_p: float,
-    repetition_penalty: float,
-    frequency_penalty: float,
-    presence_penalty: float,
-    no_repeat_ngram_size: int,
-) -> int:
+def sample_next_token(logits, generated_ids, temperature, top_k, top_p,
+                      repetition_penalty, frequency_penalty, no_repeat_ngram_size):
     logits = logits.float().clone()
     logits = apply_repetition_penalty(logits, generated_ids, repetition_penalty)
-    logits = apply_frequency_presence_penalty(logits, generated_ids, frequency_penalty, presence_penalty)
+    logits = apply_frequency_presence_penalty(logits, generated_ids, frequency_penalty, 0.0)
     logits = apply_no_repeat_ngram(logits, generated_ids, no_repeat_ngram_size)
+
     if temperature <= 0:
         return int(torch.argmax(logits).item())
+
     logits = logits / temperature
     logits = apply_top_k(logits, top_k)
     logits = apply_top_p(logits, top_p)
+
     probs = F.softmax(logits, dim=-1)
     if torch.isnan(probs).any() or torch.isinf(probs).any() or probs.sum() <= 0:
         return int(torch.argmax(logits).item())
     return int(torch.multinomial(probs, num_samples=1).item())
 
 
-def model_forward_logits(input_ids: torch.Tensor) -> torch.Tensor:
-    out = model(input_ids=input_ids)
-    if hasattr(out, "logits"):
-        return out.logits
-    if isinstance(out, tuple):
-        return out[0]
-    raise RuntimeError("Could not get logits from model output.")
-
-
-@torch.no_grad()
-def generate_manual(
-    input_ids: torch.Tensor,
-    max_new_tokens: int,
-    min_new_tokens: int,
-    temperature: float,
-    top_k: int,
-    top_p: float,
-    repetition_penalty: float,
-    frequency_penalty: float,
-    presence_penalty: float,
-    no_repeat_ngram_size: int,
-    ctx_len: int,
-    stop_strings: List[str],
-) -> List[int]:
-    idx = input_ids
-    generated_after_prompt: List[int] = []
-    stop_sequences = build_stop_sequences(stop_strings)
-    eos_id = tokenizer.eos_token_id
+def build_stop_sequences(stop_strings) -> List[List[int]]:
+    out = []
+    for s in stop_strings:
+        ids = tokenizer.encode(s, add_special_tokens=False)
+        if ids:
+            out.append(ids)
+    return out
 
-    for step in range(max_new_tokens):
-        idx_cond = idx[:, -ctx_len:] if ctx_len > 0 else idx
-        logits = model_forward_logits(idx_cond)
-        logits = logits[:, -1, :][0]
-
-        if step < min_new_tokens:
-            if eos_id is not None and 0 <= eos_id < logits.numel():
-                logits[eos_id] = -float("inf")
-            for seq in stop_sequences:
-                if len(seq) == 1:
-                    tid = seq[0]
-                    if 0 <= tid < logits.numel():
-                        logits[tid] = -float("inf")
-
-        next_id = sample_next_token(
-            logits=logits,
-            generated_ids=generated_after_prompt,
-            temperature=temperature,
-            top_k=top_k,
-            top_p=top_p,
-            repetition_penalty=repetition_penalty,
-            frequency_penalty=frequency_penalty,
-            presence_penalty=presence_penalty,
-            no_repeat_ngram_size=no_repeat_ngram_size,
-        )
 
-        next_tensor = torch.tensor([[next_id]], dtype=torch.long, device=idx.device)
-        idx = torch.cat([idx, next_tensor], dim=1)
-        generated_after_prompt.append(next_id)
-
-        if step >= min_new_tokens:
-            if eos_id is not None and next_id == eos_id:
-                break
-            full_ids = idx[0].tolist()
-            should_stop = False
-            for seq in stop_sequences:
-                if endswith_sequence(full_ids, seq):
-                    should_stop = True
-                    break
-            if should_stop:
-                break
+def endswith_sequence(ids, suffix) -> bool:
+    if not suffix or len(ids) < len(suffix):
+        return False
+    return ids[-len(suffix):] == suffix
 
-    return idx[0].tolist()
 
+def strip_after_stop_text(text, stop_strings) -> str:
+    best = None
+    for s in stop_strings:
+        if not s:
+            continue
+        pos = text.find(s)
+        if pos != -1 and (best is None or pos < best):
+            best = pos
+    return text if best is None else text[:best]
 
-def build_prompt(expression: str, use_no_think: bool) -> str:
-    """Build the prompt for the model.
 
-    Args:
-        expression: arithmetic expression (e.g. '(10 + 28) * 3 =').
-        use_no_think: if True, wrap in the /no think chat template for production use.
-    """
-    if use_no_think:
+def build_prompt(expression: str, use_think_format: bool) -> str:
+    if use_think_format:
         return (
             f"{IM_START}user\n"
             f"{expression} {NO_THINK}"
@@ -277,228 +179,165 @@ def build_prompt(expression: str, use_no_think: bool) -> str:
     return expression
 
 
-def extract_completion(full_text: str, prompt: str) -> str:
-    """Extract the completion (model output) from the full generated text."""
-    if full_text.startswith(prompt):
-        return full_text[len(prompt):]
-    pos = full_text.rfind(prompt)
-    if pos != -1:
-        return full_text[pos + len(prompt):]
-    return full_text
-
-
-def strip_after_stop_text(text: str, stop_strings: List[str]) -> str:
-    """Strip everything after the first stop string."""
-    best = None
-    for s in stop_strings:
-        if not s:
-            continue
-        pos = text.find(s)
-        if pos != -1:
-            if best is None or pos < best:
-                best = pos
-    if best is None:
-        return text
-    return text[:best]
-
-
-# ============================================================
-# Gradio inference function
-# ============================================================
-
 @spaces.GPU(duration=30)
 def solve(
     expression: str,
-    use_no_think: bool,
-    max_new_tokens: int,
-    temperature: float,
-    top_k: int,
-    top_p: float,
-    repetition_penalty: float,
-    frequency_penalty: float,
-    no_repeat_ngram_size: int,
-    seed: int,
-):
-    """Solve an arithmetic expression using the Arithmetic-SLM model.
+    use_think_format: bool = False,
+    temperature: float = 0.5,
+    top_k: int = 40,
+    top_p: float = 0.95,
+    max_new_tokens: int = 48,
+    seed: int = -1,
+) -> str:
+    """Solve an arithmetic expression with the Arithmetic-SLM model.
 
     Args:
-        expression: an arithmetic expression ending with '=' (e.g. '59 + 45 =').
-        use_no_think: use the /no think chat template (recommended for production).
-        max_new_tokens: maximum number of tokens to generate.
-        temperature: sampling temperature (lower = more deterministic).
-        top_k: top-k sampling filter.
-        top_p: nucleus sampling filter.
-        repetition_penalty: penalty for repeated tokens.
-        frequency_penalty: penalty based on token frequency.
-        no_repeat_ngram_size: ban repeating n-grams of this size.
-        seed: random seed (-1 for random).
+        expression: An arithmetic expression ending in '=', e.g. '(10 + 28) * 3 ='.
+        use_think_format: Use the production [IM_START]/[IM_END] chat template with a '/no think' tag.
+        temperature: Sampling temperature (lower = more deterministic).
+        top_k: Top-k sampling cutoff.
+        top_p: Nucleus (top-p) sampling cutoff.
+        max_new_tokens: Maximum number of tokens to generate.
+        seed: RNG seed; -1 for random.
+
+    Returns:
+        The model's completion of the expression (typically the solved result).
     """
-    if not expression or not expression.strip():
-        return "Please enter an arithmetic expression ending with '='.", ""
+    expression = (expression or "").strip()
+    if not expression:
+        return "Please enter an arithmetic expression, e.g. '(10 + 28) * 3 ='."
 
-    # Set seed
-    if seed is not None and seed >= 0:
-        random.seed(seed)
-        torch.manual_seed(seed)
-        if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(seed)
-    else:
-        seed = random.randint(0, 2**31 - 1)
-        random.seed(seed)
-        torch.manual_seed(seed)
+    if seed is not None and int(seed) >= 0:
+        random.seed(int(seed))
+        torch.manual_seed(int(seed))
         if torch.cuda.is_available():
-            torch.cuda.manual_seed_all(seed)
+            torch.cuda.manual_seed_all(int(seed))
 
-    prompt = build_prompt(expression.strip(), use_no_think)
-    stop_strings = [IM_END, IM_START]
+    repetition_penalty = 1.05
+    frequency_penalty = 0.10
+    no_repeat_ngram_size = 4
+    min_new_tokens = 1
+
+    prompt = build_prompt(expression, use_think_format)
 
     encoded = tokenizer(prompt, return_tensors="pt", add_special_tokens=False)
     encoded.pop("token_type_ids", None)
-    input_ids = encoded["input_ids"].to("cuda")
-
-    output_ids = generate_manual(
-        input_ids=input_ids,
-        max_new_tokens=max_new_tokens,
-        min_new_tokens=1,
-        temperature=temperature,
-        top_k=top_k,
-        top_p=top_p,
-        repetition_penalty=repetition_penalty,
-        frequency_penalty=frequency_penalty,
-        presence_penalty=0.0,
-        no_repeat_ngram_size=no_repeat_ngram_size,
-        ctx_len=2048,
-        stop_strings=stop_strings,
-    )
-
-    full_text = tokenizer.decode(output_ids, skip_special_tokens=False)
-    completion = extract_completion(full_text, prompt)
-    completion = strip_after_stop_text(completion, stop_strings)
-    completion = completion.strip()
-
-    # The model often repeats the input expression in its output.
-    # If the completion starts with the expression, extract just the answer.
-    expr_stripped = expression.strip()
-    if completion.startswith(expr_stripped):
-        answer = completion[len(expr_stripped):].strip()
-    else:
-        answer = completion
-
-    # Build the result display
-    full_display = f"{expr_stripped} {answer}".strip()
-
-    if use_no_think:
-        display = f"**Prompt (no-think template):**\n```\n{prompt}\n```\n\n**Model output:**\n```\n{completion}\n```\n\n**Answer:** `{answer}`"
-    else:
-        display = f"**Prompt:** `{prompt}`\n\n**Model output:** `{completion}`\n\n**Answer:** `{answer}`"
-
-    return full_display, display
-
-
-# ============================================================
-# Gradio UI
-# ============================================================
+    idx = encoded["input_ids"].to("cuda")
+
+    stop_sequences = build_stop_sequences(STOP_STRINGS)
+    eos_id = tokenizer.eos_token_id
+    generated: List[int] = []
+
+    with torch.no_grad():
+        for step in range(int(max_new_tokens)):
+            idx_cond = idx[:, -CTX_LEN:]
+            out = model(input_ids=idx_cond)
+            logits = out.logits[:, -1, :][0]
+
+            if step < min_new_tokens:
+                if eos_id is not None and 0 <= eos_id < logits.numel():
+                    logits[eos_id] = -float("inf")
+                for seq in stop_sequences:
+                    if len(seq) == 1 and 0 <= seq[0] < logits.numel():
+                        logits[seq[0]] = -float("inf")
+
+            next_id = sample_next_token(
+                logits, generated, float(temperature), int(top_k), float(top_p),
+                repetition_penalty, frequency_penalty, no_repeat_ngram_size,
+            )
+            idx = torch.cat(
+                [idx, torch.tensor([[next_id]], dtype=torch.long, device=idx.device)], dim=1
+            )
+            generated.append(next_id)
+
+            if step >= min_new_tokens:
+                if eos_id is not None and next_id == eos_id:
+                    break
+                full_ids = idx[0].tolist()
+                if any(endswith_sequence(full_ids, seq) for seq in stop_sequences):
+                    break
+
+    full_text = tokenizer.decode(idx[0].tolist(), skip_special_tokens=False)
+
+    if use_think_format:
+        # Show the completion after the prompt, cleaned of control markers.
+        if full_text.startswith(prompt):
+            completion = full_text[len(prompt):]
+        else:
+            pos = full_text.rfind(prompt)
+            completion = full_text[pos + len(prompt):] if pos != -1 else full_text
+        completion = strip_after_stop_text(completion, STOP_STRINGS)
+        return completion.strip()
+
+    # Raw mode: return the full continued expression.
+    completion = strip_after_stop_text(full_text, STOP_STRINGS)
+    return completion.strip()
+
 
 CSS = """
-#col-container { max-width: 900px; margin: 0 auto; }
+#col-container { max-width: 820px; margin: 0 auto; }
 .dark .gradio-container { color: var(--body-text-color); }
 """
 
-with gr.Blocks() as demo:
-    gr.Markdown(
-        "# 🧮 Arithmetic-SLM Playground\n"
-        "**WhirlwindAI/Arithmetic-SLM** — a tiny 31.7M-parameter model specialized for arithmetic. "
-        "Enter an expression ending with `=` and the model completes it with the answer.\n\n"
-        "Supports `+`, `-`, `*`, `/`, parentheses, and decimals. "
-        "[Model card](https://huggingface.co/WhirlwindAI/Arithmetic-SLM)"
-    )
-
+with gr.Blocks(theme=gr.themes.Citrus(), css=CSS) as demo:
     with gr.Column(elem_id="col-container"):
+        gr.Markdown(
+            """
+            # 🧮 Arithmetic-SLM
+
+            A tiny (31.7M parameter) specialized language model that **completes arithmetic
+            expressions** — it learned to do math token by token, not with a calculator.
+            Handles operator precedence, parentheses, and decimals.
+
+            Enter an expression ending in `=` and let the model finish it.
+
+            [Model card](https://huggingface.co/WhirlwindAI/Arithmetic-SLM)
+            """
+        )
         with gr.Row():
-            expression_input = gr.Textbox(
+            expression = gr.Textbox(
                 label="Arithmetic expression",
-                placeholder="e.g. (10 + 28) * 3 =",
+                placeholder="(10 + 28) * 3 =",
+                value="(10 + 28) * 3 =",
                 scale=4,
-                show_label=False,
             )
-            solve_btn = gr.Button("Solve", variant="primary", scale=1)
+            run = gr.Button("Solve", variant="primary", scale=1)
 
-        output_text = gr.Textbox(label="Result", interactive=False)
-        output_detail = gr.Textbox(label="Details", interactive=False, visible=True)
+        output = gr.Textbox(label="Model output", lines=3)
 
         with gr.Accordion("Advanced settings", open=False):
-            use_no_think = gr.Checkbox(
-                label="Use /no think template (recommended)",
-                value=True,
-                info="Wraps the expression in the model's production chat template for cleaner outputs.",
+            use_think_format = gr.Checkbox(
+                label="Use production /no think chat template",
+                value=False,
+                info="Wraps the input in the [IM_START]/[IM_END] template with a <think> block.",
             )
-            with gr.Row():
-                max_new_tokens = gr.Slider(1, 128, value=48, step=1, label="Max new tokens")
-                temperature = gr.Slider(0.0, 2.0, value=0.6, step=0.05, label="Temperature")
-            with gr.Row():
-                top_k = gr.Slider(0, 100, value=50, step=1, label="Top-k")
-                top_p = gr.Slider(0.01, 1.0, value=0.97, step=0.01, label="Top-p")
-            with gr.Row():
-                repetition_penalty = gr.Slider(0.1, 2.0, value=1.0, step=0.05, label="Repetition penalty")
-                frequency_penalty = gr.Slider(0.0, 2.0, value=0.0, step=0.05, label="Frequency penalty")
-            with gr.Row():
-                no_repeat_ngram_size = gr.Slider(0, 10, value=4, step=1, label="No-repeat n-gram size")
-                seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
+            temperature = gr.Slider(0.0, 1.0, value=0.5, step=0.05, label="Temperature")
+            top_k = gr.Slider(0, 100, value=40, step=1, label="Top-k")
+            top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.01, label="Top-p")
+            max_new_tokens = gr.Slider(8, 128, value=48, step=1, label="Max new tokens")
+            seed = gr.Number(value=-1, precision=0, label="Seed (-1 = random)")
 
         gr.Examples(
             examples=[
-                ["59 + 45 =", True],
-                ["(10 + 28) * 3 =", True],
-                ["16 + 4 * 3 =", True],
-                ["3 * 9 + 12 / 1 =", True],
-                ["(132 / 12) + (46 - 15) =", True],
-                ["0.5 * 0.5 =", True],
-                ["8 * 5 + 4 / 4 =", True],
-                ["(85 - 45) + 56 =", True],
+                ["59 + 45 ="],
+                ["16 + 4 * 3 ="],
+                ["(16 / 4) + 44 ="],
+                ["3 * 9 + 12 / 1 ="],
+                ["(132 / 12) + (46 - 15) ="],
+                ["0.5 * 0.5 ="],
+                ["8 * 5 + 4 / 4 ="],
+                ["(85 - 45) + 56 ="],
             ],
-            inputs=[expression_input, use_no_think],
-            outputs=[output_text, output_detail],
+            inputs=[expression],
+            outputs=output,
             fn=solve,
             cache_examples=True,
             cache_mode="lazy",
         )
 
-    solve_btn.click(
-        fn=solve,
-        inputs=[
-            expression_input,
-            use_no_think,
-            max_new_tokens,
-            temperature,
-            top_k,
-            top_p,
-            repetition_penalty,
-            frequency_penalty,
-            no_repeat_ngram_size,
-            seed,
-        ],
-        outputs=[output_text, output_detail],
-        api_name="solve",
-    )
-
-    expression_input.submit(
-        fn=solve,
-        inputs=[
-            expression_input,
-            use_no_think,
-            max_new_tokens,
-            temperature,
-            top_k,
-            top_p,
-            repetition_penalty,
-            frequency_penalty,
-            no_repeat_ngram_size,
-            seed,
-        ],
-        outputs=[output_text, output_detail],
-        api_name="solve_submit",
-    )
-
+    inputs = [expression, use_think_format, temperature, top_k, top_p, max_new_tokens, seed]
+    run.click(solve, inputs=inputs, outputs=output, api_name="solve")
+    expression.submit(solve, inputs=inputs, outputs=output, api_name=False)
 
 if __name__ == "__main__":
-    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)
+    demo.launch(mcp_server=True)

requirements.txt

@@ -1,2 +1,2 @@
 transformers
-safetensors
+safetensors