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") }