Create app.py

app.py

@@ -0,0 +1,361 @@
+import gradio as gr
+from flwr_datasets import FederatedDataset
+from flwr_datasets.partitioner import (
+    DirichletPartitioner,
+    IidPartitioner,
+    PathologicalPartitioner,
+    ShardPartitioner,
+    LinearPartitioner,
+    SquarePartitioner,
+    ExponentialPartitioner,
+    NaturalIdPartitioner
+)
+from flwr_datasets.visualization import plot_label_distributions
+import matplotlib.pyplot as plt
+
+partitioner_types = {
+    "DirichletPartitioner": DirichletPartitioner,
+    "IidPartitioner": IidPartitioner,
+    "PathologicalPartitioner": PathologicalPartitioner,
+    "ShardPartitioner": ShardPartitioner,
+    "LinearPartitioner": LinearPartitioner,
+    "SquarePartitioner": SquarePartitioner,
+    "ExponentialPartitioner": ExponentialPartitioner,
+    "NaturalIdPartitioner": NaturalIdPartitioner,
+}
+
+partitioner_parameters = {
+    "DirichletPartitioner": ["num_partitions", "alpha", "partition_by", "min_partition_size", "self_balancing"],
+    "IidPartitioner": ["num_partitions"],
+    "PathologicalPartitioner": ["num_partitions", "partition_by", "num_classes_per_partition", "class_assignment_mode"],
+    "ShardPartitioner": ["num_partitions", "partition_by", "num_shards_per_partition", "shard_size", "keep_incomplete_shard"],
+    "NaturalIdPartitioner": ["partition_by"],
+    "LinearPartitioner": ["num_partitions"],
+    "SquarePartitioner": ["num_partitions"],
+    "ExponentialPartitioner": ["num_partitions"],
+}
+
+def update_parameter_visibility(partitioner_type):
+    print("calling update_parameter_visibility")
+    print(partitioner_type)
+    required_params = partitioner_parameters.get(partitioner_type, [])
+    updates = []
+    # For num_partitions_input
+    if "num_partitions" in required_params:
+        updates.append(gr.update(visible=True))
+    else:
+        updates.append(gr.update(visible=False))
+    # For alpha_input
+    if "alpha" in required_params:
+        updates.append(gr.update(visible=True))
+    else:
+        updates.append(gr.update(visible=False))
+    # For partition_by_input
+    if "partition_by" in required_params:
+        updates.append(gr.update(visible=True))
+    else:
+        updates.append(gr.update(visible=False))
+    # For min_partition_size_input
+    if "min_partition_size" in required_params:
+        updates.append(gr.update(visible=True))
+    else:
+        updates.append(gr.update(visible=False))
+    # For self_balancing_input
+    if "self_balancing" in required_params:
+        updates.append(gr.update(visible=True))
+    else:
+        updates.append(gr.update(visible=False))
+    # For num_classes_per_partition_input
+    if "num_classes_per_partition" in required_params:
+        updates.append(gr.update(visible=True))
+    else:
+        updates.append(gr.update(visible=False))
+    # For class_assignment_mode_input
+    if "class_assignment_mode" in required_params:
+        updates.append(gr.update(visible=True))
+    else:
+        updates.append(gr.update(visible=False))
+    # For num_shards_per_partition_input
+    if "num_shards_per_partition" in required_params:
+        updates.append(gr.update(visible=True))
+    else:
+        updates.append(gr.update(visible=False))
+    # For shard_size_input
+    if "shard_size" in required_params:
+        updates.append(gr.update(visible=True))
+    else:
+        updates.append(gr.update(visible=False))
+    # For keep_incomplete_shard_input
+    if "keep_incomplete_shard" in required_params:
+        updates.append(gr.update(visible=True))
+    else:
+        updates.append(gr.update(visible=False))
+    return updates
+
+def partition_and_plot(
+    dataset,
+    partitioner_type,
+    num_partitions,
+    alpha,
+    partition_by,
+    min_partition_size,
+    self_balancing,
+    num_classes_per_partition,
+    class_assignment_mode,
+    num_shards_per_partition,
+    shard_size,
+    keep_incomplete_shard,
+    label_name,
+    title,
+    legend,
+    verbose_labels,
+    size_unit,
+    partition_id_axis,
+):
+    partitioner_params = {}
+    try:
+        if partitioner_type == "DirichletPartitioner":
+            partitioner_params = {
+                "num_partitions": int(num_partitions),
+                "partition_by": partition_by,
+                "alpha": float(alpha),
+                "min_partition_size": int(min_partition_size),
+                "self_balancing": self_balancing,
+            }
+        elif partitioner_type == "IidPartitioner":
+            partitioner_params = {
+                "num_partitions": int(num_partitions),
+            }
+        elif partitioner_type == "PathologicalPartitioner":
+            partitioner_params = {
+                "num_partitions": int(num_partitions),
+                "partition_by": partition_by,
+                "num_classes_per_partition": int(num_classes_per_partition),
+                "class_assignment_mode": class_assignment_mode,
+            }
+        elif partitioner_type == "ShardPartitioner":
+            partitioner_params = {
+                "num_partitions": int(num_partitions),
+                "partition_by": partition_by,
+                "num_shards_per_partition": int(num_shards_per_partition),
+                "shard_size": int(shard_size),
+                "keep_incomplete_shard": keep_incomplete_shard == "True",
+            }
+        elif partitioner_type == "NaturalIdPartitioner":
+            partitioner_params = {
+                "partition_by": partition_by,
+            }
+        elif partitioner_type in ["LinearPartitioner", "SquarePartitioner", "ExponentialPartitioner"]:
+            partitioner_params = {
+                "num_partitions": int(num_partitions),
+            }
+
+        partitioner_class = partitioner_types[partitioner_type]
+        partitioner = partitioner_class(**partitioner_params)
+        fds = FederatedDataset(
+            dataset=dataset,
+            partitioners={
+                "train": partitioner,
+            },
+            trust_remote_code=True,
+        )
+        partitioner = fds.partitioners["train"]
+        figure, axis, dataframe = plot_label_distributions(
+            partitioner=partitioner,
+            label_name=label_name,
+            title=title,
+            legend=legend,
+            verbose_labels=verbose_labels,
+            size_unit=size_unit,
+            partition_id_axis=partition_id_axis,
+        )
+
+        # Save plot to a file
+        plot_filename = "label_distribution.png"
+        figure.savefig(plot_filename, bbox_inches='tight')
+
+        # Generate the code
+        partitioner_params_str = "\n"
+        n_params = len(partitioner_params)
+        i = 0
+        for k, v in partitioner_params.items():
+            if isinstance(v, str):
+                v = f'"{v}"'
+            if i != (n_params - 1):
+                partitioner_params_str = partitioner_params_str + f"\t{k} = {v},\n"    
+            else:
+                partitioner_params_str = partitioner_params_str + f"\t{k} = {v}\n" 
+            i +=1
+
+        code = f"""
+from flwr_datasets import FederatedDataset
+from flwr_datasets.partitioner import {partitioner_type}
+from flwr_datasets.visualization import plot_label_distributions
+
+partitioner = {partitioner_type}({partitioner_params_str})
+fds = FederatedDataset(
+    dataset="{dataset}",
+    partitioners={{
+        "train": partitioner,
+    }},
+    trust_remote_code=True,
+)
+partitioner = fds.partitioners["train"]
+figure, axis, dataframe = plot_label_distributions(
+    partitioner=partitioner,
+    label_name="label",
+    title="{title}",
+    legend={legend},
+    verbose_labels={verbose_labels},
+    size_unit="{size_unit}",
+    partition_id_axis="{partition_id_axis}",
+)
+    """
+        return plot_filename, code#, plot_filename # with df: plot_filename, code, dataframe, plot_filename
+    except Exception as e:
+        # Return error messages
+        error_message = str(e)
+        return None, f"Error: {error_message}", None, None
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Federated Dataset: Partitioning Visualization")
+    gr.Markdown("See partitioned datasets for Federated Learning experiments. The partitioning and visualization was created using `flwr-datasets`.")
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            # gr.Markdown("## Federated Dataset Parameters")
+            with gr.Accordion("Federated Dataset Parameters", open=True):
+                dataset_input = gr.Textbox(label="Dataset", value="cifar10")
+                partitioner_type_input = gr.Dropdown(label="Partitioner", choices=list(partitioner_types.keys()), value="DirichletPartitioner")
+                num_partitions_input = gr.Number(label="num_partitions", value=10, visible=True)
+                alpha_input = gr.Number(label="alpha", value=0.3, visible=True)
+                partition_by_input = gr.Textbox(label="partition_by", value="label", visible=True)
+                min_partition_size_input = gr.Number(label="min_partition_size", value=0, visible=True)
+                self_balancing_input = gr.Radio(label="self_balancing", choices=[True, False], value=False, visible=True)
+    
+                num_classes_per_partition_input = gr.Number(label="num_classes_per_partition", value=2, visible=False)
+                class_assignment_mode_input = gr.Dropdown(label="class_assignment_mode", choices=["random", "first-deterministic", "deterministic"], value="first-deterministic", visible=False)
+                num_shards_per_partition_input = gr.Number(label="num_shards_per_partition", value=2, visible=False)
+                shard_size_input = gr.Number(label="shard_size", value=0, visible=False)
+                keep_incomplete_shard_input = gr.Radio(label="keep_incomplete_shard", choices=["True", "False"], value="True", visible=False)
+            with gr.Accordion("Plot Parameters", open=False):
+                label_name = gr.Textbox(label="label_name", value="label")
+                title = gr.Textbox(label="title", value="Per Partition Label Distribution")
+                # legend_title = gr.Textbox(label="legend_title", value=None)
+                legend = gr.Radio(label="legend", choices=[True, False], value=True)
+                verbose_labels = gr.Radio(label="verbose_labels", choices=[True, False], value=True)
+                size_unit = gr.Radio(label="size_unit", choices=["absolute", "percent"], value="absolute")
+                partition_id_axis = gr.Radio(label="partition_id_axis", choices=["x", "y"], value="x")
+                
+            
+            
+            # Update parameter visibility when partitioner_type_input changes
+            partitioner_type_input.change(
+                fn=update_parameter_visibility,
+                inputs=[partitioner_type_input],
+                outputs=[
+                    num_partitions_input,
+                    alpha_input,
+                    partition_by_input,
+                    min_partition_size_input,
+                    self_balancing_input,
+                    num_classes_per_partition_input,
+                    class_assignment_mode_input,
+                    num_shards_per_partition_input,
+                    shard_size_input,
+                    keep_incomplete_shard_input
+                ]
+            )    
+        with gr.Column(scale=3):
+            gr.Markdown("## Label Distribution Plot")
+            plot_output = gr.Image(label="Label Distribution Plot")
+            submit_button = gr.Button("Partition and Plot", variant="primary")
+            # download_button = gr.DownloadButton(label="Download Plot", value="label_distribution.png")
+            gr.Markdown("## Code")
+            code_output = gr.Code(label="Code", language="python")
+            # Uncomment to show dataframe (note that it only works with header that is of type "string")
+            # gr.Markdown("## Partitioning DataFrame")
+            # dataframe_output = gr.Dataframe(label="Partitioning DataFrame")
+    size_skew_examples = gr.Examples(
+        examples=[
+            ["cifar10", "IidPartitioner", 10],
+            ["cifar10", "LinearPartitioner", 10],
+            ["cifar10", "SquarePartitioner", 10],
+            ["cifar10", "ExponentialPartitioner", 10],
+        ],
+        inputs=[
+            dataset_input,
+            partitioner_type_input,
+            num_partitions_input,
+        ],
+        label="Size Skew Examples",
+    )
+    
+    dirichlet_examples = gr.Examples(
+        examples=[
+            ["cifar10", "DirichletPartitioner", 10, 0.1, "label", 0, False, "absolute"],
+            ["cifar10", "DirichletPartitioner", 10, 0.1, "label", 0, False, "percent"],
+        ],
+        inputs=[
+            dataset_input,
+            partitioner_type_input,
+            num_partitions_input,
+            alpha_input,
+            partition_by_input,
+            min_partition_size_input,
+            self_balancing_input,
+            size_unit,
+        ],
+        label="Dirichlet Examples",
+    )
+
+    pathological_examples = gr.Examples(
+        examples=[
+            ["cifar10", "PathologicalPartitioner", 10, 2, "first-deterministic", "label"],
+            ["cifar10", "PathologicalPartitioner", 10, 3, "deterministic", "label"],
+        ],
+        inputs=[
+            dataset_input,
+            partitioner_type_input,
+            num_partitions_input,
+            num_classes_per_partition_input,
+            class_assignment_mode_input,
+            partition_by_input,
+        ],
+        label="Pathological Examples",
+    )
+    markdown = gr.Markdown("See more tutorial, examples and documentation on [https://flower.ai/docs/datasets/index.html](https://flower.ai/docs/datasets/index.html).")
+
+    # Set up the event handler for the submit_button
+    submit_button.click(
+        fn=partition_and_plot,
+        inputs=[
+            dataset_input,
+            partitioner_type_input,
+            num_partitions_input,
+            alpha_input,
+            partition_by_input,
+            min_partition_size_input,
+            self_balancing_input,
+            num_classes_per_partition_input,
+            class_assignment_mode_input,
+            num_shards_per_partition_input,
+            shard_size_input,
+            keep_incomplete_shard_input,
+            label_name,
+            title,
+            legend,
+            verbose_labels,
+            size_unit,
+            partition_id_axis,
+        ],
+        outputs=[
+            plot_output,
+            code_output,
+            # dataframe_output,
+            # download_button
+        ]
+    )
+
+if __name__ == "__main__":
+    demo.launch()