Numerical error fix for NaN confidence values.
diff --git a/src/main/scala/org/template/textclassification/NBAlgorithm.scala b/src/main/scala/org/template/textclassification/NBAlgorithm.scala
index 1c1c31d..9190bd7 100644
--- a/src/main/scala/org/template/textclassification/NBAlgorithm.scala
+++ b/src/main/scala/org/template/textclassification/NBAlgorithm.scala
@@ -77,7 +77,10 @@
// Vectorize query,
val x: Vector = pd.transform(doc)
- normalize(scoreArray.map(e => exp(innerProduct(e._2, x.toArray) + e._1)))
+ val z = scoreArray
+ .map(e => innerProduct(e._2, x.toArray) + e._1)
+
+ normalize((0 until z.size).map(k => exp(z(k) - z.max)).toArray)
}
// 4. Implement predict method for our model using
