Spark MLlib 特征工程系列—特征提取Word2Vec

Spark MLlib 特征工程系列—特征提取Word2Vec

Word2Vec 是一种流行的词嵌入方法，用于将词语转换为低维的向量表示，这些向量能够捕捉词语之间的语义关系。Word2Vec 最初由 Google 的研究团队开发，其核心思想是通过训练模型来学习词语的分布式表示，使得语义相似的词语在向量空间中靠得更近。

案例演示

在 Spark 中，Word2Vec 也作为一个工具类存在，可以将文本数据中的词汇转化为低维向量。

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