实现一个CountVectorizer

最近在一个nlp问题中使用了sklearn的CountVectorizer库进行分词，目的是对一个多值离散型特征进行编码并转换成稀疏矩阵（csr_matrix），使用过程中发现CountVectorizer的速度非常慢，相当的耗时，因此决定提取最核心的功能，实现一个自己的版本，只要能实现相同的功能并更加节省时间即可。

提取CountVectorizer的功能

我需要实现的只有两个函数，fit和transform。为了准确实现自己的版本，我首先需要分析CountVectorizer中这两个函数的功能。
对于一个由字符串构成的数组，每个元素可能是一个以空格分割的句子（sentence），CountVectorizer.fit的功能是将它们分割，为每一个单词（word）编码，在这个过程中会自动滤除停止词（stop words），例如英文句子中的”a”，”.”之类的长度为1的字符串。CountVectorizer.transform的功能则是将输入的数组中每个元素进行分割，然后使用fit中生成的编码字典，将原单词转化成编码，数据以csr_matrix的形式返回。

自己实现的版本

# coding:utf-8
import numpy as np
from scipy import sparse

class MyCountVectorizer():
    def __init__(self, pass_stop=True):
        self.pass_stop = pass_stop # 提供停止词滤除功能，可禁止

    def fit(self, data):
        data = map(lambda x:str(x).split(" "), data)
        self.elements_ = set()
        for line in data:
            for x in line:
                if self.pass_stop:
                    if len(x)==1:
                        continue
                self.elements_.add(x)
        # 原元素
        self.elements_ = np.sort(list(self.elements_))
        # 编码
        self.labels_ = np.arange(len(self.elements_)).astype(int)
        # 生成字典
        self.dict_ = {}
        for i in range(len(self.elements_)):
            self.dict_[str(self.elements_[i])] = self.labels_[i]

    def transform(self, data):
        rows = []
        cols = []
        data = map(lambda x:str(x).split(" "), data)
        for i in range(len(data)):
            for x in data[i]:
                if self.pass_stop:
                    if len(x)==1:
                        continue
                rows.append(i)
                cols.append(self.dict_[x])
        vals = np.ones((len(rows),)).astype(int)

        return sparse.csr_matrix((vals, (rows, cols)), shape=(len(data), len(self.labels_)))

为什么`vals`可以直接用`ones`

可能有人会有疑问，如果一个句子中有两个相同的单词怎么办？例如”one two one”，实际上sparse.csc_matrix遇到相同两个相同位置的记录，会自动将值进行累加。为了解释这个问题，我再做一个实验，如下所示，这是一个对角矩阵的稀疏表达，在对角线上的值原本都是1，但是我把坐标(2,2)上的值重复了一次。

from scipy import sparse

rows = [0,1,2,2]
cols = [0,1,2,2]
vals = [1,1,1,1]
mat = sparse.csr_matrix((vals,(rows,cols)), shape=(3,3))
print mat.toarray()

结果如下，可见对位置重复的记录，sparse.csr_matrix会自动累加到数值上。

[[1 0 0]
 [0 1 0]
 [0 0 2]]

时间性能测试

接下来应该测试一下我的版本是否能达到要求，即既能准确计算，又能节省时间。test.csv是我伪造的一个数据集，对它简单进行重复，分别统计两种方法的时间，同时获取两者生成的数据结果，做差保证两者一致。

from sklearn.feature_extraction.text import CountVectorizer
import pandas as pd
from scipy import sparse
from MyCountVectorizer import MyCountVectorizer
import time
import numpy as np
from matplotlib import pyplot as plt

data = pd.read_csv("data/test.csv")

records = []
for n in range(100, 10100, 100):
    seq = np.tile(data['interest'].values, (n,1)).ravel()

    # CountVectorizer in sklearn
    start1 = time.time()
    cv = CountVectorizer()
    cv.fit(seq)
    res1 = cv.transform(seq).toarray()
    end1 = time.time()

    # MyCountVectorizer
    start2 = time.time()
    mcv = MyCountVectorizer()
    mcv.fit(seq)
    res2 = mcv.transform(seq).toarray()
    end2 = time.time()

    # 保证两种方法生成的结果一致
    assert np.sum(res1-res2) == 0

    records.append((n, end1-start1, end2-start2))
records = np.array(records)

plt.plot(records[:,0], records[:,1], label="CountVectorizer")
plt.plot(records[:,0], records[:,2], label="MyCountVectorizer")
plt.legend()
plt.savefig('pictures/result.png')

这里写图片描述

可见，我的版本既节省时间，计算得到的结果也与sklearn的版本一致。

完整的代码在：https://github.com/SongDark/MyCountVectorizer

参考资料

【API】scipy.sparse.csr_matrix
【API】sklearn.feature_extraction.text.CountVectorizer
【StackOverFlow】CountVectorizer: AttributeError: ‘numpy.ndarray’ object has no attribute ‘lower’