Update app.py

app.py

@@ -1,11 +1,9 @@
 import re
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
+from optimum.onnxruntime import ORTModelForTokenClassification
 import gradio as gr
 
-# Load a pre-trained multilingual NER model for entity recognition
-ner_model = pipeline("ner", model="dbmdz/bert-base-multilingual-cased", aggregation_strategy="simple")
-
-# Define categories and their associated keywords
+# Define categories and their keywords
 CATEGORIES = {
     "Need": {
         "Utilities": ["dien", "nuoc", "gas", "internet", "dienthoai"],
@@ -19,7 +17,7 @@ CATEGORIES = {
     },
     "Want": {
         "Dining Out": ["nha hang", "quan an", "cafe", "tra sua"],
-        "Entertainment": ["phim", "karaoke", "game", "nhac"],
+        "Entertainment": ["phim", "karaoke", "game", "nhac", "do choi", "bup be"],
         "Travel": ["du lich", "ve may bay", "khach san"],
         "Fitness": ["gym", "yoga", "the thao"],
         "Shopping": ["quan ao", "phu kien", "dien thoai", "luxury"],
@@ -39,9 +37,19 @@ CATEGORIES = {
 
 # Normalize Vietnamese input (remove accents)
 def normalize_vietnamese(text):
-    return re.sub(r'[àáạảãâầấậẩẫăằắặẳẵèéẹẻẽêềếệểễìíịỉĩòóọỏõôồốộổỗơờớợởỡùúụủũưừứựửữỳýỵỷỹđ]', '', text).replace("đ", "d")
+    return re.sub(
+        r'[àáạảãâầấậẩẫăằắặẳẵèéẹẻẽêềếệểễìíịỉĩòóọỏõôồốộổỗơờớợởỡùúụủũưừứựửữỳýỵỷỹđ]', '', text
+    ).replace("đ", "d")
+
+# Load and quantize the model
+model_name = "distilbert-base-multilingual-cased"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+quantized_model = ORTModelForTokenClassification.from_pretrained(model_name, from_transformers=True)
+
+# Create the NER pipeline with the quantized model
+ner_model = pipeline("ner", model=quantized_model, tokenizer=tokenizer, aggregation_strategy="simple")
 
-# Extract entities and classify
+# Classify input
 def classify_and_extract(user_input):
     normalized_input = normalize_vietnamese(user_input.lower())
 
@@ -49,10 +57,7 @@ def classify_and_extract(user_input):
     amount_match = re.search(r"(\d+(\.\d{1,2})?)", normalized_input)
     amount = amount_match.group(0) if amount_match else "Unknown"
 
-    # Run the NER model to detect entities
-    ner_results = ner_model(user_input)
-
-    # Match keywords for categories
+    # Rule-based matching for categories
     for main_category, subcategories in CATEGORIES.items():
         for subcategory, keywords in subcategories.items():
             if any(keyword in normalized_input for keyword in keywords):
@@ -60,10 +65,11 @@ def classify_and_extract(user_input):
                     "Main Category": main_category,
                     "Sub Category": subcategory,
                     "Amount": amount,
-                    "Entities": ner_results,
+                    "Entities": []  # Skip NER if matched via rules
                 }
 
-    # Default response if no match
+    # Fallback to NER model for unmatched cases
+    ner_results = ner_model(user_input)
     return {
         "Main Category": "Uncategorized",
         "Sub Category": "Unknown",
@@ -89,4 +95,4 @@ iface = gr.Interface(
     description="Classify expenditures into main and subcategories (Need, Want, Saving/Investment) and extract amounts."
 )
 
-iface.launch()
+iface.launch()