word2vec.scala

import org.apache.spark._
import org.apache.spark.mllib.feature.{Word2Vec, Word2VecModel}

object Word2VecPerfTest extends App {

  val start = System.currentTimeMillis()

  val conf = new SparkConf().setMaster("local[*]").setAppName("word2vec")

  val sc = new SparkContext(conf)

  val input = sc.textFile("text8").map(line => line.split(" ").toSeq)

  val word2vec = new Word2Vec().setNumPartitions(20)

  val model = word2vec.fit(input)

  val synonyms = model.findSynonyms("china", 40)

  for((synonym, cosineSimilarity) <- synonyms) {
    println(s"$synonym $cosineSimilarity")
  }

  // Save and load model
  model.save(sc, "result")
  val sameModel = Word2VecModel.load(sc, "result")

  val end = System.currentTimeMillis()

  val time = (end - start) / 1000

  println(time + " seconds")
}