Fix sprinkle

aclai-lab · Dec 4, 2023 · c7b10bf · c7b10bf
1 parent 7cc4d57
commit c7b10bf
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/interfaces/MLJ.jl b/src/interfaces/MLJ.jl
@@ -118,9 +118,12 @@ function MMI.fit(m::SymbolicModel, verbosity::Integer, X, y, var_grouping, class
     cache  = nothing
     report = (
         printmodel                  = printer,
-        sprinkle                    = (Xnew, ynew)->begin
+        sprinkle                    = (Xnew, ynew; simplify = false)->begin
             (Xnew, ynew, var_grouping, classes_seen, w) = MMI.reformat(m, Xnew, ynew; passive_mode = true)
             preds, sprinkledmodel = ModalDecisionTrees.sprinkle(model, Xnew, ynew)
+            if simplify
+                sprinkledmodel = MDT.prune(model; simplify = true)
+            end
             preds, translate_function(sprinkledmodel)
         end,
         # TODO remove redundancy?