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

MODEL_NAME = "Amey9766/llama32-1b-maintenance-classifier"

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Load model on CPU to avoid meta tensor issues
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float32,
    device_map=None
)
model.to("cpu")

def classify(text):
    # Strong classification prompt to force label output
    prompt = (
        "You are a maintenance request classifier. "
        "Your job is to output ONLY ONE WORD: urgent, routine, or cosmetic.\n\n"
        f"Request: {text}\n"
        "Category:"
    )

    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")

    outputs = model.generate(
        **inputs,
        max_new_tokens=3,
        temperature=0.0,
        do_sample=False,
        eos_token_id=tokenizer.eos_token_id
    )

    raw = tokenizer.decode(outputs[0], skip_special_tokens=True).lower()

    # Extract only the part after "category:"
    if "category:" in raw:
        raw = raw.split("category:")[-1].strip()

    # Match labels
    if "urgent" in raw:
        return "🔴 URGENT"
    if "routine" in raw:
        return "🟡 ROUTINE"
    if "cosmetic" in raw:
        return "🟢 COSMETIC"

    return f"Unrecognized output: {raw}"

# Gradio UI
demo = gr.Interface(
    fn=classify,
    inputs=gr.Textbox(label="Enter maintenance request"),
    outputs=gr.Textbox(label="Predicted Category"),
    title="Maintenance Request Classifier",
    description="Predicts whether a maintenance request is urgent, routine, or cosmetic."
)

if __name__ == "__main__":
    demo.launch()