Fixed error with categorical and text features - fixes #22

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.3.8 (unreleased)
+
+- Fixed error with categorical and text features
+
 ## 0.3.7 (2020-11-23)
 
 - Fixed error with LightGBM summary

lib/eps/lightgbm.rb

@@ -71,7 +71,7 @@ def _train(verbose: nil, early_stopping: nil, learning_rate: 0.1)
       end
 
       # create datasets
-      categorical_idx = @features.values.map.with_index.select { |type, _| type == "categorical" }.map(&:last)
+      categorical_idx = train_set.columns.keys.map.with_index.select { |k, _| @features[k] == "categorical" }.map(&:last)
       train_ds = ::LightGBM::Dataset.new(train_set.map_rows(&:to_a), label: train_set.label, weight: train_set.weight, categorical_feature: categorical_idx, params: params)
       validation_ds = ::LightGBM::Dataset.new(validation_set.map_rows(&:to_a), label: validation_set.label, weight: validation_set.weight, categorical_feature: categorical_idx, params: params, reference: train_ds) if validation_set
 

test/lightgbm_test.rb

@@ -218,6 +218,14 @@ def test_text_features_classification
     assert_equal ["ham", "spam"], model.predict(test_data)
   end
 
+  def test_text_features_classification_categorical
+    data = [
+      {message: "This is the first document.", category: "A", tag: "ham"},
+      {message: "Hello, this is the second document.", category: "B", tag: "spam"},
+    ]
+    Eps::LightGBM.new(data, target: :tag)
+  end
+
   def test_learning_rate
     data = mpg_data
     model = Eps::LightGBM.new(data, target: :hwy, split: false, learning_rate: 1)