ViperEkura
5ba21f4eb3
- Add tests/helpers.py with shared config, dataset, tokenizer, executor, and assertion helpers - Replace 15 copies of device one-liner with session-scoped fixture - Collapse 5 near-identical Dataset subclasses into RandomTokenDataset - Remove duplicate _make_config/_make_model/_make_frozen and FakeTokenizer/FakeExecutor definitions - Make test_callbacks and test_early_stopping use existing train_config_factory - Replace 6 duplicate meta.json read blocks with load_shard_meta - Fix mkdtemp leaks in test_lora.py with TemporaryDirectory
88 lines
2.6 KiB
Python
88 lines
2.6 KiB
Python
import json
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
|
|
import pytest
|
|
import torch
|
|
from tokenizers import Tokenizer, models, pre_tokenizers, trainers
|
|
|
|
from astrai.model.transformer import AutoRegressiveLM
|
|
from astrai.tokenize import AutoTokenizer
|
|
from tests.helpers import TINY_CONFIG, RandomTokenDataset, make_tiny_config
|
|
|
|
|
|
def pytest_configure(config):
|
|
config.addinivalue_line("markers", "slow: marks tests as slow")
|
|
config.addinivalue_line("markers", "integration: integration tests")
|
|
config.addinivalue_line("markers", "unit: fast unit tests")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def device():
|
|
"""Session-scoped device string (``"cuda"`` if available, else ``"cpu"``)."""
|
|
return "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
|
|
def create_test_tokenizer(vocab_size: int = 1000) -> AutoTokenizer:
|
|
"""Create a simple tokenizer for testing purposes."""
|
|
tokenizer = Tokenizer(models.BPE())
|
|
tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()
|
|
trainer = trainers.BpeTrainer(
|
|
vocab_size=vocab_size, min_frequency=1, special_tokens=["<unk>", "<pad>"]
|
|
)
|
|
tokenizer.train_from_iterator([chr(i) for i in range(256)], trainer)
|
|
auto_tokenizer = AutoTokenizer()
|
|
auto_tokenizer._tokenizer = tokenizer
|
|
auto_tokenizer._special_token_map = {"unk_token": "<unk>", "pad_token": "<pad>"}
|
|
return auto_tokenizer
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def test_tokenizer():
|
|
"""Session-scoped tokenizer, created once for the entire test run."""
|
|
return create_test_tokenizer()
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def test_model(device):
|
|
"""Session-scoped small AutoRegressiveLM model, created once."""
|
|
config = make_tiny_config()
|
|
model = AutoRegressiveLM(config).to(device=device)
|
|
return {"model": model, "device": device, "config": config}
|
|
|
|
|
|
@pytest.fixture
|
|
def base_test_env(test_model, test_tokenizer):
|
|
"""Function-scoped test environment with isolated temp directory."""
|
|
test_dir = tempfile.mkdtemp()
|
|
config_path = os.path.join(test_dir, "config.json")
|
|
with open(config_path, "w") as f:
|
|
json.dump(TINY_CONFIG, f)
|
|
|
|
yield {
|
|
"device": test_model["device"],
|
|
"test_dir": str(test_dir),
|
|
"config_path": config_path,
|
|
"transformer_config": test_model["config"],
|
|
"model": test_model["model"],
|
|
"tokenizer": test_tokenizer,
|
|
}
|
|
|
|
shutil.rmtree(test_dir)
|
|
|
|
|
|
@pytest.fixture
|
|
def random_dataset():
|
|
return RandomTokenDataset(length=None)
|
|
|
|
|
|
@pytest.fixture
|
|
def multi_turn_dataset():
|
|
return RandomTokenDataset(length=None, with_loss_mask=True)
|
|
|
|
|
|
@pytest.fixture
|
|
def early_stopping_dataset():
|
|
return RandomTokenDataset(length=10, stop_after=5)
|