ai4math 0.1.0__py3-none-any.whl

ai4math/__init__.py

@@ -0,0 +1,130 @@
+"""ai4math — tools for AI-assisted mathematics.
+
+A small, dependency-light toolkit for the mechanics that every AI-for-mathematics
+project ends up rewriting: extracting a final answer from a model's reasoning,
+deciding whether it matches the reference, scoring sampled attempts, and calling
+a computer algebra system as a tool.
+
+The library covers both halves of the field as described by Yang et al. (2024,
+arXiv:2412.16075). The informal, natural-language half is implemented here and
+needs only SymPy. The formal half — proving theorems in an interactive theorem
+prover — is defined as a protocol with an optional Lean backend, because a proof
+assistant toolchain does not belong in a small package's dependencies.
+
+Quick start::
+
+    from ai4math import extract_answer, verify
+
+    text = "Adding gives 4, so the answer is \\\\boxed{4}."
+    answer = extract_answer(text)                # '4'
+    assert verify("4", answer)                   # graded correct
+
+Grading a model over a dataset::
+
+    from ai4math import evaluate, BUILTIN_EXAMPLES
+
+    report = evaluate(BUILTIN_EXAMPLES, my_model_fn, n_samples=4)
+    print(report.as_dict())                      # pass@1, pass@4, maj@n
+
+Layout:
+
+* :mod:`ai4math.answers` — extract and normalize final answers
+* :mod:`ai4math.verify` — decide answer equivalence
+* :mod:`ai4math.metrics` — ``pass@k``, majority vote, best-of-n
+* :mod:`ai4math.symbolic` — a guarded SymPy tool surface for models
+* :mod:`ai4math.prompts` — prompt templates with matching answer parsers
+* :mod:`ai4math.datasets` — problem records and loaders
+* :mod:`ai4math.formal` — theorem-prover protocol, Lean adapters, and a
+  generate-compile-repair proof loop
+* :mod:`ai4math.models` — adapters turning an LLM into a plain callable
+* :mod:`ai4math.pipeline` — end-to-end solve, grade and aggregate
+* :mod:`ai4math.demo` — ``python -m ai4math.demo`` for a Lean-verified proof
+
+The formal half is reached through :mod:`ai4math.formal` rather than re-exported
+here, keeping the top-level namespace to what works with SymPy alone.
+"""
+
+from __future__ import annotations
+
+from ai4math.answers import (
+    extract_all_boxed,
+    extract_answer,
+    extract_boxed,
+    extract_final_answer,
+    normalize_answer,
+    normalize_latex,
+)
+from ai4math.datasets import BUILTIN_EXAMPLES, MathProblem, load_json, load_jsonl
+from ai4math.metrics import (
+    EvalSummary,
+    accuracy,
+    best_of_n,
+    majority_vote,
+    pass_at_k,
+)
+from ai4math.pipeline import (
+    EvalReport,
+    ModelFn,
+    SolveResult,
+    evaluate,
+    solve,
+    solve_and_verify,
+)
+from ai4math.prompts import TEMPLATES, PromptTemplate, get_template
+from ai4math.symbolic import (
+    SymbolicError,
+    differentiate,
+    evaluate_expr,
+    factor_expr,
+    integrate_expr,
+    simplify_expr,
+    solve_equation,
+)
+from ai4math.verify import VerificationResult, verify, verify_equivalence
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "__version__",
+    # answers
+    "extract_answer",
+    "extract_boxed",
+    "extract_all_boxed",
+    "extract_final_answer",
+    "normalize_answer",
+    "normalize_latex",
+    # verify
+    "verify",
+    "verify_equivalence",
+    "VerificationResult",
+    # metrics
+    "accuracy",
+    "pass_at_k",
+    "majority_vote",
+    "best_of_n",
+    "EvalSummary",
+    # symbolic
+    "simplify_expr",
+    "solve_equation",
+    "differentiate",
+    "integrate_expr",
+    "factor_expr",
+    "evaluate_expr",
+    "SymbolicError",
+    # prompts
+    "PromptTemplate",
+    "TEMPLATES",
+    "get_template",
+    # datasets
+    "MathProblem",
+    "load_jsonl",
+    "load_json",
+    "BUILTIN_EXAMPLES",
+    # pipeline
+    "solve",
+    "solve_and_verify",
+    "evaluate",
+    "SolveResult",
+    "EvalReport",
+    "ModelFn",
+]

ai4math/answers/__init__.py

@@ -0,0 +1,33 @@
+"""Extracting and normalizing final answers from mathematical solutions."""
+
+from ai4math.answers.extract import (
+    BoxedSpan,
+    extract_all_boxed,
+    extract_answer,
+    extract_boxed,
+    extract_final_answer,
+    extract_last_number,
+    find_boxed_spans,
+    remove_boxed,
+)
+from ai4math.answers.normalize import (
+    normalize_answer,
+    normalize_latex,
+    split_tuple,
+    strip_units,
+)
+
+__all__ = [
+    "BoxedSpan",
+    "extract_all_boxed",
+    "extract_answer",
+    "extract_boxed",
+    "extract_final_answer",
+    "extract_last_number",
+    "find_boxed_spans",
+    "normalize_answer",
+    "normalize_latex",
+    "remove_boxed",
+    "split_tuple",
+    "strip_units",
+]

ai4math/answers/extract.py

@@ -0,0 +1,299 @@
+"""Extract final answers from natural-language mathematical solutions.
+
+Model outputs rarely arrive in a machine-readable form: the answer is buried at
+the end of a chain of reasoning, wrapped in ``\\boxed{...}``, announced with a
+phrase like "the final answer is", or simply left as the last number in the
+text. Every evaluation harness reimplements this extraction step, usually
+slightly differently. This module collects the strategies that recur across
+OpenAI's PRM800K grader, EleutherAI's ``lm-evaluation-harness`` (Minerva-style
+prompting) and HuggingFace's ``math-verify``, behind one explicit ladder.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable, Sequence
+from typing import NamedTuple
+
+__all__ = [
+    "FINAL_ANSWER_PATTERNS",
+    "BoxedSpan",
+    "extract_all_boxed",
+    "extract_answer",
+    "extract_boxed",
+    "extract_final_answer",
+    "extract_last_number",
+    "find_boxed_spans",
+    "remove_boxed",
+]
+
+#: LaTeX macros that conventionally mark a final answer.
+_BOX_MACROS = ("\\boxed", "\\fbox", "\\framebox")
+
+#: Phrases that introduce a final answer, most explicit first. The Minerva
+#: template ("Final Answer: The final answer is $X$. I hope it is correct.")
+#: comes first because matching it is unambiguous.
+#: Each pattern stops its capture at a sentence boundary (``. `` followed by
+#: whitespace) rather than running to end of line, so a solution that states an
+#: answer and then keeps talking yields just the answer — and so that
+#: ``finditer`` can find a *later* restatement in the same line.
+FINAL_ANSWER_PATTERNS: tuple[str, ...] = (
+    r"final answer is\s*:?\s*(?P<ans>.+?)(?:\.\s*I hope it is correct\.?|\.\s|\.$|$)",
+    r"final answer\s*:\s*(?P<ans>.+?)(?:\.\s|\.$|$)",
+    r"the answer is\s*:?\s*(?P<ans>.+?)(?:\.\s|\.$|$)",
+    r"answer\s*:\s*(?P<ans>.+?)(?:\.\s|\.$|$)",
+)
+
+#: A clause boundary: a comma followed by a word, as in "3, that is ...".
+#: Tuple answers such as "1, 2" are unaffected because the tail is a digit, and
+#: bracketed answers are exempted wholesale.
+_TRAILING_CLAUSE = re.compile(r",\s+(?=[A-Za-z])")
+
+_BRACKET_PAIRS = (("(", ")"), ("[", "]"), ("{", "}"))
+
+_TRAILING_JUNK = re.compile(r"(?:\.\s*I hope it is correct\.?|[.\s]+)$", re.IGNORECASE)
+
+# A signed number: optional thousands separators, decimals, exponent.
+_NUMBER = re.compile(
+    r"-?\d+(?:[,\s]\d{3})*(?:\.\d+)?(?:[eE][-+]?\d+)?|-?\.\d+"
+)
+
+
+class BoxedSpan(NamedTuple):
+    """Location and contents of one ``\\boxed{...}``-style span."""
+
+    start: int
+    """Index of the first character of the macro (e.g. of ``\\`` in ``\\boxed``)."""
+    end: int
+    """Index one past the last character of the span."""
+    content: str
+    """The text inside the braces, with surrounding whitespace stripped."""
+
+
+def _matching_brace(text: str, open_idx: int) -> int | None:
+    """Index of the ``}`` matching the ``{`` at *open_idx*, or None.
+
+    Counts nesting depth while treating any backslash as escaping the next
+    character, so ``\\{`` and ``\\}`` do not affect the balance.
+    """
+    if open_idx >= len(text) or text[open_idx] != "{":
+        return None
+    depth = 0
+    i = open_idx
+    while i < len(text):
+        char = text[i]
+        if char == "\\":
+            i += 2
+            continue
+        if char == "{":
+            depth += 1
+        elif char == "}":
+            depth -= 1
+            if depth == 0:
+                return i
+        i += 1
+    return None
+
+
+def find_boxed_spans(text: str) -> list[BoxedSpan]:
+    """Find every ``\\boxed``/``\\fbox`` span in *text*, in order of appearance.
+
+    Handles both the braced form ``\\boxed{42}`` and the brace-less form
+    ``\\boxed 42`` that appears in some datasets, where the answer runs to the
+    next ``$`` or end of line.
+
+    >>> find_boxed_spans(r"so \\boxed{42}")[0].content
+    '42'
+    >>> find_boxed_spans(r"\\boxed{\\frac{1}{2}}")[0].content
+    '\\\\frac{1}{2}'
+    """
+    spans: list[BoxedSpan] = []
+    for macro in _BOX_MACROS:
+        start = text.find(macro)
+        while start != -1:
+            cursor = start + len(macro)
+            while cursor < len(text) and text[cursor] in " \t":
+                cursor += 1
+            if cursor < len(text) and text[cursor] == "{":
+                close = _matching_brace(text, cursor)
+                if close is not None:
+                    spans.append(
+                        BoxedSpan(start, close + 1, text[cursor + 1 : close].strip())
+                    )
+                    start = text.find(macro, close + 1)
+                    continue
+            else:
+                # Brace-less form: run to the next '$' or newline.
+                stop = len(text)
+                for terminator in ("$", "\n"):
+                    found = text.find(terminator, cursor)
+                    if found != -1:
+                        stop = min(stop, found)
+                content = text[cursor:stop].strip()
+                if content:
+                    spans.append(BoxedSpan(start, stop, content))
+                    start = text.find(macro, stop)
+                    continue
+            start = text.find(macro, cursor)
+    spans.sort(key=lambda span: span.start)
+    return spans
+
+
+def extract_all_boxed(text: str) -> list[str]:
+    """Contents of every boxed span in *text*, in order of appearance."""
+    return [span.content for span in find_boxed_spans(text)]
+
+
+def extract_boxed(text: str, which: str = "last") -> str | None:
+    """Contents of one boxed span, or None if there is none.
+
+    Args:
+        text: Solution text to search.
+        which: ``"last"`` (the default, and the convention used by MATH-style
+            evaluation, since models often box intermediate results) or
+            ``"first"``.
+
+    >>> extract_boxed(r"first \\boxed{1} then \\boxed{2}")
+    '2'
+    >>> extract_boxed("no box here") is None
+    True
+    """
+    if which not in ("first", "last"):
+        raise ValueError(f"which must be 'first' or 'last', got {which!r}")
+    spans = find_boxed_spans(text)
+    if not spans:
+        return None
+    return spans[0].content if which == "first" else spans[-1].content
+
+
+def remove_boxed(text: str) -> str:
+    """Strip a ``\\boxed{...}`` wrapper, returning *text* unchanged if absent.
+
+    >>> remove_boxed(r"\\boxed{42}")
+    '42'
+    >>> remove_boxed("42")
+    '42'
+    """
+    spans = find_boxed_spans(text)
+    if len(spans) == 1 and spans[0].start == 0 and spans[0].end == len(text.rstrip()):
+        return spans[0].content
+    return text
+
+
+def extract_final_answer(
+    text: str, patterns: Sequence[str] | None = None
+) -> str | None:
+    """Extract an answer announced by a phrase such as "the final answer is".
+
+    Tries *patterns* in order and, within each, uses the **last** match — a
+    solution's concluding statement rather than an earlier, revised attempt.
+    Trailing prose and punctuation are trimmed from the captured span.
+
+    >>> extract_final_answer("... so the final answer is 42.")
+    '42'
+    """
+    active = FINAL_ANSWER_PATTERNS if patterns is None else patterns
+    for pattern in active:
+        matches = list(re.finditer(pattern, text, re.IGNORECASE | re.MULTILINE))
+        if matches:
+            candidate = _trim_prose(matches[-1].group("ans"))
+            if candidate:
+                return candidate
+    return None
+
+
+def _strip_math_delimiters(text: str) -> str:
+    r"""Remove enclosing ``$...$``, ``\(...\)`` or ``\[...\]`` delimiters.
+
+    These mark where mathematics starts, and are not part of the answer.
+
+    >>> _strip_math_delimiters("$12$")
+    '12'
+    """
+    result = text.strip()
+    for opener, closer in (("$$", "$$"), ("$", "$"), ("\\(", "\\)"), ("\\[", "\\]")):
+        while (
+            result.startswith(opener)
+            and result.endswith(closer)
+            and len(result) > len(opener) + len(closer) - 1
+        ):
+            candidate = result[len(opener) : -len(closer)].strip()
+            if not candidate:
+                break
+            result = candidate
+    return result
+
+
+def _trim_prose(candidate: str) -> str:
+    """Strip trailing prose, punctuation and math delimiters from a capture."""
+    trimmed = _TRAILING_JUNK.sub("", candidate.strip())
+    stripped = trimmed.strip()
+    bracketed = any(
+        stripped.startswith(left) and stripped.endswith(right)
+        for left, right in _BRACKET_PAIRS
+    )
+    if not bracketed:
+        trimmed = _TRAILING_CLAUSE.split(trimmed, maxsplit=1)[0]
+    return _strip_math_delimiters(_TRAILING_JUNK.sub("", trimmed.strip()))
+
+
+def extract_last_number(text: str) -> str | None:
+    """Last numeric literal in *text*, with thousands separators removed.
+
+    The weakest strategy, and the last resort of the extraction ladder: it is
+    right for GSM8K-style word problems whose answers are always integers, and
+    wrong for anything symbolic.
+
+    >>> extract_last_number("he sold 12 apples for $1,500 total")
+    '1500'
+    """
+    matches = _NUMBER.findall(text)
+    if not matches:
+        return None
+    return re.sub(r"[,\s]", "", matches[-1])
+
+
+def extract_answer(
+    text: str,
+    strategies: Iterable[str] = ("boxed", "phrase", "last_number"),
+) -> str | None:
+    """Extract a final answer by trying *strategies* in order.
+
+    This is the entry point most callers want. The default ladder runs from
+    most to least reliable: an explicit ``\\boxed{}``, then a concluding phrase,
+    then the last number in the text. The first strategy that yields something
+    non-empty wins.
+
+    Args:
+        text: Full model output or reference solution.
+        strategies: Any of ``"boxed"``, ``"phrase"``, ``"last_number"``.
+
+    Returns:
+        The extracted answer string, or None if every strategy failed.
+
+    Raises:
+        ValueError: If *strategies* names an unknown strategy.
+
+    >>> extract_answer(r"Thus the answer is $\\boxed{7}$.")
+    '7'
+    >>> extract_answer("I count 3 then 5 apples")
+    '5'
+    """
+    if not text:
+        return None
+    runners = {
+        "boxed": extract_boxed,
+        "phrase": extract_final_answer,
+        "last_number": extract_last_number,
+    }
+    for name in strategies:
+        try:
+            runner = runners[name]
+        except KeyError:
+            raise ValueError(
+                f"unknown strategy {name!r}; expected one of {sorted(runners)}"
+            ) from None
+        found = runner(text)
+        if found:
+            return found
+    return None