Fix mobile sidebar

etowahadams · Oct 31, 2024 · a81dfa1 · a81dfa1
1 parent 8ac15d4
commit a81dfa1
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 37 deletions.
diff --git a/viz/src/SAEConfigs.ts b/viz/src/SAEConfigs.ts
@@ -1,13 +1,19 @@
-export const SAE_CONFIGS: {
-  [key: string]: {
-    baseUrl: string;
-    numHiddenDims: number;
-    plmLayer: number;
-    curated?: { name: string; dim: number; desc: string }[];
-    defaultDim: number;
-    supportsCustomSequence?: boolean;
-  };
-} = {
+export type CuratedFeature = {
+  name: string;
+  dim: number;
+  desc: string;
+};
+
+export type SAEConfig = {
+  baseUrl: string;
+  numHiddenDims: number;
+  plmLayer: number;
+  curated?: CuratedFeature[];
+  defaultDim: number;
+  supportsCustomSequence?: boolean;
+};
+
+export const SAE_CONFIGS: Record<string, SAEConfig> = {
   "SAE4096-L24": {
     baseUrl:
       "https://raw.githubusercontent.com/liambai/plm-interp-viz-data/refs/heads/main/esm2_plm1280_l24_sae4096_100Kseqs/",

diff --git a/viz/src/SAEVisualizer.tsx b/viz/src/SAEVisualizer.tsx
@@ -1,9 +1,8 @@
 import { useEffect, useState } from "react";
-// import MolstarViewer from "./components/MolstarViewer";
-import MolstarMulti from './components/MolstarMulti';
+import MolstarMulti from "./components/MolstarMulti";
 import SeqViewer, { SingleSeq } from "./components/SeqViewer";
 import CustomSeqPlayground from "./components/CustomSeqPlayground";
-import { SAE_CONFIGS } from "./SAEConfigs";
+import { CuratedFeature, SAE_CONFIGS } from "./SAEConfigs";
 import {
   Select,
   SelectContent,
@@ -17,6 +16,7 @@ import {
   SidebarProvider,
   SidebarHeader,
   SidebarTrigger,
+  useSidebar,
 } from "@/components/ui/sidebar";
 import { Separator } from "@/components/ui/separator";
 import HomeNavigator from "@/components/HomeNavigator";
@@ -25,6 +25,51 @@ import { Toggle } from "./components/ui/toggle";
 
 const NUM_SEQS_TO_DISPLAY = 9;
 
+interface FeatureListProps {
+  config: {
+    curated?: CuratedFeature[];
+    numHiddenDims: number;
+  };
+  feature: number;
+  setFeature: (feature: number) => void;
+}
+
+function FeatureList({ config, feature, setFeature }: FeatureListProps) {
+  const { setOpenMobile } = useSidebar();
+
+  const handleFeatureChange = (feature: number) => {
+    setFeature(feature);
+    setOpenMobile(false);
+  };
+
+  return (
+    <ul className="space-y-2 font-medium">
+      {config.curated?.map((c) => (
+        <Toggle
+          key={`feature-${c.dim}`}
+          style={{ width: "100%", paddingLeft: 20, textAlign: "left" }}
+          className="justify-start"
+          pressed={feature === c.dim}
+          onPressedChange={() => handleFeatureChange(c.dim)}
+        >
+          {c.name}
+        </Toggle>
+      ))}
+      {Array.from({ length: config.numHiddenDims }, (_, i) => i).map((i) => (
+        <Toggle
+          key={`feature-${i}`}
+          style={{ width: "100%", paddingLeft: 20 }}
+          className="justify-start"
+          pressed={feature === i}
+          onPressedChange={() => handleFeatureChange(i)}
+        >
+          {i}
+        </Toggle>
+      ))}
+    </ul>
+  );
+}
+
 function SAEVisualizer() {
   const [selectedModel, setSelectedModel] = useState(() => {
     const params = new URLSearchParams(window.location.search);
@@ -91,30 +136,7 @@ function SAEVisualizer() {
           <Separator />
         </SidebarHeader>
         <SidebarContent>
-          <ul className="space-y-2 font-medium">
-            {config.curated?.map((i) => (
-              <Toggle
-                key={`feature-${i.dim}`}
-                style={{ width: "100%", paddingLeft: 20, textAlign: "left" }}
-                className="justify-start"
-                pressed={feature === i.dim}
-                onPressedChange={() => setFeature(i.dim)}
-              >
-                {i.name}
-              </Toggle>
-            ))}
-            {Array.from({ length: config.numHiddenDims }, (_, i) => i).map((i) => (
-              <Toggle
-                key={`feature-${i}`}
-                style={{ width: "100%", paddingLeft: 20 }}
-                className="justify-start"
-                pressed={feature === i}
-                onPressedChange={() => setFeature(i)}
-              >
-                {i}
-              </Toggle>
-            ))}
-          </ul>
+          <FeatureList config={config} feature={feature} setFeature={setFeature} />
         </SidebarContent>
       </Sidebar>
       <main className="text-left max-w-full overflow-x-auto">