add tests and Word2VecModelWrapper

Author: Ishiihara

mllib/src/main/scala/org/apache/spark/mllib/api/python/PythonMLLibAPI.scala

@@ -29,6 +29,8 @@ import org.apache.spark.annotation.DeveloperApi
 import org.apache.spark.api.java.{JavaRDD, JavaSparkContext}
 import org.apache.spark.mllib.classification._
 import org.apache.spark.mllib.clustering._
+import org.apache.spark.mllib.feature.Word2Vec
+import org.apache.spark.mllib.feature.Word2VecModel
 import org.apache.spark.mllib.optimization._
 import org.apache.spark.mllib.linalg._
 import org.apache.spark.mllib.random.{RandomRDDs => RG}
@@ -40,8 +42,6 @@ import org.apache.spark.mllib.tree.impurity._
 import org.apache.spark.mllib.tree.model.DecisionTreeModel
 import org.apache.spark.mllib.stat.{MultivariateStatisticalSummary, Statistics}
 import org.apache.spark.mllib.stat.correlation.CorrelationNames
-import org.apache.spark.mllib.feature.Word2Vec
-import org.apache.spark.mllib.feature.Word2VecModel
 import org.apache.spark.mllib.util.MLUtils
 import org.apache.spark.rdd.RDD
 import org.apache.spark.util.Utils
@@ -290,72 +290,34 @@ class PythonMLLibAPI extends Serializable {
    * Extra care needs to be taken in the Python code to ensure it gets freed on
    * exit; see the Py4J documentation.
    * @param dataJRDD Input JavaRDD
-   * @return A handle to java Word2VecModel instance at python side
+   * @return A handle to java Word2VecModelWrapper instance at python side
    */
-  def trainWord2Vec(
-    dataJRDD: JavaRDD[java.util.ArrayList[String]]
-    ): Word2VecModel = {
-    val data = dataJRDD.rdd.map(_.toArray(new Array[String](0)).toSeq).cache()
+  def trainWord2Vec(dataJRDD: JavaRDD[java.util.ArrayList[String]]): Word2VecModelWrapper = {
+    val data = dataJRDD.rdd.cache()
     val word2vec = new Word2Vec()
     val model = word2vec.fit(data)
-    model
+    new Word2VecModelWrapper(model)
   }
 
-  /**
-   * Java stub for Python mllib Word2VecModel transform
-   * @param model Word2VecModel instance
-   * @param word a word
-   * @return serialized vector representation of word
-   */
-  def Word2VecModelTransform(
-    model: Word2VecModel,
-    word: String
-    ): Vector = {
-    model.transform(word)
-  }
+  private[python] class Word2VecModelWrapper(model: Word2VecModel) {
+    def transform(word: String): Vector = {
+      model.transform(word)
+    }
 
-  /**
-   * Java stub for Python mllib Word2VecModel findSynonyms
-   * @param model Word2VecModel instance
-   * @param word a word
-   * @param num number of synonyms to find
-   * @return a java LinkedList containing serialized version of
-   * synonyms and similarities
-   */
-  def Word2VecModelSynonyms(
-    model: Word2VecModel,
-    word: String,
-    num: Int
-    ): java.util.List[java.lang.Object] = {
-    val result = model.findSynonyms(word, num)
-    val similarity = Vectors.dense(result.map(_._2))
-    val words = result.map(_._1)
-    val ret = new java.util.LinkedList[java.lang.Object]()
-    ret.add(words)
-    ret.add(similarity)
-    ret
-  }
+    def findSynonyms(word: String, num: Int): java.util.List[java.lang.Object] = {
+      val vec = transform(word)
+      findSynonyms(vec, num)
+    }
 
-  /**
-   * Java stub for Python mllib Word2VecModel findSynonyms
-   * @param model Word2VecModel instance
-   * @param vecBytes serialization of vector representation of words
-   * @param num number of synonyms to find
-   * @return a java LinkedList containing serialized version of
-   * synonyms and similarities
-   */
-  def Word2VecModelSynonyms(
-    model: Word2VecModel,
-    vec: Vector,
-    num: Int
-    ): java.util.List[java.lang.Object] = {
-    val result = model.findSynonyms(vec, num)
-    val similarity = Vectors.dense(result.map(_._2))
-    val words = result.map(_._1)
-    val ret = new java.util.LinkedList[java.lang.Object]()
-    ret.add(words)
-    ret.add(similarity)
-    ret
+    def findSynonyms(vector: Vector, num: Int): java.util.List[java.lang.Object] = {
+      val result = model.findSynonyms(vector, num)
+      val similarity = Vectors.dense(result.map(_._2))
+      val words = result.map(_._1)
+      val ret = new java.util.LinkedList[java.lang.Object]()
+      ret.add(words)
+      ret.add(similarity)
+      ret
+    }
   }
 
   /**

python/pyspark/mllib/Word2Vec.py

@@ -19,8 +19,6 @@
 Python package for Word2Vec in MLlib.
 """
 
-from functools import wraps
-
 from pyspark import PickleSerializer
 
 from pyspark.mllib.linalg import _convert_to_vector
@@ -44,21 +42,13 @@ def __del__(self):
         self._sc._gateway.detach(self._java_model)
 
     def transform(self, word):
-        pythonAPI = self._sc._jvm.PythonMLLibAPI()
-        result = pythonAPI.Word2VecModelTransform(self._java_model, word)
+        result = self._java_model.transform(word)
         return PickleSerializer().loads(str(self._sc._jvm.SerDe.dumps(result)))
 
     def findSynonyms(self, x, num):
-        SerDe = self._sc._jvm.SerDe
-        ser = PickleSerializer()
-        pythonAPI = self._sc._jvm.PythonMLLibAPI()
-        if type(x) == str:
-            jlist = pythonAPI.Word2VecModelSynonyms(self._java_model, x, num)
-        else:
-            bytes = bytearray(ser.dumps(_convert_to_vector(x)))
-            vec = self._sc._jvm.SerDe.loads(bytes)
-            jlist = pythonAPI.Word2VecModelSynonyms(self._java_model, vec, num)
-        return PickleSerializer().loads(str(self._sc._jvm.SerDe.dumps(jlist)))
+        jlist = self._java_model.findSynonyms(x, num)
+        words, similarity = PickleSerializer().loads(str(self._sc._jvm.SerDe.dumps(jlist)))
+        return zip(words, similarity)
 
 
 class Word2Vec(object):
@@ -77,12 +67,22 @@ class Word2Vec(object):
     Efficient Estimation of Word Representations in Vector Space
     and
     Distributed Representations of Words and Phrases and their Compositionality.
+    >>> sentence = "a b " * 100 + "a c " * 10
+    >>> localDoc = [sentence, sentence]
+    >>> doc = sc.parallelize(localDoc).map(lambda line: line.split(" "))
+    >>> model = Word2Vec().setVectorSize(10).setSeed(42L).fit(doc)
+    >>> syms = model.findSynonyms("a", 2)
+    >>> str(syms[0][0])
+    'b'
+    >>> str(syms[1][0])
+    'c'
     """
     def __init__(self):
         self.vectorSize = 100
         self.startingAlpha = 0.025
         self.numPartitions = 1
         self.numIterations = 1
+        self.seed = 42L
 
     def setVectorSize(self, vectorSize):
         self.vectorSize = vectorSize
@@ -100,10 +100,11 @@ def setNumIterations(self, numIterations):
         self.numIterations = numIterations
         return self
 
+    def setSeed(self, seed):
+        self.seed = seed
+        return self
+
     def fit(self, data):
-        """
-        :param data: Input RDD
-        """
         sc = data.context
         model = sc._jvm.PythonMLLibAPI().trainWord2Vec(data._to_java_object_rdd())
         return Word2VecModel(sc, model)