File size: 1,482 Bytes
a311e94 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | """Shared fixtures + helpers for the Agron zoo test-suite."""
from __future__ import annotations
import pytest
import torch
from agron_nets.architectures import build_architecture
TINY = dict(
vocab_size=64,
hidden_size=32,
intermediate_size=64,
num_hidden_layers=2,
num_attention_heads=4,
num_key_value_heads=2,
max_position_embeddings=32,
)
def make_model(arch: str = "transformer", **overrides):
cfg = dict(TINY)
cfg.update(overrides)
model, config = build_architecture(arch, **cfg)
return model, config
def make_batch(batch_size: int = 2, seq_len: int = 8, vocab_size: int = 64, device="cpu"):
input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=device)
return {
"input_ids": input_ids,
"labels": input_ids.clone(),
"attention_mask": torch.ones_like(input_ids),
}
class ListLoader:
"""Minimal iterable dataloader yielding pre-made batches."""
def __init__(self, batches):
self.batches = list(batches)
def __iter__(self):
return iter(self.batches)
def __len__(self):
return len(self.batches)
def make_loader(steps: int, batch_size: int = 2, seq_len: int = 8, vocab_size: int = 64):
return ListLoader([make_batch(batch_size, seq_len, vocab_size) for _ in range(steps)])
@pytest.fixture
def tiny_model():
model, _ = make_model("transformer")
return model
@pytest.fixture
def loader():
return make_loader(8)
|