Update app.py

app.py

@@ -1,61 +1,70 @@
-import os, re, torch, gradio as gr
+import os
+import re
+import json
+import torch
+import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-MODEL_ID = "Amey9766/llama3.2-hotel-review-triage"
-HF_TOKEN = os.getenv("HF_TOKEN")
+MODEL_ID = "Amey9766/llama32B-hospitality-review-triage"
+HF_TOKEN = os.getenv("HF_TOKEN")  # optional secret
 
-@gr.cache_resource
-def load_model():
-    tok = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        token=HF_TOKEN,
-        device_map="auto",
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-    )
-    model.eval()
-    return tok, model
+# ---- Load once at startup (works even on older Gradio) ----
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
 
-tokenizer, model = load_model()
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    token=HF_TOKEN,
+    device_map="auto",
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+)
+model.eval()
 
 def extract_json(text: str) -> str:
     m = re.search(r"\{.*\}", text, flags=re.S)
     return m.group(0) if m else text
 
-def triage(review, max_new_tokens=256, temperature=0.0):
-    if not review.strip():
+def triage_review(review, max_new_tokens=256, temperature=0.0):
+    if not review or not review.strip():
         return "Please enter a hotel review."
 
     messages = [
-        {"role":"system", "content":"You are a hospitality review triage assistant. Output ONLY valid JSON."},
-        {"role":"user", "content": review.strip()}
+        {"role": "system", "content": "You are a hospitality review triage assistant. Output ONLY valid JSON."},
+        {"role": "user", "content": review.strip()},
     ]
+
     prompt = tokenizer.apply_chat_template(messages, tokenize=False)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     with torch.no_grad():
-        out = model.generate(
+        output = model.generate(
             **inputs,
             max_new_tokens=int(max_new_tokens),
-            do_sample=temperature > 0,
-            temperature=temperature if temperature > 0 else None,
+            do_sample=float(temperature) > 0,
+            temperature=float(temperature) if float(temperature) > 0 else None,
             pad_token_id=tokenizer.eos_token_id
         )
 
-    decoded = tokenizer.decode(out[0], skip_special_tokens=True)
-    return extract_json(decoded)
+    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
+    json_text = extract_json(decoded)
+
+    # Try to pretty print JSON
+    try:
+        obj = json.loads(json_text)
+        return json.dumps(obj, indent=2)
+    except Exception:
+        return json_text
 
-with gr.Blocks(title="Hotel Review Triage") as demo:
-    gr.Markdown("# 🏨 Hotel Review Triage Demo")
-    gr.Markdown("Paste a hotel review. The model returns **JSON only** for triage.")
+with gr.Blocks(title="Hospitality Review Triage") as demo:
+    gr.Markdown("# 🏨 Hospitality Review Triage Demo")
+    gr.Markdown("Paste a review and get **JSON-only** triage output (category, severity, department, sentiment).")
 
     review = gr.Textbox(label="Hotel Review", lines=5, placeholder="The room was dirty and the AC didn’t work.")
     with gr.Row():
-        max_new_tokens = gr.Slider(64, 512, value=256, step=32, label="Max new tokens")
-        temperature = gr.Slider(0.0, 1.0, value=0.0, step=0.1, label="Temperature")
+        max_tokens = gr.Slider(64, 512, value=256, step=32, label="Max new tokens")
+        temp = gr.Slider(0.0, 1.0, value=0.0, step=0.1, label="Temperature")
 
     out = gr.Code(label="Output (JSON)", language="json")
     btn = gr.Button("Generate JSON", variant="primary")
-    btn.click(triage, inputs=[review, max_new_tokens, temperature], outputs=out)
+    btn.click(triage_review, inputs=[review, max_tokens, temp], outputs=out)
 
 demo.launch()