Update app.py

app.py

@@ -1,76 +1,71 @@
-import spaces
-import gradio as gr
-import torch
+import re, spaces, gradio as gr, torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# Load tokenizer globally (CPU operation)
-model_name = "loocorez/reverse-text-warmup"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+MODEL_NAME = "loocorez/reverse-text-warmup"
 
-# Load model globally but keep on CPU initially
+SYSTEM_PROMPT = """Reverse the given text.
+
+Respond in the following format:
+<response>
+  <think>...</think>
+  <answer>...</answer>
+</response>"""
+
+# ---------- load once ----------
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+tokenizer.pad_token = tokenizer.eos_token  # safety
 model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16  # Use half precision for memory efficiency
+    MODEL_NAME,
+    torch_dtype=torch.float16
 )
+model.eval()  # turn off dropout etc.
 
-@spaces.GPU(duration=60)  # Reserve GPU for 60 seconds
-def reverse_text(input_text):
-    # Move model to GPU only when needed
-    model.to("cuda")
-    
+@spaces.GPU(duration=60)   # reserve GPU only while we run
+def reverse_text(user_text: str) -> str:
     try:
-        # Tokenize and move to GPU
-        inputs = tokenizer(
-            input_text, 
-            return_tensors="pt",
-            max_length=512,
-            truncation=True
-        ).to("cuda")
-        
-        # Generate
+        model.to("cuda")
+
+        # full prompt = system + user
+        prompt = f"{SYSTEM_PROMPT}\n\n{user_text.strip()}"
+        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+
         with torch.no_grad():
-            outputs = model.generate(
+            out = model.generate(
                 **inputs,
-                max_new_tokens=1000,
-                do_sample=True,
-                temperature=0.7,
+                max_new_tokens=256,
+                do_sample=False,        # <- deterministic
+                temperature=0.0,
                 pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.eos_token_id
+                eos_token_id=tokenizer.convert_tokens_to_ids("</response>")
             )
-        
-        # Decode result
-        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        generated_text = result[len(input_text):].strip()
-        
-        return generated_text
-        
+
+        # decode ONLY the newly-generated tokens
+        generated = tokenizer.decode(out[0][inputs["input_ids"].size(1):],
+                                     skip_special_tokens=True)
+
+        # pull out <answer>...</answer>
+        m = re.search(r"<answer>(.*?)</answer>", generated, re.S)
+        return (m.group(1).strip() if m else generated.strip())
+
     except Exception as e:
-        return f"Error: {str(e)}"
-    
+        return f"Error: {e}"
+
     finally:
-        # Move model back to CPU to free GPU memory
         model.to("cpu")
         torch.cuda.empty_cache()
 
-# Create interface
 demo = gr.Interface(
     fn=reverse_text,
-    inputs=gr.Textbox(
-        label="Input Text",
-        placeholder="Enter text to process...",
-        lines=3
-    ),
-    outputs=gr.Textbox(
-        label="Generated Text",
-        lines=3
-    ),
+    inputs=gr.Textbox(label="Input Text", lines=3, placeholder="Enter text…"),
+    outputs=gr.Textbox(label="Reversed Text", lines=3),
     title="🔄 Reverse Text Model Demo",
-    description="Test your custom reverse-text-warmup model using ZeroGPU",
+    description="Give it a sentence; it returns the reversed version.",
     examples=[
         ["Hello world"],
-        ["The quick brown fox"],
+        ["The quick brown fox jumps over the lazy dog."],
         ["Machine learning is amazing"]
     ]
 )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()