Spaces:
Running on Zero
Running on Zero
File size: 2,279 Bytes
b3d11b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | from __future__ import annotations
import csv
from collections import defaultdict
from pathlib import Path
def load_feature_sets(path: str | Path, max_size: int) -> dict[str, dict]:
"""
Select a same-layer feature set for each concept from a feature catalog.
The layer whose best training feature is strongest is selected first. The top
distinct features from that same SAE dictionary are then returned. Keeping a
set within one layer is essential because decoder directions from different
residual spaces should not be summed into one intervention.
"""
with Path(path).open(newline='', encoding='utf-8') as handle:
rows = list(csv.DictReader(handle))
grouped: dict[str, list[dict]] = defaultdict(list)
for row in rows:
item = dict(row)
item['layer'] = int(row['layer'])
item['feature_id'] = int(row['feature_id'])
item['train_auroc'] = float(row['train_auroc'])
item['activation_contrast'] = (
float(row['activation_rate_pos']) - float(row['activation_rate_neg'])
)
grouped[row['concept']].append(item)
result: dict[str, dict] = {}
for concept, concept_rows in grouped.items():
best_by_layer: dict[int, tuple[float, float]] = {}
for row in concept_rows:
key = (row['train_auroc'], row['activation_contrast'])
best_by_layer[row['layer']] = max(
best_by_layer.get(row['layer'], (-1.0, -1.0)),
key,
)
chosen_layer = max(best_by_layer, key=best_by_layer.get)
layer_rows = [row for row in concept_rows if row['layer'] == chosen_layer]
layer_rows.sort(
key=lambda row: (row['train_auroc'], row['activation_contrast']),
reverse=True,
)
seen: set[int] = set()
feature_ids: list[int] = []
for row in layer_rows:
feature_id = int(row['feature_id'])
if feature_id not in seen:
seen.add(feature_id)
feature_ids.append(feature_id)
if len(feature_ids) >= int(max_size):
break
result[concept] = {
'layer': int(chosen_layer),
'feature_ids': feature_ids,
}
return result
|