机器学习之特征抽取–CountVectorizer

一、概念

CountVectorizer 旨在通过计数来将一个文档转换为向量。当不存在先验字典时，Countvectorizer作为Estimator提取词汇进行训练，并生成一个CountVectorizerModel用于存储相应的词汇向量空间。该模型产生文档关于词语的稀疏表示，其表示可以传递给其他算法，例如LDA（ Latent Dirichlet Allocation 隐含狄利克雷分布）。

在CountVectorizerModel的训练过程中，CountVectorizer将根据语料库中的词频排序从高到低进行选择，词汇表的最大含量由vocabsize超参数来指定，超参数minDF则指定词汇表中的词语至少要在多少个不同文档中出现。

二、代码实现

2.1、构造文档集合

import java.util.Arrays;
import java.util.List;
import org.apache.spark.ml.feature.CountVectorizer;
import org.apache.spark.ml.feature.CountVectorizerModel;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.RowFactory;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.types.ArrayType;
import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.sql.types.Metadata;
import org.apache.spark.sql.types.StructField;
import org.apache.spark.sql.types.StructType;

 //获取spark
SparkSession spark = SparkSession.builder().appName("CountVectorizerModel").master("local").getOrCreate();


//获取数据 DataFrame
List<Row> rawData = Arrays.asList(RowFactory.create(0, new String[] {"a", "b", "c"}),
                        RowFactory.create(1, new String[] {"a", "b", "b", "c", "a"}));
StructType schema = new StructType(new StructField[] {
        new StructField("id",DataTypes.IntegerType,false,Metadata.empty()),
        new StructField("words",new ArrayType(DataTypes.StringType,true),false,Metadata.empty())
});
Dataset<Row> data = spark.createDataFrame(rawData, schema);
data.show(false);

输出结果：

+---+---------------+
|id |words          |
+---+---------------+
|0  |[a, b, c]      |
|1  |[a, b, b, c, a]|
+---+---------------+

2.2、设定参数，训练模型

通过CountVectorizer设定超参数，训练一个CountVectorizerModel，这里设定词汇表的最大量为3，设定词汇表中的词至少要在2个文档中出现过，以过滤那些偶然出现的词汇。

CountVectorizerModel cvModel = new  CountVectorizer().setInputCol("words")
                                                        .setOutputCol("features")
                                                        .setVocabSize(3)
                                                        .setMinDF(2)
                                                        .fit(data);
String[] vocabulary =  cvModel.vocabulary();

在训练结束后，可以通过CountVectorizerModel的vocabulary成员获得到模型的词汇表。

2.3、获取文档向量

使用这一模型对DataFrame进行变换，可以得到文档的向量化表示：

cvModel.transform(data).show(false);

输出结果：

+---+---------------+-------------------------+
|id |words          |features                  |
+---+---------------+-------------------------+
|0  |[a, b, c]       |(3,[0,1,2],[1.0,1.0,1.0])|
|1  |[a, b, b, c,  a]|(3,[0,1,2],[2.0,2.0,1.0])|
+---+---------------+-------------------------+

和其他Transformer不同，CountVectorizerModel可以通过指定一个先验词汇表来直接生成，如以下例子，直接指定词汇表的成员是“a”，“b”两个个词：

CountVectorizerModel cvm = new CountVectorizerModel(new String[] {"a", "b"}).setInputCol("words")
                                                                            .setOutputCol("features");
cvm.transform(data).show(false);

输出结果：

+---+---------------+-------------------+
|id |words          |features           |
+---+---------------+-------------------+
|0  |[a, b, c]      |(2,[0,1],[1.0,1.0])|
|1  |[a, b, b, c, a]|(2,[0,1],[2.0,2.0])|
+---+---------------+-------------------+

参考资料： http://spark.apache.org/docs/latest/ml-features.html#countvectorizer