search

ndif-team / nnsight

The nnsight package enables interpreting and manipulating the internals of deep learned models.
https://nnsight.net/
MIT License
360 stars 34 forks source link

Unified #65

Closed cadentj closed 7 months ago

cadentj commented 7 months ago

Test snippet.

from nnsight.models.UnifiedTransformer import UnifiedTransformer

device = "cuda:0"

# Pass in a model name from the TransformerLens library to load a HookedTransformer. 
unified_model = UnifiedTransformer("gpt2", device=device)

# kwargs are passed to from_pretrained to process the TransformerLens model.
unified_model = UnifiedTransformer("gpt2", fold_ln=True, device=device)

# Pass process=False to skip default TransformerLens processing.
unified_model = UnifiedTransformer("gpt2", processing=False, device=device)

with unified_model.invoke("Hello, my name is") as invoker:
    pass

# You can also make use of TransformerLens methods by calling `.local_model`
unified_model.local_model.to_str_tokens("Hello, my name is")