ViperEkura
115192c67c
- Remove H5Store, H5Writer, save_h5/load_h5 and h5py dependency - MmapStore (bin) is the sole pre-tokenized storage backend - Move setup_logging after imports to fix E402 in __init__.py - Clean up unused imports across test files - Move inline test imports to file top
61 lines
1.5 KiB
Python
61 lines
1.5 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from astrai.config import TrainConfig
|
|
from astrai.trainer.schedule import SchedulerFactory
|
|
from tests.helpers import RandomTokenDataset
|
|
|
|
|
|
def create_train_config(
|
|
model_fn,
|
|
dataset,
|
|
test_dir: str,
|
|
device: str,
|
|
strategy: str = "seq",
|
|
n_epoch: int = 1,
|
|
batch_per_device: int = 2,
|
|
grad_accum_steps: int = 1,
|
|
max_grad_norm: float = 1.0,
|
|
ckpt_interval: int = 5,
|
|
random_seed: int = 42,
|
|
**kwargs,
|
|
):
|
|
"""Factory function to create common TrainConfig for tests."""
|
|
|
|
def optimizer_fn(m):
|
|
return torch.optim.AdamW(m.parameters(), lr=0.001)
|
|
|
|
def scheduler_fn(optim):
|
|
return SchedulerFactory.create(
|
|
"cosine", optim, warmup_steps=10, lr_decay_steps=10, min_rate=0.05
|
|
)
|
|
|
|
return TrainConfig(
|
|
strategy=strategy,
|
|
model_fn=model_fn,
|
|
dataset=dataset,
|
|
optimizer_fn=optimizer_fn,
|
|
scheduler_fn=scheduler_fn,
|
|
ckpt_dir=test_dir,
|
|
n_epoch=n_epoch,
|
|
batch_per_device=batch_per_device,
|
|
ckpt_interval=ckpt_interval,
|
|
grad_accum_steps=grad_accum_steps,
|
|
max_grad_norm=max_grad_norm,
|
|
random_seed=random_seed,
|
|
device_type=device,
|
|
**kwargs,
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def train_config_factory():
|
|
"""Fixture providing the ``create_train_config`` factory function."""
|
|
return create_train_config
|
|
|
|
|
|
@pytest.fixture
|
|
def trainer_dataset():
|
|
"""Fixture providing a dataset for trainer tests."""
|
|
return RandomTokenDataset()
|