特征抽取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]]