Functions
Path\|None	resolve_prompts_path (str\|Path path)
Path	write_probes_jsonl (Path path, list[dict[str, Any]] probes)
str\|None	_balance_key (dict[str, Any] probe, str\|None group=None)
tuple[list[dict[str, Any]], dict[str, Any]]	balance_probes (list[dict[str, Any]] probes, *, str\|None group=None)
list[dict[str, Any]]	generate_llm_probes (str model_id, int count, *, str topic="general knowledge, reasoning, and instructions")
tuple[list[dict[str, Any]], dict[str, Any]]	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]]	load_probes (str\|Path\|None path, *, list[str]\|None fallback=None)
list[int]	parse_layers (str\|None spec, int n_layers)
tuple[str, ModelMode]	resolve_capture_model_id (str model_id)
int	llm_layer_count (str model_id, str\|Path\|None checkpoint_path=None)
torch.Tensor	_pool_activation (torch.Tensor tensor, Position position)
dict[int, torch.Tensor]	_forward_llm_layers (Any model, str text, list[int] layers, *, Position position, int max_chars=512)
tuple[dict[int, torch.Tensor], list[str]]	_forward_llm_layers_token (Any model, str text, list[int] layers, *, int max_chars=512)
dict[str, Any]	_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]	_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]	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	PROBE_TEXT_KEYS = ("instruction", "prompt", "text", "content", "response")
	RESERVED_PROBE_KEYS = frozenset({"id", *PROBE_TEXT_KEYS})
	Position = Literal["last", "mean"]
	Granularity = Literal["prompt", "token"]
	ModelMode = Literal["llm"]
int	MAX_PROBE_COUNT = 64
tuple	BALANCE_PRIORITY = ("label", "stressor", "lang", "group")

Function Documentation

◆ _balance_key()

str \| None _balance_key	(	dict[str, Any]	probe,
		str \| None	group = None )

protected

Definition at line 53 of file activation_capture.py.

Referenced by balance_probes().

◆ _forward_llm_layers()

dict[int, torch.Tensor] _forward_llm_layers	(	Any	model,
		str	text,
		list[int]	layers,
		*	,
		Position	position,
		int	max_chars = 512 )

protected

Definition at line 297 of file activation_capture.py.

References _pool_activation().

Referenced by _run_capture_llm().

◆ _forward_llm_layers_token()

tuple[dict[int, torch.Tensor], list[str]] _forward_llm_layers_token	(	Any	model,
		str	text,
		list[int]	layers,
		*	,
		int	max_chars = 512 )

protected

Definition at line 321 of file activation_capture.py.

Referenced by _run_capture_llm().

◆ _pool_activation()

torch.Tensor _pool_activation	(	torch.Tensor	tensor,
		Position	position )

protected

tensor: (seq, d_model) -> (d_model,)

Definition at line 290 of file activation_capture.py.

Referenced by _forward_llm_layers().

◆ _run_capture_llm()

dict[str, Any] _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" )

protected

Definition at line 474 of file activation_capture.py.

References _forward_llm_layers(), _forward_llm_layers_token(), and _write_capture_artifacts().

Referenced by run_capture_activations().

◆ _write_capture_artifacts()

dict[str, Any] _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 )

protected

Definition at line 347 of file activation_capture.py.

Referenced by _run_capture_llm().

◆ balance_probes()

tuple[list[dict[str, Any]], dict[str, Any]] balance_probes	(	list[dict[str, Any]]	probes,
		*	,
		str \| None	group = None )

Definition at line 67 of file activation_capture.py.

References _balance_key().

Referenced by resolve_probes_for_capture().

◆ generate_llm_probes()

list[dict[str, Any]] generate_llm_probes	(	str	model_id,
		int	count,
		*	,
		str	topic = "general knowledge, reasoning, and instructions" )

Use the loaded LLM to generate diverse probe strings.

Definition at line 106 of file activation_capture.py.

Referenced by resolve_probes_for_capture().

◆ llm_layer_count()

int llm_layer_count	(	str	model_id,
		str \| Path \| None	checkpoint_path = None )

Catalog layer count :  never load weights just to read n_layers (OOM on T4).

Definition at line 281 of file activation_capture.py.

◆ load_probes()

list[dict[str, Any]] load_probes	(	str \| Path \| None	path,
		*	,
		list[str] \| None	fallback = None )

Load probes with optional metadata from JSON/JSONL.

Definition at line 215 of file activation_capture.py.

References resolve_prompts_path().

Referenced by resolve_probes_for_capture().

◆ parse_layers()

list[int] parse_layers	(	str \| None	spec,
		int	n_layers )

Definition at line 258 of file activation_capture.py.

◆ resolve_capture_model_id()

tuple[str, ModelMode] resolve_capture_model_id ( str model_id )

Definition at line 275 of file activation_capture.py.

Referenced by run_capture_activations().

◆ resolve_probes_for_capture()

tuple[list[dict[str, Any]], dict[str, Any]] 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 )

Load probes from file or generate when --prompts omitted.

Definition at line 165 of file activation_capture.py.

References balance_probes(), generate_llm_probes(), load_probes(), resolve_prompts_path(), and write_probes_jsonl().

◆ resolve_prompts_path()

Path | None resolve_prompts_path ( str | Path path )

Resolve prompts file from cwd or repo parent (e.g. when run from cli/).

Definition at line 28 of file activation_capture.py.

Referenced by load_probes(), and resolve_probes_for_capture().

◆ run_capture_activations()

dict[str, Any] 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" )

Definition at line 587 of file activation_capture.py.

References _run_capture_llm(), and resolve_capture_model_id().

◆ write_probes_jsonl()

Path write_probes_jsonl	(	Path	path,
		list[dict[str, Any]]	probes )

Definition at line 44 of file activation_capture.py.

Referenced by resolve_probes_for_capture().

Variable Documentation

◆ BALANCE_PRIORITY

tuple aquin.compute.activation_capture.BALANCE_PRIORITY = ("label", "stressor", "lang", "group")

Definition at line 25 of file activation_capture.py.

◆ Granularity

aquin.compute.activation_capture.Granularity = Literal["prompt", "token"]

Definition at line 22 of file activation_capture.py.

◆ MAX_PROBE_COUNT

int aquin.compute.activation_capture.MAX_PROBE_COUNT = 64

Definition at line 24 of file activation_capture.py.

◆ ModelMode

aquin.compute.activation_capture.ModelMode = Literal["llm"]

Definition at line 23 of file activation_capture.py.

◆ Position

aquin.compute.activation_capture.Position = Literal["last", "mean"]

Definition at line 21 of file activation_capture.py.

◆ PROBE_TEXT_KEYS

tuple aquin.compute.activation_capture.PROBE_TEXT_KEYS = ("instruction", "prompt", "text", "content", "response")

Definition at line 18 of file activation_capture.py.

◆ RESERVED_PROBE_KEYS

aquin.compute.activation_capture.RESERVED_PROBE_KEYS = frozenset({"id", *PROBE_TEXT_KEYS})

Definition at line 19 of file activation_capture.py.

Functions

Variables

Function Documentation

◆ _balance_key()

◆ _forward_llm_layers()

◆ _forward_llm_layers_token()

◆ _pool_activation()

◆ _run_capture_llm()

◆ _write_capture_artifacts()

◆ balance_probes()

◆ generate_llm_probes()

◆ llm_layer_count()

◆ load_probes()

◆ parse_layers()

◆ resolve_capture_model_id()

◆ resolve_probes_for_capture()

◆ resolve_prompts_path()

◆ run_capture_activations()

◆ write_probes_jsonl()

Variable Documentation

◆ BALANCE_PRIORITY

◆ Granularity

◆ MAX_PROBE_COUNT

◆ ModelMode

◆ Position

◆ PROBE_TEXT_KEYS

◆ RESERVED_PROBE_KEYS