Functions
Iterator[str]	_iter_corpus_text (str\|Path\|None corpus_path)
None	collect_activations (str model_id, int layer, Path acts_dir, int n_tokens, *, str\|Path\|None checkpoint_path=None, str\|Path\|None corpus_path=None)
Path	train_sae (str model_id, int layer, Path output_path, *, str\|Path\|None checkpoint_path=None, str\|Path\|None corpus_path=None, bool quick=False, int\|None d_model=None, int\|None n_features=None, int\|None max_steps=None, int\|None max_epochs=None, str\|Path\|None activations_dir=None, bool balance=False, str\|None balance_group=None)
Path	default_user_sae_path (str model_id, str name, int layer)

Variables
int	QUICK_TOKENS = 100_000
int	FULL_TOKENS = 2_000_000
int	QUICK_MAX_STEPS = 3_000
int	FULL_MAX_STEPS = 50_000
int	QUICK_MAX_EPOCHS = 10
int	FULL_MAX_EPOCHS = 200
int	SEQ_LEN = 64
int	COLLECT_BATCH = 32
int	CHUNK_SIZE = 50_000
int	BATCH_SIZE = 4096
int	LR = 1e-4
float	L1_COEFF = 10.0
str	USER_SAE_ROOT = Path.home() / ".aquin" / "sae" / "user"

Function Documentation

◆ _iter_corpus_text()

Iterator[str] _iter_corpus_text ( str | Path | None corpus_path )

protected

Definition at line 46 of file sae_train.py.

Referenced by collect_activations().

◆ collect_activations()

None collect_activations	(	str	model_id,
		int	layer,
		Path	acts_dir,
		int	n_tokens,
		*	,
		str \| Path \| None	checkpoint_path = None,
		str \| Path \| None	corpus_path = None )

Definition at line 83 of file sae_train.py.

References _iter_corpus_text().

Referenced by train_sae().

◆ default_user_sae_path()

Path default_user_sae_path	(	str	model_id,
		str	name,
		int	layer )

Definition at line 303 of file sae_train.py.

◆ train_sae()

Path train_sae	(	str	model_id,
		int	layer,
		Path	output_path,
		*	,
		str \| Path \| None	checkpoint_path = None,
		str \| Path \| None	corpus_path = None,
		bool	quick = False,
		int \| None	d_model = None,
		int \| None	n_features = None,
		int \| None	max_steps = None,
		int \| None	max_epochs = None,
		str \| Path \| None	activations_dir = None,
		bool	balance = False,
		str \| None	balance_group = None )

Definition at line 149 of file sae_train.py.

References collect_activations().

Variable Documentation

◆ BATCH_SIZE

int aquin.compute.sae_train.BATCH_SIZE = 4096

Definition at line 40 of file sae_train.py.

◆ CHUNK_SIZE

int aquin.compute.sae_train.CHUNK_SIZE = 50_000

Definition at line 39 of file sae_train.py.

◆ COLLECT_BATCH

int aquin.compute.sae_train.COLLECT_BATCH = 32

Definition at line 38 of file sae_train.py.

◆ FULL_MAX_EPOCHS

int aquin.compute.sae_train.FULL_MAX_EPOCHS = 200

Definition at line 36 of file sae_train.py.

◆ FULL_MAX_STEPS

int aquin.compute.sae_train.FULL_MAX_STEPS = 50_000

Definition at line 34 of file sae_train.py.

◆ FULL_TOKENS

int aquin.compute.sae_train.FULL_TOKENS = 2_000_000

Definition at line 32 of file sae_train.py.

◆ L1_COEFF

float aquin.compute.sae_train.L1_COEFF = 10.0

Definition at line 42 of file sae_train.py.

◆ LR

int aquin.compute.sae_train.LR = 1e-4

Definition at line 41 of file sae_train.py.

◆ QUICK_MAX_EPOCHS

int aquin.compute.sae_train.QUICK_MAX_EPOCHS = 10

Definition at line 35 of file sae_train.py.

◆ QUICK_MAX_STEPS

int aquin.compute.sae_train.QUICK_MAX_STEPS = 3_000

Definition at line 33 of file sae_train.py.

◆ QUICK_TOKENS

int aquin.compute.sae_train.QUICK_TOKENS = 100_000

Definition at line 31 of file sae_train.py.

◆ SEQ_LEN

int aquin.compute.sae_train.SEQ_LEN = 64

Definition at line 37 of file sae_train.py.

◆ USER_SAE_ROOT

str aquin.compute.sae_train.USER_SAE_ROOT = Path.home() / ".aquin" / "sae" / "user"

Definition at line 43 of file sae_train.py.

Functions

Variables