fixed model call so that uses type argument

rekhajoshm · Mar 5, 2015 · 900b586 · 900b586
1 parent ea09b28
commit 900b586
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/classification/NaiveBayes.scala b/mllib/src/main/scala/org/apache/spark/mllib/classification/NaiveBayes.scala
@@ -310,21 +310,21 @@ object NaiveBayes {
    *
    * The model type can be set to either Multinomial NB ([[http://tinyurl.com/lsdw6p]])
    * or Bernoulli NB ([[http://tinyurl.com/p7c96j6]]). The Multinomial NB can handle
-   * discrete count data and can be called by setting the model type to "Multinomial".
+   * discrete count data and can be called by setting the model type to "multinomial".
    * For example, it can be used with word counts or TF_IDF vectors of documents.
    * The Bernoulli model fits presence or absence (0-1) counts. By making every vector a
-   * 0-1 vector and setting the model type to "Bernoulli", the  fits and predicts as
+   * 0-1 vector and setting the model type to "bernoulli", the  fits and predicts as
    * Bernoulli NB.
    *
    * @param input RDD of `(label, array of features)` pairs.  Every vector should be a frequency
    *              vector or a count vector.
    * @param lambda The smoothing parameter
    *
    * @param modelType The type of NB model to fit from the enumeration NaiveBayesModels, can be
-   *              Multinomial or Bernoulli
+   *              multinomial or bernoulli
    */
   def train(input: RDD[LabeledPoint], lambda: Double, modelType: String): NaiveBayesModel = {
-    new NaiveBayes(lambda, Multinomial).run(input)
+    new NaiveBayes(lambda, MODELTYPE.fromString(modelType)).run(input)
   }
 
 

diff --git a/mllib/src/test/scala/org/apache/spark/mllib/classification/NaiveBayesSuite.scala b/mllib/src/test/scala/org/apache/spark/mllib/classification/NaiveBayesSuite.scala
@@ -124,7 +124,7 @@ class NaiveBayesSuite extends FunSuite with MLlibTestSparkContext {
     val testRDD = sc.parallelize(testData, 2)
     testRDD.cache()
 
-    val model = NaiveBayes.train(testRDD, 1.0, "Multinomial")
+    val model = NaiveBayes.train(testRDD, 1.0, "multinomial")
     validateModelFit(pi, theta, model)
 
     val validationData = NaiveBayesSuite.generateNaiveBayesInput(
@@ -161,7 +161,7 @@ class NaiveBayesSuite extends FunSuite with MLlibTestSparkContext {
     val testRDD = sc.parallelize(testData, 2)
     testRDD.cache()
 
-    val model = NaiveBayes.train(testRDD, 1.0, "Bernoulli")
+    val model = NaiveBayes.train(testRDD, 1.0, "bernoulli")
     validateModelFit(pi, theta, model)
 
     val validationData = NaiveBayesSuite.generateNaiveBayesInput(