lib/eps/evaluators/naive_bayes.rb in eps-0.3.3 vs lib/eps/evaluators/naive_bayes.rb in eps-0.3.4

@@ -8,17 +8,18 @@
         @features = features
         @derived = derived
         @legacy = legacy
       end
 
-      def predict(x)
+      def predict(x, probabilities: false)
         probs = calculate_class_probabilities(x)
         probs.map do |xp|
-          # convert probabilities
-          # not needed when just returning label
-          # sum = xp.values.map { |v| Math.exp(v) }.sum.to_f
-          # p xp.map { |k, v| [k, Math.exp(v) / sum] }.to_h
-          xp.sort_by { |k, v| [-v, k] }[0][0]
+          if probabilities
+            sum = xp.values.map { |v| Math.exp(v) }.sum.to_f
+            xp.map { |k, v| [k, Math.exp(v) / sum] }.to_h
+          else
+            xp.sort_by { |k, v| [-v, k] }[0][0]
+          end
         end
       end
 
       # use log to prevent underflow
       # https://www.antoniomallia.it/lets-implement-a-gaussian-naive-bayes-classifier-in-python.html