Update app.py

app.py

@@ -1,27 +1,122 @@
+# app.py
 import os
+import re
+import json
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
+# ✅ Your model repo (exact)
 MODEL_ID = "Amey9766/llama32B-hospitality-review-triage"
-HF_TOKEN = os.getenv("HF_TOKEN")
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+# If your model is private/gated, add HF_TOKEN as a Space Secret
+HF_TOKEN = os.getenv("HF_TOKEN", None)
 
-config = AutoConfig.from_pretrained(MODEL_ID, token=HF_TOKEN)
 
-# IMPORTANT: remove quantization_config cleanly
-if hasattr(config, "quantization_config"):
+def extract_first_json(text: str) -> str:
+    """
+    Extract the first JSON object from a string. Falls back to raw text.
+    """
+    match = re.search(r"\{.*\}", text, flags=re.S)
+    return match.group(0) if match else text
+
+
+def load_model():
+    """
+    Load tokenizer/model once. Keep it simple: no 4-bit, no bitsandbytes.
+    This avoids the bitsandbytes / quantization crashes on Spaces.
+    """
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+
+    # Ensure pad token is set
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_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 tokenizer, model
+
+
+tokenizer, model = load_model()
+
+
+def triage(review_text: str, max_new_tokens: int = 256, temperature: float = 0.0):
+    if not review_text or not review_text.strip():
+        return "Please enter a hotel review."
+
+    messages = [
+        {
+            "role": "system",
+            "content": (
+                "You are a hospitality review triage assistant. "
+                "Output ONLY valid JSON (no extra text)."
+            ),
+        },
+        {"role": "user", "content": review_text.strip()},
+    ]
+
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    do_sample = float(temperature) > 0.0
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=int(max_new_tokens),
+            do_sample=do_sample,
+            temperature=float(temperature) if do_sample else None,
+            pad_token_id=tokenizer.eos_token_id,
+        )
+
+    decoded = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+    # Try to isolate JSON and pretty print it
+    json_text = extract_first_json(decoded)
     try:
-        delattr(config, "quantization_config")   # best option
+        obj = json.loads(json_text)
+        return json.dumps(obj, indent=2)
     except Exception:
-        config.quantization_config = {}          # fallback
-
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    token=HF_TOKEN,
-    config=config,
-    quantization_config=None,                    # prevents auto-quant logic
-    device_map="auto",
-    dtype=torch.float16 if torch.cuda.is_available() else torch.float32
-)
-model.eval()
+        return json_text
+
+
+# ✅ Gradio UI
+with gr.Blocks(title="Hospitality Review Triage") as demo:
+    gr.Markdown("# 🏨 Hospitality Review → JSON Triage")
+    gr.Markdown(
+        "Paste a guest review and get structured JSON for routing/triage. "
+        "The model is instructed to output **JSON only**."
+    )
+
+    review_in = gr.Textbox(
+        label="Guest Review",
+        lines=6,
+        placeholder="Example: The room was dirty and the AC didn’t work. Front desk didn’t respond.",
+    )
+
+    with gr.Row():
+        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")
+
+    output = gr.Code(label="JSON Output", language="json")
+
+    btn = gr.Button("Generate JSON", variant="primary")
+    btn.click(triage, inputs=[review_in, max_tokens, temp], outputs=output)
+
+    gr.Examples(
+        examples=[
+            ["The room was dirty and the AC didn’t work. I called twice and no one came."],
+            ["Great location and staff were friendly, but breakfast was overpriced and slow."],
+            ["I found bugs in the bathroom. This is unacceptable and I want a refund."],
+            ["Noise from the hallway kept us awake all night. The bed was uncomfortable."],
+        ],
+        inputs=review_in,
+        label="Try examples",
+    )
+
+demo.launch()