Add callbacks

lib/xgboost.rb

@@ -44,8 +44,12 @@ class << self
   autoload :FFI, "xgboost/ffi"
 
   class << self
-    def train(params, dtrain, num_boost_round: 10, evals: nil, early_stopping_rounds: nil, verbose_eval: true)
+    def train(params, dtrain, num_boost_round: 10, evals: nil, early_stopping_rounds: nil, verbose_eval: true,
+              callbacks: [])
       booster = Booster.new(params: params)
+      callbacks.each do |callback|
+        callback.before_training(model: booster)
+      end
       num_feature = dtrain.num_col
       booster.set_param("num_feature", num_feature)
       booster.feature_names = dtrain.feature_names
@@ -59,6 +63,9 @@ def train(params, dtrain, num_boost_round: 10, evals: nil, early_stopping_rounds
       end
 
       num_boost_round.times do |iteration|
+        callbacks.each do |callback|
+          callback.before_iteration(model: booster, epoch: iteration, evals: evals)
+        end
         booster.update(dtrain, iteration)
 
         if evals.any?
@@ -73,17 +80,23 @@ def train(params, dtrain, num_boost_round: 10, evals: nil, early_stopping_rounds
           puts message if verbose_eval
           score = res[-1][1]
 
-          # TODO handle larger better
-          if best_score.nil? || score < best_score
-            best_score = score
-            best_iter = iteration
-            best_message = message
-          elsif early_stopping_rounds && iteration - best_iter >= early_stopping_rounds
-            booster.best_iteration = best_iter
-            puts "Stopping. Best iteration:\n#{best_message}" if verbose_eval
-            break
-          end
+        # TODO: handle larger better
+        if best_score.nil? || score < best_score
+          best_score = score
+          best_iter = iteration
+          best_message = message
+        elsif early_stopping_rounds && iteration - best_iter >= early_stopping_rounds
+          booster.best_iteration = best_iter
+          booster.best_score = best_score
+          puts "Stopping. Best iteration:\n#{best_message}" if verbose_eval
+          break
         end
+        callbacks.each do |callback|
+          callback.after_iteration(model: booster, epoch: iteration, evals: evals, res: res)
+        end
+      end
+      callbacks.each do |callback|
+        callback.after_training(model: booster)
       end
 
       booster

lib/xgboost/booster.rb

@@ -1,6 +1,6 @@
 module XGBoost
   class Booster
-    attr_accessor :best_iteration, :feature_names, :feature_types
+    attr_accessor :best_iteration, :feature_names, :feature_types, :best_score
 
     def initialize(params: nil, model_file: nil)
       @handle = ::FFI::MemoryPointer.new(:pointer)

test/train_test.rb

@@ -55,6 +55,67 @@ def test_feature_names_and_types
     assert_nil model.feature_types
   end
 
+  class MockCallback
+    attr_reader :before_training_count, :after_training_count, :before_iteration_count, :after_iteration_count,
+                :before_training_args, :after_training_args, :before_iteration_args, :after_iteration_args
+
+    def initialize
+      @before_training_count = 0
+      @after_training_count = 0
+      @before_iteration_count = 0
+      @after_iteration_count = 0
+      @before_training_args = []
+      @after_training_args = []
+      @before_iteration_args = []
+      @after_iteration_args = []
+    end
+
+    def before_training(env)
+      @before_training_count += 1
+      @before_training_args << env
+    end
+
+    def after_training(env)
+      @after_training_count += 1
+      @after_training_args << env
+    end
+
+    def before_iteration(env)
+      @before_iteration_count += 1
+      @before_iteration_args << env
+    end
+
+    def after_iteration(env)
+      @after_iteration_count += 1
+      @after_iteration_args << env
+    end
+  end
+
+  def test_callback
+    callback = MockCallback.new
+
+    XGBoost.train(
+      regression_params,
+      regression_train,
+      num_boost_round: 10,
+      callbacks: [callback],
+      evals: [[regression_train, 'train'], [regression_test, 'eval']]
+    )
+
+    assert_equal 1, callback.before_training_count
+    assert_equal 1, callback.after_training_count
+    assert_equal 10, callback.before_iteration_count
+    assert_equal 10, callback.after_iteration_count
+
+    # Verify arguments
+    assert_kind_of XGBoost::Booster, callback.before_training_args.first[:model]
+    assert_kind_of XGBoost::Booster, callback.after_training_args.first[:model]
+    assert_kind_of XGBoost::Booster, callback.before_iteration_args.first[:model]
+    assert_kind_of XGBoost::Booster, callback.after_iteration_args.first[:model]
+    assert_equal 0, callback.before_iteration_args.first[:epoch]
+    assert_equal 9, callback.after_iteration_args.last[:epoch]
+  end
+
   private
 
   def rsme(y_true, y_pred)