Collaboration diagram for aquin.compute.hf_llm_shim.HfLlmShim:

Public Member Functions
None	__init__ (self, Any hf_model, Any tokenizer, *, str hf_name, int n_layers, int d_model, int n_heads)
nn.Module	ln_final (self)
_UnembedShim	unembed (self)
HfLlmShim	from_pretrained (cls, str hf_name, dict[str, Any] cfg, *, torch.dtype dtype, str device, dict[str, Any]\|None hf_kwargs=None)
torch.Tensor	W_E (self)
torch.Tensor	W_U (self)
list[Any]	blocks (self)
HfLlmShim	eval (self)
	parameters (self, bool recurse=True)
	named_parameters (self, str prefix="", bool recurse=True)
HfLlmShim	to (self, str\|torch.device device)
torch.Tensor	to_tokens (self, str\|torch.Tensor text, bool prepend_bos=False)
str	to_string (self, Any tokens)
torch.Tensor	__call__ (self, torch.Tensor tokens)
tuple[torch.Tensor, dict[str, torch.Tensor]]	run_with_cache (self, torch.Tensor tokens, Any names_filter=None, str\|None return_type="logits")
torch.Tensor	run_with_hooks (self, torch.Tensor tokens, list[tuple[str, Callable[..., torch.Tensor]]]\|None fwd_hooks=None, str\|None return_type="logits")

Public Attributes
	hf_model = hf_model
	tokenizer = tokenizer
	cfg
	device = next(hf_model.parameters()).device

Protected Member Functions
torch.Tensor	_forward_logits (self, torch.Tensor tokens)
tuple[set[int], set[int], bool]	_hooks_from_filter (self, Any names_filter)

Protected Attributes
	_backbone
	_layers
	_embed
	_lm_head = _resolve_modules(hf_model)
tuple	_final_norm
	_unembed = _UnembedShim(self.W_U, head_bias)

Detailed Description

Enough of HookedTransformer for SAE capture, steering, and eval logits.

Definition at line 135 of file hf_llm_shim.py.

Constructor & Destructor Documentation

◆ init()

None aquin.compute.hf_llm_shim.HfLlmShim.__init__	(		self,
		Any	hf_model,
		Any	tokenizer,
		*	,
		str	hf_name,
		int	n_layers,
		int	d_model,
		int	n_heads )

Definition at line 138 of file hf_llm_shim.py.

Member Function Documentation

◆ call()

torch.Tensor aquin.compute.hf_llm_shim.HfLlmShim.__call__	(		self,
		torch.Tensor	tokens )

Definition at line 262 of file hf_llm_shim.py.

References _forward_logits().

◆ _forward_logits()

torch.Tensor aquin.compute.hf_llm_shim.HfLlmShim._forward_logits	(		self,
		torch.Tensor	tokens )

protected

Definition at line 265 of file hf_llm_shim.py.

References hf_model.

Referenced by __call__(), run_with_cache(), and run_with_hooks().

◆ _hooks_from_filter()

tuple[set[int], set[int], bool] aquin.compute.hf_llm_shim.HfLlmShim._hooks_from_filter	(		self,
		Any	names_filter )

protected

Definition at line 270 of file hf_llm_shim.py.

References aquin.compute.hf_llm_shim._names_filter_match(), and cfg.

Referenced by run_with_cache().

◆ blocks()

list[Any] aquin.compute.hf_llm_shim.HfLlmShim.blocks ( self )

Definition at line 225 of file hf_llm_shim.py.

References _layers.

◆ eval()

HfLlmShim aquin.compute.hf_llm_shim.HfLlmShim.eval ( self )

Definition at line 228 of file hf_llm_shim.py.

References eval(), and hf_model.

Referenced by eval().

◆ from_pretrained()

HfLlmShim aquin.compute.hf_llm_shim.HfLlmShim.from_pretrained	(		cls,
		str	hf_name,
		dict[str, Any]	cfg,
		*	,
		torch.dtype	dtype,
		str	device,
		dict[str, Any] \| None	hf_kwargs = None )

Definition at line 175 of file hf_llm_shim.py.

References aquin.compute.hf_llm_shim._hf_load_device_map().

◆ ln_final()

nn.Module aquin.compute.hf_llm_shim.HfLlmShim.ln_final ( self )

Definition at line 167 of file hf_llm_shim.py.

References _final_norm.

◆ named_parameters()

aquin.compute.hf_llm_shim.HfLlmShim.named_parameters	(		self,
		str	prefix = "",
		bool	recurse = True )

Definition at line 235 of file hf_llm_shim.py.

References hf_model, and named_parameters().

Referenced by named_parameters().

◆ parameters()

aquin.compute.hf_llm_shim.HfLlmShim.parameters	(		self,
		bool	recurse = True )

Definition at line 232 of file hf_llm_shim.py.

References hf_model, and parameters().

Referenced by parameters(), and to().

◆ run_with_cache()

tuple[torch.Tensor, dict[str, torch.Tensor]] aquin.compute.hf_llm_shim.HfLlmShim.run_with_cache	(		self,
		torch.Tensor	tokens,
		Any	names_filter = None,
		str \| None	return_type = "logits" )

Definition at line 284 of file hf_llm_shim.py.

References _embed, _forward_logits(), _hooks_from_filter(), _layers, and hf_model.

◆ run_with_hooks()

torch.Tensor aquin.compute.hf_llm_shim.HfLlmShim.run_with_hooks	(		self,
		torch.Tensor	tokens,
		list[tuple[str, Callable[..., torch.Tensor]]] \| None	fwd_hooks = None,
		str \| None	return_type = "logits" )

Definition at line 336 of file hf_llm_shim.py.

References _embed, _forward_logits(), and _layers.

◆ to()

HfLlmShim aquin.compute.hf_llm_shim.HfLlmShim.to	(		self,
		str \| torch.device	device )

Definition at line 238 of file hf_llm_shim.py.

References aquin.compute.hf_llm_shim._model_has_device_map(), device, hf_model, parameters(), and to().

Referenced by to().

◆ to_string()

str aquin.compute.hf_llm_shim.HfLlmShim.to_string	(		self,
		Any	tokens )

Definition at line 255 of file hf_llm_shim.py.

References tokenizer.

◆ to_tokens()

torch.Tensor aquin.compute.hf_llm_shim.HfLlmShim.to_tokens	(		self,
		str \| torch.Tensor	text,
		bool	prepend_bos = False )

Definition at line 245 of file hf_llm_shim.py.

References device, and tokenizer.

◆ unembed()

_UnembedShim aquin.compute.hf_llm_shim.HfLlmShim.unembed ( self )

Definition at line 171 of file hf_llm_shim.py.

References _unembed.

◆ W_E()

torch.Tensor aquin.compute.hf_llm_shim.HfLlmShim.W_E ( self )

Definition at line 215 of file hf_llm_shim.py.

References _embed.

◆ W_U()

torch.Tensor aquin.compute.hf_llm_shim.HfLlmShim.W_U ( self )

Definition at line 219 of file hf_llm_shim.py.

References _lm_head.

Member Data Documentation

◆ _backbone

aquin.compute.hf_llm_shim.HfLlmShim._backbone

protected

Definition at line 156 of file hf_llm_shim.py.

◆ _embed

aquin.compute.hf_llm_shim.HfLlmShim._embed

protected

Definition at line 156 of file hf_llm_shim.py.

Referenced by run_with_cache(), run_with_hooks(), and W_E().

◆ _final_norm

tuple aquin.compute.hf_llm_shim.HfLlmShim._final_norm

protected

Initial value:

=  (
            getattr(self._backbone, "norm", None)
            or getattr(self._backbone, "ln_f", None)
            or getattr(self._backbone, "final_layer_norm", None)
        )

Definition at line 158 of file hf_llm_shim.py.

Referenced by ln_final().

◆ _layers

aquin.compute.hf_llm_shim.HfLlmShim._layers

protected

Definition at line 156 of file hf_llm_shim.py.

Referenced by blocks(), run_with_cache(), and run_with_hooks().

◆ _lm_head

aquin.compute.hf_llm_shim.HfLlmShim._lm_head = _resolve_modules(hf_model)

protected

Definition at line 156 of file hf_llm_shim.py.

Referenced by W_U().

◆ _unembed

aquin.compute.hf_llm_shim.HfLlmShim._unembed = _UnembedShim(self.W_U, head_bias)

protected

Definition at line 164 of file hf_llm_shim.py.

Referenced by unembed().

◆ cfg

aquin.compute.hf_llm_shim.HfLlmShim.cfg

Initial value:

=  _ShimCfg(
            n_layers=n_layers,
            d_model=d_model,
            n_heads=n_heads,
            model_name=hf_name,
        )

Definition at line 150 of file hf_llm_shim.py.

Referenced by _hooks_from_filter().

◆ device

aquin.compute.hf_llm_shim.HfLlmShim.device = next(hf_model.parameters()).device

Definition at line 157 of file hf_llm_shim.py.

Referenced by to(), and to_tokens().

◆ hf_model

aquin.compute.hf_llm_shim.HfLlmShim.hf_model = hf_model

Definition at line 148 of file hf_llm_shim.py.

Referenced by _forward_logits(), eval(), named_parameters(), parameters(), run_with_cache(), and to().

◆ tokenizer

aquin.compute.hf_llm_shim.HfLlmShim.tokenizer = tokenizer

Definition at line 149 of file hf_llm_shim.py.

Referenced by to_string(), and to_tokens().

The documentation for this class was generated from the following file:

aquin/compute/hf_llm_shim.py

Public Member Functions

Public Attributes

Protected Member Functions

Protected Attributes

Detailed Description

Constructor & Destructor Documentation

◆ __init__()

Member Function Documentation

◆ __call__()

◆ _forward_logits()

◆ _hooks_from_filter()

◆ blocks()

◆ eval()

◆ from_pretrained()

◆ ln_final()

◆ named_parameters()

◆ parameters()

◆ run_with_cache()

◆ run_with_hooks()

◆ to()

◆ to_string()

◆ to_tokens()

◆ unembed()

◆ W_E()

◆ W_U()

Member Data Documentation

◆ _backbone

◆ _embed

◆ _final_norm

◆ _layers

◆ _lm_head

◆ _unembed

◆ cfg

◆ device

◆ hf_model

◆ tokenizer

◆ init()

◆ call()