[SPARK-10299][ML] word2vec should allow users to specify the window size

mllib/src/main/scala/org/apache/spark/ml/feature/Word2Vec.scala

@@ -48,6 +48,17 @@ private[feature] trait Word2VecBase extends Params
   /** @group getParam */
   def getVectorSize: Int = $(vectorSize)
 
+  /**
+   * The window size (context words from [-window, window]).
+   * @group param
+   */
+  final val windowSize = new IntParam(
+    this, "windowSize", "the window size (context words from [-window, window])")
+  setDefault(windowSize -> 5)
+
+  /** @group getParam */
+  def getWindowSize: Int = $(windowSize)
+
   /**
    * Number of partitions for sentences of words.
    * @group param
@@ -102,6 +113,9 @@ final class Word2Vec(override val uid: String) extends Estimator[Word2VecModel]
   /** @group setParam */
   def setVectorSize(value: Int): this.type = set(vectorSize, value)
 
+  /** @group setParam */
+  def setWindowSize(value: Int): this.type = set(windowSize, value)
+
   /** @group setParam */
   def setStepSize(value: Double): this.type = set(stepSize, value)
 
@@ -127,6 +141,7 @@ final class Word2Vec(override val uid: String) extends Estimator[Word2VecModel]
       .setNumPartitions($(numPartitions))
       .setSeed($(seed))
       .setVectorSize($(vectorSize))
+      .setWindowSize($(windowSize))
       .fit(input)
     copyValues(new Word2VecModel(uid, wordVectors).setParent(this))
   }

mllib/src/main/scala/org/apache/spark/mllib/feature/Word2Vec.scala

@@ -128,6 +128,15 @@ class Word2Vec extends Serializable with Logging {
     this
   }
 
+  /**
+   * Sets the window of words (default: 5)
+   */
+  @Since("1.6.0")
+  def setWindowSize(window: Int): this.type = {
+    this.window = window
+    this
+  }
+
   /**
    * Sets minCount, the minimum number of times a token must appear to be included in the word2vec
    * model's vocabulary (default: 5).
@@ -144,7 +153,7 @@ class Word2Vec extends Serializable with Logging {
   private val MAX_SENTENCE_LENGTH = 1000
 
   /** context words from [-window, window] */
-  private val window = 5
+  private var window = 5
 
   private var trainWordsCount = 0
   private var vocabSize = 0

mllib/src/test/scala/org/apache/spark/ml/feature/Word2VecSuite.scala

@@ -131,7 +131,42 @@ class Word2VecSuite extends SparkFunSuite with MLlibTestSparkContext {
     expectedSimilarity.zip(similarity).map {
       case (expected, actual) => assert(math.abs((expected - actual) / expected) < 1E-5)
     }
+  }
+
+  test("window size") {
+
+    val sqlContext = new SQLContext(sc)
+    import sqlContext.implicits._
+
+    val sentence = "a q s t q s t b b b s t m s t m q " * 100 + "a c " * 10
+    val doc = sc.parallelize(Seq(sentence, sentence)).map(line => line.split(" "))
+    val docDF = doc.zip(doc).toDF("text", "alsotext")
 
+    val model = new Word2Vec()
+      .setVectorSize(3)
+      .setWindowSize(2)
+      .setInputCol("text")
+      .setOutputCol("result")
+      .setSeed(42L)
+      .fit(docDF)
+
+    val (synonyms, similarity) = model.findSynonyms("a", 6).map {
+      case Row(w: String, sim: Double) => (w, sim)
+    }.collect().unzip
+
+    // Increase the window size
+    val biggerModel = new Word2Vec()
+      .setVectorSize(3)
+      .setInputCol("text")
+      .setOutputCol("result")
+      .setSeed(42L)
+      .setWindowSize(10)
+      .fit(docDF)
+
+    val (synonymsLarger, similarityLarger) = model.findSynonyms("a", 6).map {
+      case Row(w: String, sim: Double) => (w, sim)
+    }.collect().unzip
+    // The similarity score should be very different with the larger window
+    assert(math.abs(similarity(5) - similarityLarger(5) / similarity(5)) > 1E-5)
   }
 }
-