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

@@ -9,11 +9,13 @@
         @labels = labels
         @features = features
         @text_features = text_features
       end
 
-      def predict(data)
+      def predict(data, probabilities: false)
+        raise "Probabilities not supported" if probabilities && @objective == "regression"
+
         rows = data.map(&:to_h)
 
         # sparse matrix
         @text_features.each do |k, v|
           encoder = TextEncoder.new(**v)
@@ -36,21 +38,32 @@
 
         case @objective
         when "regression"
           sum_trees(rows, @trees)
         when "binary"
-          sum_trees(rows, @trees).map { |s| @labels[sigmoid(s) > 0.5 ? 1 : 0] }
+          prob = sum_trees(rows, @trees).map { |s| sigmoid(s) }
+          if probabilities
+            prob.map { |v| @labels.zip([1 - v, v]).to_h }
+          else
+            prob.map { |v| @labels[v > 0.5 ? 1 : 0] }
+          end
         else
           tree_scores = []
           num_trees = @trees.size / @labels.size
           @trees.each_slice(num_trees).each do |trees|
             tree_scores << sum_trees(rows, trees)
           end
           data.size.times.map do |i|
             v = tree_scores.map { |s| s[i] }
-            idx = v.map.with_index.max_by { |v2, _| v2 }.last
-            @labels[idx]
+            if probabilities
+              exp = v.map { |vi| Math.exp(vi) }
+              sum = exp.sum
+              @labels.zip(exp.map { |e| e / sum }).to_h
+            else
+              idx = v.map.with_index.max_by { |v2, _| v2 }.last
+              @labels[idx]
+            end
           end
         end
       end
 
       private
@@ -107,10 +120,10 @@
           nil
         end
       end
 
       def sigmoid(x)
-        1.0 / (1 + Math::E**(-x))
+        1.0 / (1 + Math.exp(-x))
       end
     end
   end
 end