-
Notifications
You must be signed in to change notification settings - Fork 0
/
core.py
43 lines (30 loc) · 1.62 KB
/
core.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
import findspark
findspark.init()
from pyspark.sql import SparkSession
import pyspark.sql.functions as F
spark = SparkSession.builder.appName('SMSSpamCollection').getOrCreate()
# load preprocessor
from pyspark.ml import PipelineModel
spam_cleaner_loaded = PipelineModel.load("spam_cleaner_pipeline")
# load model
from pyspark.ml.classification import NaiveBayesModel
predictor_loaded = NaiveBayesModel.load("spam_classifier_model")
def predict(text):
'''
Input:
text: a string that contains the text to be predicted whether it is spam or not.
Output:
prediction: an integer value that represents the predicted class label of the input text.
0 indicates that the input text is not spam, while 1 indicates that it is spam.
'''
# preprocess the text
# text = "URGENT! You have won a 1 week FREE membership in our £100,000 Prize Jackpot! Txt the word: CLAIM to No: 81010"
# input_str = "We would like to invite you and your families to the annual company picnic that will be held on Saturday, June 15th, at the park near the office. The picnic will start at 11 am and end at 3 pm."
input_df = spark.createDataFrame([(text,)], ["text"])
input_df = input_df.withColumn("textLength", F.length("text"))
clean_text = spam_cleaner_loaded.transform(input_df)
# predict the class label
prediction = predictor_loaded.transform(clean_text)
prediction = int(prediction.select("prediction").first()[0])
return prediction ## 0 = Not Spam // 1 = Spam
# predict("URGENT! You have won a 1 week FREE membership in our £100,000 Prize Jackpot! Txt the word: CLAIM to No: 81010")