Update app.py

app.py

@@ -1,7 +1,6 @@
-# app.py
 import re, spaces, gradio as gr, torch
-from transformers import AutoTokenizer
-from peft import AutoPeftModelForCausalLM  # ← handles LoRA if present
+from transformers import AutoTokenizer, AutoModelForCausalLM   # ⟵ swap import
+# from peft import AutoPeftModelForCausalLM   (remove)
 
 MODEL_NAME = "loocorez/reverse-text-warmup"
 
@@ -13,20 +12,21 @@ Respond in the following format:
   <answer>...</answer>
 </response>"""
 
-# ---------- load once ----------
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 tokenizer.pad_token = tokenizer.eos_token
 
-# 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()
+# load the **full model** directly
+model = AutoModelForCausalLM.from_pretrained(             # ⟵ use AutoModel
+    MODEL_NAME,
+    torch_dtype=torch.float16
+).eval()
 
-# helper: id of "</response>"
 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):
+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")
 
@@ -34,29 +34,24 @@ def reverse_text(user_text: str, temperature: float = 0.0, max_tokens: int = 256
         out = model.generate(
             **inputs,
             max_new_tokens=max_tokens,
-            do_sample=temperature > 0.0,
-            temperature=max(temperature, 1e-6),  # avoid 0 with sampling
+            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,              # tames loops
+            repetition_penalty=1.1,
         )
 
-    # decode only the newly-generated part
-    generated = tokenizer.decode(out[0][inputs["input_ids"].size(1):],
-                                 skip_special_tokens=True)
-
-    # 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()
+    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()
 
 demo = gr.Interface(
     fn=reverse_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")
+        gr.Textbox(label="Input Text", lines=3),
+        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",