adding feature importance for new model types (#128)

salesforce · Sep 13, 2018 · fdd780e · fdd780e
1 parent 078ea62
commit fdd780e
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/core/src/main/scala/com/salesforce/op/ModelInsights.scala b/core/src/main/scala/com/salesforce/op/ModelInsights.scala
@@ -636,9 +636,13 @@ case object ModelInsights {
         case Some(m: RandomForestClassificationModel) => Seq(m.featureImportances.toArray.toSeq)
         case Some(m: NaiveBayesModel) => m.theta.rowIter.toSeq.map(_.toArray.toSeq)
         case Some(m: DecisionTreeClassificationModel) => Seq(m.featureImportances.toArray.toSeq)
+        case Some(m: GBTClassificationModel) => Seq(m.featureImportances.toArray.toSeq)
+        case Some(m: LinearSVCModel) => Seq(m.coefficients.toArray.toSeq)
         case Some(m: LinearRegressionModel) => Seq(m.coefficients.toArray.toSeq)
         case Some(m: DecisionTreeRegressionModel) => Seq(m.featureImportances.toArray.toSeq)
         case Some(m: RandomForestRegressionModel) => Seq(m.featureImportances.toArray.toSeq)
+        case Some(m: GBTRegressionModel) => Seq(m.featureImportances.toArray.toSeq)
+        case Some(m: GeneralizedLinearRegressionModel) => Seq(m.coefficients.toArray.toSeq)
         case _ => Seq.empty[Seq[Double]]
       }
       case _ => Seq.empty[Seq[Double]]