import gradio as gr
from transformers import pipeline
import json

# Load a zero-shot classification pipeline
classifier = pipeline("zero-shot-classification", 
                       model="facebook/bart-large-mnli")

# Define expense categories
expense_categories = [
    "Groceries", "Restaurants", "Coffee Shops", "Transportation",
    "Utilities", "Entertainment", "Shopping", "Health", "Travel", 
    "Education", "Home Improvement", "Personal Care", "Gifts"
]

def categorize_expense(merchant_name, item_description=""):
    """Categorize an expense based on merchant name and optional item description"""
    # Combine inputs for better context
    input_text = f"{merchant_name} {item_description}".strip()
    
    # Run zero-shot classification
    result = classifier(
        input_text, 
        expense_categories,
        multi_label=False
    )
    
    # Get top 3 categories with their scores
    top_categories = []
    for category, score in zip(result['labels'][:3], result['scores'][:3]):
        top_categories.append({"category": category, "confidence": float(score)})
    
    return json.dumps(top_categories)

# Create interface
iface = gr.Interface(
    fn=categorize_expense,
    inputs=[
        gr.Textbox(label="Merchant Name"),
        gr.Textbox(label="Item Description (Optional)")
    ],
    outputs=gr.Textbox(label="Categories"),
    title="Reciply Expense Categorizer",
    description="Categorize expenses based on merchant name and item description"
)

iface.launch()