-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathexpense_classifier.py
More file actions
282 lines (228 loc) · 9.19 KB
/
Copy pathexpense_classifier.py
File metadata and controls
282 lines (228 loc) · 9.19 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
"""
Expense Classification using sklearn
Automatically classifies expense titles into categories: food, travel, shopping,
entertainment, utilities, or other.
"""
import pickle
import os
import logging
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Global classifier (loaded once, reused for all classifications)
_classifier = None
_MODEL_PATH = 'expense_classifier_model.pkl'
def get_classifier():
"""
Get or load the classifier model.
Uses lazy loading - only loads model when first needed.
"""
global _classifier
if _classifier is None:
try:
if not os.path.exists(_MODEL_PATH):
logger.error(f"Model file not found: {_MODEL_PATH}")
logger.info("Please run 'python3 train_classifier.py' first to train the model.")
raise FileNotFoundError(f"Model file not found: {_MODEL_PATH}")
logger.info(f"Loading classifier model from {_MODEL_PATH}...")
with open(_MODEL_PATH, 'rb') as f:
_classifier = pickle.load(f)
logger.info("Model loaded successfully!")
except Exception as e:
logger.error(f"Error loading model: {e}")
raise
return _classifier
def classify_expense(title: str) -> str:
"""
Classify an expense title into a category.
Args:
title: The expense title (e.g., "McDonald", "Uber ride", "Tesco shopping")
Returns:
Category string: "food", "travel", "shopping", "entertainment", "utilities", or "other"
Examples:
>>> classify_expense("McDonald")
'food'
>>> classify_expense("Uber ride to airport")
'travel'
>>> classify_expense("Tesco grocery shopping")
'shopping'
"""
if not title or not title.strip():
return "other"
try:
classifier = get_classifier()
# Predict category
prediction = classifier.predict([title])[0]
# Get confidence score
proba = classifier.predict_proba([title])[0]
max_proba = max(proba)
predicted_idx = list(proba).index(max_proba)
# Get all class probabilities
classes = classifier.classes_
predicted_category = classes[predicted_idx]
logger.debug(f"Classified '{title}' as '{predicted_category}' "
f"(confidence: {max_proba:.2%})")
# If confidence is very low, default to "other"
if max_proba < 0.3:
logger.warning(f"Low confidence ({max_proba:.2%}) for '{title}', "
f"using 'other' instead of '{predicted_category}'")
return "other"
return predicted_category
except FileNotFoundError:
# Model not trained yet, use fallback
logger.warning("Model not found, using fallback classification")
return _fallback_classify(title)
except Exception as e:
logger.error(f"Error classifying '{title}': {e}")
# Fallback to simple rule-based if model fails
return _fallback_classify(title)
def _fallback_classify(title: str) -> str:
"""
Fallback classification using simple keyword matching.
Used if sklearn model fails to load or classify.
Includes EU/Ireland-specific keywords.
"""
title_lower = title.lower()
# Food keywords (including EU/Ireland specific)
food_keywords = [
'mcdonald', 'starbucks', 'pizza', 'restaurant', 'food', 'cafe', 'café',
'burger', 'taco', 'subway', 'kfc', 'chipotle', 'dinner', 'lunch',
'breakfast', 'coffee', 'deliveroo', 'just eat', 'ubereats', 'domino',
'supermac', 'spar', 'centra', 'dunnes', 'marks & spencer', 'm&s'
]
# Travel keywords (including EU/Ireland specific)
travel_keywords = [
'uber', 'lyft', 'taxi', 'flight', 'hotel', 'airbnb', 'train', 'bus',
'metro', 'parking', 'gas', 'petrol', 'toll', 'rental', 'scooter',
'bike', 'luas', 'dart', 'ireland rail', 'aer lingus', 'ryanair',
'dublin bus', 'bus eireann', 'train ticket', 'flight ticket'
]
# Shopping keywords (including EU/Ireland specific)
shopping_keywords = [
'tesco', 'dunnes', 'supervalu', 'aldi', 'lidl', 'penneys', 'primark',
'target', 'walmart', 'amazon', 'grocery', 'shopping', 'store',
'costco', 'pharmacy', 'boots', 'superdrug', 'argos', 'ikea',
'clothing', 'electronics', 'bookstore', 'convenience'
]
# Entertainment keywords
entertainment_keywords = [
'movie', 'cinema', 'concert', 'theater', 'theatre', 'museum',
'amusement', 'bowling', 'karaoke', 'arcade', 'sports', 'game',
'festival', 'club', 'bar', 'pub', 'croke park', 'aviva stadium',
'ticket', 'event', 'show'
]
# Utilities keywords
utilities_keywords = [
'electricity', 'esb', 'water', 'internet', 'eir', 'vodafone',
'phone', 'mobile', 'cable', 'sky', 'virgin media', 'streaming',
'netflix', 'spotify', 'insurance', 'rent', 'mortgage', 'bill',
'subscription', 'service'
]
# Check categories in order of specificity
if any(kw in title_lower for kw in food_keywords):
return 'food'
elif any(kw in title_lower for kw in travel_keywords):
return 'travel'
elif any(kw in title_lower for kw in shopping_keywords):
return 'shopping'
elif any(kw in title_lower for kw in entertainment_keywords):
return 'entertainment'
elif any(kw in title_lower for kw in utilities_keywords):
return 'utilities'
else:
return 'other'
def get_top_tags(title: str, top_n: int = 3) -> list:
"""
Get top N tag suggestions for an expense title with confidence scores.
Args:
title: The expense title
top_n: Number of top suggestions to return (default: 3)
Returns:
List of dicts with 'tag' and 'confidence' keys, sorted by confidence (highest first)
Example:
>>> get_top_tags("McDonald")
[{'tag': 'food', 'confidence': 0.85},
{'tag': 'shopping', 'confidence': 0.10},
{'tag': 'other', 'confidence': 0.05}]
"""
if not title or not title.strip():
return [{"tag": "other", "confidence": 1.0}]
try:
classifier = get_classifier()
# Get probabilities for all classes
proba = classifier.predict_proba([title])[0]
classes = classifier.classes_
# Create list of (tag, confidence) tuples
tag_scores = [(classes[i], float(proba[i])) for i in range(len(classes))]
# Sort by confidence (highest first)
tag_scores.sort(key=lambda x: x[1], reverse=True)
# Return top N
top_tags = [
{"tag": tag, "confidence": round(confidence, 4)}
for tag, confidence in tag_scores[:top_n]
]
return top_tags
except FileNotFoundError:
logger.warning("Model not found, using fallback classification")
# Return fallback suggestions
fallback_tag = _fallback_classify(title)
return [{"tag": fallback_tag, "confidence": 0.5}]
except Exception as e:
logger.error(f"Error getting top tags for '{title}': {e}")
fallback_tag = _fallback_classify(title)
return [{"tag": fallback_tag, "confidence": 0.5}]
def classify_batch(titles: list) -> list:
"""
Classify multiple expense titles at once.
More efficient than calling classify_expense() multiple times.
Args:
titles: List of expense titles
Returns:
List of category strings
"""
if not titles:
return []
try:
classifier = get_classifier()
predictions = classifier.predict(titles)
return list(predictions)
except FileNotFoundError:
logger.warning("Model not found, using fallback classification")
return [_fallback_classify(title) for title in titles]
except Exception as e:
logger.error(f"Error in batch classification: {e}")
# Fallback to individual fallback classification
return [_fallback_classify(title) for title in titles]
# Test function
if __name__ == "__main__":
print("="*60)
print("Expense Classifier Test")
print("="*60)
# Check if model exists
if not os.path.exists(_MODEL_PATH):
print(f"\n Model file not found: {_MODEL_PATH}")
print("Please run 'python3 train_classifier.py' first to train the model.")
print("\nUsing fallback classification for testing...")
print("-" * 60)
else:
print("\nTesting classification...")
print("-" * 60)
test_cases = [
"McDonald",
"Uber ride to airport",
"Tesco grocery shopping",
"Luas ticket",
"Netflix subscription",
"Movie tickets",
"Dinner at Supermac's",
"Dublin Bus fare",
"ESB electricity bill",
"Concert at 3Arena",
"Amazon order",
"Coffee at Starbucks",
]
for title in test_cases:
category = classify_expense(title)
print(f"{title:30} → {category}")
print("\n" + "="*60)
print("Test complete!")