Add plotting and dataset options

dataset_options.py

@@ -0,0 +1,257 @@
+import gradio as gr
+import numpy as np
+import numexpr
+import pandas as pd
+import time
+
+
+NUMEXPR_CONSTANTS = {
+    'pi': np.pi,
+    'PI': np.pi,
+    'e': np.e,
+}
+
+
+def get_function(function, xlim=(-1, 1), nsample=100):
+    x = np.linspace(xlim[0], xlim[1], nsample)
+    y = numexpr.evaluate(function, local_dict={'x': x, **NUMEXPR_CONSTANTS})
+    x = x.reshape(-1, 1)
+    return x, y
+
+
+
+def get_data_points(function, xlim=(-1, 1), nsample=10, sigma=0, seed=0):
+    num_points_to_generate = 100
+    if nsample > num_points_to_generate:
+        raise ValueError(f"nsample too large, limit to {num_points_to_generate}")
+
+    rng = np.random.default_rng(seed)
+    x = rng.uniform(xlim[0], xlim[1], size=num_points_to_generate)
+    x = x[:nsample]
+    x = np.sort(x)
+
+    rng = np.random.default_rng(seed)
+    noise = sigma * rng.standard_normal(nsample)
+    y = numexpr.evaluate(function, local_dict={'x': x, **NUMEXPR_CONSTANTS}) + noise
+
+    x = x.reshape(-1, 1)
+    return x, y
+
+
+class DatasetOptions:
+    def __init__(
+        self,
+        mode: str = "generate",
+        function: str = "x ** 2",
+        xmin: float = -1.0,
+        xmax: float = 1.0,
+        nsample: int = 30,
+        sigma: float = 0.0,
+        seed: int = 0,
+        csv_path: str = None,
+    ):
+        self.mode = mode
+
+        self.function = function
+        self.xmin = xmin
+        self.xmax = xmax
+        self.nsample = nsample
+        self.sigma = sigma
+        self.seed = seed
+
+        self.csv_path = csv_path
+
+        self.x, self.y = self._get_data()
+
+    def _get_data(self):
+        if self.mode == "generate":
+            return get_data_points(
+                function=self.function,
+                xlim=(self.xmin, self.xmax),
+                nsample=self.nsample,
+                sigma=self.sigma,
+                seed=self.seed,
+            )
+
+        elif self.mode == "csv":
+            if self.csv_path is None:
+                return np.array([]), np.array([])
+
+            df = pd.read_csv(self.csv_path)
+            if df.shape[1] != 2:
+                raise ValueError("CSV file must have exactly two columns")
+
+            x = df.iloc[:, 0].values.reshape(-1, 1)
+            y = df.iloc[:, 1].values
+            return x, y
+
+        else:
+            raise ValueError(f"Unknown dataset mode: {self.mode}")
+
+    def update(self, **kwargs):
+        return DatasetOptions(
+            mode=kwargs.get("mode", self.mode),
+            function=kwargs.get("function", self.function),
+            xmin=kwargs.get("xmin", self.xmin),
+            xmax=kwargs.get("xmax", self.xmax),
+            nsample=kwargs.get("nsample", self.nsample),
+            sigma=kwargs.get("sigma", self.sigma),
+            seed=kwargs.get("seed", self.seed),
+            csv_path=kwargs.get("csv_path", self.csv_path),
+        )
+
+    def _safe_hash(self, val: int) -> int | tuple[int, str]:
+        # special handling for -1 (same hash number as -2)
+        if val == -1:
+            return (-1, "special")
+        return val
+
+    def __hash__(self):
+        return hash(
+            (
+                self.mode,
+                self.function,
+                self._safe_hash(self.xmin),
+                self._safe_hash(self.xmax),
+                self.nsample,
+                self.sigma,
+                self.seed,
+                self.csv_path,
+            )
+        )
+
+
+class DatasetOptionsView:
+    def update_mode(self, mode: str, state: gr.State):
+        state = state.update(mode=mode)
+
+        if mode == "generate":
+            return (
+                state,
+                gr.update(visible=True),  # function
+                gr.update(visible=True),  # xmin
+                gr.update(visible=True),  # xmax 
+                gr.update(visible=True),  # sigma
+                gr.update(visible=True),  # nsample
+                gr.update(visible=True),  # regenerate
+                gr.update(visible=False),  # csv upload
+            )
+        elif mode == "csv":
+            return (
+                state,
+                gr.update(visible=False),  # function
+                gr.update(visible=False),  # xmin
+                gr.update(visible=False),  # xmax 
+                gr.update(visible=False),  # sigma
+                gr.update(visible=False),  # nsample
+                gr.update(visible=False),  # regenerate
+                gr.update(visible=True),  # csv upload
+            )
+        else:
+            raise ValueError(f"Unknown mode: {mode}")
+
+    def upload_csv(self, file, state):
+        try:
+            state = state.update(
+                mode="csv",
+                csv_path=file.name,
+            )
+
+        except Exception as e:
+            gr.Info(f"⚠️   {e}")
+
+        return state
+
+    def regenerate_data(self, state: gr.State):
+        seed = int(time.time() * 1000) % (2 ** 32)
+        state = state.update(seed=seed)
+        return state
+
+    def build(self, state: gr.State):
+        options = state.value
+
+        with gr.Column():
+            mode = gr.Radio(
+                label="Dataset",
+                choices=["generate", "csv"],
+                value="generate",
+            )
+
+            function = gr.Textbox(
+                label="Function (in terms of x)", 
+                value=options.function,
+            )
+            with gr.Row():
+                xmin = gr.Number(
+                    label="X min", 
+                    value=options.xmin,
+                )
+                xmax = gr.Number(
+                    label="X max", 
+                    value=options.xmax,
+                )
+            sigma = gr.Number(
+                label="Gaussian noise standard deviation",
+                value=options.sigma,
+            )
+            nsample = gr.Slider(
+                label="Number of samples", 
+                minimum=1, 
+                maximum=100, 
+                step=1, 
+                value=options.nsample,
+            )
+            regenerate = gr.Button("Regenerate Data")
+
+            csv_upload = gr.File(
+                label="Upload CSV file", 
+                file_types=['.csv'],
+                visible=False,  # function mode is default
+            )
+
+        mode.change(
+            fn=self.update_mode,
+            inputs=[mode, state],
+            outputs=[state, function, xmin, xmax, sigma, nsample, regenerate, csv_upload],
+        )
+
+        # function
+        function.submit(
+            lambda f, s: s.update(function=f),
+            inputs=[function, state],
+            outputs=[state],
+        )
+        xmin.submit(
+            lambda xmn, s: s.update(xmin=xmn),
+            inputs=[xmin, state],
+            outputs=[state],
+        )
+        xmax.submit(
+            lambda xmx, s: s.update(xmax=xmx),
+            inputs=[xmax, state],
+            outputs=[state],
+        )
+        sigma.submit(
+            lambda sig, s: s.update(sigma=sig),
+            inputs=[sigma, state],
+            outputs=[state],
+        )
+        nsample.change(
+            lambda n, s: s.update(nsample=n),
+            inputs=[nsample, state],
+            outputs=[state],
+        )
+        regenerate.click(
+            fn=self.regenerate_data,
+            inputs=[state],
+            outputs=[state],
+        )
+
+        # csv upload
+        csv_upload.upload(
+            self.upload_csv,
+            inputs=[csv_upload, state],
+            outputs=[state],
+        )
+
+

mlp_visualizer.py

@@ -1,4 +1,5 @@
 from collections import deque
+from dataclasses import dataclass, replace
 import functools
 from pathlib import Path
 import pickle
@@ -28,190 +29,14 @@ logging.basicConfig(
 )
 logger = logging.getLogger("ELVIS")
 
-
-NUMEXPR_CONSTANTS = {
-    'pi': np.pi,
-    'PI': np.pi,
-    'e': np.e,
-}
-
-
-def get_function(function, xlim=(-1, 1), nsample=100):
-    x = np.linspace(xlim[0], xlim[1], nsample)
-    y = numexpr.evaluate(function, local_dict={'x': x, **NUMEXPR_CONSTANTS})
-    x = x.reshape(-1, 1)
-    return x, y
-
-
-def get_data_points(function, xlim=(-1, 1), nsample=10, sigma=0, seed=0):
-    num_points_to_generate = 100
-    if nsample > num_points_to_generate:
-        raise ValueError(f"nsample too large, limit to {num_points_to_generate}")
-
-    rng = np.random.default_rng(seed)
-    x = rng.uniform(xlim[0], xlim[1], size=num_points_to_generate)
-    x = x[:nsample]
-    x = np.sort(x)
-
-    rng = np.random.default_rng(seed)
-    noise = sigma * rng.standard_normal(nsample)
-    y = numexpr.evaluate(function, local_dict={'x': x, **NUMEXPR_CONSTANTS}) + noise
-
-    x = x.reshape(-1, 1)
-    return x, y
-
-
-class HiddenLayerBox:
-    def __init__(self, initially_visible=False):
-        with gr.Row():
-            self.hidden_units = gr.Number(label="Hidden units", value=64, visible=initially_visible)
-            self.activation = gr.Textbox(label="Activation", value="ReLU", visible=initially_visible)
-
-    def set_visibility(self, visible):
-        return [
-            gr.update(visible=visible),
-            gr.update(visible=visible),
-        ]
-
-    def get_values(self):
-        return [self.hidden_units, self.activation]
-
-
-class ArchitectureComponent:
-    def __init__(self, update_architecture_callback, canvas, max_layers=5):
-        self.num_show = 2
-        self.components = []
-        for i in range(max_layers):
-            comp = HiddenLayerBox(initially_visible=(i < self.num_show))
-            self.components.append(comp)
-
-        self.update_architecture_callback = update_architecture_callback
-        self.canvas = canvas
-
-    def update_architecture(self, *values):
-        # values come as [hidden1, act1, hidden2, act2, ...]
-        hidden_layers = []
-        activations = []
-        for i in range(0, self.num_show * 2, 2):
-            if values[i] != "" or values[i + 1] != "":
-                hidden_layers.append(values[i])
-                activations.append(values[i + 1])
-        return self.update_architecture_callback(hidden_layers, activations)
-
-    def build(self):
-        with gr.Row():
-            add_btn = gr.Button("Add layer")
-            remove_btn = gr.Button("Remove layer")
-
-        with gr.Row():
-            gr.Number(label="Output units", value=1, interactive=False)
-            gr.Textbox(label="Activation", value="Identity", interactive=False)
-
-        # Collect all subcomponents
-        all_outputs = []
-        for comp in self.components:
-            all_outputs += [comp.hidden_units, comp.activation]
-
-        def on_add():
-            self.num_show = min(self.num_show + 1, len(self.components))
-            updates = []
-            for i, comp in enumerate(self.components):
-                updates += comp.set_visibility(i < self.num_show)
-
-            updates += [gr.update(value=self.num_show)]
-            return updates
-
-        def on_remove():
-            self.num_show = max(self.num_show - 1, 0)
-            updates = []
-            for i, comp in enumerate(self.components):
-                updates += comp.set_visibility(i < self.num_show)
-
-            updates += [gr.update(value=self.num_show)]
-            return updates
-
-        hidden_counter = gr.Number(value=self.num_show, visible=False)
-
-        add_btn.click(on_add, outputs=[*all_outputs, hidden_counter] )
-        remove_btn.click(on_remove, outputs=[*all_outputs, hidden_counter] )
-
-        for output in all_outputs:
-            output.submit(
-                fn=self.update_architecture,
-                inputs=all_outputs,
-                outputs=[self.canvas],
-            )
-        hidden_counter.change(
-            fn=self.update_architecture,
-            inputs=all_outputs,
-            outputs=[self.canvas],
-        )
+from dataset_options import DatasetOptions, DatasetOptionsView, get_function
 
 
 class MlpVisualizer:
-    DEFAULT_FUNCTION = "sin(2 * pi * x)"
-
-    DEFAULT_OPTIMIZER = "SGD"
-    DEFAULT_LEARNING_RATE = 0.01
-
-    DEFAULT_OPTIMIZER_HPARAMS = {
-        "SGD": {
-            "learning_rate": 0.1,
-            "momentum": 0.0,
-        },
-        "Adam": {
-            "learning_rate": 0.01,
-            "beta1": 0.9,
-            "beta2": 0.999,
-            "eps": 1e-8,
-        },
-    }
-
-    def _init_state(self):
-        self.data_options = {
-            "function": self.DEFAULT_FUNCTION,
-            "nsample": 30,
-            "sigma": 0,
-            "seed": 0,
-            "x_min": -1,
-            "x_max": 1,
-        }
-        self.x_train, self.y_train = self.generate_data()
-
-        self.architecture_options = {
-            "hidden_layers": [64, 64],
-            "activations": ["ReLU", "ReLU"],
-        }
-        self.basic_train_hparams = {
-            "batch_size": self.x_train.shape[0],
-            "optimizer": self.DEFAULT_OPTIMIZER,
-        }
-
-        # important to copy dict
-        self.optimizer_hparams = {}
-        for opt, params in self.DEFAULT_OPTIMIZER_HPARAMS.items():
-            self.optimizer_hparams[opt] = params.copy()
-
-        # do not initialise here, otherwise gradio will make it not work
-        # self.param_components = {}
-
-        self.criterion = nn.MSELoss()
-        self.model, self.optimizer, self.train_loss = self.init_model()
-        self.num_steps_trained = 0
-
-
-        self.plot_options = {
-            "show_training_data": True,
-            "show_true_function": True,
-            "show_predictions": True,
-        }
-
     def __init__(self, width, height):
         self.canvas_width = width
         self.canvas_height = height
 
-        self._init_state()
-
         self.plot_cmap = plt.get_cmap("tab20")
 
         self.css = """
@@ -219,97 +44,38 @@ class MlpVisualizer:
     display: none;
 }"""
 
-    def on_load(self):
-        self._init_state()
-
-    def generate_data(self):
-        function = self.data_options["function"]
-        nsample = self.data_options["nsample"]
-        sigma = self.data_options["sigma"]
-        x_min = self.data_options["x_min"]
-        x_max = self.data_options["x_max"]
-
-        return get_data_points(function, xlim=(x_min, x_max), nsample=nsample, sigma=sigma, seed=self.data_options["seed"])
-
-    def init_model(self):
-        print(self.architecture_options)
-        layers = []
-        input_size = 1
-        for hidden_units, activation in zip(self.architecture_options["hidden_layers"], self.architecture_options["activations"]):
-            layers.append(nn.Linear(input_size, hidden_units))
-            if activation == "ReLU":
-                layers.append(nn.ReLU())
-            elif activation == "Sigmoid":
-                layers.append(nn.Sigmoid())
-            elif activation == "Tanh":
-                layers.append(nn.Tanh())
-            elif activation == "LeakyReLU":
-                layers.append(nn.LeakyReLU())
-            elif activation == "Identity":
-                layers.append(nn.Identity())
-            else:
-                raise ValueError(f"Unsupported activation: {activation}")
-            input_size = hidden_units
-
-        output_layer = nn.Linear(input_size, 1)
-        model = nn.Sequential(*layers, output_layer)
-
-        if self.basic_train_hparams["optimizer"] == "Adam":
-            optimizer = torch.optim.Adam(
-                model.parameters(), 
-                lr=self.optimizer_hparams["Adam"]["learning_rate"],
-                betas=(self.optimizer_hparams["Adam"]["beta1"], self.optimizer_hparams["Adam"]["beta2"]),
-                eps=self.optimizer_hparams["Adam"]["eps"],
-            )
-        elif self.basic_train_hparams["optimizer"] == "SGD":
-            optimizer = torch.optim.SGD(
-                model.parameters(), 
-                lr=self.optimizer_hparams["SGD"]["learning_rate"],
-                momentum=self.optimizer_hparams["SGD"]["momentum"],
-            )
-        else:
-            raise ValueError(f"Unsupported optimizer: {self.basic_train_hparams['optimizer']}")
-
-        self.num_steps_trained = 0
-
-        # compute initial train loss
-        model.eval()
-        inputs = torch.from_numpy(self.x_train).float()
-        targets = torch.from_numpy(self.y_train).float().unsqueeze(1)
-        with torch.no_grad():
-            outputs = model(inputs)
-            train_loss = self.criterion(outputs, targets).item()
-
-        return model, optimizer, train_loss
-
-    def plot(self):
-        '''
-        '''
+    def plot(self, dataset_options: DatasetOptions):
+        print("Plotting")
         t1 = time.time()
-        logger.info("Initializing figure")
-        fig = plt.figure(figsize=(self.canvas_width/100., self.canvas_height/100.0), dpi=100)
+        fig = plt.figure(figsize=(self.canvas_width / 100., self.canvas_height / 100.0), dpi=100)
         # set entire figure to be the canvas to allow simple conversion of mouse
         # position to coordinates in the figure
         ax = fig.add_axes([0., 0., 1., 1.]) # 
         ax.margins(x=0, y=0) # no padding in both directions
 
-        x_test, y_test = get_function(self.data_options["function"], xlim=(-2, 2), nsample=100)
-        y_pred = self.model(torch.from_numpy(x_test).float()).detach().numpy()
+        if dataset_options.mode == "generate":
+            x_test, y_test = get_function(dataset_options.function, xlim=(-2, 2), nsample=100)
+
+        # y_pred = self.model(torch.from_numpy(x_test).float()).detach().numpy()
 
         # plot
         fig, ax = plt.subplots(figsize=(8, 8))
         ax.set_title("")
         ax.set_xlabel("x")
         ax.set_ylabel("y")
-        ax.set_ylim(y_test.min() - 1, y_test.max() + 1)
 
-        if self.plot_options["show_training_data"]:
-            plt.scatter(self.x_train.flatten(), self.y_train, label='training data', color=self.plot_cmap(0))
+        if dataset_options.mode == "generate":
+            ax.set_ylim(y_test.min() - 1, y_test.max() + 1)
+
+        x_train = dataset_options.x
+        y_train = dataset_options.y
+        if True:
+            plt.scatter(x_train.flatten(), y_train, label='training data', color=self.plot_cmap(0))
 
-        if self.plot_options["show_true_function"]:
+        if dataset_options.mode == "generate":
             plt.plot(x_test.flatten(), y_test, label='true function', color=self.plot_cmap(1))
 
-        if self.plot_options["show_predictions"]:
+        if False:
             plt.plot(x_test.flatten(), y_pred, linestyle="--", label='prediction', color=self.plot_cmap(2))
 
         plt.legend()
@@ -325,335 +91,44 @@ class MlpVisualizer:
 
         return img
 
-    def _update_data_seed(self):
-        self.data_options["seed"] += 1
-        self.x_train, self.y_train = self.generate_data()
-        self.reset_model()
-        return self.plot(), self.num_steps_trained, self.train_loss
-
-    def reset_model(self):
-        self.model, self.optimizer, self.train_loss = self.init_model()
-        return self.plot(), self.num_steps_trained, self.train_loss
-
-    def update_data_options(self, **kwargs):
-        for key, value in kwargs.items():
-            if key in self.data_options:
-
-                # if function - test if valid
-                if key == "function":
-                    try:
-                        x = np.linspace(-1, 1, 10)
-                        y = numexpr.evaluate(value, local_dict={'x': x, **NUMEXPR_CONSTANTS})
-                    except Exception as e:
-                        raise ValueError(f"Invalid function: {e}")
-
-                self.data_options[key] = value
-
-        # reset data and model
-        self.x_train, self.y_train = self.generate_data()
-        self.reset_model()
-
-        if "nsample" in kwargs:
-            slider_update = gr.update(maximum=self.x_train.shape[0], value=min(self.basic_train_hparams["batch_size"], self.x_train.shape[0]))
-            return self.plot(), slider_update, self.num_steps_trained, self.train_loss
-
-        return self.plot(), self.num_steps_trained, self.train_loss
-
-    def update_plot_options(self, **kwargs):
-        for key, value in kwargs.items():
-            if key in self.plot_options:
-                self.plot_options[key] = value
-        return self.plot()
-
-    def update_architecture(self, hidden_layers, activations):
-        self.architecture_options["hidden_layers"] = hidden_layers
-        self.architecture_options["activations"] = activations
-
-        # reset model
-        self.model, self.optimizer, self.train_loss = self.init_model()
-
-        return self.plot(), self.num_steps_trained, self.train_loss
-
-    def update_basic_train_hparams(self, **kwargs):
-        for key, value in kwargs.items():
-            if key in self.basic_train_hparams:
-                self.basic_train_hparams[key] = value
-
-        # reset model
-        self.model, self.optimizer, self.train_loss = self.init_model()
-
-        return self.plot(), self.num_steps_trained, self.train_loss
-
-    def update_optimizer(self, optimizer_name):
-        self.basic_train_hparams["optimizer"] = optimizer_name
-        # reset optimizer hyperparameters to default
-        self.optimizer_hparams[optimizer_name] = self.DEFAULT_OPTIMIZER_HPARAMS[optimizer_name].copy()
-
-        updates = []
-        for opt_name, params in self.param_components.items():
-            is_visible = (opt_name == optimizer_name)
-            for _ in params.values():
-                updates.append(gr.update(visible=is_visible))
-
-        # reset model
-        self.model, self.optimizer, self.train_loss = self.init_model()
-
-        return updates + [self.plot(), self.num_steps_trained, self.train_loss]
-
-    def build_optimizer_components(self):
-        self.param_components = {}
-        for opt_name, params in self.DEFAULT_OPTIMIZER_HPARAMS.items():
-            opt_dict = {}
-            for param_name, param_value in params.items():
-                opt_dict[param_name] = gr.Number(
-                    label=f"{param_name}",
-                    value=param_value,
-                    visible=(opt_name == self.DEFAULT_OPTIMIZER),
-                    interactive=True,
-                )
-            self.param_components[opt_name] = opt_dict
-
-        all_param_components = [
-            comp for opt in self.param_components.values() for comp in opt.values()
-        ]
-        return all_param_components
-
-    def update_hparam(self, value, optimizer_name, param_name):
-        self.optimizer_hparams[optimizer_name][param_name] = value
-
-        # reset model and plot
-        self.model, self.optimizer, self.train_loss = self.init_model()
-        return self.plot(), self.num_steps_trained, self.train_loss
-    
-    def train_step(self):
-        self.model.train()
-
-        inputs = torch.from_numpy(self.x_train).float()
-        targets = torch.from_numpy(self.y_train).float().unsqueeze(1)
-        outputs = self.model(inputs)
-        loss = self.criterion(outputs, targets)
-
-        self.optimizer.zero_grad()
-        loss.backward()
-        self.optimizer.step()
-
-        self.num_steps_trained += 1
-
-        # update train loss
-        self.model.eval()
-        with torch.no_grad():
-            outputs = self.model(inputs)
-            self.train_loss = self.criterion(outputs, targets).item()
-
-        return self.plot(), self.num_steps_trained, self.train_loss
-
     def launch(self):
         # build the Gradio interface
         with gr.Blocks(css=self.css) as demo:
             # app title
             gr.HTML("<div style='text-align:left; font-size:40px; font-weight: bold;'>MLP Training Visualizer</div>")
 
+            # states
+            dataset_options = gr.State(DatasetOptions())
+
             # GUI elements and layout 
             with gr.Row():
                 with gr.Column(scale=2):
-                    self.canvas = gr.Image(
-                        value=self.plot(), 
+                    canvas = gr.Image(
+                        value=self.plot(dataset_options.value),
                         show_download_button=False,
                         container=True,
                     )
 
                 with gr.Column(scale=1): 
                     with gr.Tab("Dataset"):
-                        dataset_radio = gr.Radio(
-                            ["Generate", "Upload"], 
-                            value="Generate", 
-                            label="Dataset", 
+                        dataset_view = DatasetOptionsView()
+                        dataset_view.build(state=dataset_options)
+                        dataset_options.change(
+                            fn=self.plot,
+                            inputs=[dataset_options],
+                            outputs=[canvas],
                         )
 
-                        with gr.Column():
-                            function_box = gr.Textbox(
-                                label="Function",
-                                placeholder="function of x",
-                                value=self.DEFAULT_FUNCTION,
-                                interactive=True,
-                            )
-                            with gr.Row():
-                                x_min = gr.Number(
-                                    label="Min x",
-                                    value=-1,
-                                    interactive=True,
-                                )
-                                x_max = gr.Number(
-                                    label="Max x",
-                                    value=1,
-                                    interactive=True,
-                                )
-                            with gr.Row():
-                                noise_value = gr.Number(
-                                    label="Gaussian noise standard deviation",
-                                    value=0,
-                                    interactive=True,
-                                )
-                                num_points_slider = gr.Slider(
-                                    label="Number of data points",
-                                    minimum=0,
-                                    maximum=100,
-                                    step=1,
-                                    value=30,
-                                    interactive=True,
-                                )
-
-                            regenerate_button = gr.Button("Regenerate Data")
-
-                        # upload data
-                        file_chooser = gr.File(label="Choose a file", visible=False, elem_id="rowheight")
-                        self.file_chooser = file_chooser
-
                     with gr.Tab("Architecture"):
-                        self.architecture_component = ArchitectureComponent(self.update_architecture, self.canvas)
-                        self.architecture_component.build()
-
+                        gr.Markdown("HI")
                     with gr.Tab("Train"):
-                        optimizer_radio = gr.Radio(
-                            ["SGD", "Adam"], 
-                            value=self.DEFAULT_OPTIMIZER,
-                            label="Optimizer", 
-                        )
-
-                        all_param_components = self.build_optimizer_components()
-                        self.temp = all_param_components
-
-                        batch_size_slider = gr.Slider(
-                            label="Batch Size",
-                            minimum=1,
-                            maximum=self.x_train.shape[0],
-                            step=1,
-                            value=self.x_train.shape[0],
-                            interactive=True,
-                        )
-
-                        with gr.Row():
-                            train_step_counter = gr.Number(
-                                label="Train steps",
-                                value=0,
-                                interactive=False,
-                            )
-                            train_loss_display = gr.Number(
-                                label="Train loss",
-                                value=self.train_loss,
-                                interactive=False,
-                            )
-
-                        train_button = gr.Button("Train Step")
-                        reset_model_button = gr.Button("Reset Model")
-
+                        gr.Markdown("HI")
                     with gr.Tab("Plot"):
-                        # plot show options
-                        with gr.Column():
-                            with gr.Row():
-                                show_training_data = gr.Checkbox(label="Show training data", value=True)
-                                show_true_function = gr.Checkbox(label="Show true function", value=True)
-                            with gr.Row():
-                                show_predictions = gr.Checkbox(label="Show mean prediction", value=True)
-
-                            #gr.Markdown(''.join(open('kernel_examples.md', 'r').readlines()))
-
+                        gr.Markdown("HI")
                     with gr.Tab("Export"):
-                        # use hidden download button to generate files on the fly
-                        # https://github.com/gradio-app/gradio/issues/9230#issuecomment-2323771634
-
-                        btn_export_data = gr.Button("Data")
-                        btn_export_data_hidden = gr.DownloadButton(label="You should not see this", elem_id="btn_export_data_hidden", elem_classes="hidden-button")
-                        
-                        btn_export_model = gr.Button('Model')
-                        btn_export_model_hidden = gr.DownloadButton(label="You should not see this", elem_id="btn_export_model_hidden", elem_classes="hidden-button")
-
-                        btn_export_code = gr.Button('Code')
-                        btn_export_code_hidden = gr.DownloadButton(label="You should not see this", elem_id="btn_export_code_hidden", elem_classes="hidden-button")
-
+                        gr.Markdown("HI")
                     with gr.Tab("Usage"):
-                        gr.Markdown(''.join(open('usage.md', 'r').readlines()))
-
-                    # data options
-                    function_box.submit(
-                        fn=lambda function: self.update_data_options(function=function),
-                        inputs=function_box,
-                        outputs=[self.canvas, train_step_counter, train_loss_display],
-                    )
-                    x_min.submit(
-                        fn=lambda xmin: self.update_data_options(x_min=xmin),
-                        inputs=x_min,
-                        outputs=[self.canvas, train_step_counter, train_loss_display],
-                    )
-                    x_max.submit(
-                        fn=lambda xmax: self.update_data_options(x_max=xmax),
-                        inputs=x_max,
-                        outputs=[self.canvas, train_step_counter, train_loss_display],
-                    )
-                    num_points_slider.change(
-                        fn=lambda nsample: self.update_data_options(nsample=nsample),
-                        inputs=num_points_slider,
-                        outputs=[self.canvas, batch_size_slider, train_step_counter, train_loss_display],
-                    )
-                    noise_value.submit(
-                        fn=lambda sigma: self.update_data_options(sigma=sigma),
-                        inputs=noise_value,
-                        outputs=[self.canvas, train_step_counter, train_loss_display],
-                    )
-                    regenerate_button.click(
-                        fn=self._update_data_seed,
-                        outputs=[self.canvas, train_step_counter, train_loss_display],
-                    )
-            
-                    # train options
-                    optimizer_radio.change(
-                        fn=self.update_optimizer,
-                        inputs=optimizer_radio,
-                        outputs=[*all_param_components, self.canvas, train_step_counter, train_loss_display],
-                    )
-                    batch_size_slider.change(
-                        fn=lambda batch_size: self.update_basic_train_hparams(batch_size=batch_size),
-                        inputs=batch_size_slider,
-                        outputs=[self.canvas, train_step_counter, train_loss_display],
-                    )
-                    train_button.click(
-                        fn=self.train_step,
-                        outputs=[self.canvas, train_step_counter, train_loss_display],
-                        show_progress="hidden",
-                    )
-                    reset_model_button.click(
-                        fn=self.reset_model,
-                        outputs=[self.canvas, train_step_counter, train_loss_display],
-                    )
-                    for opt_name, params in self.param_components.items():
-                        for param_name, comp in params.items():
-                            comp.submit(
-                                fn=functools.partial(self.update_hparam, optimizer_name=opt_name, param_name=param_name),
-                                inputs=[comp],
-                                outputs=[self.canvas, train_step_counter, train_loss_display],
-                            )
-
-                    # plot options
-                    show_training_data.change(
-                        fn=lambda show: self.update_plot_options(show_training_data=show),
-                        inputs=show_training_data,
-                        outputs=[self.canvas],
-                        show_progress="hidden",
-                    )
-                    show_true_function.change(
-                        fn=lambda show: self.update_plot_options(show_true_function=show),
-                        inputs=show_true_function,
-                        outputs=[self.canvas],
-                        show_progress="hidden",
-                    )
-                    show_predictions.change(
-                        fn=lambda show: self.update_plot_options(show_predictions=show),
-                        inputs=show_predictions,
-                        outputs=[self.canvas],
-                        show_progress="hidden",
-                    )
-
-            demo.load(self.on_load)
+                        gr.Markdown("HI")
 
         demo.launch()
 

mlp_visualizer_old.py

@@ -0,0 +1,662 @@
+from collections import deque
+import functools
+from pathlib import Path
+import pickle
+import time
+
+import gradio as gr
+import inspect
+import io
+from jinja2 import Template
+import matplotlib.pyplot as plt
+import matplotlib.lines as mlines
+import numpy as np
+import numexpr
+import pandas as pd
+from PIL import Image
+import plotly.graph_objects as go
+import torch
+import torch.nn as nn
+
+import traceback
+import yaml
+
+import logging
+logging.basicConfig(
+    level=logging.INFO,  # set minimum level to capture (DEBUG, INFO, WARNING, ERROR, CRITICAL)
+    format="%(asctime)s [%(levelname)s] %(message)s",  # log format
+)
+logger = logging.getLogger("ELVIS")
+
+
+NUMEXPR_CONSTANTS = {
+    'pi': np.pi,
+    'PI': np.pi,
+    'e': np.e,
+}
+
+
+def get_function(function, xlim=(-1, 1), nsample=100):
+    x = np.linspace(xlim[0], xlim[1], nsample)
+    y = numexpr.evaluate(function, local_dict={'x': x, **NUMEXPR_CONSTANTS})
+    x = x.reshape(-1, 1)
+    return x, y
+
+
+def get_data_points(function, xlim=(-1, 1), nsample=10, sigma=0, seed=0):
+    num_points_to_generate = 100
+    if nsample > num_points_to_generate:
+        raise ValueError(f"nsample too large, limit to {num_points_to_generate}")
+
+    rng = np.random.default_rng(seed)
+    x = rng.uniform(xlim[0], xlim[1], size=num_points_to_generate)
+    x = x[:nsample]
+    x = np.sort(x)
+
+    rng = np.random.default_rng(seed)
+    noise = sigma * rng.standard_normal(nsample)
+    y = numexpr.evaluate(function, local_dict={'x': x, **NUMEXPR_CONSTANTS}) + noise
+
+    x = x.reshape(-1, 1)
+    return x, y
+
+
+class HiddenLayerBox:
+    def __init__(self, initially_visible=False):
+        with gr.Row():
+            self.hidden_units = gr.Number(label="Hidden units", value=64, visible=initially_visible)
+            self.activation = gr.Textbox(label="Activation", value="ReLU", visible=initially_visible)
+
+    def set_visibility(self, visible):
+        return [
+            gr.update(visible=visible),
+            gr.update(visible=visible),
+        ]
+
+    def get_values(self):
+        return [self.hidden_units, self.activation]
+
+
+class ArchitectureComponent:
+    def __init__(self, update_architecture_callback, canvas, max_layers=5):
+        self.num_show = 2
+        self.components = []
+        for i in range(max_layers):
+            comp = HiddenLayerBox(initially_visible=(i < self.num_show))
+            self.components.append(comp)
+
+        self.update_architecture_callback = update_architecture_callback
+        self.canvas = canvas
+
+    def update_architecture(self, *values):
+        # values come as [hidden1, act1, hidden2, act2, ...]
+        hidden_layers = []
+        activations = []
+        for i in range(0, self.num_show * 2, 2):
+            if values[i] != "" or values[i + 1] != "":
+                hidden_layers.append(values[i])
+                activations.append(values[i + 1])
+        return self.update_architecture_callback(hidden_layers, activations)
+
+    def build(self):
+        with gr.Row():
+            add_btn = gr.Button("Add layer")
+            remove_btn = gr.Button("Remove layer")
+
+        with gr.Row():
+            gr.Number(label="Output units", value=1, interactive=False)
+            gr.Textbox(label="Activation", value="Identity", interactive=False)
+
+        # Collect all subcomponents
+        all_outputs = []
+        for comp in self.components:
+            all_outputs += [comp.hidden_units, comp.activation]
+
+        def on_add():
+            self.num_show = min(self.num_show + 1, len(self.components))
+            updates = []
+            for i, comp in enumerate(self.components):
+                updates += comp.set_visibility(i < self.num_show)
+
+            updates += [gr.update(value=self.num_show)]
+            return updates
+
+        def on_remove():
+            self.num_show = max(self.num_show - 1, 0)
+            updates = []
+            for i, comp in enumerate(self.components):
+                updates += comp.set_visibility(i < self.num_show)
+
+            updates += [gr.update(value=self.num_show)]
+            return updates
+
+        hidden_counter = gr.Number(value=self.num_show, visible=False)
+
+        add_btn.click(on_add, outputs=[*all_outputs, hidden_counter] )
+        remove_btn.click(on_remove, outputs=[*all_outputs, hidden_counter] )
+
+        for output in all_outputs:
+            output.submit(
+                fn=self.update_architecture,
+                inputs=all_outputs,
+                outputs=[self.canvas],
+            )
+        hidden_counter.change(
+            fn=self.update_architecture,
+            inputs=all_outputs,
+            outputs=[self.canvas],
+        )
+
+
+class MlpVisualizer:
+    DEFAULT_FUNCTION = "sin(2 * pi * x)"
+
+    DEFAULT_OPTIMIZER = "SGD"
+    DEFAULT_LEARNING_RATE = 0.01
+
+    DEFAULT_OPTIMIZER_HPARAMS = {
+        "SGD": {
+            "learning_rate": 0.1,
+            "momentum": 0.0,
+        },
+        "Adam": {
+            "learning_rate": 0.01,
+            "beta1": 0.9,
+            "beta2": 0.999,
+            "eps": 1e-8,
+        },
+    }
+
+    def _init_state(self):
+        self.data_options = {
+            "function": self.DEFAULT_FUNCTION,
+            "nsample": 30,
+            "sigma": 0,
+            "seed": 0,
+            "x_min": -1,
+            "x_max": 1,
+        }
+        self.x_train, self.y_train = self.generate_data()
+
+        self.architecture_options = {
+            "hidden_layers": [64, 64],
+            "activations": ["ReLU", "ReLU"],
+        }
+        self.basic_train_hparams = {
+            "batch_size": self.x_train.shape[0],
+            "optimizer": self.DEFAULT_OPTIMIZER,
+        }
+
+        # important to copy dict
+        self.optimizer_hparams = {}
+        for opt, params in self.DEFAULT_OPTIMIZER_HPARAMS.items():
+            self.optimizer_hparams[opt] = params.copy()
+
+        # do not initialise here, otherwise gradio will make it not work
+        # self.param_components = {}
+
+        self.criterion = nn.MSELoss()
+        self.model, self.optimizer, self.train_loss = self.init_model()
+        self.num_steps_trained = 0
+
+
+        self.plot_options = {
+            "show_training_data": True,
+            "show_true_function": True,
+            "show_predictions": True,
+        }
+
+    def __init__(self, width, height):
+        self.canvas_width = width
+        self.canvas_height = height
+
+        self._init_state()
+
+        self.plot_cmap = plt.get_cmap("tab20")
+
+        self.css = """
+.hidden-button {
+    display: none;
+}"""
+
+    def on_load(self):
+        self._init_state()
+
+    def generate_data(self):
+        function = self.data_options["function"]
+        nsample = self.data_options["nsample"]
+        sigma = self.data_options["sigma"]
+        x_min = self.data_options["x_min"]
+        x_max = self.data_options["x_max"]
+
+        return get_data_points(function, xlim=(x_min, x_max), nsample=nsample, sigma=sigma, seed=self.data_options["seed"])
+
+    def init_model(self):
+        print(self.architecture_options)
+        layers = []
+        input_size = 1
+        for hidden_units, activation in zip(self.architecture_options["hidden_layers"], self.architecture_options["activations"]):
+            layers.append(nn.Linear(input_size, hidden_units))
+            if activation == "ReLU":
+                layers.append(nn.ReLU())
+            elif activation == "Sigmoid":
+                layers.append(nn.Sigmoid())
+            elif activation == "Tanh":
+                layers.append(nn.Tanh())
+            elif activation == "LeakyReLU":
+                layers.append(nn.LeakyReLU())
+            elif activation == "Identity":
+                layers.append(nn.Identity())
+            else:
+                raise ValueError(f"Unsupported activation: {activation}")
+            input_size = hidden_units
+
+        output_layer = nn.Linear(input_size, 1)
+        model = nn.Sequential(*layers, output_layer)
+
+        if self.basic_train_hparams["optimizer"] == "Adam":
+            optimizer = torch.optim.Adam(
+                model.parameters(), 
+                lr=self.optimizer_hparams["Adam"]["learning_rate"],
+                betas=(self.optimizer_hparams["Adam"]["beta1"], self.optimizer_hparams["Adam"]["beta2"]),
+                eps=self.optimizer_hparams["Adam"]["eps"],
+            )
+        elif self.basic_train_hparams["optimizer"] == "SGD":
+            optimizer = torch.optim.SGD(
+                model.parameters(), 
+                lr=self.optimizer_hparams["SGD"]["learning_rate"],
+                momentum=self.optimizer_hparams["SGD"]["momentum"],
+            )
+        else:
+            raise ValueError(f"Unsupported optimizer: {self.basic_train_hparams['optimizer']}")
+
+        self.num_steps_trained = 0
+
+        # compute initial train loss
+        model.eval()
+        inputs = torch.from_numpy(self.x_train).float()
+        targets = torch.from_numpy(self.y_train).float().unsqueeze(1)
+        with torch.no_grad():
+            outputs = model(inputs)
+            train_loss = self.criterion(outputs, targets).item()
+
+        return model, optimizer, train_loss
+
+    def plot(self):
+        '''
+        '''
+        t1 = time.time()
+        logger.info("Initializing figure")
+        fig = plt.figure(figsize=(self.canvas_width/100., self.canvas_height/100.0), dpi=100)
+        # set entire figure to be the canvas to allow simple conversion of mouse
+        # position to coordinates in the figure
+        ax = fig.add_axes([0., 0., 1., 1.]) # 
+        ax.margins(x=0, y=0) # no padding in both directions
+
+        x_test, y_test = get_function(self.data_options["function"], xlim=(-2, 2), nsample=100)
+        y_pred = self.model(torch.from_numpy(x_test).float()).detach().numpy()
+
+        # plot
+        fig, ax = plt.subplots(figsize=(8, 8))
+        ax.set_title("")
+        ax.set_xlabel("x")
+        ax.set_ylabel("y")
+        ax.set_ylim(y_test.min() - 1, y_test.max() + 1)
+
+        if self.plot_options["show_training_data"]:
+            plt.scatter(self.x_train.flatten(), self.y_train, label='training data', color=self.plot_cmap(0))
+
+        if self.plot_options["show_true_function"]:
+            plt.plot(x_test.flatten(), y_test, label='true function', color=self.plot_cmap(1))
+
+        if self.plot_options["show_predictions"]:
+            plt.plot(x_test.flatten(), y_pred, linestyle="--", label='prediction', color=self.plot_cmap(2))
+
+        plt.legend()
+
+        buf = io.BytesIO()
+        fig.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
+        plt.close(fig)
+        buf.seek(0)
+        img = Image.open(buf)
+
+        t2 = time.time()
+        logger.info(f"Plotting took {t2 - t1:.4f} seconds")
+
+        return img
+
+    def _update_data_seed(self):
+        self.data_options["seed"] += 1
+        self.x_train, self.y_train = self.generate_data()
+        self.reset_model()
+        return self.plot(), self.num_steps_trained, self.train_loss
+
+    def reset_model(self):
+        self.model, self.optimizer, self.train_loss = self.init_model()
+        return self.plot(), self.num_steps_trained, self.train_loss
+
+    def update_data_options(self, **kwargs):
+        for key, value in kwargs.items():
+            if key in self.data_options:
+
+                # if function - test if valid
+                if key == "function":
+                    try:
+                        x = np.linspace(-1, 1, 10)
+                        y = numexpr.evaluate(value, local_dict={'x': x, **NUMEXPR_CONSTANTS})
+                    except Exception as e:
+                        raise ValueError(f"Invalid function: {e}")
+
+                self.data_options[key] = value
+
+        # reset data and model
+        self.x_train, self.y_train = self.generate_data()
+        self.reset_model()
+
+        if "nsample" in kwargs:
+            slider_update = gr.update(maximum=self.x_train.shape[0], value=min(self.basic_train_hparams["batch_size"], self.x_train.shape[0]))
+            return self.plot(), slider_update, self.num_steps_trained, self.train_loss
+
+        return self.plot(), self.num_steps_trained, self.train_loss
+
+    def update_plot_options(self, **kwargs):
+        for key, value in kwargs.items():
+            if key in self.plot_options:
+                self.plot_options[key] = value
+        return self.plot()
+
+    def update_architecture(self, hidden_layers, activations):
+        self.architecture_options["hidden_layers"] = hidden_layers
+        self.architecture_options["activations"] = activations
+
+        # reset model
+        self.model, self.optimizer, self.train_loss = self.init_model()
+
+        return self.plot(), self.num_steps_trained, self.train_loss
+
+    def update_basic_train_hparams(self, **kwargs):
+        for key, value in kwargs.items():
+            if key in self.basic_train_hparams:
+                self.basic_train_hparams[key] = value
+
+        # reset model
+        self.model, self.optimizer, self.train_loss = self.init_model()
+
+        return self.plot(), self.num_steps_trained, self.train_loss
+
+    def update_optimizer(self, optimizer_name):
+        self.basic_train_hparams["optimizer"] = optimizer_name
+        # reset optimizer hyperparameters to default
+        self.optimizer_hparams[optimizer_name] = self.DEFAULT_OPTIMIZER_HPARAMS[optimizer_name].copy()
+
+        updates = []
+        for opt_name, params in self.param_components.items():
+            is_visible = (opt_name == optimizer_name)
+            for _ in params.values():
+                updates.append(gr.update(visible=is_visible))
+
+        # reset model
+        self.model, self.optimizer, self.train_loss = self.init_model()
+
+        return updates + [self.plot(), self.num_steps_trained, self.train_loss]
+
+    def build_optimizer_components(self):
+        self.param_components = {}
+        for opt_name, params in self.DEFAULT_OPTIMIZER_HPARAMS.items():
+            opt_dict = {}
+            for param_name, param_value in params.items():
+                opt_dict[param_name] = gr.Number(
+                    label=f"{param_name}",
+                    value=param_value,
+                    visible=(opt_name == self.DEFAULT_OPTIMIZER),
+                    interactive=True,
+                )
+            self.param_components[opt_name] = opt_dict
+
+        all_param_components = [
+            comp for opt in self.param_components.values() for comp in opt.values()
+        ]
+        return all_param_components
+
+    def update_hparam(self, value, optimizer_name, param_name):
+        self.optimizer_hparams[optimizer_name][param_name] = value
+
+        # reset model and plot
+        self.model, self.optimizer, self.train_loss = self.init_model()
+        return self.plot(), self.num_steps_trained, self.train_loss
+    
+    def train_step(self):
+        self.model.train()
+
+        inputs = torch.from_numpy(self.x_train).float()
+        targets = torch.from_numpy(self.y_train).float().unsqueeze(1)
+        outputs = self.model(inputs)
+        loss = self.criterion(outputs, targets)
+
+        self.optimizer.zero_grad()
+        loss.backward()
+        self.optimizer.step()
+
+        self.num_steps_trained += 1
+
+        # update train loss
+        self.model.eval()
+        with torch.no_grad():
+            outputs = self.model(inputs)
+            self.train_loss = self.criterion(outputs, targets).item()
+
+        return self.plot(), self.num_steps_trained, self.train_loss
+
+    def launch(self):
+        # build the Gradio interface
+        with gr.Blocks(css=self.css) as demo:
+            # app title
+            gr.HTML("<div style='text-align:left; font-size:40px; font-weight: bold;'>MLP Training Visualizer</div>")
+
+            # GUI elements and layout 
+            with gr.Row():
+                with gr.Column(scale=2):
+                    self.canvas = gr.Image(
+                        value=self.plot(), 
+                        show_download_button=False,
+                        container=True,
+                    )
+
+                with gr.Column(scale=1): 
+                    with gr.Tab("Dataset"):
+                        dataset_radio = gr.Radio(
+                            ["Generate", "Upload"], 
+                            value="Generate", 
+                            label="Dataset", 
+                        )
+
+                        with gr.Column():
+                            function_box = gr.Textbox(
+                                label="Function",
+                                placeholder="function of x",
+                                value=self.DEFAULT_FUNCTION,
+                                interactive=True,
+                            )
+                            with gr.Row():
+                                x_min = gr.Number(
+                                    label="Min x",
+                                    value=-1,
+                                    interactive=True,
+                                )
+                                x_max = gr.Number(
+                                    label="Max x",
+                                    value=1,
+                                    interactive=True,
+                                )
+                            with gr.Row():
+                                noise_value = gr.Number(
+                                    label="Gaussian noise standard deviation",
+                                    value=0,
+                                    interactive=True,
+                                )
+                                num_points_slider = gr.Slider(
+                                    label="Number of data points",
+                                    minimum=0,
+                                    maximum=100,
+                                    step=1,
+                                    value=30,
+                                    interactive=True,
+                                )
+
+                            regenerate_button = gr.Button("Regenerate Data")
+
+                        # upload data
+                        file_chooser = gr.File(label="Choose a file", visible=False, elem_id="rowheight")
+                        self.file_chooser = file_chooser
+
+                    with gr.Tab("Architecture"):
+                        self.architecture_component = ArchitectureComponent(self.update_architecture, self.canvas)
+                        self.architecture_component.build()
+
+                    with gr.Tab("Train"):
+                        optimizer_radio = gr.Radio(
+                            ["SGD", "Adam"], 
+                            value=self.DEFAULT_OPTIMIZER,
+                            label="Optimizer", 
+                        )
+
+                        all_param_components = self.build_optimizer_components()
+                        self.temp = all_param_components
+
+                        batch_size_slider = gr.Slider(
+                            label="Batch Size",
+                            minimum=1,
+                            maximum=self.x_train.shape[0],
+                            step=1,
+                            value=self.x_train.shape[0],
+                            interactive=True,
+                        )
+
+                        with gr.Row():
+                            train_step_counter = gr.Number(
+                                label="Train steps",
+                                value=0,
+                                interactive=False,
+                            )
+                            train_loss_display = gr.Number(
+                                label="Train loss",
+                                value=self.train_loss,
+                                interactive=False,
+                            )
+
+                        train_button = gr.Button("Train Step")
+                        reset_model_button = gr.Button("Reset Model")
+
+                    with gr.Tab("Plot"):
+                        # plot show options
+                        with gr.Column():
+                            with gr.Row():
+                                show_training_data = gr.Checkbox(label="Show training data", value=True)
+                                show_true_function = gr.Checkbox(label="Show true function", value=True)
+                            with gr.Row():
+                                show_predictions = gr.Checkbox(label="Show mean prediction", value=True)
+
+                            #gr.Markdown(''.join(open('kernel_examples.md', 'r').readlines()))
+
+                    with gr.Tab("Export"):
+                        # use hidden download button to generate files on the fly
+                        # https://github.com/gradio-app/gradio/issues/9230#issuecomment-2323771634
+
+                        btn_export_data = gr.Button("Data")
+                        btn_export_data_hidden = gr.DownloadButton(label="You should not see this", elem_id="btn_export_data_hidden", elem_classes="hidden-button")
+                        
+                        btn_export_model = gr.Button('Model')
+                        btn_export_model_hidden = gr.DownloadButton(label="You should not see this", elem_id="btn_export_model_hidden", elem_classes="hidden-button")
+
+                        btn_export_code = gr.Button('Code')
+                        btn_export_code_hidden = gr.DownloadButton(label="You should not see this", elem_id="btn_export_code_hidden", elem_classes="hidden-button")
+
+                    with gr.Tab("Usage"):
+                        gr.Markdown(''.join(open('usage.md', 'r').readlines()))
+
+                    # data options
+                    function_box.submit(
+                        fn=lambda function: self.update_data_options(function=function),
+                        inputs=function_box,
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
+                    )
+                    x_min.submit(
+                        fn=lambda xmin: self.update_data_options(x_min=xmin),
+                        inputs=x_min,
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
+                    )
+                    x_max.submit(
+                        fn=lambda xmax: self.update_data_options(x_max=xmax),
+                        inputs=x_max,
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
+                    )
+                    num_points_slider.change(
+                        fn=lambda nsample: self.update_data_options(nsample=nsample),
+                        inputs=num_points_slider,
+                        outputs=[self.canvas, batch_size_slider, train_step_counter, train_loss_display],
+                    )
+                    noise_value.submit(
+                        fn=lambda sigma: self.update_data_options(sigma=sigma),
+                        inputs=noise_value,
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
+                    )
+                    regenerate_button.click(
+                        fn=self._update_data_seed,
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
+                    )
+            
+                    # train options
+                    optimizer_radio.change(
+                        fn=self.update_optimizer,
+                        inputs=optimizer_radio,
+                        outputs=[*all_param_components, self.canvas, train_step_counter, train_loss_display],
+                    )
+                    batch_size_slider.change(
+                        fn=lambda batch_size: self.update_basic_train_hparams(batch_size=batch_size),
+                        inputs=batch_size_slider,
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
+                    )
+                    train_button.click(
+                        fn=self.train_step,
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
+                        show_progress="hidden",
+                    )
+                    reset_model_button.click(
+                        fn=self.reset_model,
+                        outputs=[self.canvas, train_step_counter, train_loss_display],
+                    )
+                    for opt_name, params in self.param_components.items():
+                        for param_name, comp in params.items():
+                            comp.submit(
+                                fn=functools.partial(self.update_hparam, optimizer_name=opt_name, param_name=param_name),
+                                inputs=[comp],
+                                outputs=[self.canvas, train_step_counter, train_loss_display],
+                            )
+
+                    # plot options
+                    show_training_data.change(
+                        fn=lambda show: self.update_plot_options(show_training_data=show),
+                        inputs=show_training_data,
+                        outputs=[self.canvas],
+                        show_progress="hidden",
+                    )
+                    show_true_function.change(
+                        fn=lambda show: self.update_plot_options(show_true_function=show),
+                        inputs=show_true_function,
+                        outputs=[self.canvas],
+                        show_progress="hidden",
+                    )
+                    show_predictions.change(
+                        fn=lambda show: self.update_plot_options(show_predictions=show),
+                        inputs=show_predictions,
+                        outputs=[self.canvas],
+                        show_progress="hidden",
+                    )
+
+            demo.load(self.on_load)
+
+        demo.launch()
+
+visualizer = MlpVisualizer(width=1200, height=900)
+visualizer.launch()
+