add mod preds; todo: add legend

app.py

@@ -313,7 +313,7 @@ def get_orig_preds(i: int) -> dict[str, object]:
 
 
 @beartype.beartype
-def unscaled(x: float, max_obs: float) -> float:
+def unscaled(x: float, max_obs: float | int) -> float:
     """Scale from [-10, 10] to [10 * -max_obs, 10 * max_obs]."""
     return map_range(x, (-10.0, 10.0), (-10.0 * max_obs, 10.0 * max_obs))
 
@@ -333,12 +333,17 @@ def map_range(
 
 @beartype.beartype
 @torch.inference_mode
-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)
+def get_mod_preds(i: int, latents: dict[str, int | float]) -> dict[str, object]:
+    latents = {int(k): float(v) for k, v in latents.items()}
+    img = data.get_img(i)
 
+    split_vit, vit_transform = modeling.load_vit(DEVICE)
+    sae = load_sae(DEVICE)
+    _, top_values, _ = load_tensors()
+    clf = load_clf()
+
+    x_BCWH = vit_transform(img)[None, ...].to(DEVICE)
+    x_BPD = split_vit.forward_start(x_BCWH)
     x_hat_BPD, f_x_BPS, _ = sae(x_BPD)
 
     err_BPD = x_BPD - x_hat_BPD
@@ -346,18 +351,14 @@ def get_mod_preds(i: int, latents: dict[int, float]) -> dict[str, object]:
     values = torch.tensor(
         [
             unscaled(float(value), top_values[latent].max().item())
-            for value, latent in [
-                (value1, latent1),
-                (value2, latent2),
-                (value3, latent3),
-            ]
+            for latent, value in latents.items()
         ],
-        device=device,
+        device=DEVICE,
     )
-    f_x_BPS[..., torch.tensor([latent1, latent2, latent3], device=device)] = values
+    f_x_BPS[..., torch.tensor(list(latents.keys()), device=DEVICE)] = values
 
     # Reproduce the SAE forward pass after f_x
-    modified_x_hat_BPD = (
+    mod_x_hat_BPD = (
         einops.einsum(
             f_x_BPS,
             sae.W_dec,
@@ -365,14 +366,19 @@ def get_mod_preds(i: int, latents: dict[int, float]) -> dict[str, object]:
         )
         + sae.b_dec
     )
-    modified_BPD = err_BPD + modified_x_hat_BPD
+    mod_BPD = err_BPD + mod_x_hat_BPD
 
-    modified_BPD = rest_of_vit.forward_end(modified_BPD)
+    mod_BPD = split_vit.forward_end(mod_BPD)
+    mod_WHD = einops.rearrange(mod_BPD, "() (w h) dim -> w h dim", w=16, h=16)
 
-    logits_BPC = head(modified_BPD)
-    pred_P = logits_BPC[0].argmax(axis=-1)
-    pred_WH = einops.rearrange(pred_P, "(w h) -> w h", w=16, h=16)
-    return seg_to_img(upsample(pred_WH)), pred_P.tolist()
+    logits_WHC = clf(mod_WHD)
+    pred_WH = logits_WHC.argmax(axis=-1)
+    # pred_WH = einops.rearrange(pred_P, "(w h) -> w h", w=16, h=16)
+    return {
+        "index": i,
+        "orig_url": data.img_to_base64(data.to_sized(img)),
+        "seg_url": data.img_to_base64(data.u8_to_img(upsample(pred_WH))),
+    }
 
 
 @jaxtyped(typechecker=beartype.beartype)