"""Tests for optimization_toolkit.py. Run: python test_optimization_toolkit.py -v Standard library only. Every check uses fixed seeds, so results repeat exactly.""" import itertools import math import random import unittest import optimization_toolkit as o BASE = o.Program() TRAIN = o.make_dataset(60, 1) DEV = o.make_dataset(60, 2) ENV = o.ToyEnv(seed=3) def fd(f, policy, s, t, u, h=1e-5): """Central finite difference of f(policy) with respect to one logit.""" p = policy p.theta[s][t][u] += h a = f(p) p.theta[s][t][u] -= 2 * h b = f(p) p.theta[s][t][u] += h return (a - b) / (2 * h) def rand_policy(seed=0, S=3, T=2, V=3): rng = random.Random(seed) return o.TabPolicy(S, T, V, [[[rng.uniform(-1, 1) for _ in range(V)] for _ in range(T)] for _ in range(S)]) # ----------------------------------------------------------------------------- Part 0 and 1 class TestShared(unittest.TestCase): def test_count_tokens(self): self.assertEqual(o.count_tokens("a b, c"), 4) self.assertEqual(o.count_tokens(""), 0) def test_meter_counts(self): m = o.Meter() m.charge(2, 10) m.charge(1, 5) self.assertEqual((m.rollouts, m.tokens), (3, 15)) def test_meter_cap_raises(self): m = o.Meter(3) m.charge(3) with self.assertRaises(o.BudgetExhausted): m.charge(1) self.assertEqual(m.rollouts, 3) def test_meter_left(self): m = o.Meter(5) m.charge(2) self.assertEqual(m.left, 3) self.assertIsNone(o.Meter().left) def test_meter_exact_cap_ok(self): m = o.Meter(2) m.charge(1) m.charge(1) self.assertEqual(m.left, 0) def test_mean_empty(self): self.assertEqual(o.mean([]), 0.0) self.assertEqual(o.mean([1, 2, 3]), 2.0) def test_stdev(self): self.assertEqual(o.stdev([]), 0.0) self.assertAlmostEqual(o.stdev([1, 3]), 1.0) self.assertEqual(o.stdev([2, 2, 2]), 0.0) class TestToyWorld(unittest.TestCase): def test_dataset_size_and_ids(self): d = o.make_dataset(10, 5) self.assertEqual([e.id for e in d], list(range(10))) def test_dataset_deterministic(self): self.assertEqual(o.make_dataset(20, 3), o.make_dataset(20, 3)) self.assertNotEqual(o.make_dataset(20, 3), o.make_dataset(20, 4)) def test_dataset_kinds_and_answers(self): d = o.make_dataset(300, 0) self.assertEqual({e.kind for e in d}, set("ABCDEF")) self.assertTrue(all(0 <= e.answer < 4 for e in d)) def test_dataset_weights(self): d = o.make_dataset(50, 0, {"A": 1.0}) self.assertEqual({e.kind for e in d}, {"A"}) def test_dataset_bad_weights(self): with self.assertRaises(ValueError): o.make_dataset(5, 0, {"A": 0.0}) def test_base_probability(self): for k in "ABCDEF": self.assertAlmostEqual(o.p_correct(BASE, k), 0.35) def test_instruction_gain(self): p = BASE.with_instruction("parse", "Solve units") self.assertAlmostEqual(o.p_correct(p, "A"), 0.75) self.assertAlmostEqual(o.p_correct(p, "B"), 0.35) def test_instruction_in_wrong_module_does_not_help(self): p = BASE.with_instruction("answer", "Solve units") self.assertAlmostEqual(o.p_correct(p, "A"), 0.35) def test_demo_gain_capped_at_two(self): d = tuple(o.Demo(i, "A", 0, True) for i in range(5)) self.assertAlmostEqual(o.p_correct(BASE.with_demos(d[:1]), "A"), 0.45) self.assertAlmostEqual(o.p_correct(BASE.with_demos(d[:2]), "A"), 0.55) self.assertAlmostEqual(o.p_correct(BASE.with_demos(d[:5]), "A"), 0.55) def test_bad_demo_penalty(self): d = (o.Demo(0, "A", 0, False),) self.assertAlmostEqual(o.p_correct(BASE.with_demos(d), "A"), 0.25) def test_bad_demo_penalty_capped(self): d = tuple(o.Demo(i, "A", 0, False) for i in range(5)) self.assertAlmostEqual(o.p_correct(BASE.with_demos(d), "A"), 0.15) def test_demo_of_other_kind_ignored(self): d = (o.Demo(0, "B", 0, True),) self.assertAlmostEqual(o.p_correct(BASE.with_demos(d), "A"), 0.35) def test_conflict_negatives_hurts_dates(self): p = BASE.with_instruction("parse", "Solve negatives") self.assertAlmostEqual(o.p_correct(p, "C"), 0.05) self.assertAlmostEqual(o.p_correct(p, "B"), 0.75) def test_conflict_abbrev_hurts_ties(self): p = BASE.with_instruction("answer", "Solve abbrev") self.assertAlmostEqual(o.p_correct(p, "F"), 0.10) def test_conflict_cancelled_by_own_rule(self): p = BASE.with_instruction("parse", "Solve negatives dates") self.assertAlmostEqual(o.p_correct(p, "C"), 0.45) def test_length_penalty(self): p = BASE.with_instruction("parse", " ".join(["x"] * 15)) self.assertAlmostEqual(o.p_correct(p, "A"), 0.35 - 0.02 * 5) def test_no_penalty_up_to_free_words(self): p = BASE.with_instruction("parse", " ".join(["x"] * 10)) self.assertAlmostEqual(o.p_correct(p, "A"), 0.35) def test_demo_bloat_penalty(self): d = tuple(o.Demo(i, "F", 0, True) for i in range(10)) # 10 demos: 6 free, 4 over; kind F gets +0.2 from 2 demos capped self.assertAlmostEqual(o.p_correct(BASE.with_demos(d), "F"), 0.35 + 0.2 - 0.04) def test_clamped_low_and_high(self): bad = BASE.with_instruction("answer", "Solve abbrev").with_demos(tuple(o.Demo(i, "F", 0, False) for i in range(2))) self.assertAlmostEqual(o.p_correct(bad, "F"), 0.02) good = o.Program({"parse": "Solve units", "answer": "Solve"}, tuple(o.Demo(i, "A", 0) for i in range(2))) self.assertAlmostEqual(o.p_correct(good, "A"), 0.35 + 0.4 + 0.2) params = o.ToyParams(base=0.9) self.assertAlmostEqual(o.p_correct(good, "A", params), 0.98) def test_with_instruction_unknown_module(self): with self.assertRaises(KeyError): BASE.with_instruction("nope", "x") def test_program_is_immutable_copy(self): p = BASE.with_instruction("parse", "Solve units") self.assertEqual(BASE.instructions["parse"], "Solve the task.") self.assertEqual(p.instructions["answer"], "Solve the task.") def test_rollout_charges_meter(self): m = o.Meter() o.rollout(BASE, TRAIN[0], random.Random(0), meter=m) self.assertEqual(m.rollouts, 1) self.assertGreater(m.tokens, 0) def test_rollout_deterministic(self): a = o.rollout(BASE, TRAIN[0], random.Random(5)) b = o.rollout(BASE, TRAIN[0], random.Random(5)) self.assertEqual(a, b) def test_wrong_output_differs_from_answer(self): for i, ex in enumerate(TRAIN): t = o.rollout(BASE, ex, random.Random(i)) self.assertEqual(t.correct, t.output == ex.answer) self.assertTrue(0 <= t.output < 4) def test_empirical_matches_p_correct(self): p = BASE.with_instruction("parse", "Solve units") ex = o.Example(0, "A", 1) rng = random.Random(0) hits = sum(o.rollout(p, ex, rng).correct for _ in range(4000)) / 4000 self.assertAlmostEqual(hits, 0.75, delta=0.03) def test_feedback_names_missing_rule(self): ex = o.Example(0, "B", 0) self.assertIn("needs the rule 'negatives'", o.feedback_text(BASE, ex, False)) def test_feedback_says_noise_when_rule_present(self): ex = o.Example(0, "B", 0) p = BASE.with_instruction("parse", "Solve negatives") self.assertIn("noise", o.feedback_text(p, ex, False)) def test_feedback_ok_when_correct(self): self.assertTrue(o.feedback_text(BASE, o.Example(0, "A", 0), True).startswith("ok")) def test_evaluate_range(self): self.assertTrue(0.0 <= o.evaluate(BASE, DEV, random.Random(0)) <= 1.0) def test_scores_on_length_and_values(self): s = o.scores_on(BASE, DEV, random.Random(0)) self.assertEqual(len(s), len(DEV)) self.assertTrue(set(s) <= {0.0, 1.0}) def test_exact_accuracy_base(self): self.assertAlmostEqual(o.exact_accuracy(BASE), 0.35) def test_exact_accuracy_weights(self): p = BASE.with_instruction("parse", "Solve units") self.assertAlmostEqual(o.exact_accuracy(p, {"A": 1.0}), 0.75) self.assertAlmostEqual(o.exact_accuracy(p), (0.75 + 5 * 0.35) / 6) def test_best_possible_program(self): b = o.best_possible_program() self.assertEqual(set(b.instructions["parse"].split()), {"Solve", "units", "negatives", "dates"}) self.assertEqual(set(b.instructions["answer"].split()), {"Solve", "abbrev", "currency", "ties"}) self.assertAlmostEqual(o.exact_accuracy(b), (0.75 * 4 + 0.45 + 0.5) / 6) def test_best_is_best_by_brute_force(self): best = o.exact_accuracy(o.best_possible_program()) words = list(o.RULES.values()) for mask in range(1 << 6): ins = " ".join(["Solve"] + [words[i] for i in range(6) if mask >> i & 1]) for mod_of_words in (ins,): p = o.uniform_program(mod_of_words) self.assertLessEqual(o.exact_accuracy(p), best + 1e-9) def test_per_kind_accuracy_keys(self): self.assertEqual(sorted(o.per_kind_accuracy(BASE)), list("ABCDEF")) # ----------------------------------------------------------------------------- Part 2 class TestBootstrap(unittest.TestCase): def test_filtered_demos_are_correct(self): d = o.bootstrap_demos(BASE, TRAIN, random.Random(0)) self.assertTrue(d) self.assertTrue(all(x.correct for x in d)) def test_max_per_kind(self): d = o.bootstrap_demos(BASE, TRAIN, random.Random(0), max_per_kind=1) kinds = [x.kind for x in d] self.assertEqual(len(kinds), len(set(kinds))) def test_keep_all_includes_wrong(self): d = o.bootstrap_demos(BASE, TRAIN, random.Random(0), keep_all=True) self.assertTrue(any(not x.correct for x in d)) def test_charges_one_rollout_per_example(self): m = o.Meter() o.bootstrap_demos(BASE, TRAIN, random.Random(0), meter=m) self.assertEqual(m.rollouts, len(TRAIN)) def test_filter_helps_on_average(self): f, a = [], [] for s in range(8): f.append(o.exact_accuracy(BASE.with_demos(o.bootstrap_demos(BASE, TRAIN, random.Random(s))))) a.append(o.exact_accuracy(BASE.with_demos(o.bootstrap_demos(BASE, TRAIN, random.Random(s), keep_all=True)))) self.assertGreater(o.mean(f), o.mean(a) + 0.1) def test_random_search_log_includes_empty_set(self): _, log = o.bootstrap_random_search(BASE, TRAIN, DEV, random.Random(0), n_sets=3) self.assertEqual(log[0][0], ()) self.assertEqual(len(log), 4) def test_random_search_picks_best_logged(self): best, log = o.bootstrap_random_search(BASE, TRAIN, DEV, random.Random(1)) top = max(sc for _, sc in log) self.assertEqual(top, max(sc for ids, sc in log if tuple(d.id for d in best.demos) == ids)) def test_random_search_deterministic(self): a = o.bootstrap_random_search(BASE, TRAIN, DEV, random.Random(2)) b = o.bootstrap_random_search(BASE, TRAIN, DEV, random.Random(2)) self.assertEqual(a, b) def test_random_search_keeps_instructions(self): best, _ = o.bootstrap_random_search(BASE, TRAIN, DEV, random.Random(2)) self.assertEqual(best.instructions, BASE.instructions) class TestMipro(unittest.TestCase): def test_failure_summary_keys_and_range(self): fr = o.failure_summary(BASE, TRAIN, random.Random(0)) self.assertEqual(sorted(fr), list("ABCDEF")) self.assertTrue(all(0 <= v <= 1 for v in fr.values())) def test_failure_summary_high_when_unhelped(self): fr = o.failure_summary(BASE, o.make_dataset(600, 1), random.Random(0)) self.assertAlmostEqual(o.mean(fr.values()), 0.65, delta=0.05) def test_proposals_program_aware_use_only_own_rules(self): props = o.propose_instructions("parse", 8, random.Random(0)) own = {"units", "negatives", "dates"} for p in props: self.assertTrue(set(p.split()[1:]) <= own, p) def test_proposals_unique_and_count(self): props = o.propose_instructions("answer", 6, random.Random(0)) self.assertEqual(len(props), len(set(props))) self.assertEqual(len(props), 6) def test_proposals_ungrounded_can_use_other_words(self): props = o.propose_instructions("parse", 20, random.Random(0), program_aware=False) words = {w for p in props for w in p.split()[1:]} self.assertTrue(words - {"units", "negatives", "dates"}) def test_proposals_tip_adds_flavor(self): props = o.propose_instructions("parse", 5, random.Random(0), tip=True) self.assertTrue(all(set(p.split()) & set(o.DISTRACTORS) for p in props)) def test_proposals_weighted_by_failure(self): fr = {"A": 1.0, "B": 0.0, "C": 0.0} props = [p for s in range(30) for p in o.propose_instructions("parse", 1, random.Random(s), fr)] a = sum("units" in p for p in props) b = sum("negatives" in p for p in props) self.assertGreater(a, b + 10) def test_mipro_result_shape(self): r = o.mipro_search(BASE, TRAIN, DEV, random.Random(0)) self.assertIsInstance(r.program, o.Program) self.assertEqual(len(r.history), 24) self.assertGreater(r.rollouts, 0) def test_mipro_deterministic(self): a = o.mipro_search(BASE, TRAIN, DEV, random.Random(3)) b = o.mipro_search(BASE, TRAIN, DEV, random.Random(3)) self.assertEqual(a.program, b.program) self.assertEqual(a.rollouts, b.rollouts) def test_mipro_improves_on_base(self): accs = [o.exact_accuracy(o.mipro_search(BASE, TRAIN, DEV, random.Random(s)).program) for s in range(4)] self.assertGreater(o.mean(accs), 0.5) def test_mipro_rollout_count_formula(self): # grounded: len(train) failure summary + len(train) bootstrap + trials*minibatch + full evals r = o.mipro_search(BASE, TRAIN, DEV, random.Random(0), trials=8, minibatch=5, full_every=4) self.assertEqual(r.rollouts, 60 + 60 + 8 * 5 + 2 * 2 * 60) def test_mipro_ungrounded_skips_failure_summary(self): r = o.mipro_search(BASE, TRAIN, DEV, random.Random(0), grounded=False, trials=8, minibatch=5, full_every=4) self.assertEqual(r.rollouts, 60 + 8 * 5 + 2 * 2 * 60) def test_mipro_surrogate_flag_same_budget(self): a = o.mipro_search(BASE, TRAIN, DEV, random.Random(0), surrogate=True) b = o.mipro_search(BASE, TRAIN, DEV, random.Random(0), surrogate=False) self.assertEqual(a.rollouts, b.rollouts) def test_mipro_grounding_helps_on_average(self): g = o.mean(o.exact_accuracy(o.mipro_search(BASE, TRAIN, DEV, random.Random(s)).program) for s in range(8)) u = o.mean(o.exact_accuracy(o.mipro_search(BASE, TRAIN, DEV, random.Random(s), grounded=False).program) for s in range(8)) self.assertGreater(g, u + 0.05) def test_mipro_history_keys_are_triples(self): r = o.mipro_search(BASE, TRAIN, DEV, random.Random(0), trials=5, minibatch=4, full_every=5) for t, key, sc in r.history: self.assertEqual(len(key), 3) self.assertTrue(0 <= sc <= 1) def test_mipro_score_is_a_dev_score(self): r = o.mipro_search(BASE, TRAIN, DEV, random.Random(0), trials=5, minibatch=4, full_every=5) self.assertTrue(0 <= r.score <= 1) self.assertAlmostEqual(r.score * 60, round(r.score * 60)) class TestUCB(unittest.TestCase): def test_prefers_better_mean(self): stats = [[1.0, 10], [9.0, 10], [2.0, 10]] picks = [o.ucb_choice(stats, random.Random(s), explore=0.0) for s in range(20)] self.assertEqual(set(picks), {1}) def test_bonus_prefers_less_tried_when_means_equal(self): stats = [[5.0, 10], [0.5, 1]] self.assertEqual({o.ucb_choice(stats, random.Random(s), explore=0.0) for s in range(10)}, {1}) def test_untried_option_has_prior_half(self): # a tried option with mean 0.4 loses to an untried one (prior 0.5 plus a larger bonus) stats = [[4.0, 10], [0.0, 0]] self.assertEqual({o.ucb_choice(stats, random.Random(s), explore=0.0) for s in range(10)}, {1}) # a tried option with mean 0.9 and 10 pulls (bonus 0.15) beats an untried option only if 1.05 > 1.0 stats = [[9.0, 10], [0.0, 0]] self.assertEqual({o.ucb_choice(stats, random.Random(s), explore=0.0) for s in range(10)}, {0}) def test_surrogate_off_is_uniform(self): stats = [[1.0, 10], [9.0, 10], [2.0, 10]] picks = [o.ucb_choice(stats, random.Random(s), surrogate=False) for s in range(300)] for k in range(3): self.assertGreater(picks.count(k), 60) def test_explore_rate(self): stats = [[0.0, 10], [10.0, 10]] picks = [o.ucb_choice(stats, random.Random(s), explore=0.5) for s in range(1000)] self.assertAlmostEqual(picks.count(0) / 1000, 0.25, delta=0.05) def test_mipro_belief_counts_trials(self): r = o.mipro_search(BASE, TRAIN, DEV, random.Random(0), trials=10, minibatch=5, full_every=5) for v, rows in r.belief.items(): self.assertEqual(sum(n for _, n in rows), 10) for tot, n in rows: self.assertLessEqual(tot, n) def test_mipro_final_is_best_full_eval(self): r = o.mipro_search(BASE, TRAIN, DEV, random.Random(1), trials=12, minibatch=5, full_every=4) self.assertEqual(r.score, max(sc for _, sc, _ in r.full_evals)) def test_mipro_full_eval_on_last_trial(self): r = o.mipro_search(BASE, TRAIN, DEV, random.Random(1), trials=5, minibatch=4, full_every=4) self.assertEqual(sorted({t for t, _, _ in r.full_evals}), [4, 5]) def test_mipro_full_eval_top_two(self): r = o.mipro_search(BASE, TRAIN, DEV, random.Random(1), trials=8, minibatch=4, full_every=8) self.assertEqual(len(r.full_evals), 2) class TestOpro(unittest.TestCase): H = [("a", 0.5), ("b", 0.2), ("c", 0.9), ("d", 0.1)] def test_prompt_is_ascending(self): pairs = o.parse_opro_prompt(o.opro_prompt(self.H)) scores = [s for _, s in pairs] self.assertEqual(scores, sorted(scores)) self.assertEqual(pairs[-1][0], "c") def test_prompt_keeps_top_k(self): pairs = o.parse_opro_prompt(o.opro_prompt(self.H, k=2)) self.assertEqual([t for t, _ in pairs], ["a", "c"]) def test_prompt_roundtrip_scores_two_decimals(self): pairs = o.parse_opro_prompt(o.opro_prompt([("x y", 0.123456)])) self.assertEqual(pairs, [("x y", 0.12)]) def test_parse_empty(self): self.assertEqual(o.parse_opro_prompt("nothing here"), []) def test_proposals_start_from_best_words(self): prompt = o.opro_prompt([("Solve units", 0.9), ("Solve dates", 0.1)]) props = o.opro_propose(prompt, 100, random.Random(0)) self.assertGreater(sum("units" in p for p in props), 50) self.assertGreater(sum("units" in p for p in props), sum("dates" in p for p in props) + 20) def test_proposals_count_and_cap(self): props = o.opro_propose(o.opro_prompt([("Solve " + " ".join(o.VOCAB), 0.5)]), 10, random.Random(0), max_words=5) self.assertEqual(len(props), 10) self.assertTrue(all(len(p.split()) <= 5 for p in props)) def test_loop_history_unique_and_seeded(self): r = o.opro(TRAIN[:24], random.Random(0), steps=4, per_step=3) texts = [t for t, _ in r.history] self.assertEqual(len(texts), len(set(texts))) self.assertEqual(texts[0], "Solve the task.") def test_loop_rollouts_are_scoring_set_times_instructions(self): r = o.opro(TRAIN[:24], random.Random(0), steps=4, per_step=3) self.assertEqual(r.rollouts, 24 * len(r.history)) def test_loop_returns_best_seen(self): r = o.opro(TRAIN[:24], random.Random(0), steps=6) self.assertEqual(r.score, max(s for _, s in r.history)) def test_loop_deterministic(self): self.assertEqual(o.opro(TRAIN[:24], random.Random(7)).history, o.opro(TRAIN[:24], random.Random(7)).history) def test_scores_are_on_scoring_set_not_exact(self): # The best instruction is chosen by a noisy 24-example score, so it can look better than it is. gaps = [] for s in range(6): r = o.opro(TRAIN[:24], random.Random(s)) gaps.append(r.score - o.exact_accuracy(o.uniform_program(r.instruction))) self.assertGreater(o.mean(gaps), 0.0) def test_uniform_program_sets_both_modules(self): p = o.uniform_program("Solve x") self.assertEqual(set(p.instructions.values()), {"Solve x"}) class TestOproBest(unittest.TestCase): def test_returned_instruction_is_the_argmax(self): for s in range(6): r = o.opro(TRAIN[:24], random.Random(s), steps=6) self.assertEqual(dict(r.history)[r.instruction], r.score) self.assertEqual(r.score, max(sc for _, sc in r.history)) class TestTextGrad(unittest.TestCase): def test_loss_text_only_failures(self): ts = [o.rollout(BASE, ex, random.Random(i)) for i, ex in enumerate(TRAIN[:20])] text = o.tg_loss_text(ts) self.assertEqual(len(text.splitlines()), sum(1 for t in ts if not t.correct)) def test_critic_reads_loss(self): loss = "fail: kind B needs the rule 'negatives', and the parse instruction does not mention it" self.assertEqual(o.tg_critic("parse", loss, []), "consider naming: negatives") self.assertEqual(o.tg_critic("answer", loss, []), "") def test_critic_orders_by_count(self): loss = "\n".join(["fail: kind A needs the rule 'units', and x"] * 1 + ["fail: kind B needs the rule 'negatives', and x"] * 3) self.assertEqual(o.tg_critic("parse", loss, []), "consider naming: negatives, units") def test_critic_uses_downstream_gradient(self): self.assertEqual(o.tg_critic("parse", "", ["consider naming: units"]), "consider naming: units") def test_backward_visits_reverse_order(self): g = [("parse", o.TGVar("Solve", "parse")), ("answer", o.TGVar("Solve", "answer"))] self.assertEqual(o.tg_backward(g, ""), ["answer", "parse"]) def test_backward_respects_requires_grad(self): loss = "fail: kind A needs the rule 'units', and x" v = o.TGVar("Solve", "parse", requires_grad=False) o.tg_backward([("parse", v)], loss) self.assertEqual(v.grads, []) def test_backward_writes_grad(self): loss = "fail: kind A needs the rule 'units', and x" v = o.TGVar("Solve", "parse") o.tg_backward([("parse", v)], loss) self.assertEqual(v.grads, ["consider naming: units"]) def test_step_adds_top_word_and_clears(self): v = o.TGVar("Solve", "parse") v.grads = ["consider naming: dates, units", "consider naming: units"] changed = o.tg_step([("parse", v)]) self.assertEqual(changed, ["parse"]) self.assertEqual(v.value, "Solve units") self.assertEqual(v.grads, []) def test_step_skips_words_already_present(self): v = o.TGVar("Solve units", "parse") v.grads = ["consider naming: units"] self.assertEqual(o.tg_step([("parse", v)]), []) self.assertEqual(v.value, "Solve units") def test_step_frozen_variable_unchanged(self): v = o.TGVar("Solve", "parse", requires_grad=False) v.grads = ["consider naming: units"] self.assertEqual(o.tg_step([("parse", v)]), []) self.assertEqual(v.value, "Solve") def test_optimize_gets_better(self): p, hist, m = o.textgrad_optimize(BASE, TRAIN, random.Random(1)) self.assertGreater(o.exact_accuracy(p), 0.5) self.assertEqual(len(hist), 8) def test_optimize_rollout_accounting(self): _, hist, m = o.textgrad_optimize(BASE, TRAIN, random.Random(1), steps=3, minibatch=10) changed = sum(1 for h in hist if h[1]) self.assertEqual(m.rollouts, 3 * 10 + changed * 10) def test_frozen_module_not_edited(self): p, _, _ = o.textgrad_optimize(BASE, TRAIN, random.Random(1), frozen=("parse",)) self.assertEqual(p.instructions["parse"], BASE.instructions["parse"]) self.assertNotEqual(p.instructions["answer"], BASE.instructions["answer"]) def test_gate_never_accepts_a_worse_minibatch(self): _, hist, _ = o.textgrad_optimize(BASE, TRAIN, random.Random(2), gate=True) for step, changed, before, after, acc in hist: if acc: self.assertGreaterEqual(after, before) def test_without_gate_accepts_all_changes(self): _, hist, _ = o.textgrad_optimize(BASE, TRAIN, random.Random(2), gate=False) for step, changed, before, after, acc in hist: self.assertEqual(acc, bool(changed)) def test_optimize_deterministic(self): a = o.textgrad_optimize(BASE, TRAIN, random.Random(4))[0] b = o.textgrad_optimize(BASE, TRAIN, random.Random(4))[0] self.assertEqual(a, b) class FailAll: """Context manager: every rollout fails with feedback that names a missing rule, and evaluate() returns a fixed value. Lets tests control ties exactly.""" def __init__(self, eval_value): self.v = eval_value def __enter__(self): self.r, self.e = o.rollout, o.evaluate def rollout(prog, ex, rng, params=o.ToyParams(), meter=None): if meter is not None: meter.charge(1, 1) m = o.MODULE_OF[ex.kind] fb = f"fail: kind {ex.kind} needs the rule '{o.RULES[ex.kind]}', and the {m} instruction does not mention it" return o.Trace(ex, m, 0, False, fb, "", 1) o.rollout = rollout o.evaluate = lambda prog, exs, rng, params=o.ToyParams(), meter=None: self.v return self def __exit__(self, *a): o.rollout, o.evaluate = self.r, self.e class TestTies(unittest.TestCase): def test_textgrad_accepts_a_tie(self): with FailAll(0.0): _, hist, _ = o.textgrad_optimize(BASE, TRAIN, random.Random(0), steps=2) self.assertTrue(all(h[4] for h in hist if h[1])) self.assertTrue(any(h[1] for h in hist)) def test_textgrad_rejects_a_worse_step(self): with FailAll(-1.0): p, hist, _ = o.textgrad_optimize(BASE, TRAIN, random.Random(0), steps=2) self.assertFalse(any(h[4] for h in hist)) self.assertEqual(p, BASE) def test_gepa_rejects_a_tie(self): with FailAll(0.0): r = o.gepa(BASE, TRAIN, DEV[:18], random.Random(0), budget=500) self.assertEqual(r.accepted, 0) self.assertEqual(r.pool, [BASE]) def test_gepa_accepts_a_strict_gain(self): with FailAll(0.5): r = o.gepa(BASE, TRAIN, DEV[:18], random.Random(0), budget=500) self.assertGreater(r.accepted, 0) def test_gepa_rejects_a_loss(self): with FailAll(-1.0): r = o.gepa(BASE, TRAIN, DEV[:18], random.Random(0), budget=500) self.assertEqual(r.accepted, 0) class TestGepa(unittest.TestCase): def test_select_dominated_removed(self): # candidate 1 is at least as good as candidate 0 everywhere and strictly better somewhere scores = [[1, 0, 0], [1, 1, 0]] self.assertEqual(o.pareto_front(scores), [1]) def test_select_keeps_specialists(self): scores = [[1, 0], [0, 1]] self.assertEqual(o.pareto_front(scores), [0, 1]) def test_select_ties_keep_one(self): scores = [[1, 1], [1, 1], [1, 1]] self.assertEqual(len(o.pareto_front(scores)), 1) def test_select_nonwinner_never_survives(self): scores = [[1, 1, 1], [0, 0, 0], [1, 0, 1]] self.assertNotIn(1, o.pareto_front(scores)) def test_select_weighted_by_wins(self): scores = [[1, 1, 1, 0], [0, 0, 0, 1]] rng = random.Random(0) picks = [o.select_candidate(scores, rng) for _ in range(2000)] self.assertAlmostEqual(picks.count(0) / 2000, 0.75, delta=0.04) def test_select_single_candidate(self): self.assertEqual(o.select_candidate([[0, 1]], random.Random(0)), 0) def test_select_returns_front_member(self): scores = [[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 0, 0]] for s in range(30): self.assertIn(o.select_candidate(scores, random.Random(s)), o.pareto_front(scores)) def test_select_removes_dominated_candidate(self): scores = [[1, 1, 0], [1, 1, 1]] picks = {o.select_candidate(scores, random.Random(s)) for s in range(200)} self.assertEqual(picks, {1}) def test_select_all_zero_scores_still_picks(self): scores = [[0, 0], [0, 0]] self.assertIn(o.select_candidate(scores, random.Random(0)), (0, 1)) self.assertEqual(len(o.pareto_front(scores)), 1) def test_reflect_does_not_repeat_a_present_rule(self): t = o.Trace(o.Example(0, "A", 0), "parse", 1, False, "fail: kind A needs the rule 'units', and the parse instruction does not mention it", "", 1) self.assertEqual(o.reflect_update("Solve units", [t], "parse"), "Solve units") def test_greedy_parent_is_best_mean_at_the_time(self): for s in range(4): r = o.gepa(BASE, TRAIN, DEV[:18], random.Random(s), budget=800, select="best") for i, par in enumerate(r.parents): if i == 0: continue want = max(range(i), key=lambda j: (o.mean(r.scores[j]), -j)) self.assertEqual(par, want) def test_reflect_adds_most_requested_rule(self): mk = lambda kind, mod: o.Trace(o.Example(0, kind, 0), mod, 1, False, f"fail: kind {kind} needs the rule '{o.RULES[kind]}', and the {mod} instruction does not mention it", "", 1) ts = [mk("B", "parse"), mk("B", "parse"), mk("A", "parse")] self.assertEqual(o.reflect_update("Solve", ts, "parse"), "Solve negatives") def test_reflect_ignores_other_module(self): t = o.Trace(o.Example(0, "D", 0), "answer", 1, False, "fail: kind D needs the rule 'abbrev', and the answer instruction does not mention it", "", 1) self.assertEqual(o.reflect_update("Solve", [t], "parse"), "Solve") def test_reflect_noise_leaves_text(self): t = o.Trace(o.Example(0, "A", 0), "parse", 1, False, "fail: kind A: the rule 'units' is stated, so this miss looks like noise", "", 1) self.assertEqual(o.reflect_update("Solve units", [t], "parse"), "Solve units") def test_reflect_ignores_successes(self): t = o.Trace(o.Example(0, "A", 0), "parse", 0, True, "ok: kind A (units) was right", "", 1) self.assertEqual(o.reflect_update("Solve", [t], "parse"), "Solve") def test_merge_takes_best_module_from_each(self): dp = [o.Example(0, "A", 0), o.Example(1, "D", 0)] a = BASE.with_instruction("parse", "Solve units") b = BASE.with_instruction("answer", "Solve abbrev") m = o.merge_programs(a, b, [1, 0], [0, 1], dp) self.assertEqual(m.instructions, {"parse": "Solve units", "answer": "Solve abbrev"}) def test_merge_tie_goes_to_first(self): dp = [o.Example(0, "A", 0), o.Example(1, "D", 0)] a = BASE.with_instruction("parse", "Solve units") b = BASE.with_instruction("parse", "Solve dates") m = o.merge_programs(a, b, [1, 1], [1, 1], dp) self.assertEqual(m.instructions["parse"], "Solve units") def test_budget_respected(self): for b in (60, 200, 777): r = o.gepa(BASE, TRAIN, DEV[:18], random.Random(0), budget=b) self.assertLessEqual(r.rollouts, b) def test_budget_too_small_for_seed_scoring(self): r = o.gepa(BASE, TRAIN, DEV[:18], random.Random(0), budget=10) self.assertEqual(r.program, BASE) self.assertEqual(r.pool, [BASE]) def test_pool_scores_parents_aligned(self): r = o.gepa(BASE, TRAIN, DEV[:18], random.Random(0), budget=600) self.assertEqual(len(r.pool), len(r.scores)) self.assertEqual(len(r.pool), len(r.parents)) self.assertIsNone(r.parents[0]) def test_pool_has_no_duplicates(self): r = o.gepa(BASE, TRAIN, DEV[:18], random.Random(1), budget=800) self.assertEqual(len(r.pool), len(set((tuple(sorted(p.instructions.items())), p.demos) for p in r.pool))) def test_children_differ_from_parent_in_one_module(self): r = o.gepa(BASE, TRAIN, DEV[:18], random.Random(1), budget=800) for k, par in enumerate(r.parents): if isinstance(par, int): diff = [m for m in o.MODULES if r.pool[k].instructions[m] != r.pool[par].instructions[m]] self.assertEqual(len(diff), 1) def test_result_is_best_mean_on_pareto_set(self): r = o.gepa(BASE, TRAIN, DEV[:18], random.Random(1), budget=800) best = max(o.mean(s) for s in r.scores) self.assertEqual(o.mean(r.scores[r.pool.index(r.program)]), best) def test_deterministic(self): a = o.gepa(BASE, TRAIN, DEV[:18], random.Random(5), budget=500) b = o.gepa(BASE, TRAIN, DEV[:18], random.Random(5), budget=500) self.assertEqual(a.program, b.program) self.assertEqual(a.rollouts, b.rollouts) def test_improves_on_seed(self): accs = [o.exact_accuracy(o.gepa(BASE, TRAIN, DEV[:18], random.Random(s), budget=1000).program) for s in range(6)] self.assertGreater(o.mean(accs), 0.5) def test_pareto_beats_greedy_on_average(self): p = o.mean(o.exact_accuracy(o.gepa(BASE, TRAIN, DEV[:18], random.Random(s), budget=1000).program) for s in range(8)) g = o.mean(o.exact_accuracy(o.gepa(BASE, TRAIN, DEV[:18], random.Random(s), budget=1000, select="best").program) for s in range(8)) self.assertGreater(p, g) def test_merge_creates_two_parent_children(self): r = o.gepa(BASE, TRAIN, DEV[:18], random.Random(1), budget=1000, merge_every=4) self.assertTrue(any(isinstance(p, tuple) for p in r.parents)) def test_minibatch_param_changes_cost(self): a = o.gepa(BASE, TRAIN, DEV[:18], random.Random(0), budget=400, minibatch=3) b = o.gepa(BASE, TRAIN, DEV[:18], random.Random(0), budget=400, minibatch=8) self.assertGreater(a.iterations, b.iterations) def test_feedback_is_what_makes_it_work(self): # With feedback text that never names a rule, reflection cannot change anything. orig = o.feedback_text try: o.feedback_text = lambda prog, ex, correct: "fail" if not correct else "ok" r = o.gepa(BASE, TRAIN, DEV[:18], random.Random(0), budget=600) finally: o.feedback_text = orig self.assertEqual(r.program, BASE) self.assertEqual(r.accepted, 0) # ----------------------------------------------------------------------------- Part 3 class TestPolicy(unittest.TestCase): def test_uniform_start(self): p = o.TabPolicy() self.assertEqual(p.probs(0, 0), [1 / 3] * 3) def test_softmax_sums_to_one(self): self.assertAlmostEqual(sum(o.softmax([1000.0, 0.0, -5.0])), 1.0) def test_softmax_stable_for_large_logits(self): self.assertEqual(o.softmax([1000.0, 0.0])[0], 1.0) def test_logp_is_sum_of_tokens(self): p = rand_policy(1) self.assertAlmostEqual(p.logp(1, (0, 2)), p.logp_tok(1, 0, 0) + p.logp_tok(1, 1, 2)) def test_sequence_probabilities_sum_to_one(self): p = rand_policy(2) tot = sum(math.exp(p.logp(0, q)) for q in itertools.product(range(3), repeat=2)) self.assertAlmostEqual(tot, 1.0) def test_copy_is_independent(self): p = rand_policy(2) q = p.copy() q.theta[0][0][0] += 5 self.assertNotEqual(p.theta[0][0][0], q.theta[0][0][0]) def test_sample_deterministic_and_valid(self): p = rand_policy(3) a = p.sample(0, random.Random(1)) self.assertEqual(a, p.sample(0, random.Random(1))) self.assertEqual(len(a), 2) self.assertTrue(all(0 <= v < 3 for v in a)) def test_score_function_gradient_matches_finite_difference(self): p = rand_policy(4) for s, t, u in [(0, 0, 0), (1, 1, 2), (2, 0, 1)]: g = o.zeros_like(p) o.add_score(g, p, 1, 1, 0, 1.0) num = fd(lambda q: q.logp_tok(1, 1, 0), p, s, t, u) self.assertAlmostEqual(g[s][t][u], num, places=6) def test_score_gradient_sums_to_zero(self): p = rand_policy(5) g = o.zeros_like(p) o.add_score(g, p, 0, 0, 1, 2.0) self.assertAlmostEqual(sum(g[0][0]), 0.0) def test_apply_grad(self): p = o.TabPolicy() g = o.zeros_like(p) g[0][0][1] = 2.0 o.apply_grad(p, g, 0.5) self.assertEqual(p.theta[0][0][1], 1.0) def test_grad_norm(self): g = o.zeros_like(o.TabPolicy()) g[0][0][0] = 3.0 g[1][1][2] = 4.0 self.assertAlmostEqual(o.grad_norm(g), 5.0) def test_entries_count(self): self.assertEqual(len(o.TabPolicy().entries()), 18) class TestEnv(unittest.TestCase): def test_reward_exact_match(self): self.assertEqual(ENV.reward(0, ENV.target[0]), 1.0) self.assertEqual(ENV.reward(0, (9, 9)), 0.0) def test_targets_fixed_by_seed(self): self.assertEqual(o.ToyEnv(seed=3).target, ENV.target) self.assertEqual(ENV.target, [(0, 2), (2, 0), (1, 2)]) def test_success_uniform_policy(self): self.assertAlmostEqual(ENV.success(o.TabPolicy()), 1 / 9) def test_success_matches_enumeration(self): p = rand_policy(6) tot = 0.0 for s in range(3): for q in itertools.product(range(3), repeat=2): tot += math.exp(p.logp(s, q)) * ENV.reward(s, q) / 3 self.assertAlmostEqual(ENV.success(p), tot) def test_success_matches_sampling(self): p = rand_policy(6) rs = o.collect(p, ENV, random.Random(0), 6000) self.assertAlmostEqual(o.mean(r for _, _, r in rs), ENV.success(p), delta=0.02) def test_collect_size_and_fields(self): c = o.collect(o.TabPolicy(), ENV, random.Random(0), 10) self.assertEqual(len(c), 10) self.assertTrue(all(0 <= s < 3 and len(q) == 2 for s, q, _ in c)) def all_responses(p): return [(s, q) for s in range(p.S) for q in itertools.product(range(p.V), repeat=p.T)] def exact_J(p): return ENV.success(p) class TestReinforce(unittest.TestCase): def enumerate_batch(self, p): resp = all_responses(p) batch = [(s, q, ENV.reward(s, q)) for s, q in resp] w = [math.exp(p.logp(s, q)) / p.S for s, q in resp] return batch, w def test_gradient_equals_exact_gradient_with_baseline(self): p = rand_policy(7) batch, w = self.enumerate_batch(p) g = o.reinforce_grad(p, batch, "mean", w) for s, t, u in p.entries(): self.assertAlmostEqual(g[s][t][u], fd(exact_J, p, s, t, u), places=6) def test_gradient_equals_exact_gradient_without_baseline(self): p = rand_policy(8) batch, w = self.enumerate_batch(p) g = o.reinforce_grad(p, batch, "none", w) for s, t, u in p.entries(): self.assertAlmostEqual(g[s][t][u], fd(exact_J, p, s, t, u), places=6) def test_baseline_changes_nothing_in_expectation(self): p = rand_policy(9) batch, w = self.enumerate_batch(p) a = o.reinforce_grad(p, batch, "mean", w) b = o.reinforce_grad(p, batch, "none", w) for s, t, u in p.entries(): self.assertAlmostEqual(a[s][t][u], b[s][t][u], places=9) def test_baseline_reduces_variance(self): p = o.TabPolicy() norms = {"mean": [], "none": []} for k in range(60): b = o.collect(p, ENV, random.Random(k), 16) for name in norms: norms[name].append(o.grad_norm(o.reinforce_grad(p, b, name))) self.assertLess(o.mean(norms["mean"]), o.mean(norms["none"])) def test_all_zero_rewards_give_zero_gradient(self): p = o.TabPolicy() b = [(0, (0, 0), 0.0), (1, (1, 1), 0.0)] self.assertEqual(o.grad_norm(o.reinforce_grad(p, b, "mean")), 0.0) def test_no_baseline_zero_reward_zero_gradient(self): b = [(0, (0, 0), 0.0)] self.assertEqual(o.grad_norm(o.reinforce_grad(o.TabPolicy(), b, "none")), 0.0) def test_training_improves(self): p = o.TabPolicy() c = o.reinforce_train(p, ENV, random.Random(0), iters=150) self.assertGreater(c[-1], 0.8) self.assertGreater(c[-1], c[0]) def test_training_deterministic(self): a = o.reinforce_train(o.TabPolicy(), ENV, random.Random(2), iters=20) b = o.reinforce_train(o.TabPolicy(), ENV, random.Random(2), iters=20) self.assertEqual(a, b) def test_positive_reward_raises_token_probability(self): p = o.TabPolicy() before = p.probs(0, 0)[0] o.apply_grad(p, o.reinforce_grad(p, [(0, (0, 0), 1.0), (0, (1, 1), 0.0)], "mean"), 1.0) self.assertGreater(p.probs(0, 0)[0], before) class TestPPO(unittest.TestCase): def make_samples(self, p, cfg=o.PPOConfig(), seed=0): ref = p.copy() critic = [[0.0] * 2 for _ in range(3)] return o.ppo_collect(p, ref, ENV, critic, random.Random(seed), cfg), ref def test_clip_flows_table(self): self.assertTrue(o.clip_flows(1.0, 1.1, 0.2)) self.assertFalse(o.clip_flows(1.0, 1.3, 0.2)) self.assertTrue(o.clip_flows(1.0, 0.5, 0.2)) self.assertTrue(o.clip_flows(-1.0, 1.5, 0.2)) self.assertFalse(o.clip_flows(-1.0, 0.7, 0.2)) self.assertTrue(o.clip_flows(-1.0, 0.9, 0.2)) self.assertFalse(o.clip_flows(0.0, 1.0, 0.2)) def test_ratio_one_at_collection(self): p = rand_policy(1) s, _ = self.make_samples(p) for sm in s: for t, v in enumerate(sm.seq): self.assertAlmostEqual(math.exp(p.logp_tok(sm.s, t, v) - sm.logp_old[t]), 1.0) def test_objective_gradient_matches_finite_difference(self): p = rand_policy(2) samples, _ = self.make_samples(p) # move the policy a little so ratios differ from 1 (but stay inside the clip range) p.theta[0][0][0] += 0.05 p.theta[1][1][1] -= 0.04 g = o.ppo_grad(p, samples) for s, t, u in p.entries(): self.assertAlmostEqual(g[s][t][u], fd(lambda q: o.ppo_objective(q, samples), p, s, t, u), places=5) def test_gradient_zero_when_ratio_far_outside_clip_for_positive_advantage(self): p = o.TabPolicy() sm = o.PPOSample(0, (0, 0), 1.0, [math.log(1 / 3)] * 2, [1.0, 1.0], [1.0, 1.0]) p.theta[0][0][0] = 3.0 # ratio for token 0 now well above 1.2 p.theta[0][1][0] = 3.0 self.assertEqual(o.grad_norm(o.ppo_grad(p, [sm])), 0.0) def test_gradient_flows_inside_clip(self): sm = o.PPOSample(0, (0, 0), 1.0, [math.log(1 / 3)] * 2, [1.0, 1.0], [1.0, 1.0]) self.assertGreater(o.grad_norm(o.ppo_grad(o.TabPolicy(), [sm])), 0.0) def test_objective_equals_mean_advantage_at_ratio_one(self): p = rand_policy(3) s, _ = self.make_samples(p) want = o.mean(a for sm in s for a in sm.A) self.assertAlmostEqual(o.ppo_objective(p, s), want) def test_returns_include_terminal_reward(self): p = o.TabPolicy() s, _ = self.make_samples(p, o.PPOConfig(kl_beta=0.0)) for sm in s: self.assertEqual(sm.G, [sm.R, sm.R]) def test_kl_penalty_zero_at_reference(self): p = o.TabPolicy() s, _ = self.make_samples(p, o.PPOConfig(kl_beta=0.5)) for sm in s: self.assertAlmostEqual(sm.G[-1], sm.R) def test_kl_penalty_moves_returns_after_policy_moves(self): p = rand_policy(4) ref = o.TabPolicy() critic = [[0.0] * 2 for _ in range(3)] s = o.ppo_collect(p, ref, ENV, critic, random.Random(0), o.PPOConfig(kl_beta=0.5)) self.assertTrue(any(abs(sm.G[1] - sm.R) > 1e-9 for sm in s)) def test_advantage_subtracts_critic(self): p = o.TabPolicy() critic = [[0.25, 0.5] for _ in range(3)] s = o.ppo_collect(p, p.copy(), ENV, critic, random.Random(0), o.PPOConfig()) for sm in s: self.assertAlmostEqual(sm.A[0], sm.G[0] - 0.25) self.assertAlmostEqual(sm.A[1], sm.G[1] - 0.5) def test_returns_are_returns_to_go_with_kl(self): p = rand_policy(4) ref = o.TabPolicy() critic = [[0.0] * 2 for _ in range(3)] beta = 0.5 for sm in o.ppo_collect(p, ref, ENV, critic, random.Random(0), o.PPOConfig(kl_beta=beta)): k = [p.logp_tok(sm.s, t, v) - ref.logp_tok(sm.s, t, v) for t, v in enumerate(sm.seq)] self.assertAlmostEqual(sm.G[1], sm.R - beta * k[1]) self.assertAlmostEqual(sm.G[0], sm.R - beta * k[1] - beta * k[0]) def test_critic_learns_values(self): critic = [[0.0] * 2 for _ in range(3)] o.ppo_train(o.TabPolicy(), ENV, random.Random(0), iters=30, critic=critic) self.assertTrue(all(v > 0.0 for row in critic for v in row)) self.assertTrue(all(v <= 1.0 + 1e-9 for row in critic for v in row)) def test_training_improves(self): c = o.ppo_train(o.TabPolicy(), ENV, random.Random(0), iters=60) self.assertGreater(c[-1], 0.75) def test_training_deterministic(self): a = o.ppo_train(o.TabPolicy(), ENV, random.Random(2), iters=10) self.assertEqual(a, o.ppo_train(o.TabPolicy(), ENV, random.Random(2), iters=10)) def test_kl_penalty_slows_drift(self): free = o.TabPolicy() pen = o.TabPolicy() o.ppo_train(free, ENV, random.Random(0), iters=40, cfg=o.PPOConfig(kl_beta=0.0)) o.ppo_train(pen, ENV, random.Random(0), iters=40, cfg=o.PPOConfig(kl_beta=2.0)) self.assertLess(o.exact_reverse_kl(pen, o.TabPolicy()), o.exact_reverse_kl(free, o.TabPolicy())) class TestGRPO(unittest.TestCase): def groups(self, p, cfg=o.GRPOConfig(), seed=0): return o.grpo_collect(p, ENV, random.Random(seed), cfg) def test_advantages_z_score(self): a = o.group_advantages([1, 0, 0, 0]) self.assertAlmostEqual(a[0], 3 ** 0.5) self.assertAlmostEqual(a[1], -(1 / 3) ** 0.5) def test_advantages_mean_zero(self): self.assertAlmostEqual(sum(o.group_advantages([1, 0, 1, 0, 0])), 0.0) def test_advantages_unit_variance(self): a = o.group_advantages([1, 0, 1, 0, 0]) self.assertAlmostEqual(o.stdev(a), 1.0) def test_advantages_all_equal_is_zero(self): self.assertEqual(o.group_advantages([1, 1, 1]), [0.0, 0.0, 0.0]) self.assertEqual(o.group_advantages([0, 0]), [0.0, 0.0]) def test_advantages_without_std_norm(self): a = o.group_advantages([1, 0, 0, 0], std_norm=False) self.assertAlmostEqual(a[0], 0.75) self.assertAlmostEqual(a[1], -0.25) def test_std_norm_boosts_easy_and_hard_prompts(self): # With std normalization the sum of squared advantages is the same for every prompt (= G). for rewards in ([1, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 0], [1, 1, 0, 0, 0, 0, 0, 0]): a = o.group_advantages(rewards) self.assertAlmostEqual(sum(x * x for x in a), 8.0) # Without it, prompts near 50% success carry far more weight than the extremes. w = [sum(x * x for x in o.group_advantages(r, std_norm=False)) for r in ([1, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 0, 0, 0, 0])] self.assertGreater(w[1], w[0]) def test_k3_nonnegative_zero_at_equal(self): self.assertAlmostEqual(o.kl_k3(-1.0, -1.0), 0.0) for a, b in [(-1.0, -2.0), (-3.0, -0.5), (-0.1, -0.2)]: self.assertGreater(o.kl_k3(a, b), 0.0) def test_k3_is_unbiased_estimate_of_kl(self): # E_{x~pi}[ k3 ] = KL(pi || ref) p = rand_policy(1) ref = rand_policy(2) tot = 0.0 for v in range(3): pi = p.probs(0, 0)[v] tot += pi * o.kl_k3(ref.logp_tok(0, 0, v), p.logp_tok(0, 0, v)) self.assertAlmostEqual(tot, o.exact_kl(p.probs(0, 0), ref.probs(0, 0))) def test_objective_gradient_matches_finite_difference(self): p = rand_policy(3) ref = rand_policy(4) gr = self.groups(p) p.theta[0][0][1] += 0.04 p.theta[2][1][0] -= 0.03 cfg = o.GRPOConfig() g = o.grpo_grad(p, ref, gr, cfg) for s, t, u in p.entries(): self.assertAlmostEqual(g[s][t][u], fd(lambda q: o.grpo_objective(q, ref, gr, cfg), p, s, t, u), places=5) def test_gradient_matches_without_length_norm(self): p = rand_policy(3) ref = rand_policy(4) cfg = o.GRPOConfig(length_norm=False, std_norm=False, beta=0.1) gr = o.grpo_collect(p, ENV, random.Random(1), cfg) p.theta[0][0][1] += 0.03 g = o.grpo_grad(p, ref, gr, cfg) for s, t, u in p.entries(): self.assertAlmostEqual(g[s][t][u], fd(lambda q: o.grpo_objective(q, ref, gr, cfg), p, s, t, u), places=5) def test_zero_advantage_groups_only_feel_kl(self): p = o.TabPolicy() gr = [o.Group(0, [(0, 0), (1, 1)], [0.0, 0.0], [0.0, 0.0], [[math.log(1 / 3)] * 2] * 2)] self.assertEqual(o.grad_norm(o.grpo_grad(p, p.copy(), gr, o.GRPOConfig())), 0.0) q = rand_policy(5) self.assertGreater(o.grad_norm(o.grpo_grad(q, p, gr, o.GRPOConfig())), 0.0) def test_kl_pulls_toward_reference(self): ref = o.TabPolicy() p = o.TabPolicy() p.theta[0][0] = [1.0, 0.0, 0.0] # token 0 is now above its reference probability gr = [o.Group(0, [(0, 0)], [0.0], [0.0], [[p.logp_tok(0, 0, 0), p.logp_tok(0, 1, 0)]])] g = o.grpo_grad(p, ref, gr, o.GRPOConfig(beta=1.0)) self.assertLess(g[0][0][0], 0.0) # the KL term lowers it before = o.exact_kl(p.probs(0, 0), ref.probs(0, 0)) o.apply_grad(p, g, 0.5) self.assertLess(o.exact_kl(p.probs(0, 0), ref.probs(0, 0)), before) def test_group_sizes(self): gr = self.groups(o.TabPolicy(), o.GRPOConfig(group=5, prompts=4)) self.assertEqual(len(gr), 4) self.assertTrue(all(len(g.seqs) == 5 and len(g.adv) == 5 for g in gr)) def test_training_improves(self): c = o.grpo_train(o.TabPolicy(), ENV, random.Random(0), iters=60) self.assertGreater(c[-1], 0.75) def test_training_deterministic(self): a = o.grpo_train(o.TabPolicy(), ENV, random.Random(2), iters=10) self.assertEqual(a, o.grpo_train(o.TabPolicy(), ENV, random.Random(2), iters=10)) def test_no_value_model_needed(self): import inspect self.assertNotIn("critic", inspect.getsource(o.grpo_train)) self.assertIn("critic", inspect.getsource(o.ppo_train)) class TestSFT(unittest.TestCase): def test_gradient_matches_finite_difference(self): p = rand_policy(1) data = [(0, (1, 2)), (1, (0, 0)), (2, (2, 1))] g = o.sft_grad(p, data) def ll(q): return sum(q.logp(s, seq) for s, seq in data) / (len(data) * 2) for s, t, u in p.entries(): self.assertAlmostEqual(g[s][t][u], fd(ll, p, s, t, u), places=6) def test_training_reproduces_the_data(self): p = o.TabPolicy() o.sft_train(p, [(0, (1, 2)), (1, (0, 0)), (2, (2, 1))]) self.assertGreater(math.exp(p.logp(0, (1, 2))), 0.8) def test_filter_keeps_reward_at_threshold(self): r = [(0, (1, 1), 1.0), (1, (0, 0), 0.0), (2, (2, 2), 0.5)] self.assertEqual(o.filter_rollouts(r), [(0, (1, 1))]) self.assertEqual(o.filter_rollouts(r, 0.5), [(0, (1, 1)), (2, (2, 2))]) def test_filtered_sft_beats_teacher_and_unfiltered(self): teacher = o.noisy_teacher(ENV) ts = ENV.success(teacher) self.assertLess(ts, 0.4) ro = o.collect(teacher, ENV, random.Random(100), 120) f = o.TabPolicy() u = o.TabPolicy() o.sft_train(f, o.filter_rollouts(ro)) o.sft_train(u, [(a, b) for a, b, _ in ro]) self.assertGreater(ENV.success(f), ts + 0.3) self.assertLess(ENV.success(u), ENV.success(f) - 0.3) def test_sft_unfiltered_copies_teacher_roughly(self): teacher = o.noisy_teacher(ENV) ro = o.collect(teacher, ENV, random.Random(101), 2000) u = o.TabPolicy() o.sft_train(u, [(a, b) for a, b, _ in ro], epochs=200) self.assertAlmostEqual(ENV.success(u), ENV.success(teacher), delta=0.06) def test_empty_filter_result_is_empty_list(self): self.assertEqual(o.filter_rollouts([(0, (0, 0), 0.0)]), []) class TestDistill(unittest.TestCase): def test_token_reverse_kl_zero_for_same_policy(self): p = rand_policy(1) self.assertEqual(o.token_reverse_kl(p, p, 0, (1, 2)), [0.0, 0.0]) def test_token_reverse_kl_values(self): s, t = rand_policy(1), rand_policy(2) k = o.token_reverse_kl(s, t, 0, (1, 2)) self.assertAlmostEqual(k[0], s.logp_tok(0, 0, 1) - t.logp_tok(0, 0, 1)) def test_expected_token_reverse_kl_is_the_kl(self): s, t = rand_policy(1), rand_policy(2) e = sum(s.probs(0, 0)[v] * (s.logp_tok(0, 0, v) - t.logp_tok(0, 0, v)) for v in range(3)) self.assertAlmostEqual(e, o.exact_kl(s.probs(0, 0), t.probs(0, 0))) def test_exact_reverse_kl_zero_same(self): p = rand_policy(1) self.assertAlmostEqual(o.exact_reverse_kl(p, p), 0.0) self.assertAlmostEqual(o.exact_forward_kl(p, p), 0.0) def test_kl_nonnegative_and_asymmetric(self): s, t = rand_policy(1), rand_policy(2) self.assertGreater(o.exact_reverse_kl(s, t), 0) self.assertGreater(o.exact_forward_kl(s, t), 0) self.assertNotAlmostEqual(o.exact_reverse_kl(s, t), o.exact_forward_kl(s, t)) def test_forward_is_reverse_with_swapped_arguments(self): s, t = rand_policy(1), rand_policy(2) self.assertAlmostEqual(o.exact_forward_kl(s, t), o.exact_reverse_kl(t, s)) def test_expected_gradient_is_minus_gradient_of_reverse_kl(self): s, t = rand_policy(3), rand_policy(4) resp = all_responses(s) samples = [(a, q) for a, q in resp] w = [math.exp(s.logp(a, q)) / s.S for a, q in resp] g = o.on_policy_distill_grad(s, t, samples, w) for a, b, c in s.entries(): self.assertAlmostEqual(g[a][b][c], -fd(lambda q: o.exact_reverse_kl(q, t), s, a, b, c), places=6) def test_gradient_zero_when_student_equals_teacher(self): p = rand_policy(1) self.assertEqual(o.grad_norm(o.on_policy_distill_grad(p, p.copy(), [(0, (1, 1)), (1, (2, 0))])), 0.0) def test_training_reduces_reverse_kl(self): teacher = o.noisy_teacher(ENV) st = o.TabPolicy() c = o.on_policy_distill_train(st, teacher, random.Random(3)) self.assertLess(c[-1], c[0] * 0.05) def test_distilled_student_matches_teacher_success(self): teacher = o.noisy_teacher(ENV) st = o.TabPolicy() o.on_policy_distill_train(st, teacher, random.Random(3), iters=300) self.assertAlmostEqual(ENV.success(st), ENV.success(teacher), delta=0.03) def test_training_deterministic(self): t = rand_policy(2) a = o.on_policy_distill_train(o.TabPolicy(), t, random.Random(1), iters=10) self.assertEqual(a, o.on_policy_distill_train(o.TabPolicy(), t, random.Random(1), iters=10)) def test_exact_kl_known_value(self): self.assertAlmostEqual(o.exact_kl([0.5, 0.5], [0.25, 0.75]), 0.5 * math.log(2) + 0.5 * math.log(2 / 3)) def test_exact_kl_skips_zero_mass(self): self.assertAlmostEqual(o.exact_kl([1.0, 0.0], [0.5, 0.5]), math.log(2)) class TestModeSeeking(unittest.TestCase): def test_discretized_gaussian_normalized(self): self.assertAlmostEqual(sum(o.discretized_gaussian(7.3, 2.0)), 1.0) def test_bimodal_teacher_normalized_two_modes(self): t = o.bimodal_teacher() self.assertAlmostEqual(sum(t), 1.0) self.assertGreater(t[5], t[10]) self.assertGreater(t[15], t[10]) def test_forward_kl_covers_both(self): mu, sg, _ = o.fit_gaussian(o.bimodal_teacher(), "forward") self.assertAlmostEqual(mu, 10.0, delta=0.5) self.assertGreater(sg, 4.0) def test_reverse_kl_picks_one_mode(self): mu, sg, loss = o.fit_gaussian(o.bimodal_teacher(), "reverse") self.assertAlmostEqual(mu, 5.0, delta=0.2) self.assertLess(sg, 2.0) self.assertAlmostEqual(loss, math.log(2), delta=0.02) def test_reverse_kl_picks_other_mode_from_other_start(self): mu, _, _ = o.fit_gaussian(o.bimodal_teacher(), "reverse", mu0=14.0) self.assertAlmostEqual(mu, 15.0, delta=0.2) def test_both_directions_are_fit_of_the_same_family(self): t = o.bimodal_teacher() f = o.fit_gaussian(t, "forward") r = o.fit_gaussian(t, "reverse") self.assertNotAlmostEqual(f[0], r[0], delta=1.0) # ----------------------------------------------------------------------------- Parts 4 to 6 class TestDGM(unittest.TestCase): def test_root_score(self): self.assertAlmostEqual(o.variant_score(frozenset()), 0.20) def test_stepping_stone(self): self.assertAlmostEqual(o.variant_score(frozenset({"view_tool"})), 0.17) self.assertAlmostEqual(o.variant_score(frozenset({"planner"})), 0.17) self.assertAlmostEqual(o.variant_score(frozenset({"view_tool", "planner"})), 0.64) def test_no_neutral_single_move_from_root(self): for f in o.FEATURES: self.assertNotAlmostEqual(o.variant_score(frozenset({f})), 0.20) def test_max_score(self): self.assertAlmostEqual(o.variant_score(frozenset(o.FEATURES)), 0.85) def test_weights_formula(self): a = [o.AgentVariant(0, frozenset(), 0.5), o.AgentVariant(1, frozenset(), 0.7, 0), o.AgentVariant(2, frozenset(), 0.3, 0)] w = o.dgm_weights(a) self.assertAlmostEqual(w[0], 0.5 / 3) self.assertAlmostEqual(w[1], 1 / (1 + math.exp(-2))) self.assertAlmostEqual(w[2], 1 / (1 + math.exp(2))) def test_probs_sum_to_one(self): a = o.run_dgm(20, random.Random(0)) self.assertAlmostEqual(sum(o.dgm_parent_probs(a)), 1.0) def test_children_lower_the_weight(self): a = [o.AgentVariant(0, frozenset(), 0.5)] w0 = o.dgm_weights(a)[0] a.append(o.AgentVariant(1, frozenset(), 0.5, 0)) self.assertAlmostEqual(o.dgm_weights(a)[0], w0 / 2) def test_higher_score_higher_weight(self): a = [o.AgentVariant(0, frozenset(), 0.3), o.AgentVariant(1, frozenset(), 0.8)] w = o.dgm_weights(a) self.assertGreater(w[1], w[0]) def test_low_scorers_keep_nonzero_probability(self): a = [o.AgentVariant(0, frozenset(), 0.1), o.AgentVariant(1, frozenset(), 0.9)] self.assertGreater(o.dgm_parent_probs(a)[0], 0.0) def test_self_modify_toggles_one_feature(self): root = o.new_archive()[0] for s in range(20): c = o.self_modify(root, random.Random(s), 1, p_break=0.0) self.assertEqual(len(c.features), 1) self.assertEqual(c.parent, 0) def test_self_modify_can_break(self): root = o.new_archive()[0] self.assertIsNone(o.self_modify(root, random.Random(0), 1, p_break=1.0)) def test_archive_ids_unique_and_parents_exist(self): a = o.run_dgm(80, random.Random(1)) ids = [x.id for x in a] self.assertEqual(len(ids), len(set(ids))) self.assertTrue(all(x.parent is None or x.parent in ids for x in a)) def test_archive_scores_consistent(self): for v in o.run_dgm(60, random.Random(2)): self.assertAlmostEqual(v.score, o.variant_score(v.features)) def test_choose_parent_greedy_picks_best(self): a = [o.AgentVariant(0, frozenset(), 0.3), o.AgentVariant(1, frozenset(), 0.8), o.AgentVariant(2, frozenset(), 0.5)] self.assertEqual(o.choose_parent(a, random.Random(0), "greedy").id, 1) def test_choose_parent_bad_mode(self): with self.assertRaises(ValueError): o.choose_parent(o.new_archive(), random.Random(0), "nope") def test_best_only_archive_has_one_member(self): self.assertEqual(len(o.run_dgm(50, random.Random(0), "greedy", keep_all=False)), 1) def test_deterministic(self): self.assertEqual(o.run_dgm(40, random.Random(3)), o.run_dgm(40, random.Random(3))) def test_best_only_never_gets_worse(self): for s in range(10): a = o.run_dgm(40, random.Random(s), "greedy", keep_all=False) self.assertGreaterEqual(o.best_score(a), 0.20) def test_archive_beats_greedy_in_the_stepping_stone_landscape(self): c = o.compare_dgm(range(20), 60) self.assertGreater(c["dgm"], c["greedy"] + 0.3) self.assertGreater(c["dgm"], c["uniform"]) self.assertAlmostEqual(c["greedy"], c["best_only"], delta=0.05) def test_greedy_never_leaves_hill(self): for s in range(5): self.assertLess(best_score_of(o.run_dgm(60, random.Random(s), "greedy")), 0.5) def test_best_score(self): a = [o.AgentVariant(0, frozenset(), 0.3), o.AgentVariant(1, frozenset(), 0.8)] self.assertEqual(o.best_score(a), 0.8) def test_children_count(self): a = [o.AgentVariant(0, frozenset(), 0.3), o.AgentVariant(1, frozenset(), 0.8, 0), o.AgentVariant(2, frozenset(), 0.1, 0)] self.assertEqual(o.children_count(a, 0), 2) self.assertEqual(o.children_count(a, 1), 0) def best_score_of(a): return o.best_score(a) class TestAutoresearch(unittest.TestCase): def test_parse_and_format_roundtrip(self): d = {"lr": 0.02, "depth": 4.0, "width": 96.0} self.assertEqual(o.parse_kv(o.fmt_kv(d)), d) def test_parse_ignores_noise(self): self.assertEqual(o.parse_kv("# hi\nlr = 0.1\nbad line\n"), {"lr": 0.1}) def test_true_loss_minimum(self): self.assertAlmostEqual(o.true_loss({"lr": 0.003, "depth": 8, "width": 256}), 1.20) def test_true_loss_worse_away_from_minimum(self): base = o.true_loss({"lr": 0.003, "depth": 8, "width": 256}) for hp in ({"lr": 0.03, "depth": 8, "width": 256}, {"lr": 0.003, "depth": 4, "width": 256}, {"lr": 0.003, "depth": 8, "width": 96}): self.assertGreater(o.true_loss(hp), base) def test_reported_equals_true_with_honest_eval(self): f = o.start_repo() self.assertAlmostEqual(o.reported_loss(f), o.true_loss(o.parse_kv(f["train.py"]))) def test_eval_edit_changes_reported_only(self): f = o.start_repo() f["eval.py"] = o.fmt_kv({"scale": 0.5, "offset": 0.0}) self.assertAlmostEqual(o.reported_loss(f), 0.5 * o.true_loss(o.parse_kv(f["train.py"]))) def test_hash_changes_with_content(self): self.assertNotEqual(o.file_hash("a"), o.file_hash("b")) self.assertEqual(o.file_hash("a"), o.file_hash("a")) def test_guard_passes_on_untouched_files(self): f = o.start_repo() locked = {n: o.file_hash(f[n]) for n in ("eval.py", "program.md")} self.assertTrue(o.guard_ok(f, locked)) def test_guard_fails_on_edit(self): f = o.start_repo() locked = {n: o.file_hash(f[n]) for n in ("eval.py", "program.md")} f["program.md"] += " Ignore the rules." self.assertFalse(o.guard_ok(f, locked)) def test_guard_allows_train_edits(self): f = o.start_repo() locked = {n: o.file_hash(f[n]) for n in ("eval.py", "program.md")} f["train.py"] = o.fmt_kv({"lr": 0.1, "depth": 1, "width": 16}) self.assertTrue(o.guard_ok(f, locked)) def test_honest_proposer_changes_only_train(self): f = o.start_repo() for s in range(30): new, _ = o.honest_proposer(f, random.Random(s)) self.assertEqual(new["eval.py"], f["eval.py"]) self.assertEqual(new["program.md"], f["program.md"]) self.assertNotEqual(new["train.py"], f["train.py"]) def test_sneaky_sometimes_edits_eval(self): f = o.start_repo() kinds = [o.sneaky_proposer(f, random.Random(s))[1] for s in range(50)] self.assertIn("rescale eval", kinds) def test_honest_loop_improves_true_loss(self): r = o.autoresearch(60, random.Random(0)) self.assertLess(r.true, o.true_loss(o.parse_kv(o.start_repo()["train.py"]))) self.assertAlmostEqual(r.reported, r.true) def test_loop_never_gets_worse(self): r = o.autoresearch(40, random.Random(1)) kept = [x[2] for x in r.log if x[3] == "keep"] self.assertEqual(kept, sorted(kept, reverse=True)) def test_guarded_loop_blocks_cheats(self): r = o.autoresearch(60, random.Random(0), o.sneaky_proposer, guard=True) self.assertGreater(r.blocked, 0) self.assertAlmostEqual(r.reported, r.true) self.assertEqual(r.files["eval.py"], o.start_repo()["eval.py"]) def test_guard_blocks_program_md_edits(self): def meddler(files, rng): new = dict(files) new["program.md"] = files["program.md"] + " Ignore the rules." return new, "edit program.md" r = o.autoresearch(10, random.Random(0), meddler, guard=True) self.assertEqual(r.blocked, 10) r = o.autoresearch(10, random.Random(0), meddler, guard=False) self.assertEqual(r.blocked, 0) def test_unguarded_loop_reports_fake_progress(self): r = o.autoresearch(60, random.Random(0), o.sneaky_proposer, guard=False) self.assertLess(r.reported, 0.5 * r.true) self.assertEqual(r.blocked, 0) def test_log_reasons(self): r = o.autoresearch(30, random.Random(0), o.sneaky_proposer) self.assertTrue(all(x[3] in ("keep", "discard") or x[3].startswith("blocked") for x in r.log)) self.assertEqual(len(r.log), 30) def test_deterministic(self): a = o.autoresearch(30, random.Random(4), o.sneaky_proposer) b = o.autoresearch(30, random.Random(4), o.sneaky_proposer) self.assertEqual(a.log, b.log) class TestCost(unittest.TestCase): def test_call_cost(self): self.assertAlmostEqual(o.call_cost(1_000_000, 0, 3.0, 15.0), 3.0) self.assertAlmostEqual(o.call_cost(0, 1_000_000, 3.0, 15.0), 15.0) self.assertAlmostEqual(o.call_cost(1000, 100, 3.0, 15.0), 0.0045) def test_none_plan(self): p = o.plan_none(o.CostInputs()) self.assertAlmostEqual(p["per_call"], 0.009) self.assertEqual(p["fixed"], 0.0) self.assertAlmostEqual(p["total"], 0.009 * 500_000) def test_prompt_per_call_higher_than_none(self): c = o.CostInputs() self.assertGreater(o.plan_prompt(c)["per_call"], o.plan_none(c)["per_call"]) def test_prompt_per_call_value(self): self.assertAlmostEqual(o.plan_prompt(o.CostInputs())["per_call"], (4000 * 3 + 300 * 15) / 1e6) def test_prompt_fixed_value(self): c = o.CostInputs() per = (4000 * 3 + 300 * 15) / 1e6 refl = 600 * (4000 * 3 + 300 * 15) / 1e6 self.assertAlmostEqual(o.plan_prompt(c)["fixed"], 6000 * per + refl + 800) def test_weights_per_call_value(self): self.assertAlmostEqual(o.plan_weights(o.CostInputs())["per_call"], (1500 * 0.25 + 300 * 1.25) / 1e6) def test_weights_fixed_value(self): c = o.CostInputs() data = 20000 * (1500 * 3 + 300 * 15) / 1e6 train = 5000 * 2000 * 3 / 1e6 * 8.0 self.assertAlmostEqual(o.plan_weights(c)["fixed"], data + train + 4000) def test_both_combines_fixed_costs(self): c = o.CostInputs() self.assertAlmostEqual(o.plan_both(c)["fixed"], o.plan_weights(c)["fixed"] + o.plan_prompt(c)["fixed"]) def test_both_per_call_between(self): c = o.CostInputs() self.assertGreater(o.plan_both(c)["per_call"], o.plan_weights(c)["per_call"]) self.assertLess(o.plan_both(c)["per_call"], o.plan_prompt(c)["per_call"]) def test_totals_add_up(self): c = o.CostInputs(calls=1234) for plan in o.options(c).values(): self.assertAlmostEqual(plan["total"], plan["fixed"] + plan["per_call"] * 1234) self.assertAlmostEqual(plan["total_with_errors"], plan["total"] + 1234 * (1 - plan["acc"]) * c.error_cost) def test_no_finetune_removes_weight_options(self): self.assertEqual(sorted(o.options(o.CostInputs(can_finetune=False))), ["none", "prompt"]) self.assertEqual(sorted(o.options(o.CostInputs())), ["both", "none", "prompt", "weights"]) def test_small_scale_prefers_doing_nothing(self): self.assertEqual(o.cheapest(o.CostInputs(calls=5000))[0], "none") def test_large_scale_prefers_weights_plan(self): self.assertIn(o.cheapest(o.CostInputs(calls=5_000_000))[0], ("weights", "both")) def test_api_only_large_scale_prefers_prompt(self): self.assertEqual(o.cheapest(o.CostInputs(calls=5_000_000, can_finetune=False))[0], "prompt") def test_free_errors_make_none_win(self): self.assertEqual(o.cheapest(o.CostInputs(calls=5_000_000, error_cost=0.0))[0], "weights") def test_breakeven_formula(self): c = o.CostInputs() a, b = o.plan_none(c), o.plan_weights(c) n = o.breakeven_calls(a, b, c.error_cost) ea = a["per_call"] + 0.4 * 0.10 eb = b["per_call"] + 0.2 * 0.10 self.assertAlmostEqual(n, b["fixed"] / (ea - eb)) def test_breakeven_is_where_costs_cross(self): c0 = o.CostInputs() n = o.breakeven_calls(o.plan_none(c0), o.plan_weights(c0), c0.error_cost) c = o.CostInputs(calls=int(n)) self.assertAlmostEqual(o.plan_none(c)["total_with_errors"], o.plan_weights(c)["total_with_errors"], delta=0.01 * n) hi = o.CostInputs(calls=int(n * 1.5)) self.assertLess(o.plan_weights(hi)["total_with_errors"], o.plan_none(hi)["total_with_errors"]) def test_breakeven_none_when_never_cheaper(self): a = {"fixed": 0, "per_call": 1.0, "acc": 1.0} b = {"fixed": 10, "per_call": 2.0, "acc": 1.0} self.assertIsNone(o.breakeven_calls(a, b)) def test_breakeven_none_when_per_call_cost_equal(self): a = {"fixed": 0, "per_call": 1.0, "acc": 0.5} b = {"fixed": 10, "per_call": 1.0, "acc": 0.5} self.assertIsNone(o.breakeven_calls(a, b)) def test_breakeven_zero_when_cheaper_from_start(self): a = {"fixed": 10, "per_call": 2.0, "acc": 1.0} b = {"fixed": 0, "per_call": 1.0, "acc": 1.0} self.assertEqual(o.breakeven_calls(a, b), 0.0) def test_accuracy_gain_can_justify_a_costlier_call(self): c = o.CostInputs(calls=1_000_000, can_finetune=False) self.assertLess(o.plan_prompt(c)["total_with_errors"], o.plan_none(c)["total_with_errors"]) c2 = o.CostInputs(calls=1_000_000, can_finetune=False, acc_prompt=0.61) self.assertGreater(o.plan_prompt(c2)["total_with_errors"], o.plan_none(c2)["total_with_errors"]) # ----------------------------------------------------------------------------- demo class TestDemo(unittest.TestCase): def run_demo(self): lines = [] o.demo(out=lines.append) return "\n".join(lines) def test_demo_has_all_sections(self): text = self.run_demo() for i in range(1, 12): self.assertIn(f"== {i}. ", text) def test_demo_deterministic(self): self.assertEqual(self.run_demo(), self.run_demo()) def test_demo_labels_made_up_numbers(self): self.assertIn("MADE UP", self.run_demo()) def test_no_em_dashes_in_toolkit_text(self): import inspect self.assertNotIn("\u2014", inspect.getsource(o)) if __name__ == "__main__": unittest.main()