AQIT 0.1.0
Loading...
Searching...
No Matches
aquin.compute.sae_stats Namespace Reference

Functions

Path _sae_cache_dir (str model_id)
list[int]|None _parse_layers_wanted (str|None layers)
str _format_load_sae_commands (str model_id, list[int] layers)
str _missing_sae_error (str model_id, *, str mode, list[int] available, str|None layers_arg)
Path _expand (str path)
list[dict[str, Any]] _rows_to_probes (list[Any] rows)
list[dict[str, Any]] load_probes (str|list[Any] path)
list[int] parse_layers_arg (str|None layers, list[int] available, *, str|None model_id=None, str mode="llm")
list[int] list_llm_sae_layers (str model_id)
dict[str, Any] _feature_stats (torch.Tensor mean_acts, *, int top_k)
torch.Tensor _llm_probe_acts (Any model, Any sae, str model_id, int layer, str text)
dict[str, Any] run_llm_sae_stats (Any model, str model_id, list[dict[str, Any]] probes, *, list[int] layers, int top_k=10)
dict[str, Any] _build_payload (*, str model_id, str mode, list[dict[str, Any]] probes, list[int] layers, list[dict[str, Any]] layer_stats, list[str] heatmap_rows, list[str] heatmap_cols, list[list[float]] heatmap_values, int top_k)
dict[str, Any] run_sae_stats (dict[str, Any] args)

Variables

 DEVICE = resolve_compute_device()
tuple _PROBE_TEXT_KEYS = ("text", "prompt", "input", "query", "instruction")
tuple _META_KEYS = ("id", "stressor", "lang", "quant_run_id", "condition", "label", "group")

Function Documentation

◆ _build_payload()

dict[str, Any] _build_payload ( * ,
str model_id,
str mode,
list[dict[str, Any]] probes,
list[int] layers,
list[dict[str, Any]] layer_stats,
list[str] heatmap_rows,
list[str] heatmap_cols,
list[list[float]] heatmap_values,
int top_k )
protected

Definition at line 319 of file sae_stats.py.

Referenced by run_llm_sae_stats().

◆ _expand()

Path _expand ( str path)
protected

Definition at line 97 of file sae_stats.py.

Referenced by load_probes().

◆ _feature_stats()

dict[str, Any] _feature_stats ( torch.Tensor mean_acts,
* ,
int top_k )
protected 
mean_acts: (n_features,)

Definition at line 204 of file sae_stats.py.

Referenced by run_llm_sae_stats().

◆ _format_load_sae_commands()

str _format_load_sae_commands ( str model_id,
list[int] layers )
protected

Definition at line 50 of file sae_stats.py.

Referenced by _missing_sae_error().

◆ _llm_probe_acts()

torch.Tensor _llm_probe_acts ( Any model,
Any sae,
str model_id,
int layer,
str text )
protected

Definition at line 226 of file sae_stats.py.

Referenced by run_llm_sae_stats().

◆ _missing_sae_error()

str _missing_sae_error ( str model_id,
* ,
str mode,
list[int] available,
str | None layers_arg )
protected 
Actionable error when SAE checkpoint files are missing on disk.

Definition at line 54 of file sae_stats.py.

References _format_load_sae_commands(), _parse_layers_wanted(), and _sae_cache_dir().

Referenced by parse_layers_arg(), and run_sae_stats().

◆ _parse_layers_wanted()

list[int] | None _parse_layers_wanted ( str | None layers)
protected

Definition at line 33 of file sae_stats.py.

Referenced by _missing_sae_error(), and parse_layers_arg().

◆ _rows_to_probes()

list[dict[str, Any]] _rows_to_probes ( list[Any] rows)
protected

Definition at line 101 of file sae_stats.py.

Referenced by load_probes().

◆ _sae_cache_dir()

Path _sae_cache_dir ( str model_id)
protected

Definition at line 29 of file sae_stats.py.

Referenced by _missing_sae_error().

◆ list_llm_sae_layers()

list[int] list_llm_sae_layers ( str model_id)

Definition at line 198 of file sae_stats.py.

Referenced by run_sae_stats().

◆ load_probes()

list[dict[str, Any]] load_probes ( str | list[Any] path) 
Load probe rows from a path, inline JSON, or an in-memory list.

Definition at line 122 of file sae_stats.py.

References _expand(), and _rows_to_probes().

Referenced by run_sae_stats().

◆ parse_layers_arg()

list[int] parse_layers_arg ( str | None layers,
list[int] available,
* ,
str | None model_id = None,
str mode = "llm" )

Definition at line 166 of file sae_stats.py.

References _missing_sae_error(), and _parse_layers_wanted().

Referenced by run_sae_stats().

◆ run_llm_sae_stats()

dict[str, Any] run_llm_sae_stats ( Any model,
str model_id,
list[dict[str, Any]] probes,
* ,
list[int] layers,
int top_k = 10 )

Definition at line 244 of file sae_stats.py.

References _build_payload(), _feature_stats(), and _llm_probe_acts().

Referenced by run_sae_stats().

◆ run_sae_stats()

dict[str, Any] run_sae_stats ( dict[str, Any] args)

Definition at line 362 of file sae_stats.py.

References _missing_sae_error(), list_llm_sae_layers(), load_probes(), parse_layers_arg(), and run_llm_sae_stats().

Variable Documentation

◆ _META_KEYS

tuple aquin.compute.sae_stats._META_KEYS = ("id", "stressor", "lang", "quant_run_id", "condition", "label", "group")
protected

Definition at line 26 of file sae_stats.py.

◆ _PROBE_TEXT_KEYS

tuple aquin.compute.sae_stats._PROBE_TEXT_KEYS = ("text", "prompt", "input", "query", "instruction")
protected

Definition at line 25 of file sae_stats.py.

◆ DEVICE

aquin.compute.sae_stats.DEVICE = resolve_compute_device()

Definition at line 23 of file sae_stats.py.