partI_nextstep/notebooks/I_nextstep_lean__part4_evaluation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Neural next-step prediction | part 4: evaluation\n",
    "Tutorial on neural theorem proving\\\n",
    "Author: Sean Welleck\n",
    "\n",
    "----------------\n",
    "\n",
    "#### High-level goal\n",
    "\n",
    "To get a quantitative estimate of our model's performance, we can perform proof search on an evaluation set of theorem statements. Intuitively, a model that is good at next-step suggestion will be effective for fully proving a theorem when paired with a suitable search algorithm. Therefore, proof search performance gives some measure of how useful the model's next-step suggestions will be when integrated into an interactive suggestion tool. \n",
    "\n",
    "First, we will evaluate on a small set of manually written theorem statements:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluation_theorems = [\n",
    "    \"\"\"example (m n : Nat) (h : m.Coprime n) : Nat.gcd m n = 1 := by {}\"\"\",\n",
    "    \"\"\"theorem thm1 (a b c : Nat) : a + b = c → a ≤ c := by {}\"\"\",\n",
    "    \"\"\"theorem thm2 (x y : ℝ) : x < y → 0 < y - x := by {}\"\"\",\n",
    "    \"\"\"theorem thm3 (n : Nat) : n ≥ 0 := by {}\"\"\",\n",
    "    \"\"\"theorem thm4 (x y z : ℝ) : x ≤ y → y ≤ z → x ≤ z := by {}\"\"\",\n",
    "    \"\"\"theorem thm5 (m n : Nat) (h : m.Coprime n) : m.gcd n = 1 := by {}\"\"\",\n",
    "    \"\"\"theorem thm6: r ⊆ s → s ⊆ t → r ⊆ t := by {}\"\"\",\n",
    "    \"\"\"theorem thm7 (f : ℕ → ℕ) : Monotone f → ∀ n, f n ≤ f (n + 1) := by {}\"\"\",\n",
    "    \"\"\"theorem thm8 (c : ℝ) : Injective fun x => x + c := by {}\"\"\",\n",
    "    \"\"\"theorem thm9 (A B : Set ℕ) : A ⊆ B → ∀ n, n ∈ A → n ∈ B := by {}\"\"\",\n",
    "    \"\"\"theorem thm10 (injg : Injective g) (injf : Injective f) : Injective fun x => g (f x) := by {}\"\"\",\n",
    "]\n",
    "    \n",
    "# Shared header for the theorems above\n",
    "header = \"\"\"import Mathlib.Data.Nat.Factorization.Basic\n",
    "import Mathlib.Data.Nat.Prime\n",
    "import Mathlib.Data.Real.Basic\n",
    "\n",
    "open Function\n",
    "variable {α : Type _} (r s t : Set α)\n",
    "\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's load our model and run best-first search:\n",
    "\n",
    "> The search budget allows for running in < 10 minutes on a MacBook Pro (M3 Max, 14 CPU). With a CUDA GPU it will be faster.\n",
    "\n",
    "Feel free to study the trajectories that are printed below, which include both successes and failures:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append('../ntp-interact/')\n",
    "\n",
    "import proofsearch\n",
    "import os\n",
    "os.environ['TOKENIZERS_PARALLELISM'] = 'false'  # prevents an annoying warning\n",
    "\n",
    "MODEL = 'l3lab/ntp-mathlib-st-deepseek-coder-1.3b'\n",
    "model, tokenizer = proofsearch.load_model(MODEL)\n",
    "\n",
    "PATH_TO_REPL = '/Users/wellecks/projects/ntptutorial/partI_nextstep/ntp-interact/repl'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "transformers.set_seed(43)\n",
    "\n",
    "results = {True: [], False: []}\n",
    "for theorem in evaluation_theorems:\n",
    "    result = proofsearch.best_first_search(\n",
    "        model, tokenizer, header, theorem, \n",
    "        max_iters=16,\n",
    "        temperatures=[0.5],\n",
    "        num_samples=8,\n",
    "        verbose=True,\n",
    "        path_to_repl=PATH_TO_REPL\n",
    "    )\n",
    "    print(\"Success: %s\" % result['success'])\n",
    "    results[result['success']].append(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are the successfully closed theorems and their generated proofs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_result(result):\n",
    "    print(result['theorem_statement'].replace('{}', '') + '\\n\\t' + '\\n\\t'.join(result['proof']) + '\\n')\n",
    "\n",
    "print(\"%.3f closed\" % (len(results[True])/ (len(results[True])+len(results[False]))))\n",
    "for result in results[True]:\n",
    "    print_result(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Performance depends on model and search\n",
    "\n",
    "Our method closed a lot of the proofs! \n",
    "\n",
    "When running experiments, it is important to note that theorem proving performance is a function of the model $p_\\theta$, the search algorithm $\\mathcal{A}$, and the search budget $k$.\n",
    "\n",
    "$\\quad\\text{pass rate} = f(p_\\theta, \\mathcal{A}, k)$.\n",
    "\n",
    "In principle, we can improve theorem proving performance by improving the model, improving the search algorithm (for a fixed budget), or by increasing the budget. As a result, when comparing models it is important to account for possible performance variations that arise from the search algorithm or budget (e.g., by holding the search algorithm and budget fixed).\n",
    "\n",
    "\n",
    "Feel free to try out different temperatures, number of samples, etc. to see how performance varies."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "-----------\n",
    "\n",
    "### Evaluating neural theorem provers\n",
    "\n",
    "Above, we evaluated on hand-specified theorems. In practice, evaluation is done in two settings:\n",
    "\n",
    "\n",
    "1. Benchmarks\n",
    "2. Test split\n",
    "\n",
    "Benchmarks provide theorem statements that can characterize performance on a certain kind of theorems (e.g., competition problems or undergraduate math) and can test distribution shift for a model (e.g., competition problems for a model trained on mathlib).\n",
    "\n",
    "A test split measures performance on theorems drawn from the same distribution as the training set."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Benchmarks in neural theorem proving\n",
    "\n",
    "[**MiniF2F** [Zheng et al ICLR 2022]](https://arxiv.org/abs/2109.00110) is a widely used benchmark of 488 problems statements drawn from the AIME, AMC, and the International Mathematical Olympiad (IMO), as well as material from high-school and undergraduate mathematics courses. \n",
    "\n",
    "Specifically, given $x_\\text{formal statement}$, our model must produce a correct formal proof $y_1,\\ldots,y_{T_x}$. Below, you can look at examples:\n",
    "\n",
    "> As a reference point, we show the informal statement and informal proof, though the model is only given the formal statement. (The informal annotations were added in [Jiang et al ICLR 2023](https://arxiv.org/abs/2210.12283))."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "#### Problem: induction_divisibility_3divnto3m2n\n",
       "\n",
       "#### Formal statement \n",
       "\n",
       "```lean\n",
       "theorem induction_divisibility_3divnto3m2n\n",
       "  (n : ℕ) :\n",
       "  3 ∣ n^3 + 2 * n := sorry\n",
       "```\n",
       "\n",
       "\n",
       "#### Informal statement\n",
       "\n",
       "Show that for any natural number $n \\in \\mathbb{N}$, $3 \\mid n^3 + 2n$ .\n",
       "#### Informal proof\n",
       "\n",
       "We show the result by induction on $n$. The result is trivial for $n=0$. Let us assume it is true for $n \\geq 0$.\n",
       "We have $(n+1)^3+2(n+1) = (n^3+3n^2+3n+1) + (2n+2) = n^3+2n + 3n^2+3n+3$. From the induction hypothesis, we know that $3$ divides $n^3+2n$. Since $3$ also divides $3n^2+3n+3$, the result is also true in $n+1$ and we have by induction that the result is true for all $n$."
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython.display import Markdown\n",
    "import datasets\n",
    "\n",
    "minif2f = datasets.load_dataset('cat-searcher/minif2f-lean4')\n",
    "\n",
    "idx = 35\n",
    "\n",
    "example = minif2f['validation'][idx]\n",
    "Markdown(\n",
    "    '#### Problem: ' + example['id'] + \n",
    "    '\\n\\n#### Formal statement \\n\\n' + '```lean\\n' + example['formal_statement'] + '\\n```\\n' + \n",
    "    '\\n\\n#### Informal statement\\n\\n' + example['informal_stmt'] + \n",
    "    '\\n#### Informal proof\\n\\n' + example['informal_proof']\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The problems vary in difficulty. Some problems can be closed easily (especially when the model deploys built in tactics such as `simp`), while others require long-form reasoning that can also be difficult to formalize. Here is a success case and a failure case (at least with this search budget):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluation_theorems = [\n",
    "    \"\"\"theorem mathd_algebra_15 (s : ℕ → ℕ → ℕ) (h₀ : ∀ a b, 0 < a ∧ 0 < b → s a b = a ^ (b : ℕ) + b ^ (a : ℕ)) : s 2 6 = 100 := by {}\"\"\",\n",
    "    \"\"\"theorem imo_2001_p6 (a b c d : ℕ) (h₀ : 0 < a ∧ 0 < b ∧ 0 < c ∧ 0 < d) (h₁ : d < c) (h₂ : c < b) (h₃ : b < a) (h₄ : a * c + b * d = (b + d + a - c) * (b + d + c - a)) : ¬Nat.Prime (a * b + c * d) := by {}\"\"\"\n",
    "]\n",
    "\n",
    "for theorem in evaluation_theorems:\n",
    "    result = proofsearch.best_first_search(\n",
    "        model, tokenizer, header, theorem, \n",
    "        max_iters=16,\n",
    "        temperatures=[0.5],\n",
    "        num_samples=8,\n",
    "        verbose=True,\n",
    "        path_to_repl=PATH_TO_REPL\n",
    "    )\n",
    "    print(\"Success: %s\" % result['success'])\n",
    "    if result['success']:\n",
    "        print_result(result)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Our model\n",
    "\n",
    "We evaluated our model on miniF2F, by changing the prompt in this [llemma_formal2formal](https://github.com/wellecks/llemma_formal2formal) repo to adhere to our model's prompt. \n",
    "\n",
    "The search uses best-first search with beam search expansion, and an expansion size of 32, and a timeout of 10 minutes (these have become informal standards).\n",
    "\n",
    "**The model succeeds on 29.1% (71/244) of the proofs.** To get a sense of this performance, here is a table from the [InternLM-Math paper](https://arxiv.org/abs/2402.06332) from February 2024:\n",
    "\n",
    "<img src=\"images/internlm.png\" width=\"500px\">\n",
    "\n",
    "Not bad!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Other benchmarks\n",
    "\n",
    "[**ProofNet** [Azerbayev et al 2023]](https://arxiv.org/abs/2302.12433) is a benchmark targeting undergraduate-level mathematics. It consists of 371 problems  drawn from popular undergraduate pure mathematics textbooks, and covering topics such as real and complex analysis, linear algebra, abstract algebra, and topology. \n",
    "\n",
    "ProofNet theorems tend to depend on more background knowledge than competition problems, which means that a learned model needs to use theorems and definitions from a wider subset of mathematics. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "#### In-distribution test split\n",
    "We can also test the model on theorems from the same distribution as those it was trained on. For our model, this means splitting the mathlib4 repository into training theorems and evaluation theorems. A common splitting strategy is via uniform sampling of theorem declarations. The resulting dataset covers a range of topics."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "Several works report nontrivial performance on mathlib, suggesting that existing models may already lead to useful tools for certain kinds of proofs in this setting. For instance, [Welleck \\& Saha 2023](https://arxiv.org/pdf/2306.15626.pdf) report 50.1\\% on a mathlib4 test split, while [Yang et al 2023](https://arxiv.org/pdf/2306.15626.pdf) report 48.6\\% on a mathlib4 test split with a retrieval-augmented model and [Polu et al 2023](https://arxiv.org/abs/2202.01344) achieve over 70\\% pass@8 on mathlib3 using a variant of reinforcement learning (and a large search budget). \n",
    "\n",
    "However, there are several caveats, one of which is that uniformly sampling theorems to hold out in the test set may not resemble the actual process of theorem proving. For instance, the model may have already seen the theorem being used, which is not always the case in practice."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Next steps\n",
    "\n",
    "Next we will consider proving theorems that depend on definitions or other theorems that were not observed during training, which occurs in nearly every real-world formal theorem proving scenario.\n",
    "To do so, we will train next-step generation models that are conditioned on additional context, such as the contents of the Lean file preceding the theorem being proven, and the preceding steps of the proof."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}