Move build_website to autoencoder dir

Author: thedch

autoencoder/feature-browser/build_website.py renamed to autoencoder/build_website.py

@@ -20,17 +20,16 @@
 """
 
 import logging
+from pathlib import Path
 from tqdm.auto import trange
 from dataclasses import dataclass
 import torch
 from tensordict import TensorDict
 import os
-import sys
 from math import ceil
-from main_page import create_main_html_page
-from subpages import write_alive_feature_page, write_dead_feature_page, write_ultralow_density_feature_page
+from feature_browser.main_page import create_main_html_page
+from feature_browser.subpages import write_alive_feature_page, write_dead_feature_page, write_ultralow_density_feature_page
 
-sys.path.insert(1, '../')
 from resource_loader import ResourceLoader
 from utils.plotting_utils import make_activations_histogram, make_logits_histogram
 
@@ -399,7 +398,7 @@ def write_main_page(self):
     logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
     # -----------------------------------------------------------------------------
     config_keys = [k for k, v in globals().items() if not k.startswith('_') and isinstance(v, (int, float, bool, str))]
-    configurator = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'configurator.py')
+    configurator = Path(__file__).parent / 'configurator.py'
     exec(open(configurator).read())  # overrides from command line or config file
     config = {k: globals()[k] for k in config_keys}  # will be useful for logging
     # -----------------------------------------------------------------------------

autoencoder/feature-browser/main_page.py renamed to autoencoder/feature_browser/main_page.py

@@ -8,8 +8,7 @@
 # Extend the Python path to include the transformer subdirectory for GPT class import
 base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 sys.path.insert(0, os.path.join(base_dir, 'transformer'))
-from model import GPTConfig
-from hooked_model import HookedGPT
+from model import GPTConfig, HookedGPT
 
 
 class ResourceLoader:
@@ -163,17 +162,8 @@ def select_resampling_data(self, size: int):
         return resampling_data
 
     def load_tokenizer(self):
-        load_meta = False
         meta_path = os.path.join(self.base_dir, 'transformer', 'data', self.dataset, 'meta.pkl')
-        load_meta = os.path.exists(meta_path)
-        if load_meta:
-            print(f"Loading meta from {meta_path}...")
-            with open(meta_path, 'rb') as f:
-                meta = pickle.load(f)
-            # TODO want to make this more general to arbitrary encoder/decoder schemes
-            stoi, itos = meta['stoi'], meta['itos']
-            encode = lambda s: [stoi[c] for c in s]
-            decode = lambda l: ''.join([itos[i] for i in l])
-        else:
-            raise DeprecationWarning('must load from dataset dir')
-        return encode, decode
+        with open(meta_path, 'rb') as f:
+            meta = pickle.load(f)
+
+        return meta['encode'], meta['decode']

autoencoder/feature-browser/subpages.py renamed to autoencoder/feature_browser/subpages.py

@@ -166,6 +166,32 @@ def _init_weights(self, module):
         elif isinstance(module, nn.Embedding):
             torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
 
+    # TODO: Remove this and only use the forward method below
+    def forward(self, idx, targets=None):
+        device = idx.device
+        b, t = idx.size()
+        assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
+        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+
+        if targets is not None:
+            # if we are given some desired targets also calculate the loss
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+        else:
+            # inference-time mini-optimization: only forward the lm_head on the very last position
+            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
+            loss = None
+
+        return logits, loss
+
     def configure_optimizers(self, weight_decay, learning_rate, betas):
         # start with all of the candidate parameters
         param_dict = {pn: p for pn, p in self.named_parameters()}