formal-mas / tests /test_reference_set.py

Formal-MAS source code

dfe1826 verified about 19 hours ago

6.96 kB

"""Offline tests for the reference set Q post-processing (toy corpus and sources)."""

from __future__ import annotations

from pathlib import Path

import pytest

from formal_mas.retrieval.corpus import SourceIndex
from formal_mas.retrieval.reference_set import (
    QueryItem,
    build_candidates,
    gold_sets,
    infrastructure,
    is_concrete,
    proof_text,
    select,
    split,
    summary,
)

TOY_LINES = [
    "/-- Doc with := inside. -/",                                     # 1
    "theorem Nat.foo (n : ℕ) (h : a < b) : n = n := by",              # 2
    "  -- uses bar_lemma in a comment only",                          # 3
    "  have := h.not_ge",                                             # 4
    "  simp [Nat.baz, mul_one]",                                      # 5
    "theorem Nat.one_line : True := Nat.baz",                         # 6
    "theorem Nat.eqns : ℕ → True",                                    # 7
    "  | 0 => Nat.baz",                                               # 8
    "  | _ + 1 => trivial",                                           # 9
]
MODULE = "Mathlib.Data.Nat.Toy"


def rec(name, area="number_theory", module=MODULE, sig="(n : ℕ) : n = n", origin="manual"):
    return {"name": name, "signature": sig, "area": area, "module": module, "origin": origin}


CORPUS = {r["name"]: r for r in [
    rec("Nat.foo"), rec("Nat.one_line"), rec("Nat.eqns"),
    rec("Nat.baz"), rec("mul_one", area="algebra", module="Mathlib.Algebra.Group.Defs"),
    rec("LT.lt.not_ge", area="other", module="Mathlib.Order.Basic"),
    rec("Nat.simp_only_lemma"), rec("Nat.bar_lemma"),
    rec("Mathlib.Tactic.Ring.add_pf", area="other", module="Mathlib.Tactic.Ring.Basic"),
]}


def entry(name, consts, rng, sel):
    return {"name": name, "module": MODULE, "range": rng, "sel": sel, "consts": consts}


FOO = entry("Nat.foo",
            ["Nat.baz", "mul_one", "LT.lt.not_ge", "Nat.simp_only_lemma", "Nat.bar_lemma",
             "Mathlib.Tactic.Ring.add_pf", "Eq.mpr", "Nat.foo"],
            [1, 0, 5, 27], [2, 8, 2, 15])
ONE_LINE = entry("Nat.one_line", ["Nat.baz"], [6, 0, 6, 39], [6, 8, 6, 20])
EQNS = entry("Nat.eqns", ["Nat.baz"], [7, 0, 9, 22], [7, 8, 7, 16])


@pytest.fixture
def sources(tmp_path: Path) -> SourceIndex:
    path = tmp_path / "Mathlib" / "Data" / "Nat" / "Toy.lean"
    path.parent.mkdir(parents=True)
    path.write_text("\n".join(TOY_LINES) + "\n", encoding="utf-8")
    return SourceIndex(tmp_path)


def test_proof_text_skips_docstring_and_comments(sources):
    text = proof_text(FOO, sources)
    assert text.strip().startswith("by")
    assert "Doc with" not in text and "bar_lemma" not in text
    assert "Nat.baz" in proof_text(ONE_LINE, sources)
    assert "| 0 => Nat.baz" in proof_text(EQNS, sources)


def test_proof_text_missing_source(sources):
    assert proof_text({**FOO, "module": "Mathlib.Missing"}, sources) is None
    assert proof_text({**FOO, "range": None}, sources) is None


def test_gold_methods(sources):
    gold_c, gold_b = gold_sets(FOO, CORPUS, {"mul_one"}, sources)
    # (b): corpus constants, no self, no Mathlib.Tactic, no infrastructure, no non-corpus (Eq.mpr)
    assert gold_b == {"Nat.baz", "LT.lt.not_ge", "Nat.simp_only_lemma", "Nat.bar_lemma"}
    # (c): written in the proof (dot notation counts; comments and implicit simp lemmas do not)
    assert gold_c == {"Nat.baz", "LT.lt.not_ge"}


@pytest.mark.parametrize(
    ("const", "text", "written"),
    [
        ("Finset.ext", "by ext x; simp", False),         # bare tactic
        ("Finset.ext", "Finset.ext fun x => h x", True),
        ("Eq.symm", "by symm; exact h", False),
        ("Eq.symm", "h.symm", True),
        ("by_contra", "by by_contra h; exact h rfl", False),
        ("Nat.le_of_lt", "le_of_lt h", True),            # ordinary lemma: bare name counts
    ],
)
def test_tactic_names_need_dotted_use(const, text, written):
    from formal_mas.retrieval.reference_set import is_written, written_tokens
    assert is_written(const, *written_tokens(text)) is written


def test_infrastructure_threshold():
    entries = [{"name": f"t{i}", "consts": ["mul_one"] + (["Nat.baz"] if i == 0 else [])} for i in range(10)]
    infra = infrastructure(entries, CORPUS, 50.0)
    assert infra == {"mul_one": 100.0}
    assert infrastructure([], CORPUS, 1.0) == {}


@pytest.mark.parametrize(
    ("sig", "concrete"),
    [("(n : ℕ) : n = n", True), ("(s : Finset α) : s.card = s.card", True), ("(i : Fin n) : i = i", True),
     ("(x : ℝ) : 0 ≤ x ^ 2", True), ("[Fintype α] : True", False), ("{G : Type*} [Group G] (a : G) : a = a", False)],
)
def test_is_concrete(sig, concrete):
    assert is_concrete(sig) is concrete


def test_build_candidates_filters(sources):
    corpus = {**CORPUS, "Nat.one_line": rec("Nat.one_line", origin="to_additive")}
    cands = build_candidates([FOO, ONE_LINE, EQNS], corpus, {"mul_one"}, sources)
    assert [q.name for q in cands] == ["Nat.eqns", "Nat.foo"]  # to_additive dropped
    foo = cands[1]
    assert foo.gold == ["LT.lt.not_ge", "Nat.baz"] and foo.concrete and foo.query == "(n : ℕ) : n = n"
    assert len(foo.gold_b) == 4


def item(name, area, concrete=True, gold=("x",)):
    return QueryItem(name=name, query="q", area=area, module="M", concrete=concrete,
                     gold=list(gold), gold_b=list(gold))


def toy_pool():
    pool = []
    for area in ("algebra", "number_theory", "combinatorics"):
        pool += [item(f"{area}.c{i}", area, True, (f"g.{area}.{i}",)) for i in range(3)]
        pool += [item(f"{area}.a{i}", area, False, (f"g.{area}.a{i}",)) for i in range(5)]
    return pool


def test_select_balanced_concrete_first():
    chosen = select(toy_pool(), total=12, seed=0)
    s = summary(split(chosen, seed=0))
    assert s["by_area"] == {"algebra": 4, "number_theory": 4, "combinatorics": 4}
    # 3 concrete available per area, the 4th is filled with an abstract theorem
    assert s["by_type"] == {"concrete": 9, "abstract": 3}


def test_select_avoids_query_gold_conflicts():
    pool = [item("A", "algebra", gold=("B",)), item("B", "algebra", gold=("z",)), item("C", "algebra", gold=("y",))]
    for seed in range(5):
        chosen = select(pool, total=3, seed=seed)
        names = {q.name for q in chosen}
        assert not ({"A", "B"} <= names)  # B is A's gold: they never appear together
        assert all(not names & set(q.gold) for q in chosen)


def test_select_is_deterministic():
    assert select(toy_pool(), 9, seed=7) == select(toy_pool(), 9, seed=7)


def test_split_stratified():
    items = split(select(toy_pool(), total=24, seed=1), seed=1)
    s = summary(items)
    assert s["n"] == 24
    for area in ("algebra", "number_theory", "combinatorics"):
        dev = sum(1 for q in items if q.area == area and q.split == "dev")
        assert dev == round(2 / 3 * 8)
    assert s["by_split"] == {"dev": 15, "test": 9}
    assert split(items, seed=1) == split(items, seed=1)

	"""Offline tests for the reference set Q post-processing (toy corpus and sources)."""

	from __future__ import annotations

	from pathlib import Path

	import pytest

	from formal_mas.retrieval.corpus import SourceIndex
	from formal_mas.retrieval.reference_set import (
	QueryItem,
	build_candidates,
	gold_sets,
	infrastructure,
	is_concrete,
	proof_text,
	select,
	split,
	summary,
	)

	TOY_LINES = [
	"/-- Doc with := inside. -/", # 1
	"theorem Nat.foo (n : ℕ) (h : a < b) : n = n := by", # 2
	" -- uses bar_lemma in a comment only", # 3
	" have := h.not_ge", # 4
	" simp [Nat.baz, mul_one]", # 5
	"theorem Nat.one_line : True := Nat.baz", # 6
	"theorem Nat.eqns : ℕ → True", # 7
	" \| 0 => Nat.baz", # 8
	" \| _ + 1 => trivial", # 9
	]
	MODULE = "Mathlib.Data.Nat.Toy"


	def rec(name, area="number_theory", module=MODULE, sig="(n : ℕ) : n = n", origin="manual"):
	return {"name": name, "signature": sig, "area": area, "module": module, "origin": origin}


	CORPUS = {r["name"]: r for r in [
	rec("Nat.foo"), rec("Nat.one_line"), rec("Nat.eqns"),
	rec("Nat.baz"), rec("mul_one", area="algebra", module="Mathlib.Algebra.Group.Defs"),
	rec("LT.lt.not_ge", area="other", module="Mathlib.Order.Basic"),
	rec("Nat.simp_only_lemma"), rec("Nat.bar_lemma"),
	rec("Mathlib.Tactic.Ring.add_pf", area="other", module="Mathlib.Tactic.Ring.Basic"),
	]}


	def entry(name, consts, rng, sel):
	return {"name": name, "module": MODULE, "range": rng, "sel": sel, "consts": consts}


	FOO = entry("Nat.foo",
	["Nat.baz", "mul_one", "LT.lt.not_ge", "Nat.simp_only_lemma", "Nat.bar_lemma",
	"Mathlib.Tactic.Ring.add_pf", "Eq.mpr", "Nat.foo"],
	[1, 0, 5, 27], [2, 8, 2, 15])
	ONE_LINE = entry("Nat.one_line", ["Nat.baz"], [6, 0, 6, 39], [6, 8, 6, 20])
	EQNS = entry("Nat.eqns", ["Nat.baz"], [7, 0, 9, 22], [7, 8, 7, 16])


	@pytest.fixture
	def sources(tmp_path: Path) -> SourceIndex:
	path = tmp_path / "Mathlib" / "Data" / "Nat" / "Toy.lean"
	path.parent.mkdir(parents=True)
	path.write_text("\n".join(TOY_LINES) + "\n", encoding="utf-8")
	return SourceIndex(tmp_path)


	def test_proof_text_skips_docstring_and_comments(sources):
	text = proof_text(FOO, sources)
	assert text.strip().startswith("by")
	assert "Doc with" not in text and "bar_lemma" not in text
	assert "Nat.baz" in proof_text(ONE_LINE, sources)
	assert "\| 0 => Nat.baz" in proof_text(EQNS, sources)


	def test_proof_text_missing_source(sources):
	assert proof_text({**FOO, "module": "Mathlib.Missing"}, sources) is None
	assert proof_text({**FOO, "range": None}, sources) is None


	def test_gold_methods(sources):
	gold_c, gold_b = gold_sets(FOO, CORPUS, {"mul_one"}, sources)
	# (b): corpus constants, no self, no Mathlib.Tactic, no infrastructure, no non-corpus (Eq.mpr)
	assert gold_b == {"Nat.baz", "LT.lt.not_ge", "Nat.simp_only_lemma", "Nat.bar_lemma"}
	# (c): written in the proof (dot notation counts; comments and implicit simp lemmas do not)
	assert gold_c == {"Nat.baz", "LT.lt.not_ge"}


	@pytest.mark.parametrize(
	("const", "text", "written"),
	[
	("Finset.ext", "by ext x; simp", False), # bare tactic
	("Finset.ext", "Finset.ext fun x => h x", True),
	("Eq.symm", "by symm; exact h", False),
	("Eq.symm", "h.symm", True),
	("by_contra", "by by_contra h; exact h rfl", False),
	("Nat.le_of_lt", "le_of_lt h", True), # ordinary lemma: bare name counts
	],
	)
	def test_tactic_names_need_dotted_use(const, text, written):
	from formal_mas.retrieval.reference_set import is_written, written_tokens
	assert is_written(const, *written_tokens(text)) is written


	def test_infrastructure_threshold():
	entries = [{"name": f"t{i}", "consts": ["mul_one"] + (["Nat.baz"] if i == 0 else [])} for i in range(10)]
	infra = infrastructure(entries, CORPUS, 50.0)
	assert infra == {"mul_one": 100.0}
	assert infrastructure([], CORPUS, 1.0) == {}


	@pytest.mark.parametrize(
	("sig", "concrete"),
	[("(n : ℕ) : n = n", True), ("(s : Finset α) : s.card = s.card", True), ("(i : Fin n) : i = i", True),
	("(x : ℝ) : 0 ≤ x ^ 2", True), ("[Fintype α] : True", False), ("{G : Type*} [Group G] (a : G) : a = a", False)],
	)
	def test_is_concrete(sig, concrete):
	assert is_concrete(sig) is concrete


	def test_build_candidates_filters(sources):
	corpus = {**CORPUS, "Nat.one_line": rec("Nat.one_line", origin="to_additive")}
	cands = build_candidates([FOO, ONE_LINE, EQNS], corpus, {"mul_one"}, sources)
	assert [q.name for q in cands] == ["Nat.eqns", "Nat.foo"] # to_additive dropped
	foo = cands[1]
	assert foo.gold == ["LT.lt.not_ge", "Nat.baz"] and foo.concrete and foo.query == "(n : ℕ) : n = n"
	assert len(foo.gold_b) == 4


	def item(name, area, concrete=True, gold=("x",)):
	return QueryItem(name=name, query="q", area=area, module="M", concrete=concrete,
	gold=list(gold), gold_b=list(gold))


	def toy_pool():
	pool = []
	for area in ("algebra", "number_theory", "combinatorics"):
	pool += [item(f"{area}.c{i}", area, True, (f"g.{area}.{i}",)) for i in range(3)]
	pool += [item(f"{area}.a{i}", area, False, (f"g.{area}.a{i}",)) for i in range(5)]
	return pool


	def test_select_balanced_concrete_first():
	chosen = select(toy_pool(), total=12, seed=0)
	s = summary(split(chosen, seed=0))
	assert s["by_area"] == {"algebra": 4, "number_theory": 4, "combinatorics": 4}
	# 3 concrete available per area, the 4th is filled with an abstract theorem
	assert s["by_type"] == {"concrete": 9, "abstract": 3}


	def test_select_avoids_query_gold_conflicts():
	pool = [item("A", "algebra", gold=("B",)), item("B", "algebra", gold=("z",)), item("C", "algebra", gold=("y",))]
	for seed in range(5):
	chosen = select(pool, total=3, seed=seed)
	names = {q.name for q in chosen}
	assert not ({"A", "B"} <= names) # B is A's gold: they never appear together
	assert all(not names & set(q.gold) for q in chosen)


	def test_select_is_deterministic():
	assert select(toy_pool(), 9, seed=7) == select(toy_pool(), 9, seed=7)


	def test_split_stratified():
	items = split(select(toy_pool(), total=24, seed=1), seed=1)
	s = summary(items)
	assert s["n"] == 24
	for area in ("algebra", "number_theory", "combinatorics"):
	dev = sum(1 for q in items if q.area == area and q.split == "dev")
	assert dev == round(2 / 3 * 8)
	assert s["by_split"] == {"dev": 15, "test": 9}
	assert split(items, seed=1) == split(items, seed=1)