61 lines
2.0 KiB
Python
61 lines
2.0 KiB
Python
import pytest
|
|
import torch
|
|
|
|
from microgpt.data import CharTokenizer, get_batch
|
|
from microgpt.sampling import filter_top_k, filter_top_p, sample_next
|
|
|
|
|
|
def test_temperature_zero_is_argmax():
|
|
logits = torch.tensor([[0.1, 2.0, -1.0], [3.0, 0.0, 0.0]])
|
|
assert sample_next(logits, temperature=0).flatten().tolist() == [1, 0]
|
|
|
|
|
|
def test_negative_temperature_rejected():
|
|
with pytest.raises(ValueError):
|
|
sample_next(torch.zeros(1, 3), temperature=-1.0)
|
|
|
|
|
|
def test_top_k_keeps_exactly_k():
|
|
logits = torch.tensor([[1.0, 5.0, 3.0, 4.0, 2.0]])
|
|
kept = torch.isfinite(filter_top_k(logits, 2))
|
|
assert kept.tolist() == [[False, True, False, True, False]]
|
|
|
|
|
|
def test_top_p_keeps_smallest_nucleus():
|
|
probs = torch.tensor([[0.5, 0.3, 0.15, 0.05]])
|
|
out = filter_top_p(probs.log(), 0.7)
|
|
assert torch.isfinite(out).tolist() == [[True, True, False, False]]
|
|
|
|
|
|
def test_top_p_always_keeps_the_top_token():
|
|
out = filter_top_p(torch.tensor([[10.0, 0.0, 0.0]]), 0.01)
|
|
assert torch.isfinite(out).tolist() == [[True, False, False]]
|
|
|
|
|
|
def test_samples_respect_top_k():
|
|
g = torch.Generator().manual_seed(0)
|
|
logits = torch.tensor([[0.0, 1.0, 2.0, 3.0, 4.0]]).expand(500, -1)
|
|
draws = sample_next(logits, top_k=2, generator=g)
|
|
assert set(draws.flatten().tolist()) <= {3, 4}
|
|
|
|
|
|
def test_sampling_matches_softmax_distribution():
|
|
g = torch.Generator().manual_seed(0)
|
|
logits = torch.log(torch.tensor([[0.7, 0.2, 0.1]])).expand(20000, -1)
|
|
freq = torch.bincount(sample_next(logits, generator=g).flatten(), minlength=3) / 20000
|
|
assert torch.allclose(freq, torch.tensor([0.7, 0.2, 0.1]), atol=0.02)
|
|
|
|
|
|
def test_tokenizer_roundtrip():
|
|
text = "hello, world\n"
|
|
tok = CharTokenizer.from_text(text)
|
|
assert tok.decode(tok.encode(text)) == text
|
|
assert tok.vocab_size == len(set(text))
|
|
|
|
|
|
def test_get_batch_targets_are_shifted_inputs():
|
|
data = torch.arange(100)
|
|
x, y = get_batch(data, batch_size=4, block_size=8)
|
|
assert x.shape == y.shape == (4, 8)
|
|
assert torch.equal(y, x + 1)
|