Update app.py

app.py

@@ -1,15 +1,38 @@
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-def critique(text):
-    return "MiniCrit demo is running.\n\nYou entered:\n" + (text or "")
-
-demo = gr.Interface(
-    fn=critique,
-    inputs=gr.Textbox(lines=4, label="Input"),
-    outputs=gr.Textbox(lines=8, label="Output"),
-    title="MiniCrit Demo",
-    description="Minimal working demo. Full model loads after GPU grant."
+MODEL = "wmaousley/MiniCrit-1.5B"   # replace with your model name
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL)
+
+# use low memory 4-bit load if supported
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL,
+    torch_dtype=torch.float16,
+    device_map="auto"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+def critique(text):
+    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            do_sample=False,
+            temperature=0.0
+        )
+    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return result
+
+
+with gr.Blocks(title="MiniCrit Demo") as demo:
+    gr.Markdown("# 🔍 MiniCrit Demo\nA lightweight adversarial critic model.")
+    
+    inp = gr.Textbox(label="Input text", placeholder="Enter text for critique")
+    out = gr.Textbox(label="Model Output")
+
+    btn = gr.Button("Run Critique")
+    btn.click(critique, inp, out)
+
+demo.launch()