Update app.py

app.py

@@ -1,67 +1,82 @@
 import re
-from transformers import pipeline
 import gradio as gr
 
-# Load a lightweight model for classification
-classifier = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english", return_all_scores=True)
-
-# Define main and subcategories
+# Define main and subcategories with keywords for fast matching
 CATEGORIES = {
-    "Need": [
-        "utilities", "housing", "groceries", "transportation", "education", "medical", "insurance", "childcare"
-    ],
-    "Want": [
-        "dining out", "entertainment", "travel", "fitness", "shopping", "hobbies", "personal care"
-    ],
-    "Saving/Investment": [
-        "emergency fund", "retirement", "investments", "debt repayment", "education fund", "savings for goals", "health savings"
-    ]
-}
-
-# Predefined keywords for fast classification
-KEYWORDS = {
-    "saving": ["gui tiet kiem", "tiet kiem", "lai suat", "savings", "interest"],
-    "utilities": ["electricity", "water", "gas", "internet", "phone"],
-    "housing": ["rent", "mortgage", "property tax", "maintenance"],
-    "groceries": ["food", "beverages", "supermarket"],
-    "transportation": ["gas", "car", "vehicle", "public transit"],
-    "education": ["tuition", "books", "school", "course"],
-    "medical": ["insurance", "doctor", "prescriptions", "medicine"],
-    "dining out": ["restaurant", "cafe", "fast food", "delivery"],
-    # Add more keywords for all subcategories...
+    "Need": {
+        "Utilities": ["electricity", "water", "gas", "internet", "phone"],
+        "Housing": ["rent", "mortgage", "property tax", "maintenance"],
+        "Groceries": ["food", "beverages", "supermarket", "groceries"],
+        "Transportation": ["gasoline", "car payment", "public transit", "vehicle maintenance"],
+        "Education": ["tuition", "books", "school", "course"],
+        "Medical": ["health insurance", "doctor", "prescription", "medicine"],
+        "Insurance": ["home insurance", "auto insurance", "life insurance"],
+        "Childcare": ["daycare", "babysitter", "school fee"],
+    },
+    "Want": {
+        "Dining Out": ["restaurant", "cafe", "fast food", "delivery"],
+        "Entertainment": ["movie", "concert", "streaming", "game", "doll", "toy"],
+        "Travel": ["vacation", "flight", "hotel", "recreation"],
+        "Fitness": ["gym", "yoga", "sports", "fitness"],
+        "Shopping": ["clothing", "accessories", "gadget", "luxury"],
+        "Hobbies": ["art", "crafts", "sports gear", "book", "collectibles"],
+        "Personal Care": ["spa", "beauty", "haircut", "salon"],
+    },
+    "Saving/Investment": {
+        "Emergency Fund": ["emergency fund"],
+        "Retirement": ["retirement", "pension"],
+        "Investments": ["stocks", "bonds", "real estate", "crypto"],
+        "Debt Repayment": ["loan repayment", "credit card payment"],
+        "Education Fund": ["education fund"],
+        "Savings for Goals": ["down payment", "vacation savings", "wedding savings"],
+        "Health Savings": ["health savings account", "hsa", "fsa"],
+    }
 }
 
 # Normalize Vietnamese input (remove accents)
 def normalize_vietnamese(text):
     return re.sub(r'[àáạảãâầấậẩẫăằắặẳẵèéẹẻẽêềếệểễìíịỉĩòóọỏõôồốộổỗơờớợởỡùúụủũưừứựửữỳýỵỷỹđ]', '', text).replace("đ", "d")
 
-# Classify input into main and subcategories
+# Classify input into categories
 def classify_input(user_input):
-    # Normalize input
     normalized_input = normalize_vietnamese(user_input.lower())
 
-    # Match keywords for faster classification
-    for main_cat, subcats in CATEGORIES.items():
-        for subcat in subcats:
-            if any(keyword in normalized_input for keyword in KEYWORDS.get(subcat, [])):
-                return {"Main Category": main_cat, "Sub Category": subcat.capitalize()}
+    # Extract amount if mentioned
+    amount_match = re.search(r"(\d+(\.\d{1,2})?)", normalized_input)
+    amount = amount_match.group(0) if amount_match else "Unknown"
+
+    # Match subcategory and main category based on keywords
+    for main_category, subcategories in CATEGORIES.items():
+        for subcategory, keywords in subcategories.items():
+            if any(keyword in normalized_input for keyword in keywords):
+                return {
+                    "Main Category": main_category,
+                    "Sub Category": subcategory,
+                    "Amount": amount
+                }
 
-    # Fallback to model classification
-    result = classifier(normalized_input)
-    category = max(result, key=lambda x: x["score"])["label"]
-    return {"Main Category": category, "Sub Category": "Unknown"}
+    # Default response for unmatched cases
+    return {
+        "Main Category": "Uncategorized",
+        "Sub Category": "Unknown",
+        "Amount": amount
+    }
 
 # Define Gradio interface
 def process_user_input(user_input):
-    classification = classify_input(user_input)
-    return f"Main Category: {classification['Main Category']}\nSub Category: {classification['Sub Category']}"
+    result = classify_input(user_input)
+    return (
+        f"Main Category: {result['Main Category']}\n"
+        f"Sub Category: {result['Sub Category']}\n"
+        f"Amount: {result['Amount']}"
+    )
 
 iface = gr.Interface(
     fn=process_user_input,
     inputs="text",
     outputs="text",
     title="Expenditure Classifier",
-    description="Classify expenditures into main and subcategories (Need, Want, Saving/Investment)."
+    description="Classify expenditures into main and subcategories (Need, Want, Saving/Investment) instantly."
 )
 
 iface.launch()