✨ Class names to be used as level names for target encoding

Author: EssamWisam

src/encoders/target_encoding/interface_mlj.jl

@@ -51,6 +51,8 @@ struct TargetEncoderResult{
     y_stat_given_feat_level::Dict{A, A}
     task::S            # "Regression", "Classification" 
     num_classes::I     # num_classes in case of classification
+    y_classes::A      # y_classes in case of classification
+    
 end
 
 
@@ -76,6 +78,7 @@ function MMI.fit(transformer::TargetEncoder, verbosity::Int, X, y)
         generic_cache[:y_stat_given_feat_level],
         generic_cache[:task],
         generic_cache[:num_classes],
+        generic_cache[:y_classes],
     )
     report = (encoded_features = generic_cache[:encoded_features],)        # report only has list of encoded features
     cache = nothing
@@ -90,6 +93,7 @@ function MMI.transform(transformer::TargetEncoder, fitresult, Xnew)
             fitresult.y_stat_given_feat_level,
         :num_classes => fitresult.num_classes,
         :task => fitresult.task,
+        :y_classes => fitresult.y_classes,
     )
     Xnew_transf = target_encoder_transform(Xnew, generic_cache)
     return Xnew_transf

src/encoders/target_encoding/target_encoding.jl

@@ -215,6 +215,7 @@ function target_encoder_fit(
         :num_classes => (task == "Regression") ? -1 : length(y_classes),
         :y_stat_given_feat_level => y_stat_given_feat_level,
         :encoded_features => encoded_features,
+        :y_classes => (task == "Regression") ? nothing : y_classes,
     )
     return cache
 end
@@ -243,11 +244,13 @@ function target_encoder_transform(X, cache)
     task = cache[:task]
     y_stat_given_feat_level = cache[:y_stat_given_feat_level]
     num_classes = cache[:num_classes]
+    y_classes = cache[:y_classes]
 
     return generic_transform(
         X,
         y_stat_given_feat_level;
         single_feat = task == "Regression" || (task == "Classification" && num_classes < 3),
-    )
+        use_levelnames = true,
+        custom_levels = y_classes,)
 end
 

test/encoders/target_encoding.jl

@@ -277,22 +277,21 @@ end
     X_tr = target_encoder_transform(X, cache)
 
     enc = (col, level) -> cache[:y_stat_given_feat_level][col][level]
-
     target = (
-        A_1 = [enc(:A, X[:A][i])[1] for i in 1:10],
-        A_2 = [enc(:A, X[:A][i])[2] for i in 1:10],
-        A_3 = [enc(:A, X[:A][i])[3] for i in 1:10],
+        A_0 = [enc(:A, X[:A][i])[1] for i in 1:10],
+        A_1 = [enc(:A, X[:A][i])[2] for i in 1:10],
+        A_2 = [enc(:A, X[:A][i])[3] for i in 1:10],
         B = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0],
-        C_1 = [enc(:C, X[:C][i])[1] for i in 1:10],
-        C_2 = [enc(:C, X[:C][i])[2] for i in 1:10],
-        C_3 = [enc(:C, X[:C][i])[3] for i in 1:10],
-        D_1 = [enc(:D, X[:D][i])[1] for i in 1:10],
-        D_2 = [enc(:D, X[:D][i])[2] for i in 1:10],
-        D_3 = [enc(:D, X[:D][i])[3] for i in 1:10],
+        C_0 = [enc(:C, X[:C][i])[1] for i in 1:10],
+        C_1 = [enc(:C, X[:C][i])[2] for i in 1:10],
+        C_2 = [enc(:C, X[:C][i])[3] for i in 1:10],
+        D_0 = [enc(:D, X[:D][i])[1] for i in 1:10],
+        D_1 = [enc(:D, X[:D][i])[2] for i in 1:10],
+        D_2 = [enc(:D, X[:D][i])[3] for i in 1:10],
         E = [1, 2, 3, 4, 5, 6, 6, 3, 2, 1],
-        F_1 = [enc(:F, X[:F][i])[1] for i in 1:10],
-        F_2 = [enc(:F, X[:F][i])[2] for i in 1:10],
-        F_3 = [enc(:F, X[:F][i])[3] for i in 1:10],
+        F_0 = [enc(:F, X[:F][i])[1] for i in 1:10],
+        F_1 = [enc(:F, X[:F][i])[2] for i in 1:10],
+        F_2 = [enc(:F, X[:F][i])[3] for i in 1:10],
     )
     for col in keys(target)
         @test all(X_tr[col] .== target[col])