import json
import re
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM

MODEL_ID = "Amey9766/llama32-hotel-review-triage"

def load_model():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        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 extract_json(text):
    match = re.search(r"\{.*\}", text, re.DOTALL)
    return match.group(0) if match else text

def triage_review(review, max_tokens, temperature):
    if not review.strip():
        return "Please enter a hotel review."

    system_prompt = "You are a hospitality review triage assistant. Output ONLY valid JSON."

    messages = [
        {"role": "system", "content": system_prompt},
        {"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=max_tokens,
            do_sample=temperature > 0,
            temperature=temperature if temperature > 0 else None,
            pad_token_id=tokenizer.eos_token_id
        )

    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
    return extract_json(decoded)

with gr.Blocks(title="Hotel Review Triage") as demo:
    gr.Markdown("## 🏨 Hotel Review Triage Demo")
    gr.Markdown(
        "Paste a hotel review below. The model will return **structured JSON** "
        "for operational triage (department, severity, sentiment, summary)."
    )

    review_input = gr.Textbox(
        label="Hotel Review",
        placeholder="The room was dirty and the AC didn’t work.",
        lines=5
    )

    with gr.Row():
        max_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")

    output = gr.Code(label="Model Output (JSON)", language="json")

    run_btn = gr.Button("Generate JSON", variant="primary")
    run_btn.click(
        fn=triage_review,
        inputs=[review_input, max_tokens, temperature],
        outputs=output
    )

demo.launch()