spark机器学习从0到1特征抽取–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;块 

然后是获取spark

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

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); 

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

+-------------------------------------+ |text | +-------------------------------------+ |[Hi I heard about Spark] | |[I wish Java could use case classes] | |[Logistic regression models are neat]| +-------------------------------------+ 

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); 

我们看一下输出结果：

+-----------------------------------------+-------------------------------------------------------------------------------+ |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]| +--------------------------------------- -+-------------------------------------------------------------------------------+ 

我们可以看到，通过Word2VecModel将文档转换为向量，然后这个向量可以作为预测的特征，来计算文档相似度计算啦。