机器学习之特征抽取–Word2Vec_使用word2vec提取文本特征代码

一、概念

Word2vec是一个Estimator，它采用一系列代表文档的词语来训练word2vecmodel。该模型将每个词语映射到一个固定大小的向量。word2vecmodel使用文档中每个词语的平均数来将文档转换为向量，然后这个向量可以作为预测的特征，来计算文档相似度计算等等。

二、代码实现

2.1、引包，获取spark

首先，我们引入相关包：

import java.util.Arrays;
import java.util.List;
import org.apache.spark.ml.feature.Word2Vec;
import org.apache.spark.ml.feature.Word2VecModel;
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;块
1
2
3
4
5
6
7
8
9
10
11
12
13

然后是获取spark

SparkSession spark =  SparkSession.builder().appName("Word2VecTest").master("local").getOrCreate();
1

2.2、构建测试数据

接着呢来构建一个DataFrame，往DateFrame里加一些测试的文档信息

List<Row> rawData =  Arrays.asList(RowFactory.create(Arrays.asList("Hi I heard about Spark".split(","))),
                        RowFactory.create(Arrays.asList("I wish Java  could use case classes".split(","))),
                        RowFactory.create(Arrays.asList("Logistic  regression models are neat".split(","))));
StructType schema = new StructType(new  StructField[] {
        new StructField("text",new  ArrayType(DataTypes.StringType,false),false,Metadata.empty())
});

Dataset<Row> documentDF =  spark.createDataFrame(rawData,schema);
documentDF.show(false);
1
2
3
4
5
6
7
8
9

我们来查看一下控制台的输出结果：

+-------------------------------------+
|text                                 |
+-------------------------------------+
|[Hi I heard about Spark]             |
|[I wish Java could use case classes] |
|[Logistic regression models are neat]|
+-------------------------------------+
1
2
3
4
5
6
7

2.3、新建评估器，训练，转换得到向量

接下来我们新建一个Word2Vec的评估器，把单词和向量建立一个映射，设定输入为文本信息text,输出为追加列result,变量的大小为3，最小计数为0。建立完之后，用Word2Vec评估器对文档进行训练和转换，得到Dataset的数据集。

Word2Vec word2Vec = new  Word2Vec().setInputCol("text")
                                   .setOutputCol("result")
                                   .setVectorSize(3)
                                   .setMinCount(0);

Word2VecModel model =  word2Vec.fit(documentDF);
Dataset<Row> result =  model.transform(documentDF);
result.show(false);
1
2
3
4
5
6
7
8

我们看一下输出结果：

+-----------------------------------------+-------------------------------------------------------------------------------+
|text                                                        |result                                                                                                            |
+-----------------------------------------+-------------------------------------------------------------------------------+
|[Hi I heard about Spark]                       |[-0.12674053013324738,0.09846510738134384,-0.10375533252954483] |
|[I wish Java could use case classes]      |[-0.1633371263742447,-0.14517612755298615,0.11354436725378036]   |
|[Logistic regression models are  neat] |[-0.019123395904898643,-0.13107778131961823,0.14307855069637299]|
+--------------------------------------- -+-------------------------------------------------------------------------------+
1
2
3
4
5
6
7

我们可以看到，通过Word2VecModel将文档转换为向量，然后这个向量可以作为预测的特征，来计算文档相似度计算啦。
参考资料：http://spark.apache.org/docs/latest/ml-features.html#word2vec