import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import os

# Get Hugging Face token from environment variable
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    raise ValueError("Please set HF_TOKEN environment variable with your Hugging Face access token")

# Load model and tokenizer
model_name = "iimran/AnalyserV1"
tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
model = AutoModelForSequenceClassification.from_pretrained(model_name, token=HF_TOKEN)

def classify_complaint(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=256)
    with torch.no_grad():
        outputs = model(**inputs)
    return model.config.id2label[torch.argmax(outputs.logits).item()]

# Create Gradio interface
demo = gr.Interface(
    fn=classify_complaint,
    inputs=gr.Textbox(lines=3, placeholder="Enter your complaint here...", label="Complaint Text"),
    outputs=gr.Label(label="Predicted Category"),
    title="Complaint Category Classifier",
    description="Automatically classify community complaints into specific categories",
    examples=[
        ["I wanted to bring to your attention that a huge big truck has been parked on Main Street"],
        ["There are overgrown bushes on Oak Road that pose a fire risk"],
        ["Excessive noise from construction site during night hours"]
    ]
)

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