Some cleanup

Author: blublinsky

src/main/scala/com/lightbend/kafka/ModelProvider.scala

@@ -42,7 +42,6 @@ object ModelProvider {
         tRecord.writeTo(bos)
         sender.writeValue(ModelServingConfiguration.MODELS_TOPIC, bos.toByteArray)
         pause()
-
       })
     }
   }

src/main/scala/com/lightbend/modelServer/DataProcessor.scala

@@ -46,8 +46,8 @@ class DataProcessor extends RichCoProcessFunction[WineRecord, ModelToServe, Doub
     println(s"New model - $model")
     newModelState.update(new ModelToServeStats(model))
     newModel = model.modelType match {
-      case ModelDescriptor.ModelType.PMML => Some(PMMLModel(model.model))
-      case ModelDescriptor.ModelType.TENSORFLOW => Some(TensorFlowModel(model.model))
+      case ModelDescriptor.ModelType.PMML => PMMLModel(model.model)             // PMML
+      case ModelDescriptor.ModelType.TENSORFLOW => TensorFlowModel(model.model) // Tensorflow
       case _ => None // Not supported yet
     }
   }

src/main/scala/com/lightbend/modelServer/DataProcessorMap.scala

@@ -27,9 +27,10 @@ class DataProcessorMap extends RichCoFlatMapFunction[WineRecord, ModelToServe, D
 
   override def flatMap2(model: ModelToServe, out: Collector[Double]): Unit = {
     println(s"New model - $model")
-    newModel = model.modelType match {
-      case ModelDescriptor.ModelType.PMML => Some(new PMMLModel(model.model))
-      case ModelDescriptor.ModelType.TENSORFLOW => Some(TensorFlowModel(model.model))
+    newModel =
+      model.modelType match {
+      case ModelDescriptor.ModelType.PMML => PMMLModel(model.model)               // PMML
+      case ModelDescriptor.ModelType.TENSORFLOW => TensorFlowModel(model.model)   // Tensorflow
       case _ => None // Not supported yet
     }
   }

src/main/scala/com/lightbend/modelServer/ModelTester.scala

@@ -24,8 +24,8 @@ object ModelTester {
     val Tensormodel = TensorFlowModel(TensorflowbyteArray)
     println("PMML | Tensorflow")
     records.foreach(r => {
-      val presult = PMMLmodel.score(r.asInstanceOf[AnyVal]).asInstanceOf[Double]
-      val tresult = Tensormodel.score(r.asInstanceOf[AnyVal]).asInstanceOf[Double]
+      val presult = PMMLmodel.get.score(r.asInstanceOf[AnyVal]).asInstanceOf[Double]
+      val tresult = Tensormodel.get.score(r.asInstanceOf[AnyVal]).asInstanceOf[Double]
       println(s"$presult  | $tresult")
     })
   }

src/main/scala/com/lightbend/modelServer/model/PMMLModel.scala

@@ -71,7 +71,13 @@ class PMMLModel(inputStream: Array[Byte]) extends Model {
 
 object PMMLModel{
 
-  def apply(inputStream: Array[Byte]): PMMLModel = new PMMLModel(inputStream)
+  def apply(inputStream: Array[Byte]): Option[PMMLModel] = {
+    try {
+      Some(new PMMLModel(inputStream))
+    }catch{
+      case t: Throwable => None
+    }
+  }
   private val optimizers = Array(new ExpressionOptimizer, new FieldOptimizer, new PredicateOptimizer, new GeneralRegressionModelOptimizer, new NaiveBayesModelOptimizer, new RegressionModelOptimizer)
   def optimize(pmml : PMML) = this.synchronized {
     optimizers.foreach(opt =>

src/main/scala/com/lightbend/modelServer/model/TensorFlowModel.scala

@@ -44,9 +44,18 @@ class TensorFlowModel(inputStream: Array[Byte]) extends Model  {
     value._1.toDouble
   }
 
-  override def cleanup(): Unit = session.close
+  override def cleanup(): Unit = {
+    session.close
+    graph.close
+  }
 }
 
 object TensorFlowModel{
-  def apply(inputStream: Array[Byte]): TensorFlowModel = new TensorFlowModel(inputStream)
+  def apply(inputStream: Array[Byte]): Option[TensorFlowModel] = {
+    try {
+      Some(new TensorFlowModel(inputStream))
+    }catch{
+      case t: Throwable => None
+    }
+  }
 }