Skip to content

Commit

Permalink
Print the feature names in report.print_tree()
Browse files Browse the repository at this point in the history
This commit addresses issue JuliaAI#23 by modifying TreePrinter struct and fit function to include the feature_names parameter.
  • Loading branch information
adarshpalaskar1 committed Feb 20, 2024
1 parent 51dc062 commit b1a41d5
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions src/MLJDecisionTreeInterface.jl
Original file line number Diff line number Diff line change
Expand Up @@ -15,9 +15,10 @@ const PKG = "MLJDecisionTreeInterface"

struct TreePrinter{T}
tree::T
features::Vector{Symbol}
end
(c::TreePrinter)(depth) = DT.print_tree(c.tree, depth)
(c::TreePrinter)() = DT.print_tree(c.tree, 5)
(c::TreePrinter)(depth) = DT.print_tree(c.tree, depth, feature_names = c.features)
(c::TreePrinter)() = DT.print_tree(c.tree, 5, feature_names = c.features)

Base.show(stream::IO, c::TreePrinter) =
print(stream, "TreePrinter object (call with display depth)")
Expand Down Expand Up @@ -71,7 +72,7 @@ function MMI.fit(
cache = nothing
report = (
classes_seen=classes_seen,
print_tree=TreePrinter(tree),
print_tree=TreePrinter(tree, features),
features=features,
)
return fitresult, cache, report
Expand Down

0 comments on commit b1a41d5

Please sign in to comment.