wip: adding mod preds

app.py

@@ -46,7 +46,7 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 CWD = pathlib.Path(".")
 """Current working directory."""
 
-N_SAE_LATENTS = 3
+N_SAE_LATENTS = 2
 """Number of SAE latents to show."""
 
 N_LATENT_EXAMPLES = 4
@@ -289,7 +289,7 @@ def get_sae_latents(img_i: int, patches: list[int]) -> list[SaeActivation]:
 
 
 @torch.inference_mode
-def get_preds(i: int) -> dict[str, object]:
+def get_orig_preds(i: int) -> dict[str, object]:
     img = data.get_img(i)
     split_vit, vit_transform = modeling.load_vit(DEVICE)
 
@@ -331,16 +331,10 @@ def map_range(
     return c + (x - a) * (d - c) / (b - a)
 
 
+@beartype.beartype
 @torch.inference_mode
-def get_modified_labels(
-    i: int,
-    latent1: int,
-    latent2: int,
-    latent3: int,
-    value1: float,
-    value2: float,
-    value3: float,
-) -> list[Image.Image | list[int]]:
+def get_mod_preds(i: int, latents: dict[int, float]) -> dict[str, object]:
+    breakpoint()
     sample = vit_dataset[i]
     x = sample["image"][None, ...].to(device)
     x_BPD = rest_of_vit.forward_start(x)
@@ -429,38 +423,38 @@ with gr.Blocks() as demo:
         api_name="get-sae-latents",
     )
 
-    #############
-    # get-preds #
-    #############
+    ##################
+    # get-orig-preds #
+    ##################
 
     # Outputs
-    get_preds_out = gr.JSON(label="get_preds_out", value=[])
+    get_orig_preds_out = gr.JSON(label="get_orig_preds_out", value=[])
 
     get_pred_labels_btn = gr.Button(value="Get Predictions")
     get_pred_labels_btn.click(
-        get_preds, inputs=[img_number], outputs=[get_preds_out], api_name="get-preds"
+        get_orig_preds,
+        inputs=[img_number],
+        outputs=[get_orig_preds_out],
+        api_name="get-orig-preds",
     )
 
-    # get_true_labels_btn = gr.Button(value="Get True Label")
-    # get_true_labels_btn.click(
-    #     get_true_labels,
-    #     inputs=[img_number],
-    #     outputs=semseg_img,
-    #     api_name="get-true-labels",
-    # )
-
-    # latent_numbers = [gr.Number(label=f"Latent {i + 1}") for i in range(3)]
-    # value_sliders = [
-    #     gr.Slider(label=f"Value {i + 1}", minimum=-10, maximum=10) for i in range(3)
-    # ]
-
-    # get_modified_labels_btn = gr.Button(value="Get Modified Label")
-    # get_modified_labels_btn.click(
-    #     get_modified_labels,
-    #     inputs=[img_number] + latent_numbers + value_sliders,
-    #     outputs=[semseg_img, semseg_colors],
-    #     api_name="get-modified-labels",
-    # )
+    #################
+    # get-mod-preds #
+    #################
+
+    # Inputs
+    latents_json = gr.JSON(label="Modified Latents", value={})
+
+    # Outputs
+    get_mod_preds_out = gr.JSON(label="get_mod_preds_out", value=[])
+
+    get_pred_labels_btn = gr.Button(value="Get Predictions")
+    get_pred_labels_btn.click(
+        get_mod_preds,
+        inputs=[img_number, latents_json],
+        outputs=[get_mod_preds_out],
+        api_name="get-mod-preds",
+    )
 
 if __name__ == "__main__":
     demo.launch()