aquin.compute.causal_trace Namespace Reference

Functions
torch.Tensor	_project_unembed (model, torch.Tensor resid)
str	_format_prompt (HookedTransformer model, str prompt)
	_get_embeds (tokens, str model_id="llama-3.2-1b")
	_corrupt_all (tokens, noise_scale, str model_id="llama-3.2-1b")
	_corrupt_single_pos (tokens, pos, noise_scale, str model_id="llama-3.2-1b")
	run_trace (str prompt, str target_token, str model_id="llama-3.2-1b", float noise_scale=3.0, int n_noise_runs=10)
	run_prompt_attribution (prompt, response, prompt_tokens, response_tokens, sig_prompt_tis, sig_response_tis, str model_id="llama-3.2-1b", float noise_scale=3.0, int n_noise_runs=5)
	run_logit_lens (str prompt, str model_id="llama-3.2-1b", int top_k=5)
	_resolve_chat_model (str model_id, Any|None model=None)
torch.device	_generation_device (Any model)
str	_decode_chat_completion (Any model, torch.Tensor input_ids, torch.Tensor output_ids)
str	run_chat (str prompt, str model_id="llama-3.2-1b", int max_new_tokens=200, float temperature=0.7, Any|None model=None)
	stream_chat (str prompt, str model_id="llama-3.2-1b", int max_new_tokens=200, float temperature=0.7)

Variables
	DEVICE = resolve_compute_device()
	DTYPE = default_dtype_for_device(DEVICE)

Function Documentation

_corrupt_all()

_corrupt_all ( tokens, noise_scale, str model_id = "llama-3.2-1b" )

protected

Definition at line 57 of file causal_trace.py.

References _get_embeds().

Referenced by run_trace().

_corrupt_single_pos()

_corrupt_single_pos ( tokens, pos, noise_scale, str model_id = "llama-3.2-1b" )

protected

Definition at line 62 of file causal_trace.py.

References _get_embeds().

Referenced by run_prompt_attribution().

_decode_chat_completion()

str _decode_chat_completion ( Any model, torch.Tensor input_ids, torch.Tensor output_ids )

protected

Definition at line 267 of file causal_trace.py.

Referenced by run_chat().

_format_prompt()

str _format_prompt ( HookedTransformer model, str prompt )

protected

Definition at line 42 of file causal_trace.py.

Referenced by run_chat(), and stream_chat().

_generation_device()

torch.device _generation_device ( Any model )

protected

Definition at line 256 of file causal_trace.py.

Referenced by run_chat().

_get_embeds()

_get_embeds ( tokens, str model_id = "llama-3.2-1b" )

protected

Definition at line 51 of file causal_trace.py.

Referenced by _corrupt_all(), and _corrupt_single_pos().

_project_unembed()

torch.Tensor _project_unembed ( model, torch.Tensor resid )

protected

Map residual vector to vocab logits (handles tied vs untied lm_head layouts).

Definition at line 35 of file causal_trace.py.

Referenced by run_logit_lens().

_resolve_chat_model()

_resolve_chat_model ( str model_id, Any | None model = None )

protected

Definition at line 246 of file causal_trace.py.

Referenced by run_chat().

run_chat()

str run_chat	(	str	prompt,
		str	model_id = "llama-3.2-1b",
		int	max_new_tokens = 200,
		float	temperature = 0.7,
		Any | None	model = None )

Definition at line 273 of file causal_trace.py.

References _decode_chat_completion(), _format_prompt(), _generation_device(), and _resolve_chat_model().

run_logit_lens()

run_logit_lens	(	str	prompt,
		str	model_id = "llama-3.2-1b",
		int	top_k = 5 )

Definition at line 217 of file causal_trace.py.

References _project_unembed().

run_prompt_attribution()

run_prompt_attribution	(		prompt,
			response,
			prompt_tokens,
			response_tokens,
			sig_prompt_tis,
			sig_response_tis,
		str	model_id = "llama-3.2-1b",
		float	noise_scale = 3.0,
		int	n_noise_runs = 5 )

Definition at line 131 of file causal_trace.py.

References _corrupt_single_pos().

run_trace()

run_trace	(	str	prompt,
		str	target_token,
		str	model_id = "llama-3.2-1b",
		float	noise_scale = 3.0,
		int	n_noise_runs = 10 )

Definition at line 69 of file causal_trace.py.

References _corrupt_all().

stream_chat()

stream_chat	(	str	prompt,
		str	model_id = "llama-3.2-1b",
		int	max_new_tokens = 200,
		float	temperature = 0.7 )

Definition at line 326 of file causal_trace.py.

References _format_prompt().

Variable Documentation

DEVICE

aquin.compute.causal_trace.DEVICE = resolve_compute_device()

Definition at line 31 of file causal_trace.py.

DTYPE

aquin.compute.causal_trace.DTYPE = default_dtype_for_device(DEVICE)

Definition at line 32 of file causal_trace.py.