Upload app.py with huggingface_hub

app.py

@@ -4,9 +4,8 @@ Nawah-Reasoning-v1 — Gradio demo.
 The prompt rendering (ChatML + BOS prepend) is IDENTICAL to train_reasoning.py. A 51M model
 is very sensitive to format drift, so do not change render_prompt() without changing training.
 
-Runs on CPU: the model is ~52M parameters, so a full answer takes a couple of seconds without
-a GPU. That is the point of the model, so the Space is deliberately CPU-only — no `spaces`
-import, no ZeroGPU quota.
+Runs on ZeroGPU. The model is only ~52M parameters and works on CPU too, but ZeroGPU keeps
+responses snappy. `import spaces` must come BEFORE torch so it can patch the CUDA calls.
 
 The model emits <think>…</think> before its answer, so the stream is split live into two
 panels: the reasoning trace and the final answer.
@@ -19,6 +18,7 @@ import os
 import re
 import threading
 
+import spaces  # import BEFORE torch so it can patch CUDA calls
 import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
@@ -32,18 +32,16 @@ HF_TOKEN = os.environ.get("MODEL_HF_TOKEN") or os.environ.get("HF_TOKEN")
 
 IM_START, IM_END = "<|im_start|>", "<|im_end|>"
 THINK_OPEN, THINK_CLOSE = "<think>", "</think>"
-MAX_NEW_TOKENS_CAP = 1500
+MAX_NEW_TOKENS_CAP = 512
 
 # ── Load (once, at startup) ───────────────────────────────────────────────────
 
-torch.set_num_threads(os.cpu_count() or 2)
-
 print("[*] token env vars present:",
       [k for k in ("MODEL_HF_TOKEN", "HF_TOKEN") if os.environ.get(k)] or "NONE")
 print(f"[*] Loading {MODEL_ID} ...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.float32, token=HF_TOKEN)
-model.eval()
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16, token=HF_TOKEN)
+model.to("cuda").eval()
 
 CTX = getattr(model.config, "max_position_embeddings", 2048)
 
@@ -78,6 +76,7 @@ def split_stream(text: str):
 
 # ── Generate ──────────────────────────────────────────────────────────────────
 
+@spaces.GPU(duration=60)
 def solve(question, max_new_tokens, temperature, repetition_penalty):
     question = (question or "").strip()
     if not question:
@@ -87,7 +86,7 @@ def solve(question, max_new_tokens, temperature, repetition_penalty):
     ids = tokenizer(render_prompt(question), add_special_tokens=False)["input_ids"]
     if tokenizer.bos_token_id is not None:
         ids = [tokenizer.bos_token_id] + ids  # match training's explicit BOS
-    input_ids = torch.tensor([ids])
+    input_ids = torch.tensor([ids], device=model.device)
 
     # skip_special_tokens must stay False — <think>/</think> are real special tokens
     # in this tokenizer, and stripping them would destroy the split.
@@ -131,8 +130,8 @@ DESCRIPTION = """
 ثم يعطي الإجابة النهائية.<br>
 A ~52M-parameter Arabic reasoning model that thinks step by step inside
 <code>&lt;think&gt;</code> before answering.</p>
-<p><i>يعمل على المعالج (CPU) بدون كرت رسومات — هذا هو بيت القصيد.<br>
-Runs on CPU, no GPU needed — that is the whole point.</i></p>
+<p><i>نموذج صغير بما يكفي ليعمل حتى على المعالج (CPU).<br>
+Small enough to run on a CPU — this Space uses ZeroGPU for snappier responses.</i></p>
 </div>
 """