constants

src/constants.py

@@ -15,4 +15,7 @@ DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 MODEL_NAME = "lc0-10-4238.onnx"
 SAE_CONFIG = "debug"
 LAYER = 9
-N_FEATURES = 7680
+ACTIVATION_DIM = 256
+DICTIONARY_SIZE = 7680
+PRE_BIAS = False
+INIT_NORMALISE_DICT = None

src/global_variables.py

@@ -7,7 +7,7 @@ import gradio as gr
 from lczerolens import ModelWrapper
 import torch
 
-from src.constants import HF_TOKEN, ASSETS_FOLDER, DEVICE, MODEL_NAME, SAE_CONFIG, LAYER
+from src import constants
 from src.helpers import SparseAutoEncoder, OutputGenerator
 
 hf_api: HfApi
@@ -22,32 +22,36 @@ def setup():
     global sae
     global generator
 
-    hf_api = HfApi(token=HF_TOKEN)
+    hf_api = HfApi(token=constants.HF_TOKEN)
     hf_api.snapshot_download(
-        local_dir=f"{ASSETS_FOLDER}/models",
+        local_dir=f"{constants.ASSETS_FOLDER}/models",
         repo_id="Xmaster6y/lczero-planning-models",
         repo_type="model",
     )
     hf_api.snapshot_download(
-        local_dir=f"{ASSETS_FOLDER}/saes",
+        local_dir=f"{constants.ASSETS_FOLDER}/saes",
         repo_id="Xmaster6y/lczero-planning-saes",
         repo_type="model",
     )
 
-    wrapper = ModelWrapper.from_onnx_path(f"{ASSETS_FOLDER}/models/{MODEL_NAME}").to(DEVICE)
+    wrapper = ModelWrapper.from_onnx_path(f"{constants.ASSETS_FOLDER}/models/{constants.MODEL_NAME}").to(constants.DEVICE)
     sae_dict = torch.load(
-        f"{ASSETS_FOLDER}/saes/{SAE_CONFIG}/model.pt",
-        map_location=DEVICE,
-        weights_only=True
+        f"{constants.ASSETS_FOLDER}/saes/{constants.SAE_CONFIG}/model.pt",
+        map_location=constants.DEVICE,
+    )
+    sae = SparseAutoEncoder(
+        constants.ACTIVATION_DIM,
+        constants.DICTIONARY_SIZE,
+        pre_bias=constants.PRE_BIAS,
+        init_normalise_dict=constants.INIT_NORMALISE_DICT,
     )
-    sae = SparseAutoEncoder()
     sae.load_state_dict(
         sae_dict
     )
     generator = OutputGenerator(
         sae=sae,
         wrapper=wrapper,
-        module_exp=rf".*block{LAYER}/conv2/relu"
+        module_exp=rf".*block{constants.LAYER}/conv2/relu"
     )
 
 if gr.NO_RELOAD: