added minor tests

tests/conftest.py

@@ -0,0 +1,28 @@
+"""Pytest fixtures.
+"""
+from __future__ import annotations
+
+import pytest
+import numpy as np
+
+
+TEST_CAUSAL_LMS = [
+    "hf-internal-testing/tiny-random-gpt2",
+    # "hf-internal-testing/tiny-random-MistralForCausalLM",
+]
+
+
+@pytest.fixture(params=[42])
+def random_seed(request) -> int:
+    return request.param
+
+
+@pytest.fixture
+def rng(random_seed: int) -> np.random.Generator:
+    return np.random.default_rng(random_seed)
+
+
+@pytest.fixture(params=TEST_CAUSAL_LMS)
+def causal_lm_name_or_path(request) -> str:
+    """Name or path of the CausalLM used for testing."""
+    return request.param

tests/test_load_huggingface_model.py

@@ -0,0 +1,15 @@
+"""Test functions in folktexts.llm_utils
+"""
+
+from transformers import PreTrainedModel, PreTrainedTokenizerBase
+from folktexts.llm_utils import load_model_tokenizer
+
+
+def test_load_model_tokenizer(causal_lm_name_or_path: str):
+    model, tokenizer = load_model_tokenizer(causal_lm_name_or_path)
+
+    assert isinstance(model, PreTrainedModel), \
+        f"Expected model type `PreTrainedModel`, got {type(model)}."
+
+    assert isinstance(tokenizer, PreTrainedTokenizerBase), \
+        f"Expected tokenizer type `PreTrainedTokenizer`, got {type(tokenizer)}."