demo-app / app.py
Jomsky's picture
Update app.py
66f69d7 verified
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()