Added save_config method to Booster - #11

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 - Updated XGBoost to 2.1.1
 - Added support for callbacks
+- Added `save_config` method to `Booster`
 - Dropped support for Ruby < 3.1
 
 ## 0.8.0 (2023-09-13)

lib/xgboost/booster.rb

@@ -40,6 +40,13 @@ def []=(key_name, raw_value)
       set_attr(**{key_name => raw_value})
     end
 
+    def save_config
+      length = ::FFI::MemoryPointer.new(:int)
+      json_string = ::FFI::MemoryPointer.new(:pointer)
+      check_result FFI.XGBoosterSaveJsonConfig(handle_pointer, length, json_string)
+      json_string.read_pointer.read_string(length.read_int).force_encoding(Encoding::UTF_8)
+    end
+
     def attr(key_name)
       key = string_pointer(key_name.to_s)
       success = ::FFI::MemoryPointer.new(:int)

lib/xgboost/ffi.rb

@@ -40,6 +40,7 @@ module FFI
     attach_function :XGBoosterPredict, %i[pointer pointer int int int pointer pointer], :int
     attach_function :XGBoosterLoadModel, %i[pointer string], :int
     attach_function :XGBoosterSaveModel, %i[pointer string], :int
+    attach_function :XGBoosterSaveJsonConfig, %i[pointer pointer pointer], :int
     attach_function :XGBoosterDumpModelExWithFeatures, %i[pointer int pointer pointer int string pointer pointer], :int
     attach_function :XGBoosterGetAttr, %i[pointer pointer pointer pointer], :int
     attach_function :XGBoosterSetAttr, %i[pointer pointer pointer], :int

test/booster_test.rb

@@ -27,6 +27,12 @@ def test_dump_model_json
     assert JSON.parse(File.read(tempfile))
   end
 
+  def test_save_config
+    config = booster.save_config
+    assert_kind_of Hash, JSON.parse(config)
+    assert_equal Encoding::UTF_8, config.encoding
+  end
+
   def test_score
     expected = {"f0" => 118, "f2" => 93, "f1" => 104, "f3" => 43}
     assert_equal expected.values.sort, booster.score.values.sort