Update app.py

app.py

@@ -1,76 +1,58 @@
 # app.py
 import re, spaces, gradio as gr, torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
-
-MODEL_NAME = "loocorez/reverse-text-warmup"
-
-SYSTEM_PROMPT = """Reverse the given text.
-
-Respond in the following format:
-<response>
-  <think>...</think>
-  <answer>...</answer>
-</response>"""
-
-# ── load model & tokenizer once ────────────────────────────────────────────────
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-tokenizer.pad_token = tokenizer.eos_token
-
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.float16
-).eval()                               # keep weights on CPU until each call
-
-EOS_ID = tokenizer.encode("</response>", add_special_tokens=False)[0]
+import verifiers as vf
 
 
-@spaces.GPU(duration=60)               # ZeroGPU reserves a card only inside this fn
-def reverse_text(
-    user_text: str,
-    temperature: float = 0.0,
-    max_tokens: int = 256,
-) -> str:
-    prompt = f"{SYSTEM_PROMPT}\n\n{user_text.strip()}"
-
-    # ── move model to GPU for this request ────────────────────────────────────
-    model.to("cuda")
-
-    try:
-        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-
-        with torch.no_grad():
-            out = model.generate(
-                **inputs,
-                max_new_tokens=max_tokens,
-                do_sample=temperature > 0,
-                temperature=max(temperature, 1e-6),
-                eos_token_id=EOS_ID,
-                pad_token_id=tokenizer.eos_token_id,
-                repetition_penalty=1.1,
-            )
-
-        gen = tokenizer.decode(out[0][inputs["input_ids"].size(1):],
-                               skip_special_tokens=True)
-        m = re.search(r"<answer>(.*?)</answer>", gen, re.S)
-        return m.group(1).strip() if m else gen.strip()
-
-    finally:
-        # ── free the GPU so the next visitor can use it ───────────────────────
-        model.to("cpu")
-        torch.cuda.empty_cache()
+MODEL_NAME = "loocorez/reverse-text-warmup"
 
 
-demo = gr.Interface(
-    fn=reverse_text,
-    inputs=[
-        gr.Textbox(label="Input Text", lines=3, placeholder="Enter text…"),
-        gr.Slider(0.0, 1.0, step=0.05, value=0.0, label="Temperature"),
-        gr.Slider(32, 512, step=32, value=256, label="Max new tokens"),
-    ],
-    outputs=gr.Textbox(label="Reversed Text", lines=3),
-    title="🔄 Reverse Text Model Demo",
-    description="Give it a sentence; it returns the reversed version.",
+# ----  prompt helpers --------------------------------------------------------
+parser = vf.XMLParser(["think", "answer"])          # <think> … </think>\n<answer> … </answer>
+SYSTEM_MSG = (
+    "You are a helpful assistant that reverses the user's text. "
+    f"Respond in exactly this format: {parser.get_format_str()}"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+def build_prompt(user_msg: str, tok) -> str:
+    """Use the model’s native chat template so all special tokens are right."""
+    return tok.apply_chat_template(
+        [{"role": "system", "content": SYSTEM_MSG},
+         {"role": "user",   "content": user_msg}],
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+# ----  lazy-load model the first time a GPU is granted -----------------------
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
+model = None                        # brought into scope & moved to CUDA inside @spaces.GPU
+
+@spaces.GPU(duration=60)            # ← the “proper annotation” for ZeroGPU
+def reverse(user_msg: str) -> str:
+    global model
+    if model is None:               # cold-start: happens on the first request
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            torch_dtype=torch.float16,
+            device_map="auto"
+        )
+    prompt = build_prompt(user_msg, tokenizer)
+    with torch.inference_mode():
+        out = model.generate(
+            tokenizer(prompt, return_tensors="pt").to(model.device),
+            max_new_tokens=64,
+            do_sample=False,
+            temperature=0.0
+        )
+    full = tokenizer.decode(out[0], skip_special_tokens=True)
+    return full[len(prompt):]       # strip the prompt – return only the reply
+
+# ----  Gradio UI -------------------------------------------------------------
+with gr.Blocks() as demo:
+    gr.Markdown("### Reverse-Text demo (ZeroGPU)")
+    txt_in = gr.Textbox(label="Input")
+    txt_out = gr.Textbox(label="Model reply")
+    btn = gr.Button("Run")
+    btn.click(reverse, txt_in, txt_out)
+
+demo.queue().launch()