fixed model call so that uses type argument

leahmcguire · leahmcguire · commit 900b5864c16c · 2015-03-05T11:53:46.000-08:00
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(

Original file line number	Diff line number	Diff line change
`@@ -310,21 +310,21 @@ object NaiveBayes {`
`310`	`310`	`*`
`311`	`311`	`* The model type can be set to either Multinomial NB ([[http://tinyurl.com/lsdw6p]])`
`312`	`312`	`* or Bernoulli NB ([[http://tinyurl.com/p7c96j6]]). The Multinomial NB can handle`
`313`		`- * discrete count data and can be called by setting the model type to "Multinomial".`
	`313`	`+ * discrete count data and can be called by setting the model type to "multinomial".`
`314`	`314`	`* For example, it can be used with word counts or TF_IDF vectors of documents.`
`315`	`315`	`* The Bernoulli model fits presence or absence (0-1) counts. By making every vector a`
`316`		`- * 0-1 vector and setting the model type to "Bernoulli", the fits and predicts as`
	`316`	`+ * 0-1 vector and setting the model type to "bernoulli", the fits and predicts as`
`317`	`317`	`* Bernoulli NB.`
`318`	`318`	`*`
`319`	`319`	* @param input RDD of `(label, array of features)` pairs. Every vector should be a frequency
`320`	`320`	`* vector or a count vector.`
`321`	`321`	`* @param lambda The smoothing parameter`
`322`	`322`	`*`
`323`	`323`	`* @param modelType The type of NB model to fit from the enumeration NaiveBayesModels, can be`
`324`		`- * Multinomial or Bernoulli`
	`324`	`+ * multinomial or bernoulli`
`325`	`325`	`*/`
`326`	`326`	`def train(input: RDD[LabeledPoint], lambda: Double, modelType: String): NaiveBayesModel = {`
`327`		`- new NaiveBayes(lambda, Multinomial).run(input)`
	`327`	`+ new NaiveBayes(lambda, MODELTYPE.fromString(modelType)).run(input)`
`328`	`328`	`}`
`329`	`329`
`330`	`330`