Spark大数据-特征抽取Word2Vec(Scala版)

特征抽取Word2Vec(Scala版)

import org.apache.spark.sql.SparkSession
val spark=SparkSession.builder().master("local").appName("word2vec").getOrCreate()
import spark.implicits._
import org.apache.spark.ml.feature.Word2Vec

// 我们首先用一组文档，其中一个词语序列代表一个文档。
// 对于每一个文档，我们将其转换为一个特征向量。
// 此特征向量可以被传递到一个学习算法。
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")
// 上述数组需要Tuple1.apply，如果list的话就不需要
documentDF.show()

// 创建一个word2vec
val word2vec=new Word2Vec().
    setInputCol("text").
    setOutputCol("result").
    setVectorSize(3).//这里设置特征向量的维度为3
    setMinCount(0) //单词0次以上列入统计
val model=word2vec.fit(documentDF)

val result=model.transform(documentDF)
result.select("result").take(3).foreach(println)

结果为：
[[-0.028139343485236168,0.04554025698453188,-0.013317196490243079]]
[[0.06872416580361979,-0.02604914902310286,0.02165239889706884]]
[[0.023467857390642166,0.027799883112311366,0.0331136979162693]]