import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained("King-8/interview_statements")

classifier = pipeline("text-classification", model=model, tokenizer=tokenizer)

label_map = {
    "LABEL_0": "Education & Innovation",
    "LABEL_1": "Community & Health",
    "LABEL_2": "Family & History",
    "LABEL_3": "Faith & Spirituality",
    "LABEL_4": "Business & Finance"
}

def classify_statement(statement):
    prediction = classifier(statement)[0]
    
    label = prediction["label"]
    
    predicted_category = label_map.get(label, "Unknown Category")  # Default to "Unknown" if not found
    
    return f"Predicted Category: {predicted_category}"

gr.Interface(
    fn=classify_statement,
    inputs=gr.Textbox(label="Enter a community statement"),
    outputs=gr.Textbox(label="Predicted Category"),
    title="Interview Statement Categorizer",
    description="This app predicts the category of a community statement based on trained data from community interviews."
).launch()