Added support for pandas_categorical to predict method - resolves #8

Co-authored-by: Nuno Silva <[email protected]>

Author: ankane

CHANGELOG.md

@@ -1,6 +1,7 @@
 ## 0.4.0 (unreleased)
 
 - Added support for different prediction types
+- Added support for `pandas_categorical` to `predict` method
 - 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

lib/lightgbm.rb

@@ -1,6 +1,9 @@
 # dependencies
 require "ffi"
 
+# stdlib
+require "json"
+
 # modules
 require_relative "lightgbm/utils"
 require_relative "lightgbm/booster"

lib/lightgbm/booster.rb

@@ -12,6 +12,7 @@ def initialize(params: nil, train_set: nil, model_file: nil, model_str: nil)
         create_handle do |handle|
           safe_call FFI.LGBM_BoosterCreateFromModelfile(model_file, out_num_iterations, handle)
         end
+        @pandas_categorical = load_pandas_categorical(file_name: model_file)
       else
         params ||= {}
         set_verbosity(params)
@@ -96,6 +97,7 @@ def model_from_string(model_str)
       create_handle do |handle|
         safe_call FFI.LGBM_BoosterLoadModelFromString(model_str, out_num_iterations, handle)
       end
+      @pandas_categorical = load_pandas_categorical(model_str: model_str)
       @cached_feature_name = nil
       self
     end
@@ -235,5 +237,48 @@ def feature_importance_type_mapper(importance_type)
         -1
       end
     end
+
+    def load_pandas_categorical(file_name: nil, model_str: nil)
+      pandas_key = "pandas_categorical:"
+      offset = -pandas_key.length
+      if !file_name.nil?
+        max_offset = -File.size(file_name)
+        lines = []
+        File.open(file_name, "rb") do |f|
+          loop do
+            offset = [offset, max_offset].max
+            f.seek(offset, IO::SEEK_END)
+            lines = f.readlines
+            if lines.length >= 2
+              break
+            end
+            offset *= 2
+          end
+        end
+        last_line = lines[-1].strip
+        if !last_line.start_with?(pandas_key)
+          last_line = lines[-2].strip
+        end
+      elsif !model_str.nil?
+        idx = model_str[..offset].rindex("\n")
+        last_line = model_str[idx..].strip
+      end
+      if last_line.start_with?(pandas_key)
+        JSON.parse(last_line[pandas_key.length..])
+      end
+    end
+
+    def loaded_param
+      buffer_len = 1 << 20
+      out_len = ::FFI::MemoryPointer.new(:int64)
+      out_str = ::FFI::MemoryPointer.new(:char, buffer_len)
+      safe_call FFI.LGBM_BoosterGetLoadedParam(@handle, 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_BoosterGetLoadedParam(@handle, actual_len, out_len, out_str)
+      end
+      JSON.parse(out_str.read_string)
+    end
   end
 end

lib/lightgbm/ffi.rb

@@ -49,6 +49,7 @@ module FFI
     attach_function :LGBM_BoosterCreate, %i[pointer string pointer], :int
     attach_function :LGBM_BoosterCreateFromModelfile, %i[string pointer pointer], :int
     attach_function :LGBM_BoosterLoadModelFromString, %i[string pointer pointer], :int
+    attach_function :LGBM_BoosterGetLoadedParam, %i[pointer int64 pointer pointer], :int
     attach_function :LGBM_BoosterFree, %i[pointer], :int
     attach_function :LGBM_BoosterAddValidData, %i[pointer pointer], :int
     attach_function :LGBM_BoosterGetNumClasses, %i[pointer pointer], :int

lib/lightgbm/inner_predictor.rb

@@ -6,6 +6,7 @@ class InnerPredictor
 
     def initialize(booster, pred_parameter)
       @handle = booster.instance_variable_get(:@handle)
+      @pandas_categorical = booster.instance_variable_get(:@pandas_categorical)
       @pred_parameter = params_str(pred_parameter)
 
       # keep booster for cached_feature_name
@@ -50,6 +51,15 @@ def predict(data, start_iteration: 0, num_iteration: -1, raw_score: false, pred_
         singular = !data.first.is_a?(Array)
         data = [data] if singular
         check_2d_array(data)
+        data = data.map(&:dup) if @pandas_categorical&.any?
+      end
+
+      if @pandas_categorical&.any?
+        apply_pandas_categorical(
+          data,
+          @booster.send(:loaded_param)["categorical_feature"],
+          @pandas_categorical
+        )
       end
 
       preds, nrow =
@@ -127,5 +137,16 @@ def sorted_feature_values(input_hash)
     def cached_feature_name
       @booster.send(:cached_feature_name)
     end
+
+    def apply_pandas_categorical(data, categorical_feature, pandas_categorical)
+      (categorical_feature || []).each_with_index do |cf, i|
+        cat_codes = pandas_categorical[i].map.with_index.to_h
+        # TODO confirm column is categorical
+        data.each do |r|
+          # TODO decide how to handle missing values
+          r[cf] = cat_codes.fetch(r[cf])
+        end
+      end
+    end
   end
 end

test/booster_test.rb

@@ -130,6 +130,20 @@ def test_predict_pred_contrib
     assert_elements_in_delta expected[0], y_pred
   end
 
+  def test_predict_pandas_categorical_model_file
+    x_test = [[3.7, 1.2, 7.2, "cat9"], [7.5, 0.5, 7.9, "cat0"]]
+    booster = LightGBM::Booster.new(model_file: "test/support/categorical.txt")
+    y_pred = booster.predict(x_test)
+    assert_elements_in_delta [0.996415541144579, 1.0809369939979934], y_pred.first(2)
+  end
+
+  def test_predict_pandas_categorical_model_str
+    x_test = [[3.7, 1.2, 7.2, "cat9"], [7.5, 0.5, 7.9, "cat0"]]
+    booster = LightGBM::Booster.new(model_str: File.read("test/support/categorical.txt"))
+    y_pred = booster.predict(x_test)
+    assert_elements_in_delta [0.996415541144579, 1.0809369939979934], y_pred.first(2)
+  end
+
   def test_model_to_string
     assert booster.model_to_string
   end

test/support/categorical.py

@@ -0,0 +1,21 @@
+import lightgbm as lgb
+import pandas as pd
+
+df = pd.read_csv('test/support/data.csv')
+df['x3'] = ('cat' + df['x3'].astype(str)).astype('category')
+
+X = df.drop(columns=['y'])
+y = df['y']
+
+X_train = X[:300]
+y_train = y[:300]
+X_test = X[300:]
+y_test = y[300:]
+
+train_data = lgb.Dataset(X_train, label=y_train)
+bst = lgb.train({}, train_data, num_boost_round=5)
+bst.save_model('test/support/categorical.txt')
+
+bst = lgb.Booster(model_file='test/support/categorical.txt')
+print('x', X_train[:2].to_numpy().tolist())
+print('predict', bst.predict(X_train)[:2].tolist())