Merge pull request #284 from EleutherAI/fix-actually-removing-embeddi…

…ngs-layer

fix: remove first layer (embeddings) from layer_indices

elk/extraction/extraction.py

@@ -194,7 +194,7 @@ def extract_hiddens(
         seed=cfg.seed,
     )
 
-    layer_indices = cfg.layers or tuple(range(model.config.num_hidden_layers))
+    layer_indices = cfg.layers or tuple(range(1, model.config.num_hidden_layers))
 
     global_max_examples = cfg.max_examples[0 if split_type == "train" else 1]
 
@@ -366,12 +366,13 @@ def hidden_features(cfg: Extract) -> tuple[DatasetInfo, Features]:
     if num_dropped:
         print(f"Dropping {num_dropped} non-multiple choice templates")
 
+    layer_indices = cfg.layers or tuple(range(1, model_cfg.num_hidden_layers))
     layer_cols = {
         f"hidden_{layer}": Array3D(
             dtype="int16",
             shape=(num_variants, num_classes, model_cfg.hidden_size),
         )
-        for layer in cfg.layers or range(model_cfg.num_hidden_layers)
+        for layer in layer_indices
     }
     other_cols = {
         "variant_ids": Sequence(