Skip to content

Commit

Permalink
updated some commented tests to use the new ModelExporter object
Browse files Browse the repository at this point in the history
reordered the imports
  • Loading branch information
selvinsource committed Nov 5, 2014
1 parent df8a89e commit ae8b993
Showing 1 changed file with 7 additions and 8 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -17,14 +17,13 @@

package org.apache.spark.mllib.export.pmml

import org.scalatest.FunSuite
import org.apache.spark.mllib.linalg.Vectors
import org.apache.spark.mllib.export.ModelExportFactory
import scala.annotation.varargs
import org.apache.spark.mllib.clustering.KMeansModel
import org.apache.spark.mllib.export.ModelExportFactory
import org.apache.spark.mllib.export.ModelExportType
import org.apache.spark.mllib.linalg.Vectors
import org.scalatest.FunSuite
import org.dmg.pmml.ClusteringModel
import javax.xml.parsers.DocumentBuilderFactory
import java.io.ByteArrayOutputStream

class KMeansPMMLModelExportSuite extends FunSuite{

Expand All @@ -43,7 +42,7 @@ class KMeansPMMLModelExportSuite extends FunSuite{

//assert that the PMML format is as expected
assert(modelExport.isInstanceOf[PMMLModelExport])
var pmml = modelExport.asInstanceOf[PMMLModelExport].getPmml()
val pmml = modelExport.asInstanceOf[PMMLModelExport].getPmml()
assert(pmml.getHeader().getDescription() === "k-means clustering")
//check that the number of fields match the single vector size
assert(pmml.getDataDictionary().getNumberOfFields() === clusterCenters(0).size)
Expand All @@ -52,8 +51,8 @@ class KMeansPMMLModelExportSuite extends FunSuite{
assert(pmml.getModels().get(0).asInstanceOf[ClusteringModel].getNumberOfClusters() === clusterCenters.size)

//manual checking
//modelExport.asInstanceOf[PMMLModelExport].save(System.out)
//modelExport.asInstanceOf[PMMLModelExport].saveLocalFile("/tmp/kmeans.xml")
//ModelExporter.toPMML(kmeansModel,new StreamResult(System.out))
//ModelExporter.toPMML(kmeansModel,"/tmp/kmeans.xml")

}

Expand Down

0 comments on commit ae8b993

Please sign in to comment.