Improve UI

Code_Browser.py

@@ -20,33 +20,28 @@ st.set_page_config(
 
 st.title("Codebook Features")
 
+pretty_model_names = {
+    "TinyStories-1Layer-21M#100ksteps_vcb_mlp": "TinyStories-1L-21M-MLP",
+    "TinyStories-1Layer-21M_ccb_attn_preproj": "TinyStories-1L-21M-Attn",
+    "TinyStories-33M_ccb_attn_preproj": "TinyStories-4L-33M-Attn",
+}
+orig_model_name = {v: k for k, v in pretty_model_names.items()}
+
 base_cache_dir = "cache/"
 dirs = glob.glob(base_cache_dir + "models/*/")
 model_name_options = [d.split("/")[-2].split("_")[:-2] for d in dirs]
 model_name_options = ["_".join(m) for m in model_name_options]
 model_name_options = sorted(set(model_name_options))
-def_model_idx = ["attn" in m for m in model_name_options].index(True)
+def_model_idx = ["attn" in m.lower() for m in model_name_options].index(True)
 
-model_name = st.selectbox(
+p_model_name = st.selectbox(
     "Model",
-    model_name_options,
+    [pretty_model_names.get(m, m) for m in model_name_options],
     index=def_model_idx,
     key=webapp_utils.persist("model_name"),
 )
-
+model_name = orig_model_name.get(p_model_name, p_model_name)
 model = model_name.split("_")[0].split("#")[0]
-model_layers = {
-    "pythia-410m-deduped": 24,
-    "pythia-70m-deduped": 6,
-    "gpt2": 12,
-    "TinyStories-1Layer-21M": 1,
-}
-model_heads = {
-    "pythia-410m-deduped": 16,
-    "pythia-70m-deduped": 8,
-    "gpt2": 12,
-    "TinyStories-1Layer-21M": 16,
-}
 ccb = model_name.split("_")[1]
 ccb = "_ccb" if ccb == "ccb" else ""
 cb_at = "_".join(model_name.split("_")[2:])
@@ -59,12 +54,12 @@ dirs.sort(key=os.path.getmtime)
 
 # session states
 is_attn = "attn" in cb_at
-num_layers = model_layers[model]
-num_heads = model_heads[model]
 codes_cache_path = dirs[-1] + "/"
 
 model_info = code_search_utils.parse_model_info(codes_cache_path)
 num_codes = model_info.num_codes
+num_layers = model_info.n_layers
+num_heads = model_info.n_heads
 dataset_cache_path = base_cache_dir + f"datasets/{model_info.dataset_name}/"
 
 (
@@ -96,6 +91,12 @@ if not DEPLOY_MODE:
         st.write(metrics)
 
 st.markdown("## Demo Codes")
+demo_codes_desc = (
+    "This section contains codes that we've found to be interpretable along "
+    "with a description of the feature we think they are capturing."
+    "Click on the 🔍 search button for a code to see the tokens that code activates on."
+)
+st.write(demo_codes_desc)
 demo_file_path = codes_cache_path + "demo_codes.txt"
 
 if st.checkbox("Show Demo Codes"):