Update app.py

app.py

@@ -1,70 +1,23 @@
 import os
-import re
-import json
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
 
 MODEL_ID = "Amey9766/llama32B-hospitality-review-triage"
-HF_TOKEN = os.getenv("HF_TOKEN")  # optional secret
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-# ---- Load once at startup (works even on older Gradio) ----
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
 
+# Load config and REMOVE any quantization config
+config = AutoConfig.from_pretrained(MODEL_ID, token=HF_TOKEN)
+if hasattr(config, "quantization_config"):
+    config.quantization_config = None
+
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     token=HF_TOKEN,
+    config=config,                         # <-- important
     device_map="auto",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+    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(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()},
-    ]
-
-    prompt = tokenizer.apply_chat_template(messages, tokenize=False)
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    with torch.no_grad():
-        output = model.generate(
-            **inputs,
-            max_new_tokens=int(max_new_tokens),
-            do_sample=float(temperature) > 0,
-            temperature=float(temperature) if float(temperature) > 0 else None,
-            pad_token_id=tokenizer.eos_token_id
-        )
-
-    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="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_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_review, inputs=[review, max_tokens, temp], outputs=out)
-
-demo.launch()