Spark MLlib机器学习开发指南(5)--特征提取--Word2Vec

Spark MLlib机器学习开发指南(5)--特征提取，转换，选择--Word2Vec

翻译自Word2Vec,基于最新2.2.0版本翻译,转载注明出处 xcrossed 机器学习

Word2Vec

在spark ml中，Word2Vec是一个估计器（前面说过估计器和转换器的概念了，可以往回看具体概念），由表示文档的单词序列训练而成的一个Word2VecModel。模型映射每个单词为一个唯一固定大小的向量。Word2VecModel使用文档中所有单词的平均值将每个文档转换成一个向量，这个向量可以作为预测的特征，文档相似性计算等等。请参阅Word2VecMLlib的用户指南，以了解更多细节。

示例
在下面的代码片断，我们一个文档集合开始，每个文档由一序列的单词表示。每个文档我们将转换成一个特征向量，这个特征向量可以被传递给一个学习算法。

详细API请参考Word2Vec Scala docs

import org.apache.spark.ml.feature.Word2Vec
import org.apache.spark.ml.linalg.Vector
import org.apache.spark.sql.Row

// Input data: Each row is a bag of words from a sentence or document.
val documentDF = spark.createDataFrame(Seq(
  "Hi I heard about Spark".split(" "),
  "I wish Java could use case classes".split(" "),
  "Logistic regression models are neat".split(" ")
).map(Tuple1.apply)).toDF("text")

// Learn a mapping from words to Vectors.
val word2Vec = new Word2Vec()
  .setInputCol("text")
  .setOutputCol("result")
  .setVectorSize(3)
  .setMinCount(0)
val model = word2Vec.fit(documentDF)

val result = model.transform(documentDF)
result.collect().foreach { case Row(text: Seq[_], features: Vector) =>
  println(s"Text: [${text.mkString(", ")}] => \nVector: $features\n") }