import gradio as gr
from transformers import pipeline

classifier = pipeline("text-classification", model="kuro-08/bert-transaction-categorization")
labels = {
    0: "Utilities",
    1: "Health",
    2: "Dining",
    3: "Travel",
    4: "Education",
    5: "Subscription",
    6: "Family",
    7: "Food",
    8: "Festivals",
    9: "Culture",
    10: "Apparel",
    11: "Transportation",
    12: "Investment",
    13: "Shopping",
    14: "Groceries",
    15: "Documents",
    16: "Grooming",
    17: "Entertainment",
    18: "Social Life",
    19: "Beauty",
    20: "Rent",
    21: "Money transfer",
    22: "Salary",
    23: "Tourism",
    24: "Household",
}

def categorize(description: str):
    result = classifier(description)[0]
    label_id = int(result['label'].split('_')[1])
    category = labels.get(label_id, "Miscellaneous")
    return {"category": category}

demo = gr.Interface(
    fn=categorize,
    inputs=gr.Textbox(label="Transaction Description", placeholder="e.g., Dinner at Subway"),
    outputs=gr.JSON(label="Category"),
    title="Transaction Categorizer API"
)

demo.launch()