From d2637fb01d5001968324c16a9632cd0fe913c2af Mon Sep 17 00:00:00 2001
From: Wing Lian <wing@axolotl.ai>
Date: Wed, 16 Apr 2025 21:14:04 -0700
Subject: [PATCH] first pass at modifying tests to use llama-7m

---
 tests/conftest.py        | 6 ++++++
 tests/test_perplexity.py | 2 +-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 97c48db41..62d3ee1ba 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -496,6 +496,12 @@ def dataset_fozziethebeat_alpaca_messages_2k_dpo_test_rev_ea82cff(
     return datasets.load_from_disk(ds_path)["train"]
 
 
+@pytest.fixture(scope="session", autouse=True)
+def download_tiny_llama_7m_model():
+    # download the model
+    return snapshot_download_w_retry("axolotl-ai-internal/llama-7m", repo_type="model")
+
+
 # # pylint: disable=redefined-outer-name,unused-argument
 # def test_load_fixtures(
 #     download_smollm2_135m_model,
diff --git a/tests/test_perplexity.py b/tests/test_perplexity.py
index 9a1c9b223..51e3dea3a 100644
--- a/tests/test_perplexity.py
+++ b/tests/test_perplexity.py
@@ -8,7 +8,7 @@ from transformers.models.auto.tokenization_auto import AutoTokenizer
 
 from axolotl.utils.callbacks.perplexity import Perplexity
 
-MODEL_NAME = "HuggingFaceTB/SmolLM2-135M"
+MODEL_NAME = "axolotl-ai-internal/llama-7m"
 
 
 @fixture()