chore: add download plot button

app.py

@@ -92,15 +92,15 @@ def toggle_inputs(use_average):
     else:
         return gr.update(interactive=True, visible=True), gr.update(interactive=True, visible=True)
 
-i = 0
+
+plot_path = "plots/last_plot.pdf"
+
 
 def plot_distances(model, dataset, use_average, cluster_method, cluster_method_param, plot_fn):
     """
     Plots all languages from the distances matrix using t-SNE.
     """
 
-    global i
-
     updated_matrix, updated_languages = filter_languages_nan(model, dataset, use_average)
 
     if cluster_method == "HDBSCAN":
@@ -126,9 +126,19 @@ def plot_distances(model, dataset, use_average, cluster_method, cluster_method_p
 
     fig = plot_fn(model, dataset, use_average, filtered_matrix, filtered_languages, clusters, legends)
     fig.tight_layout()
-    fig.savefig(f"plots/plot_{i}.pdf", format="pdf")
-    i += 1
-    return fig
+    fig.savefig(plot_path, format="pdf")
+    return fig, gr.DownloadButton(label="Download Plot", value=plot_path)
+
+
+def plot_families_subfamilies(families, model, dataset, use_average, figsize_h, figsize_w):
+    updated_matrix, updated_languages = filter_languages_nan(model, dataset, use_average)
+    updated_matrix, updated_languages = filter_languages_by_families(updated_matrix, updated_languages, families)
+
+    clusters, legends = cluster_languages_by_subfamilies(updated_languages)
+    fig = plot_mst(model, dataset, use_average, updated_matrix, updated_languages, clusters, legends, fig_size=(figsize_w, figsize_h))
+    fig.tight_layout()
+    fig.savefig(plot_path, format="pdf")
+    return fig, gr.DownloadButton(label="Download Plot", value=plot_path)
 
 
 with gr.Blocks() as demo:
@@ -185,19 +195,22 @@ with gr.Blocks() as demo:
             plot_tsne_button = gr.Button("Plot t-SNE")
             plot_umap_button = gr.Button("Plot UMAP")
             plot_mst_button = gr.Button("Plot MST")
+        
+        with gr.Row():
+            download_plot_button = gr.DownloadButton("Download Plot")
 
         with gr.Row():
             plot_output = gr.Plot(label="Distance Plot")
 
         plot_tsne_button.click(fn=partial(plot_distances, plot_fn=plot_distances_tsne), 
                             inputs=[model_input, dataset_input, average_checkbox, cluster_method_input, clusters_input], 
-                            outputs=plot_output)
+                            outputs=[plot_output, download_plot_button])
         plot_umap_button.click(fn=partial(plot_distances, plot_fn=plot_distances_umap), 
                             inputs=[model_input, dataset_input, average_checkbox, cluster_method_input, clusters_input], 
-                            outputs=plot_output)
+                            outputs=[plot_output, download_plot_button])
         plot_mst_button.click(fn=partial(plot_distances, plot_fn=plot_mst), 
                             inputs=[model_input, dataset_input, average_checkbox, cluster_method_input, clusters_input], 
-                            outputs=plot_output)
+                            outputs=[plot_output, download_plot_button])
 
     with gr.Tab(label="Language Families Subplot"):
         
@@ -227,23 +240,15 @@ with gr.Blocks() as demo:
             plot_family_button = gr.Button("Plot Families")
             plot_figsize_h_input = gr.Slider(label="Figure Height", minimum=5, maximum=30, step=1, value=15)
             plot_figsize_w_input = gr.Slider(label="Figure Width", minimum=5, maximum=30, step=1, value=15)
+
+        with gr.Row():
+            download_families_plot_button = gr.DownloadButton("Download Plot", value=plot_path)
+
         plot_family_output = gr.Plot(label="Families Plot")
-        def plot_families_subfamilies(families, model, dataset, use_average, figsize_h, figsize_w):
-            global i
-
-            updated_matrix, updated_languages = filter_languages_nan(model, dataset, use_average)
-            updated_matrix, updated_languages = filter_languages_by_families(updated_matrix, updated_languages, families)
-
-            clusters, legends = cluster_languages_by_subfamilies(updated_languages)
-            fig = plot_mst(model, dataset, use_average, updated_matrix, updated_languages, clusters, legends, fig_size=(figsize_w, figsize_h))
-            fig.tight_layout()
-            fig.savefig(f"plots/plot_{i}.pdf", format="pdf")
-            i += 1
-            return fig
         
         plot_family_button.click(fn=plot_families_subfamilies,
                                 inputs=[checked_families_input, model_input, dataset_input, average_checkbox, plot_figsize_h_input, plot_figsize_w_input],
-                                outputs=plot_family_output)
+                                outputs=[plot_family_output, download_families_plot_button])
     
 
 demo.launch(share=True)