-
Notifications
You must be signed in to change notification settings - Fork 0
/
job.py
55 lines (48 loc) · 1.78 KB
/
job.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
import pandas as pd
from sklearn.feature_extraction.text import CountVectorizer
from mrjob.job import MRJob
from mrjob.protocol import JSONValueProtocol
import pickle
import re
import time
class SMSSpamDetection(MRJob):
OUTPUT_PROTOCOL = JSONValueProtocol
def preprocess_message(self, message):
message = message.lower()
message = re.sub(r'[^\w\s]', '', message)
message = re.sub(r'\s+', ' ', message)
return message
def mapper_init(self):
data = pd.read_csv('spam.csv', encoding = "ISO-8859-1")
with open('model.pkl', 'rb') as f:
self.model = pickle.load(f)
messages = data['v2']
preprocessed_messages = [self.preprocess_message(message) for message in messages]
self.vectorizer = CountVectorizer()
self.vectorizer.fit_transform(preprocessed_messages)
def mapper(self, _, message):
current_time = time.time()
preprocessed_message = self.preprocess_message(message)
features = self.vectorizer.transform([preprocessed_message])
proba = self.model.predict_proba(features).tolist()
yield "messages", [list(proba[0]), current_time, message]
def reducer(self, key, values):
output = []
for value in values:
record = value
current_time = time.time()
label = "spam"
proba = record[0]
mapTime = record[1]
if proba[0] < proba[1]:
label = "spam"
else:
label = "ham"
output.append({
"message": record[2],
"label": label,
"time": "{0:.2f}s".format(current_time - mapTime)
})
yield key, output
if __name__ == '__main__':
SMSSpamDetection.run()