lib/eps/evaluators/naive_bayes.rb in eps-0.3.3 vs lib/eps/evaluators/naive_bayes.rb in eps-0.3.4
- old
+ new
@@ -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