sae_train.py File Reference

Go to the source code of this file.

Namespaces
namespace	aquin
namespace	aquin.compute
namespace	aquin.compute.sae_train

Functions
Iterator[str]	aquin.compute.sae_train._iter_corpus_text (str|Path|None corpus_path)
None	aquin.compute.sae_train.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	aquin.compute.sae_train.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	aquin.compute.sae_train.default_user_sae_path (str model_id, str name, int layer)

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