Skip to content

Commit 3f9c0a9

Browse files
authored
Add basic SAELens code snippet (#835)
Follow-up to #826, and (currently) an intermediate PR before `.from_pretrained` support is added for any model on the HF Hub. See jbloomAus/SAELens#234 (comment) for additional discussion.
1 parent 98e32f0 commit 3f9c0a9

File tree

2 files changed

+11
-0
lines changed

2 files changed

+11
-0
lines changed

packages/tasks/src/model-libraries-snippets.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -401,6 +401,16 @@ export const timm = (model: ModelData): string[] => [
401401
model = timm.create_model("hf_hub:${model.id}", pretrained=True)`,
402402
];
403403

404+
export const saelens = (/* model: ModelData */): string[] => [
405+
`# pip install sae-lens
406+
from sae_lens import SAE
407+
408+
sae, cfg_dict, sparsity = SAE.from_pretrained(
409+
release = "RELEASE_ID", # e.g., "gpt2-small-res-jb". See other options in https://github.com/jbloomAus/SAELens/blob/main/sae_lens/pretrained_saes.yaml
410+
sae_id = "SAE_ID", # e.g., "blocks.8.hook_resid_pre". Won't always be a hook point
411+
)`,
412+
];
413+
404414
const skopsPickle = (model: ModelData, modelFile: string) => {
405415
return [
406416
`import joblib

packages/tasks/src/model-libraries.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -433,6 +433,7 @@ export const MODEL_LIBRARIES_UI_ELEMENTS = {
433433
prettyLabel: "SAELens",
434434
repoName: "SAELens",
435435
repoUrl: "https://github.com/jbloomAus/SAELens",
436+
snippets: snippets.saelens,
436437
filter: false,
437438
},
438439
"sample-factory": {

0 commit comments

Comments
 (0)