chore: black

app.py

@@ -2,19 +2,23 @@ import gradio as gr
 import pandas as pd
 import numpy as np
 import os
-from utils import (plot_distances_tsne,
-                   plot_distances_umap, 
-                   cluster_languages_hdbscan, 
-                   cluster_languages_kmeans, 
-                   plot_mst, 
-                   cluster_languages_by_families, 
-                   cluster_languages_by_subfamilies,
-                   filter_languages_by_families)
+from utils import (
+    plot_distances_tsne,
+    plot_distances_umap,
+    cluster_languages_hdbscan,
+    cluster_languages_kmeans,
+    plot_mst,
+    cluster_languages_by_families,
+    cluster_languages_by_subfamilies,
+    filter_languages_by_families,
+)
 from functools import partial
 import datasets
 
 
-dataset = datasets.load_dataset("mshamrai/language-metric-data", split="train", trust_remote_code=True)
+dataset = datasets.load_dataset(
+    "mshamrai/language-metric-data", split="train", trust_remote_code=True
+)
 
 languages = dataset["languages_list"][0]
 average_distances_matrix = np.array(dataset["average_distances_matrix"][0])
@@ -27,7 +31,7 @@ distance_matrices = {
         MODELS[j]: np.array(dataset["distances_matrices"][0]["models"][i]["matrix"][j])
         for j in range(len(MODELS))
     }
-    for i in range(len(DATASETS)) 
+    for i in range(len(DATASETS))
 }
 
 
@@ -63,6 +67,7 @@ def get_similar_languages(model, dataset, selected_language, use_average, n):
     sorted_distances["Distance"] = sorted_distances["Distance"].round(4)
     return sorted_distances.head(n)
 
+
 def update_languages(model, dataset):
     """
     Returns the language list based on the given model and dataset.
@@ -85,21 +90,29 @@ def update_language_options(model, dataset, language, use_average):
 
 def toggle_inputs(use_average):
     if use_average:
-        return gr.update(interactive=False, visible=False), gr.update(interactive=False, visible=False)
+        return gr.update(interactive=False, visible=False), gr.update(
+            interactive=False, visible=False
+        )
     else:
-        return gr.update(interactive=True, visible=True), gr.update(interactive=True, visible=True)
+        return gr.update(interactive=True, visible=True), gr.update(
+            interactive=True, visible=True
+        )
 
 
 plot_path = "plots/last_plot.pdf"
 os.makedirs("plots", exist_ok=True)
 
 
-def plot_distances(model, dataset, use_average, cluster_method, cluster_method_param, plot_fn):
+def plot_distances(
+    model, dataset, use_average, cluster_method, cluster_method_param, plot_fn
+):
     """
     Plots all languages from the distances matrix using t-SNE.
     """
 
-    updated_matrix, updated_languages = filter_languages_nan(model, dataset, use_average)
+    updated_matrix, updated_languages = filter_languages_nan(
+        model, dataset, use_average
+    )
 
     if cluster_method == "HDBSCAN":
         filtered_matrix, filtered_languages, clusters = cluster_languages_hdbscan(
@@ -122,18 +135,41 @@ def plot_distances(model, dataset, use_average, cluster_method, cluster_method_p
     else:
         raise ValueError("Invalid cluster method")
 
-    fig = plot_fn(model, dataset, use_average, filtered_matrix, filtered_languages, clusters, legends)
+    fig = plot_fn(
+        model,
+        dataset,
+        use_average,
+        filtered_matrix,
+        filtered_languages,
+        clusters,
+        legends,
+    )
     fig.tight_layout()
     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)
+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 = 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)
@@ -145,108 +181,242 @@ with gr.Blocks() as demo:
     with gr.Row():
         model_input = gr.Dropdown(label="Model", choices=MODELS, value=MODELS[0])
         dataset_input = gr.Dropdown(
-            label="Dataset",
-            choices=DATASETS,
-            value=DATASETS[0]
+            label="Dataset", choices=DATASETS, value=DATASETS[0]
         )
 
     with gr.Tab(label="Closest Languages Table"):
         with gr.Row():
-            language_input = gr.Dropdown(label="Language", choices=languages, value=languages[0])
-            top_n_input = gr.Slider(label="Top N", minimum=1, maximum=30, step=1, value=10)
-        
+            language_input = gr.Dropdown(
+                label="Language", choices=languages, value=languages[0]
+            )
+            top_n_input = gr.Slider(
+                label="Top N", minimum=1, maximum=30, step=1, value=10
+            )
+
         output_table = gr.Dataframe(label="Similar Languages")
-        
-        model_input.change(fn=update_language_options, inputs=[model_input, dataset_input, language_input, average_checkbox], outputs=language_input)
-        dataset_input.change(fn=update_language_options, inputs=[model_input, dataset_input, language_input, average_checkbox], outputs=language_input)
-        language_input.change(fn=get_similar_languages, inputs=[model_input, dataset_input, language_input, average_checkbox, top_n_input], outputs=output_table)
-        model_input.change(fn=get_similar_languages, inputs=[model_input, dataset_input, language_input, average_checkbox, top_n_input], outputs=output_table)
-        dataset_input.change(fn=get_similar_languages, inputs=[model_input, dataset_input, language_input, average_checkbox, top_n_input], outputs=output_table)
-        top_n_input.change(fn=get_similar_languages, inputs=[model_input, dataset_input, language_input, average_checkbox, top_n_input], outputs=output_table)
+
+        model_input.change(
+            fn=update_language_options,
+            inputs=[model_input, dataset_input, language_input, average_checkbox],
+            outputs=language_input,
+        )
+        dataset_input.change(
+            fn=update_language_options,
+            inputs=[model_input, dataset_input, language_input, average_checkbox],
+            outputs=language_input,
+        )
+        language_input.change(
+            fn=get_similar_languages,
+            inputs=[
+                model_input,
+                dataset_input,
+                language_input,
+                average_checkbox,
+                top_n_input,
+            ],
+            outputs=output_table,
+        )
+        model_input.change(
+            fn=get_similar_languages,
+            inputs=[
+                model_input,
+                dataset_input,
+                language_input,
+                average_checkbox,
+                top_n_input,
+            ],
+            outputs=output_table,
+        )
+        dataset_input.change(
+            fn=get_similar_languages,
+            inputs=[
+                model_input,
+                dataset_input,
+                language_input,
+                average_checkbox,
+                top_n_input,
+            ],
+            outputs=output_table,
+        )
+        top_n_input.change(
+            fn=get_similar_languages,
+            inputs=[
+                model_input,
+                dataset_input,
+                language_input,
+                average_checkbox,
+                top_n_input,
+            ],
+            outputs=output_table,
+        )
 
         average_checkbox.change(
             fn=toggle_inputs,
             inputs=[average_checkbox],
-            outputs=[model_input, dataset_input]
+            outputs=[model_input, dataset_input],
         )
 
-        average_checkbox.change(fn=update_language_options, inputs=[model_input, dataset_input, language_input, average_checkbox], outputs=language_input)
-        average_checkbox.change(fn=get_similar_languages, inputs=[model_input, dataset_input, language_input, average_checkbox, top_n_input], outputs=output_table)
-
+        average_checkbox.change(
+            fn=update_language_options,
+            inputs=[model_input, dataset_input, language_input, average_checkbox],
+            outputs=language_input,
+        )
+        average_checkbox.change(
+            fn=get_similar_languages,
+            inputs=[
+                model_input,
+                dataset_input,
+                language_input,
+                average_checkbox,
+                top_n_input,
+            ],
+            outputs=output_table,
+        )
 
     with gr.Tab(label="Distance Plot"):
         with gr.Row():
-            cluster_method_input = gr.Dropdown(label="Cluster Method", choices=["HDBSCAN", "KMeans", "Family", "Subfamily"], value="HDBSCAN")
-            clusters_input = gr.Slider(label="Minimum Elements in a Cluster", minimum=2, maximum=10, step=1, value=2)
+            cluster_method_input = gr.Dropdown(
+                label="Cluster Method",
+                choices=["HDBSCAN", "KMeans", "Family", "Subfamily"],
+                value="HDBSCAN",
+            )
+            clusters_input = gr.Slider(
+                label="Minimum Elements in a Cluster",
+                minimum=2,
+                maximum=10,
+                step=1,
+                value=2,
+            )
 
         def update_clusters_input_option(cluster_method):
             if cluster_method == "HDBSCAN":
-                return gr.Slider(label="Minimum Elements in a Cluster", minimum=2, maximum=10, step=1, value=2, visible=True, interactive=True)
+                return gr.Slider(
+                    label="Minimum Elements in a Cluster",
+                    minimum=2,
+                    maximum=10,
+                    step=1,
+                    value=2,
+                    visible=True,
+                    interactive=True,
+                )
             elif cluster_method == "KMeans":
-                return gr.Slider(label="Number of Clusters", minimum=2, maximum=20, step=1, value=2, visible=True, interactive=True)
+                return gr.Slider(
+                    label="Number of Clusters",
+                    minimum=2,
+                    maximum=20,
+                    step=1,
+                    value=2,
+                    visible=True,
+                    interactive=True,
+                )
             else:
                 return gr.update(interactive=False, visible=False)
 
-        cluster_method_input.change(fn=update_clusters_input_option, inputs=[cluster_method_input], outputs=clusters_input)
-            
+        cluster_method_input.change(
+            fn=update_clusters_input_option,
+            inputs=[cluster_method_input],
+            outputs=clusters_input,
+        )
+
         with gr.Row():
             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, 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, 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, download_plot_button])
+        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, 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, 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, download_plot_button],
+        )
 
     with gr.Tab(label="Language Families Subplot"):
-        
-        checked_families_input = gr.CheckboxGroup(label="Language Families", 
-                             choices=[
-                                'Afroasiatic',
-                                'Austroasiatic',
-                                'Austronesian',
-                                'Constructed',
-                                'Creole',
-                                'Dravidian',
-                                'Germanic',
-                                'Indo-European',
-                                'Japonic',
-                                'Kartvelian',
-                                'Koreanic',
-                                'Language Isolate',
-                                'Niger-Congo',
-                                'Northeast Caucasian',
-                                'Romance',
-                                'Sino-Tibetan',
-                                'Turkic',
-                                'Uralic'
-                            ], 
-                            value=["Indo-European"])
+
+        checked_families_input = gr.CheckboxGroup(
+            label="Language Families",
+            choices=[
+                "Afroasiatic",
+                "Austroasiatic",
+                "Austronesian",
+                "Constructed",
+                "Creole",
+                "Dravidian",
+                "Germanic",
+                "Indo-European",
+                "Japonic",
+                "Kartvelian",
+                "Koreanic",
+                "Language Isolate",
+                "Niger-Congo",
+                "Northeast Caucasian",
+                "Romance",
+                "Sino-Tibetan",
+                "Turkic",
+                "Uralic",
+            ],
+            value=["Indo-European"],
+        )
         with gr.Row():
             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)
+            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)
+            download_families_plot_button = gr.DownloadButton(
+                "Download Plot", value=plot_path
+            )
 
         plot_family_output = gr.Plot(label="Families Plot")
-        
-        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, download_families_plot_button])
-    
+
+        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, download_families_plot_button],
+        )
+
 
 demo.launch(share=True)

constants.py

@@ -104,7 +104,7 @@ language_subfamilies = {
     "Western Punjabi": "Punjabi",
     "Yoruba": "Yoruboid",
     "Esperanto": "Constructed",
-    "Crimean Tatar": "Kypchak"
+    "Crimean Tatar": "Kypchak",
 }
 
 language_families = {
@@ -213,5 +213,5 @@ language_families = {
     "Western Punjabi": "Indo-European",
     "Yoruba": "Niger-Congo",
     "Esperanto": "Constructed",
-    "Crimean Tatar": "Turkic"
+    "Crimean Tatar": "Turkic",
 }

utils.py

@@ -21,7 +21,11 @@ def filter_languages_by_families(matrix, languages, families):
     Returns:
     - filtered_languages: list of languages that belong to the specified families.
     """
-    filtered_languages = [(i, lang) for i, lang in enumerate(languages) if language_families[lang] in families]
+    filtered_languages = [
+        (i, lang)
+        for i, lang in enumerate(languages)
+        if language_families[lang] in families
+    ]
     filtered_indices = [i for i, lang in filtered_languages]
     filtered_languages = [lang for i, lang in filtered_languages]
     filtered_matrix = matrix[np.ix_(filtered_indices, filtered_indices)]
@@ -51,13 +55,25 @@ def cluster_languages_by_families(languages):
 
 
 def cluster_languages_by_subfamilies(languages):
-    labels = [language_families[lang] + f" ({language_subfamilies[lang]})" for lang in languages]
+    labels = [
+        language_families[lang] + f" ({language_subfamilies[lang]})"
+        for lang in languages
+    ]
     legend = sorted(set(labels))
     clusters = [legend.index(family) for family in labels]
     return clusters, legend
 
 
-def plot_mst(model, dataset, use_average, matrix, languages, clusters, legend=None, fig_size=(20,20)):
+def plot_mst(
+    model,
+    dataset,
+    use_average,
+    matrix,
+    languages,
+    clusters,
+    legend=None,
+    fig_size=(20, 20),
+):
     """
     Plots a Minimum Spanning Tree (MST) from a given distance matrix, node labels, and cluster assignments.
 
@@ -68,21 +84,21 @@ def plot_mst(model, dataset, use_average, matrix, languages, clusters, legend=No
     """
     # Create an empty undirected graph
     G = nx.Graph()
-    
+
     # Number of nodes
     N = len(languages)
-    
+
     # Add edges to the graph from the distance matrix.
     # Only iterate over the upper triangle of the matrix (i < j)
     for i in range(N):
         for j in range(i + 1, N):
             G.add_edge(i, j, weight=matrix[i, j])
-    
+
     # Compute the Minimum Spanning Tree using NetworkX's built-in function.
     mst = nx.minimum_spanning_tree(G)
-    
+
     # Choose a layout for the MST. Here we use Kamada-Kawai layout which considers edge weights.
-    pos = nx.kamada_kawai_layout(mst, weight='weight')
+    pos = nx.kamada_kawai_layout(mst, weight="weight")
 
     # Map each cluster to a color
     unique_clusters = sorted(set(clusters))
@@ -90,22 +106,24 @@ def plot_mst(model, dataset, use_average, matrix, languages, clusters, legend=No
     cluster_colors = {cluster: cmap[i] for i, cluster in enumerate(unique_clusters)}
 
     node_colors = [cluster_colors.get(cluster) for cluster in clusters]
-    
+
     # Create a figure for plotting.
     fig, ax = plt.subplots(figsize=fig_size)
-    
+
     # Draw the MST edges.
-    nx.draw_networkx_edges(mst, pos, edge_color='gray', ax=ax)
-    
+    nx.draw_networkx_edges(mst, pos, edge_color="gray", ax=ax)
+
     # Draw the nodes with colors corresponding to their clusters.
-    nx.draw_networkx_nodes(mst, pos, node_color=node_colors, node_size=100, ax=ax, alpha=0.7)
+    nx.draw_networkx_nodes(
+        mst, pos, node_color=node_colors, node_size=100, ax=ax, alpha=0.7
+    )
 
     # Instead of directly drawing labels, we create text objects to adjust them later
     texts = []
     for i, label in enumerate(languages):
         x, y = pos[i]
         texts.append(ax.text(x, y, label, fontsize=10))
-    
+
     # Adjust text labels to minimize overlap.
     # The arrowprops argument can draw arrows from labels to nodes if desired.
     adjust_text(texts, expand_text=(1.05, 1.2))
@@ -114,17 +132,27 @@ def plot_mst(model, dataset, use_average, matrix, languages, clusters, legend=No
     if legend is None:
         legend = {cluster: str(cluster) for cluster in unique_clusters}
     legend_handles = [
-        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor=cluster_colors[cluster], markersize=10, alpha=0.7, label=legend[cluster])
+        plt.Line2D(
+            [0],
+            [0],
+            marker="o",
+            color="w",
+            markerfacecolor=cluster_colors[cluster],
+            markersize=10,
+            alpha=0.7,
+            label=legend[cluster],
+        )
         for cluster in unique_clusters
     ]
     ax.legend(handles=legend_handles, title="Clusters", loc="best")
-    
+
     # Remove axis for clarity.
-    ax.axis('off')
+    ax.axis("off")
     # ax.set_title(f"Minimum Spanning Tree of Languages ({'Average' if use_average else f'{model}, {dataset}'})")
 
     return fig
 
+
 def cluster_languages_kmeans(dist_matrix, languages, n_clusters=5):
     """
     Clusters languages using a distance matrix and KMeans.
@@ -172,9 +200,7 @@ def cluster_languages_hdbscan(dist_matrix, languages, min_cluster_size=2):
     - clusters: list of length N containing the cluster assignment (or ID) for each language.
     """
     # Perform clustering using HDBSCAN with the precomputed distance matrix
-    clustering_model = HDBSCAN(
-        metric='precomputed', min_cluster_size=min_cluster_size
-    )
+    clustering_model = HDBSCAN(metric="precomputed", min_cluster_size=min_cluster_size)
     clusters = clustering_model.fit_predict(dist_matrix)
 
     # Filter out points belonging to cluster -1 using NumPy
@@ -185,7 +211,9 @@ def cluster_languages_hdbscan(dist_matrix, languages, min_cluster_size=2):
     return filtered_matrix, filtered_languages, filtered_clusters
 
 
-def plot_distances_tsne(model, dataset, use_average, matrix, languages, clusters, legend=None):
+def plot_distances_tsne(
+    model, dataset, use_average, matrix, languages, clusters, legend=None
+):
     """
     Plots all languages from the distances matrix using t-SNE and colors them by clusters.
     """
@@ -198,7 +226,12 @@ def plot_distances_tsne(model, dataset, use_average, matrix, languages, clusters
     cluster_colors = {cluster: cmap[i] for i, cluster in enumerate(unique_clusters)}
 
     fig, ax = plt.subplots(figsize=(16, 12))
-    scatter = ax.scatter(tsne_results[:, 0], tsne_results[:, 1], c=[cluster_colors[cluster] for cluster in clusters], alpha=0.7)
+    scatter = ax.scatter(
+        tsne_results[:, 0],
+        tsne_results[:, 1],
+        c=[cluster_colors[cluster] for cluster in clusters],
+        alpha=0.7,
+    )
 
     # for i, lang in enumerate(languages):
     #     ax.text(tsne_results[i, 0], tsne_results[i, 1], lang, fontsize=8, alpha=0.8)
@@ -208,7 +241,7 @@ def plot_distances_tsne(model, dataset, use_average, matrix, languages, clusters
     for i, label in enumerate(languages):
         x, y = tsne_results[i, 0], tsne_results[i, 1]
         texts.append(ax.text(x, y, label, fontsize=10))
-    
+
     # Adjust text labels to minimize overlap.
     # The arrowprops argument can draw arrows from labels to nodes if desired.
     adjust_text(texts, expand_text=(1.05, 1.2))
@@ -217,18 +250,30 @@ def plot_distances_tsne(model, dataset, use_average, matrix, languages, clusters
     if legend is None:
         legend = {cluster: str(cluster) for cluster in unique_clusters}
     legend_handles = [
-        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor=cluster_colors[cluster], markersize=10, label=legend[cluster])
+        plt.Line2D(
+            [0],
+            [0],
+            marker="o",
+            color="w",
+            markerfacecolor=cluster_colors[cluster],
+            markersize=10,
+            label=legend[cluster],
+        )
         for cluster in unique_clusters
     ]
     ax.legend(handles=legend_handles, title="Clusters", loc="best")
 
-    ax.set_title(f"t-SNE Visualization of Language Distances ({'Average' if use_average else f'{model}, {dataset}'})")
+    ax.set_title(
+        f"t-SNE Visualization of Language Distances ({'Average' if use_average else f'{model}, {dataset}'})"
+    )
     ax.set_xlabel("t-SNE Dimension 1")
     ax.set_ylabel("t-SNE Dimension 2")
     return fig
 
 
-def plot_distances_umap(model, dataset, use_average, matrix, languages, clusters, legend=None):
+def plot_distances_umap(
+    model, dataset, use_average, matrix, languages, clusters, legend=None
+):
     """
     Plots all languages from the distances matrix using UMAP and colors them by clusters.
     """
@@ -242,7 +287,12 @@ def plot_distances_umap(model, dataset, use_average, matrix, languages, clusters
     cluster_colors = {cluster: cmap[i] for i, cluster in enumerate(unique_clusters)}
 
     fig, ax = plt.subplots(figsize=(16, 12))
-    scatter = ax.scatter(umap_results[:, 0], umap_results[:, 1], c=[cluster_colors[cluster] for cluster in clusters], alpha=0.7)
+    scatter = ax.scatter(
+        umap_results[:, 0],
+        umap_results[:, 1],
+        c=[cluster_colors[cluster] for cluster in clusters],
+        alpha=0.7,
+    )
 
     # for i, lang in enumerate(languages):
     #     ax.text(umap_results[i, 0], umap_results[i, 1], lang, fontsize=8, alpha=0.8)
@@ -252,7 +302,7 @@ def plot_distances_umap(model, dataset, use_average, matrix, languages, clusters
     for i, label in enumerate(languages):
         x, y = umap_results[i, 0], umap_results[i, 1]
         texts.append(ax.text(x, y, label, fontsize=10))
-    
+
     # Adjust text labels to minimize overlap.
     # The arrowprops argument can draw arrows from labels to nodes if desired.
     adjust_text(texts, expand_text=(1.05, 1.2))
@@ -261,12 +311,22 @@ def plot_distances_umap(model, dataset, use_average, matrix, languages, clusters
     if legend is None:
         legend = {cluster: str(cluster) for cluster in unique_clusters}
     legend_handles = [
-        plt.Line2D([0], [0], marker='o', color='w', markerfacecolor=cluster_colors[cluster], markersize=10, label=legend[cluster])
+        plt.Line2D(
+            [0],
+            [0],
+            marker="o",
+            color="w",
+            markerfacecolor=cluster_colors[cluster],
+            markersize=10,
+            label=legend[cluster],
+        )
         for cluster in unique_clusters
     ]
     ax.legend(handles=legend_handles, title="Clusters", loc="best")
 
-    ax.set_title(f"UMAP Visualization of Language Distances ({'Average' if use_average else f'{model}, {dataset}'})")
+    ax.set_title(
+        f"UMAP Visualization of Language Distances ({'Average' if use_average else f'{model}, {dataset}'})"
+    )
     ax.set_xlabel("UMAP Dimension 1")
     ax.set_ylabel("UMAP Dimension 2")
-    return fig
+    return fig