Added importance_type option to dump_model, model_to_string, and save_model methods

Author: ankane

CHANGELOG.md

@@ -3,6 +3,7 @@
 - Added support for different prediction types
 - Added support for hashes and Rover data frames to `predict` method
 - Added support for hashes to `Dataset`
+- Added `importance_type` option to `dump_model`, `model_to_string`, and `save_model` methods
 - Changed `Dataset` to use column names for feature names with Rover and Daru
 - Changed `predict` method to match feature names with Daru
 - Dropped support for Ruby < 3.1

lib/lightgbm/booster.rb

@@ -38,17 +38,17 @@ def current_iteration
       out.read_int
     end
 
-    def dump_model(num_iteration: nil, start_iteration: 0)
+    def dump_model(num_iteration: nil, start_iteration: 0, importance_type: "split")
       num_iteration ||= best_iteration
+      importance_type_int = feature_importance_type_mapper(importance_type)
       buffer_len = 1 << 20
       out_len = ::FFI::MemoryPointer.new(:int64)
       out_str = ::FFI::MemoryPointer.new(:char, buffer_len)
-      feature_importance_type = 0 # TODO add option
-      safe_call FFI.LGBM_BoosterDumpModel(@handle, start_iteration, num_iteration, feature_importance_type, buffer_len, out_len, out_str)
+      safe_call FFI.LGBM_BoosterDumpModel(@handle, start_iteration, num_iteration, importance_type_int, buffer_len, out_len, out_str)
       actual_len = out_len.read_int64
       if actual_len > buffer_len
         out_str = ::FFI::MemoryPointer.new(:char, actual_len)
-        safe_call FFI.LGBM_BoosterDumpModel(@handle, start_iteration, num_iteration, feature_importance_type, actual_len, out_len, out_str)
+        safe_call FFI.LGBM_BoosterDumpModel(@handle, start_iteration, num_iteration, importance_type_int, actual_len, out_len, out_str)
       end
       out_str.read_string
     end
@@ -64,19 +64,10 @@ def eval_train
 
     def feature_importance(iteration: nil, importance_type: "split")
       iteration ||= best_iteration
-      importance_type =
-        case importance_type
-        when "split"
-          FFI::C_API_FEATURE_IMPORTANCE_SPLIT
-        when "gain"
-          FFI::C_API_FEATURE_IMPORTANCE_GAIN
-        else
-          -1
-        end
-
+      importance_type_int = feature_importance_type_mapper(importance_type)
       num_feature = self.num_feature
       out_result = ::FFI::MemoryPointer.new(:double, num_feature)
-      safe_call FFI.LGBM_BoosterFeatureImportance(@handle, iteration, importance_type, out_result)
+      safe_call FFI.LGBM_BoosterFeatureImportance(@handle, iteration, importance_type_int, out_result)
       out_result.read_array_of_double(num_feature).map(&:to_i)
     end
 
@@ -109,17 +100,17 @@ def model_from_string(model_str)
       self
     end
 
-    def model_to_string(num_iteration: nil, start_iteration: 0)
+    def model_to_string(num_iteration: nil, start_iteration: 0, importance_type: "split")
       num_iteration ||= best_iteration
+      importance_type_int = feature_importance_type_mapper(importance_type)
       buffer_len = 1 << 20
       out_len = ::FFI::MemoryPointer.new(:int64)
       out_str = ::FFI::MemoryPointer.new(:char, buffer_len)
-      feature_importance_type = 0 # TODO add option
-      safe_call FFI.LGBM_BoosterSaveModelToString(@handle, start_iteration, num_iteration, feature_importance_type, buffer_len, out_len, out_str)
+      safe_call FFI.LGBM_BoosterSaveModelToString(@handle, start_iteration, num_iteration, importance_type_int, buffer_len, out_len, out_str)
       actual_len = out_len.read_int64
       if actual_len > buffer_len
         out_str = ::FFI::MemoryPointer.new(:char, actual_len)
-        safe_call FFI.LGBM_BoosterSaveModelToString(@handle, start_iteration, num_iteration, feature_importance_type, actual_len, out_len, out_str)
+        safe_call FFI.LGBM_BoosterSaveModelToString(@handle, start_iteration, num_iteration, importance_type_int, actual_len, out_len, out_str)
       end
       out_str.read_string
     end
@@ -162,10 +153,10 @@ def predict(data, start_iteration: 0, num_iteration: nil, raw_score: false, pred
       )
     end
 
-    def save_model(filename, num_iteration: nil, start_iteration: 0)
+    def save_model(filename, num_iteration: nil, start_iteration: 0, importance_type: "split")
       num_iteration ||= best_iteration
-      feature_importance_type = 0 # TODO add
-      safe_call FFI.LGBM_BoosterSaveModel(@handle, start_iteration, num_iteration, feature_importance_type, filename)
+      importance_type_int = feature_importance_type_mapper(importance_type)
+      safe_call FFI.LGBM_BoosterSaveModel(@handle, start_iteration, num_iteration, importance_type_int, filename)
       self # consistent with Python API
     end
 
@@ -233,5 +224,16 @@ def num_class
     def cached_feature_name
       @cached_feature_name ||= feature_name
     end
+
+    def feature_importance_type_mapper(importance_type)
+      case importance_type
+      when "split"
+        FFI::C_API_FEATURE_IMPORTANCE_SPLIT
+      when "gain"
+        FFI::C_API_FEATURE_IMPORTANCE_GAIN
+      else
+        -1
+      end
+    end
   end
 end