Clean up code for review

Author: k8culver

assets/demo.css

@@ -460,6 +460,7 @@ progress::-webkit-progress-bar {
     transition: filter 0.1s linear, opacity 0.1s linear;
 }
 
+/* Overwriting bootstrap tooltip styling */
 .tooltip {
     --bs-tooltip-arrow-height: 1.5rem;
     --bs-tooltip-arrow-width: 1.5rem;

demo_callbacks.py

@@ -52,6 +52,10 @@
     generate_model_fig,
 )
 
+from src.utils.global_vars import (
+    STEP_4_FILE, STEP_5_FILE, LATENT_ENCODED_FILE, STEP_2_FILE, LATENT_QPU_FILE
+)
+
 
 @dash.callback(
     Output({"type": "to-collapse-class", "index": MATCH}, "className"),
@@ -116,8 +120,8 @@ def update_model_diagram_imgs(
     progress: int,
     fig_qpu: go.Figure,
     fig_encoded: go.Figure,
-    latent_mapping: list
-) -> tuple[str, str, str, str, go.Figure, go.Figure, list]:
+    latent_mapping: list,
+) -> tuple[str, str, str, go.Figure, go.Figure, list]:
     """Force refresh images to get around Dash caching. Updates image src with a incrementing
     query string.
 
@@ -138,10 +142,10 @@ def update_model_diagram_imgs(
     fig_qpu = go.Figure(fig_qpu)
     fig_encoded = go.Figure(fig_encoded)
 
-    with open("static/model_diagram/latent_qpu.json", "r") as f:
+    with open(LATENT_QPU_FILE, "r") as f:
         latent_qpu = json.load(f)
 
-    with open("static/model_diagram/latent_encoded.json", "r") as f:
+    with open(LATENT_ENCODED_FILE, "r") as f:
         latent_encoded = json.load(f)
 
     color_mapping_qpu = [GRAPH_COLORS[int(latent_qpu[i] > 0)] for i in latent_mapping]
@@ -152,9 +156,9 @@ def update_model_diagram_imgs(
 
     if GENERATE_NEW_MODEL_DIAGRAM:
         return (
-            f"static/model_diagram/step_2_encode.png?interval={progress}",
-            f"static/model_diagram/step_4_decode.png?interval={progress}",
-            f"static/model_diagram/step_5_output.png?interval={progress}",
+            f"{STEP_2_FILE}?interval={progress}",
+            f"{STEP_4_FILE}?interval={progress}",
+            f"{STEP_5_FILE}?interval={progress}",
             fig_qpu,
             fig_encoded,
             generate_latent_vector(latent_encoded[:5], latent_encoded[-1]),
@@ -163,6 +167,20 @@ def update_model_diagram_imgs(
     raise PreventUpdate
 
 
+class CheckQpuAndUpdateModelReturn(NamedTuple):
+    """Return type for the ``check_qpu_and_update_model`` callback function."""
+
+    popup_classname: str = "display-none"
+    generate_button_disabled: bool = False
+    model_details: dict = dash.no_update
+    fig_qpu_graph: go.Figure = dash.no_update
+    fig_encoded_graph: go.Figure = dash.no_update
+    latent_diagram_size: int = dash.no_update
+    latent_mapping: list[int] = dash.no_update
+    step_2_img: str = dash.no_update
+    step_4_img: str = dash.no_update
+    step_5_img: str = dash.no_update
+
 @dash.callback(
     Output("popup", "className"),
     Output("generate-button", "disabled"),
@@ -188,7 +206,7 @@ def check_qpu_and_update_model(
     n_latents: int,
     setting_tabs_value: str,
     example_image: list,
-) -> tuple[str, bool, dict, go.Figure, go.Figure, int, list[int]]:
+) -> CheckQpuAndUpdateModelReturn:
     """Checks whether user has access to QPU associated with model and updates the model details
     when model changes.
 
@@ -200,61 +218,71 @@ def check_qpu_and_update_model(
         example_image: The example image to show all the steps for in the UI.
 
     Returns:
-        popup-classname: The class name to hide the popup.
-        generate-button-disabled: Whether to disable or enable the Generate button.
-        model-details-children: The model details to display.
-        fig-qpu-graph: The QPU graph figure.
-        fig-encoded-graph: The not QPU graph figure.
-        latent-diagram-size: The dimension of the latent space.
-        latent-mapping: The mapping of the nodes to latent space indices.
-        step-2-encode-img: The src url for the encode image.
-        step-4-decode-img: The src url for the decode image.
-        step-5-output-img: The src url for the output image.
+        CheckQpuAndUpdateModelReturn named tuple:
+            popup_classname: The class name to hide the popup.
+            generate_button_disabled: Whether to disable or enable the Generate button.
+            model_details: The model details to display.
+            fig_qpu_graph: The QPU graph figure.
+            fig_encoded_graph: The not QPU graph figure.
+            latent_diagram_size: The dimension of the latent space.
+            latent_mapping: The mapping of the nodes to latent space indices.
+            step_2_img: The src url for the encode image.
+            step_4_img: The src url for the decode image.
+            step_5_img: The src url for the output image.
     """
-    model_data = None
-    model_details = None
+    switched_to_generate_tab = ctx.triggered_id == "setting-tabs" and setting_tabs_value == "generate-tab"
 
-    if not ctx.triggered_id or ctx.triggered_id == "model-file-name" or (ctx.triggered_id == "setting-tabs" and setting_tabs_value == "generate-tab"):
+    # If first load, or a new model is chosen, or the settings tab is changed to "generate"
+    if not ctx.triggered_id or ctx.triggered_id == "model-file-name" or switched_to_generate_tab:
         with open(MODEL_PATH / model_file_name / "parameters.json") as file:
             model_data = json.load(file)
 
         model_details = generate_model_data(model_data)
 
+        # If model_data has a QPU that is no longer available, show warning popup
+        if model_data["qpu"] and not (len(SOLVERS) and model_data["qpu"] in SOLVERS):
+            return CheckQpuAndUpdateModelReturn(
+                popup_classname="",
+                generate_button_disabled=True,
+                model_details=model_details,
+            )
+
         # Create model instance to generate model diagram images
         model = ModelWrapper(qpu=model_data["qpu"], n_latents=model_data["n_latents"])
         model.load(file_path=MODEL_PATH / model_file_name)
+
         # Dash converts the tensor to a list of floats, convert back to tensor
         example_image = torch.tensor(example_image, dtype=torch.float32)
         example_image = example_image.unsqueeze(0)
-
         generate_model_diagram(model, example_image)
+
+        fig_qpu, fig_encoded, latent_mapping = generate_model_fig(
+            model_data["qpu"],
+            model_data["n_latents"],
+            model_data["random_seed"],
+        )
+
         force_refresh = random.randint(1, 9999999)
 
-        if model_data["qpu"] and not (len(SOLVERS) and model_data["qpu"] in SOLVERS):
-            return (
-                "", True, model_details, dash.no_update,
-                f"static/model_diagram/step_2_encode.png?force_refresh={force_refresh}",
-                f"static/model_diagram/step_4_decode.png?force_refresh={force_refresh}",
-                f"static/model_diagram/step_5_output.png?force_refresh={force_refresh}",
-            )
+        return CheckQpuAndUpdateModelReturn(
+            model_details=model_details,
+            fig_qpu_graph=fig_qpu,
+            fig_encoded_graph=fig_encoded,
+            latent_diagram_size=model_data["n_latents"],
+            latent_mapping=latent_mapping,
+            step_2_img=f"{STEP_2_FILE}?force_refresh={force_refresh}",
+            step_4_img=f"{STEP_4_FILE}?force_refresh={force_refresh}",
+            step_5_img=f"{STEP_5_FILE}?force_refresh={force_refresh}",
+        )
 
-    fig_qpu, fig_encoded, latent_mapping = generate_model_fig(
-        model_data["qpu"] if model_data else qpu,
-        model_data["n_latents"] if model_data else n_latents,
-        model_data["random_seed"] if model_data else 4,
-    )
+    # No model data, proceed with defaults
+    fig_qpu, fig_encoded, latent_mapping = generate_model_fig(qpu, n_latents, 4)
 
-    return (
-        "display-none",
-        False,
-        model_details if model_details else dash.no_update,
-        fig_qpu,
-        fig_encoded,
-        n_latents,
-        latent_mapping,
-        dash.no_update,
-        dash.no_update,
-        dash.no_update,
+    return CheckQpuAndUpdateModelReturn(
+        fig_qpu_graph=fig_qpu,
+        fig_encoded_graph=fig_encoded,
+        latent_diagram_size=n_latents,
+        latent_mapping=latent_mapping,
     )
 
 

demo_interface.py

@@ -34,6 +34,7 @@
     THUMBNAIL,
 )
 from src.utils.callback_helpers import get_example_image
+from src.utils.global_vars import LATENT_ENCODED_FILE, STEP_1_FILE, STEP_2_FILE, STEP_4_FILE, STEP_5_FILE_DEFAULT
 
 # Initialize available QPUs
 try:
@@ -48,7 +49,7 @@
 
 # Initialize the latent diagram with either the available file or random +/- 1s
 try:
-    with open("static/model_diagram/latent_encoded.json", "r") as f:
+    with open(LATENT_ENCODED_FILE, "r") as f:
         latent_qpu = json.load(f)
 
     LATENT_DIAGRAM_START = latent_qpu[:5]
@@ -564,13 +565,13 @@ def create_interface():
                                             html.Div(
                                                 [
                                                     html.Img(
-                                                        src="static/model_diagram/step_1_input.png",
+                                                        src=STEP_1_FILE,
                                                         id="step-1-input-img",
                                                     ),
                                                     html.Div([
                                                         html.Div(className="forward-arrow"),
                                                         html.Img(
-                                                            src="static/model_diagram/step_2_encode.png",
+                                                            src=STEP_2_FILE,
                                                             id="step-2-encode-img",
                                                         ),
                                                     ], className="graph-model-itermediate-step"),
@@ -595,9 +596,9 @@ def create_interface():
                                                     ),
                                                     html.Div([
                                                         html.Div(className="forward-arrow"),
-                                                        html.Img(src="static/model_diagram/step_4_decode.png", id="step-4-decode-img"),
+                                                        html.Img(src=STEP_4_FILE, id="step-4-decode-img"),
                                                     ], className="graph-model-itermediate-step"),
-                                                    html.Img(src="static/model_diagram/step_5_output_default.png", id="step-5-output-img"),
+                                                    html.Img(src=STEP_5_FILE_DEFAULT, id="step-5-output-img"),
                                                 ],
                                                 className="graph-model-wrapper"
                                             ),

src/model_wrapper.py

@@ -17,6 +17,7 @@
 from pathlib import Path
 from typing import Optional
 
+from src.utils.global_vars import LATENT_QPU_FILE
 import numpy as np
 import plotly.express as px
 import torch
@@ -368,7 +369,7 @@ def generate_output(self, sharpen: bool = False, save_to_file: str = "") -> go.F
                 sample_params=self.sampler_kwargs,
             )
 
-        with open("static/model_diagram/latent_qpu.json", "w") as f:
+        with open(LATENT_QPU_FILE, "w") as f:
             json.dump(samples[0].tolist(), f)
 
         images = self._dvae.decoder(samples.unsqueeze(1)).squeeze(1).clip(0.0, 1.0).detach().cpu()

src/training_parameters.yaml

@@ -2,7 +2,7 @@ ANNEALING_TIME: 1
 NUM_READS: 256
 
 IMAGE_SIZE: 32
-DATASET_SIZE: 60000  # full MNIST is 60_000
+DATASET_SIZE: null  # full MNIST is 60_000
 BATCH_SIZE: 128
 RANDOM_SEED: 775321899904
 LOSS_FUNCTION: mmd