![ab443dc9a285a504e978ecc86d7593d9.png](https://i-blog.csdnimg.cn/blog_migrate/17b66ce3e7e22782229255cc3235ba86.png)
同步更新在个人网站:http://www.wangpengcufe.com/machinelearning/ml-ml14/
一、概念
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将文档转换为向量,然后这个向量可以作为预测的特征,来计算文档相似度计算啦。
参考资料:http://spark.apache.org/docs/latest/ml-features.html#word2vec