使用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) + '
')
            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.'