Update app.py

app.py

@@ -1,6 +1,6 @@
+# app.py
 import re, spaces, gradio as gr, torch
-from transformers import AutoTokenizer, AutoModelForCausalLM   # ⟵ swap import
-# from peft import AutoPeftModelForCausalLM   (remove)
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 MODEL_NAME = "loocorez/reverse-text-warmup"
 
@@ -12,50 +12,64 @@ Respond in the following format:
   <answer>...</answer>
 </response>"""
 
+# ── load model & tokenizer once ────────────────────────────────────────────────
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 tokenizer.pad_token = tokenizer.eos_token
 
-# load the **full model** directly
-model = AutoModelForCausalLM.from_pretrained(             # ⟵ use AutoModel
+model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
     torch_dtype=torch.float16
-).eval()
+).eval()                               # keep weights on CPU until each call
 
 EOS_ID = tokenizer.encode("</response>", add_special_tokens=False)[0]
 
-@spaces.GPU(duration=60)
-def reverse_text(user_text: str,
-                 temperature: float = 0.0,
-                 max_tokens: int = 256):
+
+@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()}"
-    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()
+
+    # ── 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()
+
 
 demo = gr.Interface(
     fn=reverse_text,
     inputs=[
-        gr.Textbox(label="Input Text", lines=3),
+        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")
+        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."
+    description="Give it a sentence; it returns the reversed version.",
 )
 
 if __name__ == "__main__":