MLJ Integration

.github/workflows/CI.yml

@@ -28,6 +28,7 @@ jobs:
         version:
           - '1.6'
           - '1' 
+          - 'nightly'
         os:
           - ubuntu-latest
         arch:

examples/mlj/binary.jl

@@ -6,24 +6,25 @@ using MLJBase
 using PythonCall
 
 # Initialize data
-cat_features = [0, 1]
-train_data = DataFrame([["a", "a", "c"], ["b", "b", "d"], [1, 4, 30], [4, 5, 40],
-                        [5, 6, 50], [6, 7, 60]], :auto)
-train_labels = [1, 1, -1]
-eval_data = DataFrame([["a", "a"], ["b", "d"], [2, 1], [4, 4], [6, 50], [8, 60]], :auto)
+train_data = DataFrame([coerce(["a", "a", "c"], Multiclass),
+                        coerce(["b", "b", "d"], Multiclass),
+                        coerce([0, 0, 1], OrderedFactor), [4, 5, 40], [5, 6, 50],
+                        [6, 7, 60]], :auto)
+train_labels = coerce([1, 1, -1], OrderedFactor)
+eval_data = DataFrame([coerce(["a", "a"], Multiclass), coerce(["b", "d"], Multiclass),
+                       coerce([0, 0], OrderedFactor), [4, 4], [6, 50], [8, 60]], :auto)
 
 # Initialize CatBoostClassifier
-model = CatBoostClassifier(; iterations=2, learning_rate=1, depth=2,
-                           cat_features=cat_features)
+model = CatBoostClassifier(; iterations=2, learning_rate=1, depth=2)
 mach = machine(model, train_data, train_labels)
 
 # Fit model
 MLJBase.fit!(mach)
 
 # Get predicted classes
-preds_class = MLJBase.predict(mach, eval_data)
+preds_class = MLJBase.predict_mean(mach, eval_data)
 
 # Get predicted probabilities for each class
-preds_proba = MLJBase.predict_mean(mach, eval_data)
+preds_proba = MLJBase.predict(mach, eval_data)
 
 end # module

examples/mlj/multiclass.jl

@@ -0,0 +1,31 @@
+module Multiclass
+
+using CatBoost.MLJCatBoostInterface
+using DataFrames
+using MLJBase
+using PythonCall
+
+# Initialize data
+train_data = DataFrame([coerce(["a", "a", "c"], MLJBase.Multiclass),
+                        coerce(["b", "b", "d"], MLJBase.Multiclass),
+                        coerce([0, 0, 1], OrderedFactor), [4, 5, 40], [5, 6, 50],
+                        [6, 7, 60]], :auto)
+train_labels = coerce([1, 0, -1], MLJBase.Multiclass)
+eval_data = DataFrame([coerce(["a", "a"], MLJBase.Multiclass),
+                       coerce(["b", "d"], MLJBase.Multiclass), [2, 1], [4, 4], [6, 50],
+                       [8, 60]], :auto)
+
+# Initialize CatBoostClassifier
+model = CatBoostClassifier(; iterations=2, learning_rate=1, depth=2)
+mach = machine(model, train_data, train_labels)
+
+# Fit model
+MLJBase.fit!(mach)
+
+# Get predicted classes
+preds_class = MLJBase.predict_mean(mach, eval_data)
+
+# Get predicted probabilities for each class
+preds_proba = MLJBase.predict(mach, eval_data)
+
+end # module

examples/mlj/regression.jl

@@ -6,14 +6,12 @@ using MLJBase
 using PythonCall
 
 # Initialize data
-cat_features = [0, 1]
 train_data = DataFrame([[1, 4, 30], [4, 5, 40], [5, 6, 50], [6, 7, 60]], :auto)
 train_labels = [10.0, 20.0, 30.0]
 eval_data = DataFrame([[2, 1], [4, 4], [6, 50], [8, 60]], :auto)
 
 # Initialize CatBoostClassifier
-model = CatBoostRegressor(; iterations=2, learning_rate=1, depth=2,
-                          cat_features=cat_features)
+model = CatBoostRegressor(; iterations=2, learning_rate=1, depth=2)
 mach = machine(model, train_data, train_labels)
 
 # Fit model

src/MLJCatBoostInterface.jl

@@ -28,6 +28,40 @@ function mlj_to_kwargs(model)
                             for name in fieldnames(typeof(model)))
 end
 
+"""
+    get_dtype_feature_ix(X, dtype)
+
+Get the index of the columns with a specific dtype
+
+Parameters
+----------
+- `X`: Table
+- `dtype`: DataType 
+
+Returns
+-------
+- `Vector{Int64}`
+"""
+function get_dtype_feature_ix(X, dtype)
+    return findall(eltype.(MMI.scitype.(eachcol(X))) .<: dtype)
+end
+
+"""
+Get cat features for model
+get_cat_features
+"""
+function prepare_input(X)
+    order_factor_ix = get_dtype_feature_ix(X, OrderedFactor)
+    for col in names(X)[order_factor_ix]
+        X[:, col] = MMI.int(X[:, col])
+    end
+
+    cat_features = get_dtype_feature_ix(X, Multiclass) .- 1 # convert to 0 based indexing
+    text_features = get_dtype_feature_ix(X, AbstractString) .- 1 # convert to 0 based indexing
+
+    return X, cat_features, text_features
+end
+
 include("mlj_catboostclassifier.jl")
 include("mlj_catboostregressor.jl")
 
@@ -47,16 +81,15 @@ MMI.metadata_pkg.((CatBoostClassifier, CatBoostRegressor), name="CatBoost.jl",
 
 MMI.metadata_model(CatBoostClassifier;
                    input_scitype=Union{MMI.Table(MMI.Continuous, MMI.Count,
-                                                 MMI.OrderedFactor),
+                                                 MMI.OrderedFactor, MMI.Multiclass),
                                        AbstractMatrix{MMI.Continuous}},
-                   target_scitype=Union{AbstractVector{<:MMI.Finite},
-                                        AbstractVector{<:MMI.Continuous}},
+                   target_scitype=Union{AbstractVector{<:MMI.Finite}},
                    human_name="CatBoost classifier",
                    load_path="$PKG.MLJCatBoostInterface.CatBoostClassifier")
 
 MMI.metadata_model(CatBoostRegressor;
                    input_scitype=Union{MMI.Table(MMI.Continuous, MMI.Count,
-                                                 MMI.OrderedFactor),
+                                                 MMI.OrderedFactor, MMI.Multiclass),
                                        AbstractMatrix{MMI.Continuous}},
                    target_scitype=AbstractVector{<:MMI.Continuous},
                    human_name="CatBoost regressor",

src/mlj_catboostclassifier.jl

@@ -23,7 +23,6 @@ MMI.@mlj_model mutable struct CatBoostClassifier <: MMI.Probabilistic
     thread_count = nothing
     random_seed = nothing
     use_best_model = nothing
-    verbose = nothing
     logging_level = nothing
     metric_period = nothing
     ctr_leaf_count_limit = nothing
@@ -78,7 +77,6 @@ MMI.@mlj_model mutable struct CatBoostClassifier <: MMI.Probabilistic
     data_partition = nothing
     metadata = nothing
     early_stopping_rounds = nothing
-    cat_features = nothing
     grow_policy = nothing
     min_data_in_leaf = nothing
     min_child_samples = nothing
@@ -97,25 +95,25 @@ MMI.@mlj_model mutable struct CatBoostClassifier <: MMI.Probabilistic
     diffusion_temperature = nothing
     posterior_sampling = nothing
     boost_from_average = nothing
-    text_features = nothing
     tokenizers = nothing
     dictionaries = nothing
     feature_calcers = nothing
     text_processing = nothing
 end
 
-function model_init(mlj_model::CatBoostClassifier)
-    return catboost.CatBoostClassifier(; mlj_to_kwargs(mlj_model)...)
+function model_init(mlj_model::CatBoostClassifier; kw...)
+    return catboost.CatBoostClassifier(; mlj_to_kwargs(mlj_model)..., kw...)
 end
 
 function MMI.fit(mlj_model::CatBoostClassifier, verbosity::Int, X, y)
-    silent = verbosity > 0 ? false : true
+    verbose = verbosity > 0 ? false : true
 
-    py_X = to_pandas(X)
+    X_preprocessed, cat_features, text_features = prepare_input(X)
+    py_X = to_pandas(X_preprocessed)
     py_y = numpy.array(Array(y))
 
-    model = model_init(mlj_model)
-    model.fit(py_X, py_y; silent=silent)
+    model = model_init(mlj_model; cat_features, text_features, verbose)
+    model.fit(py_X, py_y)
 
     cache = nothing
     report = (feature_importances=feature_importances(model),)

src/mlj_catboostregressor.jl

@@ -24,8 +24,6 @@ MMI.@mlj_model mutable struct CatBoostRegressor <: MMI.Deterministic
     random_seed = nothing
     use_best_model = nothing
     best_model_min_trees = nothing
-    verbose = nothing
-    silent = nothing
     logging_level = nothing
     metric_period = nothing
     ctr_leaf_count_limit = nothing
@@ -83,7 +81,6 @@ MMI.@mlj_model mutable struct CatBoostRegressor <: MMI.Deterministic
     data_partition = nothing
     metadata = nothing
     early_stopping_rounds = nothing
-    cat_features = nothing
     grow_policy = nothing
     min_data_in_leaf = nothing
     min_child_samples = nothing
@@ -103,7 +100,6 @@ MMI.@mlj_model mutable struct CatBoostRegressor <: MMI.Deterministic
     diffusion_temperature = nothing
     posterior_sampling = nothing
     boost_from_average = nothing
-    text_features = nothing
     tokenizers = nothing
     dictionaries = nothing
     feature_calcers = nothing
@@ -112,18 +108,19 @@ MMI.@mlj_model mutable struct CatBoostRegressor <: MMI.Deterministic
     eval_fraction = nothing
 end
 
-function model_init(mlj_model::CatBoostRegressor)
-    return catboost.CatBoostRegressor(; mlj_to_kwargs(mlj_model)...)
+function model_init(mlj_model::CatBoostRegressor; kw...)
+    return catboost.CatBoostRegressor(; mlj_to_kwargs(mlj_model)..., kw...)
 end
 
 function MMI.fit(mlj_model::CatBoostRegressor, verbosity::Int, X, y)
-    silent = verbosity > 0 ? false : true
+    verbose = verbosity > 0 ? false : true
 
-    py_X = to_pandas(X)
+    X_preprocessed, cat_features, text_features = prepare_input(X)
+    py_X = to_pandas(X_preprocessed)
     py_y = numpy.array(y)
 
-    model = model_init(mlj_model)
-    model.fit(py_X, py_y; silent=silent)
+    model = model_init(mlj_model; cat_features, text_features, verbose)
+    model.fit(py_X, py_y)
 
     cache = nothing
     report = (feature_importances=feature_importances(model),)

test/mlj_interface.jl

@@ -13,7 +13,7 @@
 
     @testset "CatBoostClassifier" begin
         X = DataFrame(; a=[1, 4, 5, 6], b=[4, 5, 6, 7])
-        y = [0, 0, 1, 1]
+        y = coerce([0, 0, 1, 1], OrderedFactor)
 
         # MLJ Interface
         model = CatBoostClassifier(; iterations=5)
@@ -42,15 +42,19 @@
 
     @testset "generic interface tests" begin
         @testset "CatBoostRegressor" begin
-            failures, summary = MLJTestInterface.test([CatBoostRegressor],
-                                                      MLJTestInterface.make_regression()...;
+            data = MLJTestInterface.make_regression()
+            X = DataFrame(data[1])
+            y = data[2]
+            failures, summary = MLJTestInterface.test([CatBoostRegressor], X, y;
                                                       mod=@__MODULE__, verbosity=0, # bump to debug
                                                       throw=false)
             @test isempty(failures)
         end
         @testset "CatBoostClassifier" begin
             for data in [MLJTestInterface.make_binary(), MLJTestInterface.make_multiclass()]
-                failures, summary = MLJTestInterface.test([CatBoostClassifier], data...;
+                X = DataFrame(data[1])
+                y = data[2]
+                failures, summary = MLJTestInterface.test([CatBoostClassifier], X, y;
                                                           mod=@__MODULE__, verbosity=0, # bump to debug
                                                           throw=false)
                 @test isempty(failures)