Skip to content

Commit

Permalink
Updated dtreeviz usage to fix tests.
Browse files Browse the repository at this point in the history
  • Loading branch information
gabriel-rodriguez committed Apr 9, 2024
1 parent 17ce551 commit b9563d4
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions src/marta/analyzer/classification.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay
from dtreeviz.trees import dtreeviz
from dtreeviz.trees import model
import matplotlib.pyplot as plt
import numpy as np

Expand Down Expand Up @@ -99,7 +99,7 @@ def export_graph_tree(self, output_path="") -> None:
orientation = "TD"
else:
orientation = "LR"
viz = dtreeviz(
viz = model(
self.clf,
x_data=self.data,
y_data=self.target_data,
Expand Down

0 comments on commit b9563d4

Please sign in to comment.