Update app.py

app.py

@@ -1,5 +1,7 @@
+# app.py
 import re, spaces, gradio as gr, torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer
+from peft import AutoPeftModelForCausalLM  # ← handles LoRA if present
 
 MODEL_NAME = "loocorez/reverse-text-warmup"
 
@@ -13,57 +15,52 @@ Respond in the following format:
 
 # ---------- 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
-)
-model.eval()  # turn off dropout etc.
-
-@spaces.GPU(duration=60)   # reserve GPU only while we run
-def reverse_text(user_text: str) -> str:
-    try:
-        model.to("cuda")
+tokenizer.pad_token = tokenizer.eos_token
 
-        # full prompt = system + user
-        prompt = f"{SYSTEM_PROMPT}\n\n{user_text.strip()}"
-        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+# pulls base + adapter; merge folds LoRA into the base so inference is fast
+model = AutoPeftModelForCausalLM.from_pretrained(
+    MODEL_NAME, torch_dtype=torch.float16
+).merge_and_unload().eval()
 
-        with torch.no_grad():
-            out = model.generate(
-                **inputs,
-                max_new_tokens=2048,
-                temperature=0.3,
-                pad_token_id=tokenizer.eos_token_id,
-                eos_token_id=tokenizer.convert_tokens_to_ids("</response>")
-            )
+# helper: id of "</response>"
+EOS_ID = tokenizer.encode("</response>", add_special_tokens=False)[0]
 
-        # decode ONLY the newly-generated tokens
-        generated = tokenizer.decode(out[0][inputs["input_ids"].size(1):],
-                                     skip_special_tokens=True)
+@spaces.GPU(duration=60)
+def reverse_text(user_text: str, temperature: float = 0.0, max_tokens: int = 256):
+    prompt = f"{SYSTEM_PROMPT}\n\n{user_text.strip()}"
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
 
-        # pull out <answer>...</answer>
-        m = re.search(r"<answer>(.*?)</answer>", generated, re.S)
-        return (m.group(1).strip() if m else generated.strip())
+    with torch.no_grad():
+        out = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            do_sample=temperature > 0.0,
+            temperature=max(temperature, 1e-6),  # avoid 0 with sampling
+            eos_token_id=EOS_ID,
+            pad_token_id=tokenizer.eos_token_id,
+            repetition_penalty=1.1,              # tames loops
+        )
 
-    except Exception as e:
-        return f"Error: {e}"
+    # decode only the newly-generated part
+    generated = tokenizer.decode(out[0][inputs["input_ids"].size(1):],
+                                 skip_special_tokens=True)
 
-    finally:
-        model.to("cpu")
-        torch.cuda.empty_cache()
+    # pull the answer out of the XML
+    m = re.search(r"<answer>(.*?)</answer>", generated, re.S)
+    return m.group(1).strip() if m else generated.strip()
 
 demo = gr.Interface(
     fn=reverse_text,
-    inputs=gr.Textbox(label="Input Text", lines=3, placeholder="Enter text…"),
+    inputs=[
+        gr.Textbox(label="Input Text", lines=3, placeholder="Enter text…"),
+        gr.Slider(minimum=0.0, maximum=1.0, step=0.05,
+                  value=0.0, label="Temperature"),
+        gr.Slider(minimum=32, maximum=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.",
-    examples=[
-        ["Hello world"],
-        ["The quick brown fox jumps over the lazy dog."],
-        ["Machine learning is amazing"]
-    ]
+    description="Give it a sentence; it returns the reversed version."
 )
 
 if __name__ == "__main__":