-
Notifications
You must be signed in to change notification settings - Fork 2
/
main.py
executable file
·194 lines (169 loc) · 6.69 KB
/
main.py
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
#!/usr/bin/env python3
import random
import re
import sys
import traceback
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.pipeline import Pipeline
# Expected bank extracts CSV format:
# 1 transcation per line, in the following format:
# date;<any number of fields describing a transaction>;amount
# Removing words that aren't relevant for classification, therefore
# introducing classification noise.
# These words may need to be adapted.
STOP_WORDS=('carte', 'cb', 'du', 'facture', 'paiement')
def read_lines(filename):
lines = []
with open(filename) as fd:
l = fd.readline()
while l != "":
if re.match("^\d{4}/\d{2}/\d{2};", l):
# Only consider line starting with a valid date
# This allows skipping header, and possible comments
lines.append(l.strip())
l = fd.readline()
return lines
def cleaned_training_transaction(tr):
# Keep all transaction fields except date and price, respectively
# at 1st and last position (not using these fields for
# categorization).
return ';'.join(tr.split(';')[1:-1])
class Corpus:
def __enrich_training_set(self, transaction, category):
# Keeping all information to dump into training file
self.__training_set_str.append("{};{}".format(transaction,category))
# Setting x and y training set vectors
self.__training_set_x.append(cleaned_training_transaction(transaction))
self.__training_set_y.append(category)
def __init_training_set(self, training_fname):
training_set = read_lines(training_fname)
transactions = [";".join(l.split(';')[:-1]) for l in training_set]
categories = [ l.split(';')[-1] for l in training_set]
self.__training_set_str = []
self.__training_set_x = []
self.__training_set_y = []
for t,c in zip(transactions, categories):
self.__enrich_training_set(t, c)
def __update_state(self):
self.__categories = {}
for x,y in zip(self.__corpus, self.__prediction):
c = self.__categories.setdefault(y, [])
c.append(x)
self.__overview = []
for cat,transacs in self.__categories.items():
count = len(transacs)
amount = sum([float(v.split(';')[-1].replace(' ','')) for v in transacs])
self.__overview.append((amount, count, cat))
self.__overview.sort()
def __predict(self):
# https://scikit-learn.org/stable/tutorial/text_analytics/working_with_text_data.html
self.__text_clf.fit(self.__training_set_x, self.__training_set_y)
self.__prediction = self.__text_clf.predict(self.__corpus)
self.__update_state()
def __init__(self, training_fname, corpus_fnames):
self.__init_training_set(training_fname)
# Process all corpus files
self.__corpus = []
for fname in corpus_fnames:
self.__corpus.extend(read_lines(fname))
self.__vectorizer = CountVectorizer(
stop_words=STOP_WORDS,
token_pattern= '(?u)\\b\\w[a-zA-Z0-9_\\-\\.]+\\b',
ngram_range=(1,3),
)
self.__text_clf = Pipeline([
('vect', self.__vectorizer),
('clf', MultinomialNB()),
])
self.__predict()
# Commands
##########
def c_save_prediction(self, filename):
with open(filename, "w") as fd:
for x,y in zip(self.__corpus, self.__prediction):
fd.write("{};{}\n".format(x,y))
def c_save_training(self, filename):
with open(filename, "w") as fd:
fd.write("\n".join(self.__training_set_str))
fd.write("\n")
def c_overview(self):
for amount, count, category in self.__overview:
print("{:<10} {:<4} {}".format(round(amount,2), count, category))
def c_list_category(self, category):
for item in self.__categories[category]:
# When this becomes too slow, we should use a dictionary to fetch item indexes
print("{:<5} {}".format(self.__corpus.index(item), item))
def c_categorize(self, transaction_id, category):
transaction = self.__corpus[transaction_id]
self.__enrich_training_set(transaction, category)
self.__predict()
def c_debug(self):
print('\n***** Feature names *****')
print(self.__vectorizer.get_feature_names())
print('\n***** Training set str *****')
print(self.__training_set_str)
print('\n***** Training set x *****')
print(self.__training_set_x)
print('\n***** Training set y *****')
print(self.__training_set_y)
def help():
msg="""h Display this help message
o Display an overview of the expenses / incomes per category
p prediction_fname Write corpus with predicted categories appended to prediction_fname file
t trainingset_fname Write training set to trainingset_fname file
l category List entries in category
c item_id category Classify item with id item_id to category
q Quit"""
print(msg)
if len(sys.argv) < 3:
print("Syntax is: {} <training_set> <corpus_1> [<corpus_2> ...]".format(sys.argv[0]))
exit(1)
training_fname = sys.argv[1]
corpus_fnames = sys.argv[2:]
corp = Corpus(training_fname, corpus_fnames)
cmd = 'o'
while cmd != 'q':
try:
if cmd == '':
pass
elif cmd == 'o':
print("*** Overview")
corp.c_overview()
elif cmd == 'p':
out_fname = ln[1]
print("*** Writing predictions to '{}'".format(out_fname))
corp.c_save_prediction(out_fname)
elif cmd == 'l':
category = ln[1]
print("*** Listing items in category '{}'".format(category))
corp.c_list_category(category)
elif cmd == 'c':
item = int(ln[1])
cat = ln[2]
print("*** Classifying '{}' into '{}'".format(item, cat))
corp.c_categorize(item, cat)
elif cmd == 't':
out_fname = ln[1]
print("*** Saving training file to '{}'".format(out_fname))
corp.c_save_training(out_fname)
elif cmd == 'd':
print("*** Debugging")
corp.c_debug()
elif cmd == 'h':
print("*** Help")
help()
else:
print("*** Unknown command: {}".format(cmd))
help()
except Exception as e:
print("*** Error while processing command '{}'".format(cmd))
print(traceback.format_exc())
help()
try:
ln = input("> ").split()
cmd = ln[0]
except (EOFError, KeyboardInterrupt):
cmd = 'q'
except IndexError:
cmd = ''