Repository Public Research preview
1import pytest23from wp_ai_0001.config import FSATaskConfig, MorphTaskConfig, OpSeqTaskConfig4from wp_ai_0001.datasets import fsa, morph, opseq5from wp_ai_0001.datasets.common import assert_disjoint, encode_split678def test_opseq_validator_and_determinism():9 cfg = OpSeqTaskConfig()10 a = opseq.generate(cfg, 500, seed=1, name="t")11 b = opseq.generate(cfg, 500, seed=1, name="t")12 opseq.validate(a, cfg)13 assert a.hash == b.hash and len(a) == 50014 c = opseq.generate(cfg, 500, seed=2, name="t")15 assert c.hash != a.hash161718def test_opseq_train_excludes_heldout_pairs_and_operators():19 cfg = OpSeqTaskConfig(excluded_operators=["qua"])20 s = opseq.generate(cfg, 1000, seed=3, name="train", mode="train")21 pairs = set(map(tuple, cfg.held_out_pairs))22 for e in s.examples:23 ops = list(e.tokens[2:-2])24 assert "qua" not in ops25 assert not any((x, y) in pairs for x, y in zip(ops, ops[1:]))262728def test_opseq_comp_and_require():29 cfg = OpSeqTaskConfig()30 s = opseq.generate(cfg, 300, seed=4, name="comp", mode="comp", min_len=2)31 pairs = set(map(tuple, cfg.held_out_pairs))32 assert all(any((x, y) in pairs for x, y in zip(e.tokens[2:-2], e.tokens[3:-2])) for e in s.examples)33 t = opseq.generate(cfg, 300, seed=5, name="target", mode="any", require_operator="qua")34 assert all("qua" in e.tokens for e in t.examples)35 u = opseq.generate(cfg, 300, seed=6, name="control", mode="any", forbid_operator="qua")36 assert all("qua" not in e.tokens for e in u.examples)373839def test_disjointness_enforced():40 cfg = OpSeqTaskConfig()41 a = opseq.generate(cfg, 2000, seed=7, name="a")42 b = opseq.generate(cfg, 2000, seed=8, name="b", exclude_hashes=a.item_hashes())43 assert_disjoint(a, b)44 with pytest.raises(AssertionError):45 assert_disjoint(a, a)464748def test_encode_split_positions():49 cfg = OpSeqTaskConfig()50 s = opseq.generate(cfg, 50, seed=9, name="e")51 v = opseq.vocab_for(cfg)52 enc = encode_split(s, v, 24)53 for i, e in enumerate(s.examples):54 pl = enc["prompt_len"][i]55 assert v.itos[enc["ids"][i, pl - 1]] == "="56 assert v.decode(enc["answers"][i, : enc["n_answer"][i]]) == list(e.answers)57 assert enc["loss_mask"][i].sum() == e.n_answer and enc["loss_mask"][i, pl - 1] == 158 assert v.itos[enc["ids"][i, pl + e.n_answer - 1]] == e.answer596061def test_fsa_and_morph():62 f = FSATaskConfig()63 s = fsa.generate(f, 300, seed=1, name="f")64 fsa.validate(s, f)65 m = MorphTaskConfig()66 t = morph.generate(m, 300, seed=1, name="m")67 morph.validate(t, m)68 assert len(set(e.answer for e in t.examples)) == 3697071def test_trace_stride():72 cfg = OpSeqTaskConfig(trace_stride=2, min_len=1, max_len=5)73 s = opseq.generate(cfg, 300, seed=12, name="stride")74 opseq.validate(s, cfg)75 for e in s.examples:76 k = e.meta["len"]77 assert e.n_answer == len(set(list(range(1, k, 2)) + [k - 1]))78