使用Python 2.7实现的垃圾短信识别器

　　最近参加比赛，写了一个垃圾短信识别器，在这里做一下记录。

　　官方提供的数据是csv文件，其中训练集有80万条数据，测试集有20万条数据，训练集的格式为：行号 标记(0为普通短信，1为垃圾短信) 短信内容；测试集的格式为： 行号 短信内容；要求输出的数据格式要求为： 行号 标记，以csv格式保存。

　　实现的原理可概括为以下几步：

　　　　1.读取文件，输入数据

　　　　2.对数据进行分割，将每一行数据分成行号、标记、短信内容。由于短信内容中可能存在空格，故不能简单地用split()分割字符串，应该用正则表达式模块re进行匹配分割。

　　　　3.将分割结果存入数据库(MySQL)，方便下次测试时直接从数据库读取结果，省略步骤。

　　　　4.对短信内容进行分词，这一步用到了第三方库结巴分词：https://github.com/fxsjy/jieba

　　　　5.将分词的结果用于训练模型，训练的算法为朴素贝叶斯算法，可调用第三方库Scikit-Learn：http://scikit-learn.org/stable

　　　　6.从数据库中读取测试集，进行判断，输出结果并写入文件。

　　最终实现出来一共有4个py文件：

　　　　1.ImportIntoDB.py 将数据进行预处理并导入数据库，仅在第一次使用。

　　　　2.DataHandler.py 从数据库中读取数据，进行分词，随后处理数据，训练模型。

　　　　3.Classifier.py 从数据库中读取测试集数据，利用训练好的模型进行判断，输出结果到文件中。

　　　　4.Main.py 程序的入口

　


　　最终程序每次运行耗时平均在260秒-270秒之间，附代码：

　　ImportIntoDB.py：

# -*- coding:utf-8 -*-
__author__ = 'Jz'

import MySQLdb
import codecs
import re
import time

# txt_path = 'D:/coding_file/python_file/Big Data/trash message/train80w.txt'
txt_path = 'D:/coding_file/python_file/Big Data/trash message/test20w.txt'

# use regular expression to split string into parts
# split_pattern_80w = re.compile(u'([0-9]+).*?([01])(.*)')
split_pattern_20w = re.compile(u'([0-9]+)(.*)')

txt = codecs.open(txt_path, 'r')
lines = txt.readlines()
start_time = time.time()

#connect mysql database
con = MySQLdb.connect(host = 'localhost', port = 3306, user = 'root', passwd = '*****', db = 'TrashMessage', charset = 'UTF8')
cur = con.cursor()

# insert into 'train' table
# sql = 'insert into train(sms_id, sms_type, content) values (%s, %s, %s)'
# for line in lines:
#     match = re.match(split_pattern_80w, line)
#     sms_id, sms_type, content = match.group(1), match.group(2), match.group(3).lstrip()
#     cur.execute(sql, (sms_id, sms_type, content))
#     print sms_id
# # commit transaction
# con.commit()

# insert into 'test' table
sql = 'insert into test(sms_id, content) values (%s, %s)'
for line in lines:
match = re.match(split_pattern_20w, line)
sms_id, content = match.group(1), match.group(2).lstrip()
cur.execute(sql, (sms_id, content))
print sms_id
# commit transaction
con.commit()

cur.close()
con.close()
txt.close()
end_time = time.time()
print 'time-consuming: ' + str(end_time - start_time) + 's.'

　　DataHandler.py：

# -*- coding:utf-8 -*-
__author__ = 'Jz'

import MySQLdb
import jieba
import re

class DataHandler:
def __init__(self):
try:
self.con = MySQLdb.connect(host = 'localhost', port = 3306, user = 'root', passwd = '*****', db = 'TrashMessage', charset = 'UTF8')
self.cur = self.con.cursor()
except MySQLdb.OperationalError, oe:
print 'Connection error! Details:', oe

def __del__(self):
self.cur.close()
self.con.close()

# obsolete function
# def getConnection(self):
#     return self.con

# obsolete function
# def getCursor(self):
#     return self.cur

def query(self, sql):
self.cur.execute(sql)
result_set = self.cur.fetchall()
return result_set

def resultSetTransformer(self, train, test):
# list of words divided by jieba module after de-duplication
train_division = []
test_division = []
# list of classification of each message
train_class = []
# divide messages into words
for record in train:
train_class.append(record[1])
division = jieba.cut(record[2])
filtered_division_set = set()
for word in division:
filtered_division_set.add(word + ' ')
division = list(filtered_division_set)
str_word = ''.join(division)
train_division.append(str_word)

# handle test set in a similar way as above
for record in test:
division = jieba.cut(record[1])
filtered_division_set = set()
for word in division:
filtered_division_set.add(word + ' ')
division = list(filtered_division_set)
str_word = ''.join(division)
test_division.append(str_word)

return train_division, train_class, test_division

　　Classifier.py：

# -*- coding:utf-8 -*-
__author__ = 'Jz'

from DataHandler import DataHandler
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
import time

class Classifier:
def __init__(self):
start_time = time.time()
self.data_handler = DataHandler()
# get result set
self.train = self.data_handler.query('select * from train')
self.test = self.data_handler.query('select * from test')
self.train_division, self.train_class, self.test_division = self.data_handler.resultSetTransformer(self.train, self.test)
end_time = time.time()
print 'Classifier finished initializing, time-consuming:' + str(end_time - start_time) + 's.'

def getMatrices(self):
start_time = time.time()
# convert a collection of raw documents to a matrix of TF-IDF features.
self.tfidf_vectorizer = TfidfVectorizer()
# learn vocabulary and idf, return term-document matrix [sample, feature]
self.train_count_matrix = self.tfidf_vectorizer.fit_transform(self.train_division)
# transform the count matrix of the train set to a normalized tf-idf representation
self.tfidf_transformer = TfidfTransformer()
self.train_tfidf_matrix = self.tfidf_transformer.fit_transform(self.train_count_matrix)
end_time = time.time()
print 'Classifier finished getting matrices, time-consuming:' + str(end_time - start_time) + 's.'

def classify(self):
self.getMatrices()
start_time = time.time()
# convert a collection of text documents to a matrix of token counts
# scikit-learn doesn't support chinese vocabulary
test_tfidf_vectorizer = CountVectorizer(vocabulary = self.tfidf_vectorizer.vocabulary_)
# learn the vocabulary dictionary and return term-document matrix.
test_count_matrix = test_tfidf_vectorizer.fit_transform(self.test_division)
# transform a count matrix to a normalized tf or tf-idf representation
test_tfidf_transformer = TfidfTransformer()
test_tfidf_matrix = test_tfidf_transformer.fit(self.train_count_matrix).transform(test_count_matrix)

# the multinomial Naive Bayes classifier is suitable for classification with discrete features
# e.g., word counts for text classification).
naive_bayes = MultinomialNB(alpha = 0.65)
naive_bayes.fit(self.train_tfidf_matrix, self.train_class)
prediction = naive_bayes.predict(test_tfidf_matrix)

# output result to a csv file
index = 0
csv = open('result.csv', 'w')
for sms_type in prediction:
csv.write(str(self.test[index][0]) + ',' + str(sms_type) + '\n')
index += 1
csv.close()
end_time = time.time()
print 'Classifier finished classifying, time-consuming: ' + str(end_time - start_time) + 's.'

　　Main.py：

# -*- coding:utf-8 -*-
__author__ = 'Jz'

import time
from Classifier import Classifier

start_time = time.time()
classifier = Classifier()
classifier.classify()
end_time = time.time()
print 'total time-consuming: ' + str(end_time - start_time) + 's.'