b4hsu/jacobian-lens

This folder contains pre-fitted Jacobian Lenses for specific models at target_layer<=-1. They were trained using Anthropic's Jacobian Lens library. The same library can be used to load the lens and local visualizer.

Code used:

import torch
import transformers

import jlens
from jlens import fit
from jlens.examples import load_wikitext_prompts


def save_fit(output_path, checkpoint_path, target_layer, cache_dir=None):
    hf_model = transformers.AutoModelForCausalLM.from_pretrained(
        MODEL_NAME, dtype=torch.bfloat16, cache_dir=cache_dir
    ).cuda()
    tokenizer = transformers.AutoTokenizer.from_pretrained(
        MODEL_NAME, cache_dir=cache_dir
    )
    model = jlens.from_hf(hf_model, tokenizer)
    prompts = load_wikitext_prompts(n_prompts=100)
    lens = fit(
        model,
        prompts,
        target_layer=target_layer,
        dim_batch=32,
        max_seq_len=128,
        checkpoint_path=checkpoint_path,
    )
    lens.save(output_path)
    print(f"saved lens: {lens} -> {output_path}")
    return lens

---

Pretrained Jacobian Lens (target_layer=-1) for different models can be accessed at https://huggingface.co/neuronpedia/jacobian-lens.