Update app.py

app.py

@@ -39,7 +39,7 @@ MODELS = [
 
 # ========================================
 
-def create_plot(feature_string, n_estimators, model_idx):
+def create_plot(feature_string, n_estimators, max_depth, model_idx):
     np.random.seed(SEED)
 
     feature_list = feature_string.split(',')
@@ -58,6 +58,10 @@ def create_plot(feature_string, n_estimators, model_idx):
 
     model_name = MODEL_NAMES[model_idx]
     model = MODELS[model_idx]
+
+    if model_idx != 0: model.n_estimators = n_estimators
+    if model_idx != 3: model.max_depth = max_depth
+    if model_idx == 3: model.estimator.max_depth = max_depth
     
     model.fit(X, y)
     score = round(model.score(X, y), 3)
@@ -94,11 +98,23 @@ def iter_grid(n_rows, n_cols):
                 with gr.Column():
                     yield
 
+info = '''
+## Plot the decision surfaces of ensembles of trees on the Iris dataset
+
+This plot compares the decision surfaces learned by a decision tree classifier, a random forest classifier, an extra-trees classifier, and by an AdaBoost classifier.
+
+There are in total four features in the Iris dataset. In this example you can select two features at a time for visualization purposes using the dropdown box below.
+
+You can also vary the number of estimators in the ensembles and the max depth of the trees using the sliders.
+'''
+
 with gr.Blocks() as demo:
+    gr.Markdown(info)
     selections = combinations(FEATURES, 2)
     selections = [f'{s[0]}, {s[1]}' for s in selections] 
     dd = gr.Dropdown(selections, value=selections[0], interactive=True, label="Input features")
-    slider = gr.Slider(1, 100, value=30, step=1, label='n_estimators')
+    slider_estimators = gr.Slider(1, 100, value=30, step=1, label='n_estimators')
+    slider_max_depth = gr.Slider(1, 50, value=10, step=1, label='max_depth')
 
     counter = 0
     for _ in iter_grid(2, 2):
@@ -108,10 +124,11 @@ with gr.Blocks() as demo:
         plot = gr.Plot(label=f'{MODEL_NAMES[counter]}')
         fn = partial(create_plot, model_idx=counter)
 
-        dd.change(fn, inputs=[dd, slider], outputs=[plot])
-        slider.change(fn, inputs=[dd, slider], outputs=[plot])
-        demo.load(fn, inputs=[dd, slider], outputs=[plot])
+        dd.change(fn, inputs=[dd, slider_estimators, slider_max_depth], outputs=[plot])
+        slider_estimators.change(fn, inputs=[dd, slider_estimators, slider_max_depth], outputs=[plot])
+        slider_max_depth.change(fn, inputs=[dd, slider_estimators, slider_max_depth], outputs=[plot])
+        demo.load(fn, inputs=[dd, slider_estimators, slider_max_depth], outputs=[plot])
 
         counter += 1
 
-demo.launch()
+demo.launch(share=True, debug=True)