tests/test_data.py

1import numpy as np2import pandas as pd3import pytest45from data import load_corpus, split_ids, make_windows, iterate_batches67IDS = np.array([3, 1, 4, 1, 5, 2, 6, 5, 3, 5])   # N = 108910def test_window_table_matches_hand_example():11    # T = 4, stride 1 -> W = N - T = 6 windows12    x, y = make_windows(IDS, seq_len=4)13    x, y = x.to_numpy(), y.to_numpy()14    assert x.shape == (6, 4) and y.shape == (6, 4)15    assert x[0].tolist() == [3, 1, 4, 1] and y[0].tolist() == [1, 4, 1, 5]16    assert x[5].tolist() == [2, 6, 5, 3] and y[5].tolist() == [6, 5, 3, 5]171819def test_targets_are_inputs_shifted_by_one():20    x, y = make_windows(IDS, seq_len=4)21    assert (y.to_numpy()[:, :-1] == x.to_numpy()[:, 1:]).all()222324def test_stride_skips_windows():25    x, _ = make_windows(IDS, seq_len=4, stride=2)26    assert x.to_numpy()[:, 0].tolist() == [3, 4, 5]       # starts 0, 2, 4272829def test_too_few_tokens_raises():30    with pytest.raises(ValueError):31        make_windows(np.arange(4), seq_len=4)323334def test_split_is_contiguous_and_disjoint():35    ids = np.arange(100)36    train, val = split_ids(ids, 0.1)37    assert len(train) == 90 and len(val) == 1038    assert train[-1] + 1 == val[0]                        # cut, no overlap394041def test_window_counts_for_n_3000():42    # guide: N = 3000, T = 16, rho = 0.1 -> 2684 train and 284 val windows43    train, val = split_ids(np.arange(3000), 0.1)44    assert len(make_windows(train, 16)[0]) == 268445    assert len(make_windows(val, 16)[0]) == 284464748def test_batches_cover_every_window_once():49    x, y = make_windows(np.arange(50), seq_len=4)         # 46 windows50    rng = np.random.default_rng(0)51    seen, sizes = [], []52    for xb, yb in iterate_batches(x, y, batch_size=16, rng=rng):53        assert xb.shape == yb.shape and xb.shape[1] == 454        assert (yb[:, :-1] == xb[:, 1:]).all()            # pairs stay aligned55        seen.extend(xb[:, 0].tolist())56        sizes.append(len(xb))57    assert sorted(seen) == list(range(46))58    assert sizes == [16, 16, 14]596061def test_no_shuffle_keeps_order():62    x, y = make_windows(np.arange(20), seq_len=4)63    first = next(iterate_batches(x, y, 5, np.random.default_rng(0), shuffle=False))[0]64    assert first[:, 0].tolist() == [0, 1, 2, 3, 4]656667def test_load_corpus_joins_rows_and_skips_missing(tmp_path):68    f = tmp_path / "c.csv"69    pd.DataFrame({"text": [" hello ", None, "world"]}).to_csv(f, index=False)70    assert load_corpus(str(f), "text", 10) == "hello\nworld\n"717273def test_load_corpus_respects_max_rows(tmp_path):74    f = tmp_path / "c.csv"75    pd.DataFrame({"text": ["a", "b", "c"]}).to_csv(f, index=False)76    assert load_corpus(str(f), "text", 2) == "a\nb\n"777879def test_load_corpus_wrong_column_raises(tmp_path):80    f = tmp_path / "c.csv"81    pd.DataFrame({"body": ["a"]}).to_csv(f, index=False)82    with pytest.raises(ValueError):83        load_corpus(str(f), "text", 10)