Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion tests/acceptance/test_evals.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@

@pytest.fixture(scope="module")
def model():
return HookedTransformer.from_pretrained("gpt2-small")
return HookedTransformer.from_pretrained("gpt2-small", device="cpu")


def test_basic_ioi_eval(model):
Expand Down
18 changes: 17 additions & 1 deletion tests/acceptance/test_hooked_transformer.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,19 @@

PYTHIA_MODEL_NAMES = [name for name in OFFICIAL_MODEL_NAMES if name.startswith("EleutherAI/pythia")]

model_names = [
# Small models for basic testing
PUBLIC_MODEL_NAMES = [
"attn-only-demo",
"gpt2-small",
"opt-125m",
"pythia-70m",
"tiny-stories-33M",
"microsoft/phi-1",
"google/gemma-2b",
]

# Full set of models to test
FULL_MODEL_NAMES = [
"attn-only-demo",
"gpt2-small",
"opt-125m",
Expand All @@ -42,6 +54,10 @@
"google/gemma-2b",
"google/gemma-7b",
]

# Use full model list if HF_TOKEN is available, otherwise use public models only
model_names = FULL_MODEL_NAMES if os.environ.get("HF_TOKEN", "") else PUBLIC_MODEL_NAMES

text = "Hello world!"
"""
# Code to regenerate loss store
Expand Down
14 changes: 12 additions & 2 deletions tests/acceptance/test_tokenizer_special_tokens.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,15 @@
import os

from transformers import AutoTokenizer

import transformer_lens.loading_from_pretrained as loading
from transformer_lens import HookedTransformer, HookedTransformerConfig

# Get's tedious typing these out everytime I want to sweep over all the distinct small models
MODEL_TESTING_LIST = [
# Small models for basic testing
PUBLIC_MODEL_TESTING_LIST = ["gpt2-small", "opt-125m", "pythia-70m"]

# Full set of models to test when HF_TOKEN is available
FULL_MODEL_TESTING_LIST = [
"solu-1l",
"gpt2-small",
"gpt-neo-125M",
Expand All @@ -14,6 +19,11 @@
"pythia-70m",
]

# Use full model list if HF_TOKEN is available, otherwise use public models only
MODEL_TESTING_LIST = (
FULL_MODEL_TESTING_LIST if os.environ.get("HF_TOKEN", "") else PUBLIC_MODEL_TESTING_LIST
)


def test_d_vocab_from_tokenizer():
cfg = HookedTransformerConfig(
Expand Down