diff --git a/transformer_lens/ActivationCache.py b/transformer_lens/ActivationCache.py index d40fcc109..ed16cf322 100644 --- a/transformer_lens/ActivationCache.py +++ b/transformer_lens/ActivationCache.py @@ -60,7 +60,7 @@ class ActivationCache: The :class:`ActivationCache` is at the core of Transformer Lens. It is a wrapper that stores all important activations from a forward pass of the model, and provides a variety of helper functions to investigate them. The common way to access it is to run the model with - :meth:`transformer_lens.HookedTransformer.HookedTransformer.run_with_cache`. + :meth:`transformer_lens.model_bridge.TransformerBridge.run_with_cache`. Examples: @@ -71,9 +71,9 @@ class ActivationCache: the model predicting "road". This kind of analysis commonly falls under the category of "logit attribution" or "direct logit attribution" (DLA). - >>> from transformer_lens import HookedTransformer - >>> model = HookedTransformer.from_pretrained("tiny-stories-1M") - Loaded pretrained model tiny-stories-1M into HookedTransformer + >>> from transformer_lens.model_bridge import TransformerBridge + >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M") + >>> model.enable_compatibility_mode() >>> _logits, cache = model.run_with_cache("Why did the chicken cross the") >>> residual_stream, labels = cache.decompose_resid(return_labels=True, mode="attn") @@ -101,9 +101,10 @@ class ActivationCache: Warning: :class:`ActivationCache` is designed to be used with - :class:`transformer_lens.HookedTransformer`, and will not work with other models. It's also - designed to be used with all activations of :class:`transformer_lens.HookedTransformer` being - cached, and some internal methods will break without that. + :class:`transformer_lens.HookedTransformer` or + :class:`transformer_lens.model_bridge.TransformerBridge`. Advanced helpers expect the model to + expose the TransformerLens weight-processing interface and generally expect a complete cache; + some internal methods may break with other models or partial caches. The biggest footgun and source of bugs in this code will be keeping track of indexes, dimensions, and the numbers of each. There are several kinds of activations: @@ -270,12 +271,14 @@ def keys(self): Examples: - >>> from transformer_lens import HookedTransformer - >>> model = HookedTransformer.from_pretrained("tiny-stories-1M") - Loaded pretrained model tiny-stories-1M into HookedTransformer + >>> from transformer_lens.model_bridge import TransformerBridge + >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M") + >>> model.enable_compatibility_mode() >>> _logits, cache = model.run_with_cache("Some prompt") - >>> list(cache.keys())[0:3] - ['hook_embed', 'hook_pos_embed', 'blocks.0.hook_resid_pre'] + >>> list(cache.keys())[0:8] + ['embed.hook_in', 'hook_embed', 'embed.hook_out', 'pos_embed.hook_in', + 'hook_pos_embed', 'pos_embed.hook_out', 'blocks.0.hook_in', + 'blocks.0.hook_resid_pre'] Returns: List of all keys. @@ -306,16 +309,18 @@ def __iter__(self) -> Iterator[str]: Examples: - >>> from transformer_lens import HookedTransformer - >>> model = HookedTransformer.from_pretrained("tiny-stories-1M") - Loaded pretrained model tiny-stories-1M into HookedTransformer + >>> from transformer_lens.model_bridge import TransformerBridge + >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M") + >>> model.enable_compatibility_mode() >>> _logits, cache = model.run_with_cache("Some prompt") >>> cache_interesting_names = [] >>> for key in cache: ... if not key.startswith("blocks.") or key.startswith("blocks.0"): ... cache_interesting_names.append(key) - >>> print(cache_interesting_names[0:3]) - ['hook_embed', 'hook_pos_embed', 'blocks.0.hook_resid_pre'] + >>> print(cache_interesting_names[0:8]) + ['embed.hook_in', 'hook_embed', 'embed.hook_out', 'pos_embed.hook_in', + 'hook_pos_embed', 'pos_embed.hook_out', 'blocks.0.hook_in', + 'blocks.0.hook_resid_pre'] Returns: Iterator over the cache. @@ -392,12 +397,12 @@ def accumulated_resid( Logit Lens analysis can be done as follows: - >>> from transformer_lens import HookedTransformer + >>> from transformer_lens.model_bridge import TransformerBridge >>> import torch >>> import pandas as pd - >>> model = HookedTransformer.from_pretrained("tiny-stories-1M", device="cpu", fold_ln=True) - Loaded pretrained model tiny-stories-1M into HookedTransformer + >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M", device="cpu") + >>> model.enable_compatibility_mode() >>> prompt = "Why did the chicken cross the" >>> answer = " road" diff --git a/transformer_lens/evals.py b/transformer_lens/evals.py index 8423f301e..ee093a121 100644 --- a/transformer_lens/evals.py +++ b/transformer_lens/evals.py @@ -323,10 +323,10 @@ class IOIDataset(Dataset): .. code-block:: python >>> from transformer_lens.evals import ioi_eval, IOIDataset - >>> from transformer_lens.HookedTransformer import HookedTransformer + >>> from transformer_lens.model_bridge import TransformerBridge - >>> model = HookedTransformer.from_pretrained('gpt2-small') - Loaded pretrained model gpt2-small into HookedTransformer + >>> model = TransformerBridge.boot_transformers("gpt2", device="cpu") + >>> model.enable_compatibility_mode() >>> # Evaluate on a deterministic dataset (seed makes results reproducible) >>> ds = IOIDataset(tokenizer=model.tokenizer, num_samples=100, seed=42) @@ -551,10 +551,11 @@ def mmlu_eval( .. code-block:: python - >>> from transformer_lens import HookedTransformer + >>> from transformer_lens.model_bridge import TransformerBridge >>> from transformer_lens.evals import mmlu_eval - >>> model = HookedTransformer.from_pretrained("gpt2-small") # doctest: +SKIP + >>> model = TransformerBridge.boot_transformers("gpt2") # doctest: +SKIP + >>> model.enable_compatibility_mode() # doctest: +SKIP >>> results = mmlu_eval(model, subjects="abstract_algebra", num_samples=10) # doctest: +SKIP >>> print(f"Accuracy: {results['accuracy']:.2%}") # doctest: +SKIP """ diff --git a/transformer_lens/utilities/exploratory_utils.py b/transformer_lens/utilities/exploratory_utils.py index 426aee6a0..3e69d53b9 100644 --- a/transformer_lens/utilities/exploratory_utils.py +++ b/transformer_lens/utilities/exploratory_utils.py @@ -31,9 +31,10 @@ def test_prompt( Examples: - >>> from transformer_lens import HookedTransformer, utilities - >>> model = HookedTransformer.from_pretrained("tiny-stories-1M") - Loaded pretrained model tiny-stories-1M into HookedTransformer + >>> from transformer_lens import utilities + >>> from transformer_lens.model_bridge import TransformerBridge + >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M") + >>> model.enable_compatibility_mode() >>> prompt = "Why did the elephant cross the" >>> answer = "road" @@ -156,8 +157,8 @@ def test_prompt( try: import pytest - # Note: Docstring won't be tested with PyTest (it's ignored), as it thinks this is a regular unit - # test (because its name is prefixed `test_`). + # Note: PyTest collects and runs this docstring as a doctest. The skip marker only applies to the + # accidentally collected function item (because its name is prefixed `test_`). pytest.mark.skip(test_prompt) except ModuleNotFoundError: pass # disregard if pytest not in env