Spark-MLlib 学习入门到掌握-Word2Vec分词向量[7]

word2vector 是google开源的一个生成词向量的工具，以语言模型为优化目标，迭代更新训练文本中的词向量，最终收敛获得词向量。词向量可以作为文本分析中重要的特征，在分类问题、标注问题等场景都有着重要的应用价值。

  def Word2VecTest(): Unit = {
    import org.apache.spark.ml.feature.Word2Vec
    import org.apache.spark.ml.linalg.Vector
    import org.apache.spark.sql.Row
    val spark: SparkSession = SparkSession.builder().appName("implicits").master("local[2]").getOrCreate()

    // 输入数据:每一行是一个单词数组。
    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")

    //从单词到向量的映射。
    val word2Vec = new Word2Vec()
      //输入映射列
      .setInputCol("text")
      //输出映射列
      .setOutputCol("result")
      //输入映射向量的维度
      .setVectorSize(3)
      //必须出现在word2vec模型的词汇表中最少的次数。
      .setMinCount(1)
    //拟合模型
    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")
    }
  }

运行结果