diff --git a/tests/acceptance/test_evals.py b/tests/acceptance/test_evals.py index e9b8b497c..5936466f6 100644 --- a/tests/acceptance/test_evals.py +++ b/tests/acceptance/test_evals.py @@ -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): diff --git a/tests/acceptance/test_hooked_transformer.py b/tests/acceptance/test_hooked_transformer.py index 3fd739f94..c54846787 100644 --- a/tests/acceptance/test_hooked_transformer.py +++ b/tests/acceptance/test_hooked_transformer.py @@ -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", @@ -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 diff --git a/tests/acceptance/test_tokenizer_special_tokens.py b/tests/acceptance/test_tokenizer_special_tokens.py index 6e4a93a98..840114122 100644 --- a/tests/acceptance/test_tokenizer_special_tokens.py +++ b/tests/acceptance/test_tokenizer_special_tokens.py @@ -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", @@ -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(