decision_tree.py

def create_and_visualize_tree(X_train,y_train,max_depth=3):
  decision_tree = DecisionTreeClassifier(max_depth=max_depth, min_samples_leaf=1,random_state=1)
  decision_tree = decision_tree.fit(X_train, y_train)
  tree_str = export_graphviz(decision_tree, feature_names=X_train.columns, 
                                     filled=True, out_file=None)
  graph = pydotplus.graph_from_dot_data(tree_str)  
  graph.write_png('dt.png')
  display(Image('dt.png'))
  return decision_tree