aquin.compute.activation_store Namespace Reference

Functions
list[Path]	find_chunk_paths (Path acts_dir)
dict[str, Any]|None	read_manifest (Path acts_dir)
Path	write_collect_manifest (Path acts_dir, *, str model_id, int layer, str model_mode, int d_model, str source, int n_vectors, str|None checkpoint_path=None, str|None corpus_path=None)
int	_split_save_chunks (torch.Tensor stacked, Path out_dir)
int	import_capture_to_chunks (Path capture_dir, int layer, Path out_dir)
int	import_capture_to_chunks_balanced (Path capture_dir, int layer, Path out_dir, *, str|None group=None)
Path	resolve_training_acts_dir (str|Path path, int layer, *, bool balance=False, str|None balance_group=None)
int	count_vectors (Path acts_dir)
tuple[torch.Tensor, torch.Tensor]	compute_norm (Path acts_dir, int d_model)
None	validate_acts_manifest (Path acts_dir, *, str model_id, int layer)

Variables
int	CHUNK_SIZE = 50_000
str	MANIFEST_NAME = "manifest.json"

Function Documentation

_split_save_chunks()

int _split_save_chunks ( torch.Tensor stacked, Path out_dir )

protected

Definition at line 71 of file activation_store.py.

Referenced by import_capture_to_chunks(), and import_capture_to_chunks_balanced().

compute_norm()

tuple[torch.Tensor, torch.Tensor] compute_norm	(	Path	acts_dir,
		int	d_model )

Definition at line 227 of file activation_store.py.

References find_chunk_paths().

count_vectors()

int count_vectors

(

Path

acts_dir

)

Definition at line 219 of file activation_store.py.

References find_chunk_paths().

find_chunk_paths()

list[Path] find_chunk_paths

(

Path

acts_dir

)

Definition at line 21 of file activation_store.py.

Referenced by compute_norm(), count_vectors(), resolve_training_acts_dir(), and write_collect_manifest().

import_capture_to_chunks()

int import_capture_to_chunks	(	Path	capture_dir,
		int	layer,
		Path	out_dir )

Definition at line 81 of file activation_store.py.

References _split_save_chunks(), read_manifest(), and write_collect_manifest().

Referenced by import_capture_to_chunks_balanced(), and resolve_training_acts_dir().

import_capture_to_chunks_balanced()

int import_capture_to_chunks_balanced	(	Path	capture_dir,
		int	layer,
		Path	out_dir,
		*	,
		str | None	group = None )

Definition at line 114 of file activation_store.py.

References _split_save_chunks(), import_capture_to_chunks(), read_manifest(), and write_collect_manifest().

Referenced by resolve_training_acts_dir().

read_manifest()

dict[str, Any] | None read_manifest

(

Path

acts_dir

)

Definition at line 25 of file activation_store.py.

Referenced by import_capture_to_chunks(), import_capture_to_chunks_balanced(), resolve_training_acts_dir(), and validate_acts_manifest().

resolve_training_acts_dir()

Path resolve_training_acts_dir	(	str | Path	path,
		int	layer,
		*	,
		bool	balance = False,
		str | None	balance_group = None )

Return a directory containing chunk_*.pt for training.

Accepts:
- a collect dir (chunk_*.pt + optional manifest.json)
- a prior train run dir (…/my-run/_acts_layer{N})
- a capture-activations output dir (manifest.json + layers/layer_{N}.pt)

Definition at line 172 of file activation_store.py.

References find_chunk_paths(), import_capture_to_chunks(), import_capture_to_chunks_balanced(), and read_manifest().

validate_acts_manifest()

None validate_acts_manifest	(	Path	acts_dir,
		*	,
		str	model_id,
		int	layer )

Definition at line 254 of file activation_store.py.

References read_manifest().

write_collect_manifest()

Path write_collect_manifest	(	Path	acts_dir,
		*	,
		str	model_id,
		int	layer,
		str	model_mode,
		int	d_model,
		str	source,
		int	n_vectors,
		str | None	checkpoint_path = None,
		str | None	corpus_path = None )

Definition at line 36 of file activation_store.py.

References find_chunk_paths().

Referenced by import_capture_to_chunks(), and import_capture_to_chunks_balanced().

Variable Documentation

CHUNK_SIZE

int aquin.compute.activation_store.CHUNK_SIZE = 50_000

Definition at line 17 of file activation_store.py.

MANIFEST_NAME

str aquin.compute.activation_store.MANIFEST_NAME = "manifest.json"

Definition at line 18 of file activation_store.py.