import torch
import joblib
import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from huggingface_hub import hf_hub_download

HF_MODEL_ID = "Jomsky/brgy-complaint-classifier"
LABEL_ENCODER_FILE = "label_encoder_v2.pkl"

# Download label encoder from HF repo
label_encoder_path = hf_hub_download(
    repo_id=HF_MODEL_ID,
    filename=LABEL_ENCODER_FILE
)

label_encoder = joblib.load(label_encoder_path)

# Load model & tokenizer
tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(HF_MODEL_ID)

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()

# Classification function
def classify_complaint(text):
    if not text.strip():
        return "Please enter a complaint."

    inputs = tokenizer(
        text,
        truncation=True,
        padding=True,
        return_tensors="pt"
    ).to(device)

    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        predicted_class_id = torch.argmax(logits, dim=-1).item()

    predicted_label = label_encoder.inverse_transform([predicted_class_id])[0]
    return predicted_label

# Gradio UI
demo = gr.Interface(
    fn=classify_complaint,
    inputs=gr.Textbox(lines=4, placeholder="Ano ang iyong complaint..."),
    outputs=gr.Label(label="Predicted Category"),
    title="Online Sumbong: Barangay Poblacion Complaint Classifier",
    description="Classifies barangay complaints into categories using a fine-tuned transformer model."
)

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