activation_capture.py File Reference

Go to the source code of this file.

Namespaces
namespace	aquin
namespace	aquin.compute
namespace	aquin.compute.activation_capture

Functions
Path|None	aquin.compute.activation_capture.resolve_prompts_path (str|Path path)
Path	aquin.compute.activation_capture.write_probes_jsonl (Path path, list[dict[str, Any]] probes)
str|None	aquin.compute.activation_capture._balance_key (dict[str, Any] probe, str|None group=None)
tuple[list[dict[str, Any]], dict[str, Any]]	aquin.compute.activation_capture.balance_probes (list[dict[str, Any]] probes, *, str|None group=None)
list[dict[str, Any]]	aquin.compute.activation_capture.generate_llm_probes (str model_id, int count, *, str topic="general knowledge, reasoning, and instructions")
tuple[list[dict[str, Any]], dict[str, Any]]	aquin.compute.activation_capture.resolve_probes_for_capture (*, str model_id, ModelMode model_mode, str|Path|None prompts_path, int count, str|None topic, bool balance=False, str|None balance_group=None, Path|None output_dir=None)
list[dict[str, Any]]	aquin.compute.activation_capture.load_probes (str|Path|None path, *, list[str]|None fallback=None)
list[int]	aquin.compute.activation_capture.parse_layers (str|None spec, int n_layers)
tuple[str, ModelMode]	aquin.compute.activation_capture.resolve_capture_model_id (str model_id)
int	aquin.compute.activation_capture.llm_layer_count (str model_id, str|Path|None checkpoint_path=None)
torch.Tensor	aquin.compute.activation_capture._pool_activation (torch.Tensor tensor, Position position)
dict[int, torch.Tensor]	aquin.compute.activation_capture._forward_llm_layers (Any model, str text, list[int] layers, *, Position position, int max_chars=512)
tuple[dict[int, torch.Tensor], list[str]]	aquin.compute.activation_capture._forward_llm_layers_token (Any model, str text, list[int] layers, *, int max_chars=512)
dict[str, Any]	aquin.compute.activation_capture._write_capture_artifacts (*, Path out_root, list[dict[str, Any]] probes, list[int] layer_list, dict[int, list[torch.Tensor]] per_layer, list[dict[str, Any]] summary_rows, str model_id, ModelMode model_mode, int d_model, str|Path|None checkpoint_path, str|None checkpoint_name, Position position, bool encode_sae, str|None capture_name, str|None sae_file, dict[str, Any]|None sae_features, dict[str, Any]|None manifest_extras=None, Granularity granularity="prompt", list[dict[str, Any]]|None token_spans=None)
dict[str, Any]	aquin.compute.activation_capture._run_capture_llm (str model_id, list[dict[str, Any]] probes, str|Path output_dir, *, list[int] layers, str|Path|None checkpoint_path, str|None checkpoint_name, Position position, bool encode_sae, int|None sae_layer, str|None capture_name, dict[str, Any]|None manifest_extras=None, Granularity granularity="prompt")
dict[str, Any]	aquin.compute.activation_capture.run_capture_activations (str model_id, list[dict[str, Any]] probes, str|Path output_dir, *, ModelMode|None model_mode=None, list[int]|None layers=None, str|Path|None checkpoint_path=None, str|None checkpoint_name=None, Position position="last", bool encode_sae=False, int|None sae_layer=None, str|None capture_name=None, dict[str, Any]|None manifest_extras=None, Granularity granularity="prompt")

Variables
tuple	aquin.compute.activation_capture.PROBE_TEXT_KEYS = ("instruction", "prompt", "text", "content", "response")
	aquin.compute.activation_capture.RESERVED_PROBE_KEYS = frozenset({"id", *PROBE_TEXT_KEYS})
	aquin.compute.activation_capture.Position = Literal["last", "mean"]
	aquin.compute.activation_capture.Granularity = Literal["prompt", "token"]
	aquin.compute.activation_capture.ModelMode = Literal["llm"]
int	aquin.compute.activation_capture.MAX_PROBE_COUNT = 64
tuple	aquin.compute.activation_capture.BALANCE_PRIORITY = ("label", "stressor", "lang", "group")