partI_nextstep/notebooks/I_nextstep_lean__part1_data.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Neural next-step prediction | part 1: data\n",
    "Tutorial on neural theorem proving\\\n",
    "Author: Sean Welleck\n",
    "\n",
    "----------------"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### High-level goal\n",
    "\n",
    "Our goal is to train a neural next-step prediction model, $p(y_t|x_t)$. Here $x_t$ is a _proof state_, and $y_t$ is a next-step. \n",
    "\n",
    "To do so, we will create a dataset $\\mathcal{D}=\\{(x_t,y_t)\\}$ from human-written proofs. \n",
    "\n",
    "We can then train a next-step prediction model using a next-token prediction loss on the dataset."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Simple example\n",
    "\n",
    "To see what proof states and next-steps look like, let's look at an example human-written theorem and proof:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "import Mathlib.Data.Nat.Prime\n",
      "\n",
      "theorem test_thm (m n : Nat) (h : m.Coprime n) : m.gcd n = 1 := by\n",
      "  rw [Nat.Coprime] at h\n",
      "  exact h\n"
     ]
    }
   ],
   "source": [
    "!cat ../ntp-training-data/Examples/Example0.lean"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We would like to transform this theorem and proof into a sequence of (proof_state, next_step) examples.\n",
    "\n",
    "First, notice that the proof has two steps:\n",
    "\n",
    "1. $y_1=$ `rw [Nat.Coprime] at h`\n",
    "2. $y_2=$ `exact h`"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can manually see the proof states by looking in VSCode. \n",
    "\n",
    "For example, placing the cursor before $y_1$ gives us the proof state $x_1$ (shown as \"Tactic state\"):"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"images/proof_state_1.png\" width=600px>\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That is, the image above corresponds to $(x_1,y_1)$ defined as:\n",
    "\n",
    "  $x_1$: \n",
    "  ```\n",
    "    m n : ℕ\n",
    "    h : Nat.Coprime m n\n",
    "    ⊢ Nat.gcd m n = 1\n",
    "  ```\n",
    "\n",
    "  $y_1$: `rw [Nat.coprime] at h`\n",
    "\n",
    "\n",
    "Similarly, we can get the proof state $x_2$ prior to the step $y_2$ (`exact h`):\n",
    "\n",
    "<!-- ![title](images/proof_state_2.png) -->\n",
    "\n",
    "<img src=\"images/proof_state_2.png\" width=600px>\n",
    "\n",
    "After step $y_2$, the proof is complete: the proof state $x_3$ says we have \"No goals\":\n",
    "\n",
    "<!-- ![title](images/proof_state_3.png) -->\n",
    "<img src=\"images/proof_state_3.png\" width=600px>\n",
    "\n",
    "\n",
    "In summary, it is possible to *manually* transform the theorem and proof into a sequence $[(x_1,y_1),(x_2,y_2),(x_3)]$.\n",
    "\n",
    "Formally, the steps $y_t$ are known as a \"tactics\". As a result, we will often use \"step\" and \"tactic\" interchangeably."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---------------------\n",
    "\n",
    "## Automatically extracting proof states and next-steps \n",
    "\n",
    "To scale up data collection, we need a way to *automatically* extract proof states and next-steps (tactics) from human-written proofs.\n",
    "\n",
    "The [ntp-training-data](../ntp-training-data/) directory contains Lean code that automatically extracts proof states and tactics. \n",
    "\n",
    "It is a modified version of Scott Morrison's [lean-training-data](https://github.com/semorrison/lean-training-data)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. Transform a Lean file\n",
    "\n",
    "The extraction is done by [training_data.lean](../ntp-training-data/scripts/training_data.lean), which conceptually implements:\n",
    "\n",
    "$\\quad f_{\\text{extract}}(\\text{lean file})\\rightarrow \\{(x,y,c)_i\\}$,\n",
    "\n",
    "where $\\{(x,y,c)_i\\}$ denotes a collection of (proof state, next-tactic) pairs along with additional context $c$.\n",
    "\n",
    "We can run it on the [Example0.lean](../ntp-training-data/Examples/Example0.lean) file from above:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd ../ntp-training-data && lake exe training_data Examples/Example0.lean > ../notebooks/data/example0.jsonl\n",
    "!cat data/example0.jsonl"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This output is a `.jsonl` file where each line is an example of the form:\n",
    "```json\n",
    "{\n",
    "   \"state\": \"{tactic state}\",\n",
    "   \"nextTactic\" : \"{pretty-printed next tactic}\",\n",
    "   \"srcUpToTactic\" : \"{source code in the file up to the tactic invocation}\",\n",
    "   \"declUpToTactic\" : \"{source code in the declaration up to the tactic invocation}\",\n",
    "   \"decl\": \"{declaration without proof (e.g., statement of a theorem)}\",\n",
    "   \"declId\": \"{unique identifier of the declaration}\"\n",
    "}\n",
    "```\n",
    "Here are the extracted (state, next-tactic) examples:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== State:\n",
      "m n : ℕ\n",
      "h : Nat.Coprime m n\n",
      "⊢ Nat.gcd m n = 1\n",
      "\n",
      "--- Next tactic\n",
      "rw [Nat.Coprime] at h\n",
      "\n",
      "=== State:\n",
      "m n : ℕ\n",
      "h : Nat.gcd m n = 1\n",
      "⊢ Nat.gcd m n = 1\n",
      "\n",
      "--- Next tactic\n",
      "exact h\n",
      "\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "with open('data/example0.jsonl') as f:\n",
    "    examples = [json.loads(line) for line in f.readlines()]\n",
    "\n",
    "for example in examples:\n",
    "    print(\"=== State:\", example['state'], '', sep='\\n')\n",
    "    print(\"--- Next tactic\", example['nextTactic'], '', sep='\\n')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that the proof states are the ones we saw above in VS Code.\n",
    "\n",
    "We can also view additional extracted information, such as the source up to the tactic invocation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "==========\n",
      "import Mathlib.Data.Nat.Prime\n",
      "\n",
      "theorem test_thm (m n : Nat) (h : m.Coprime n) : m.gcd n = 1 := by\n",
      "  \n",
      "/- State:\n",
      "m n : ℕ\n",
      "h : Nat.Coprime m n\n",
      "⊢ Nat.gcd m n = 1\n",
      "-/\n",
      "/- Next tactic -/\n",
      "rw [Nat.Coprime] at h\n",
      "\n",
      "==========\n",
      "import Mathlib.Data.Nat.Prime\n",
      "\n",
      "theorem test_thm (m n : Nat) (h : m.Coprime n) : m.gcd n = 1 := by\n",
      "  rw [Nat.Coprime] at h\n",
      "  \n",
      "/- State:\n",
      "m n : ℕ\n",
      "h : Nat.gcd m n = 1\n",
      "⊢ Nat.gcd m n = 1\n",
      "-/\n",
      "/- Next tactic -/\n",
      "exact h\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for example in examples:\n",
    "    print(\"==========\", example['srcUpToTactic'], sep='\\n')\n",
    "    print(\"/- State:\", example['state'], '-/', sep='\\n')\n",
    "    print(\"/- Next tactic -/\", example['nextTactic'], '', sep='\\n')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "-----------------\n",
    "\n",
    "## Scaling up data extraction\n",
    "\n",
    "We can run the above script on a full repository:\n",
    "\n",
    "$\\quad f_{\\text{extract}}(\\text{lean repository})\\rightarrow \\mathcal{D}.$\n",
    "\n",
    "Doing so involves calling the script and keeping track of files, which we prefer to do in a scripting language such as Python.\n",
    "\n",
    "Thus `ntp-training-data` uses two Python scripts: \n",
    "1. [extract_repos.py](../ntp-training-data/scripts/extract_repos.py), which reads in a configuration file with repository information. For each repository, it calls:\n",
    "2.  [run_pipeline.py](../ntp-training-data/scripts/run_pipeline.py) script, which runs `lake exe training_data` and keeps track of files."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Extracting 300k examples from mathlib\n",
    "\n",
    "We use these scripts to extract data from [Mathlib](https://github.com/leanprover-community/mathlib4), a large community-driven library of mathematics in Lean. \n",
    "\n",
    "To do so, we first specify Mathlib in [ntp-training-data/configs/config.json](../ntp-training-data/configs/config.json):\n",
    "\n",
    "```json\n",
    "    {\n",
    "        \"repo\": \"https://github.com/leanprover-community/mathlib4\",\n",
    "        \"commit\": \"cf8e23a62939ed7cc530fbb68e83539730f32f86\",\n",
    "        \"lean\": \"leanprover/lean4:v4.4.0\",\n",
    "        \"name\": \"mathlib\",\n",
    "        \"import_file\": \"Mathlib.lean\",\n",
    "        \"imports\": [\"Mathlib\"]\n",
    "    }\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we execute the [extract_repos.py](../ntp-training-data/scripts/extract_repos.py) script. \n",
    "\n",
    "On a Macbook Pro (M3 Max, 14 CPU) it takes around 2 hours."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd ../ntp-training-data && python scripts/extract_repos.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    3759\n"
     ]
    }
   ],
   "source": [
    "# Number of files\n",
    "!ls ../ntp-training-data/Examples/Mathlib/TacticPrediction/ | wc -l"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Output data\n",
    "\n",
    "The extracted data is in [llm-training-data/Examples/Mathlib/TacticPrediction](../ntp_lean/llm-training-data/Examples/Mathlib/TacticPrediction).\n",
    "\n",
    "It is organized by file. For instance, here are extracted examples for Mathlib's Algebra/AddTorsor.lean:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!head -n 2 ../ntp-training-data/Examples/Mathlib/TacticPrediction/Mathlib_Algebra_AddTorsor.jsonl"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Extracted data on Huggingface\n",
    "\n",
    "We provide extracted Mathlib data (i.e., the result of running the command above) on HuggingFace:\n",
    "- [`l3lab/ntp-mathlib`](https://huggingface.co/datasets/l3lab/ntp-mathlib)\n",
    "\n",
    "Additional repositories can be extracted by adding them to `configs/config.json`.\n",
    "\n",
    "*If you use this data or code, we kindly ask that you cite this neural theorem proving tutorial*."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "--------------\n",
    "## Fine-tuning data\n",
    "\n",
    "Finally, we would like to format the data so that we can finetune a language model with a standard finetuning script.\n",
    "\n",
    "To this end, we format the data into (prompt, completion) examples using [ntp-training-data/scripts/instruction_tuning.py](../ntp-training-data/scripts/instruction_tuning.py).\n",
    "\n",
    "Notationally, our formatted dataset is of the form:\n",
    "\n",
    "\\begin{align}\n",
    "\\mathcal{D}=\\{(f_{\\text{prompt}}(x_t), f_{\\text{completion}}(y_t))\\},\n",
    "\\end{align}\n",
    "\n",
    "where  $f_{\\text{prompt}}$ maps a state $x_t$ to a string, and $f_{\\text{completion}}$ maps a next-tactic $y_t$ to a string.\n",
    "\n",
    "Here is the prompt and completion created from the first (state, tactic) pair from our simple example above (`example0.jsonl`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=== Prompt:\n",
      "/- You are proving a theorem in Lean 4.\n",
      "You are given the following information:\n",
      "- The current proof state, inside [STATE]...[/STATE]\n",
      "\n",
      "Your task is to generate the next tactic in the proof.\n",
      "Put the next tactic inside [TAC]...[/TAC]\n",
      "-/\n",
      "[STATE]\n",
      "m n : ℕ\n",
      "h : Nat.Coprime m n\n",
      "⊢ Nat.gcd m n = 1\n",
      "[/STATE]\n",
      "[TAC]\n",
      "=== Completion:\n",
      "rw [Nat.Coprime] at h\n",
      "[/TAC]\n"
     ]
    }
   ],
   "source": [
    "import sys, json\n",
    "sys.path.append('../ntp-training-data/scripts/')\n",
    "\n",
    "from instruction_tuning import prompt_state_tactic\n",
    "\n",
    "examples = [json.loads(line) for line in open('data/example0.jsonl').readlines()]\n",
    "\n",
    "prompt, completion = prompt_state_tactic(examples[0])\n",
    "print(f'=== Prompt:\\n{prompt}', f'=== Completion:\\n{completion}', sep='')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that we added a natural language description of the task into the prompt, commonly known as an \"instruction\" (and hence we call the script [instruction_tuning.py](../ntp-training-data/scripts/instruction_tuning.py)). In this case the instruction may not be strictly necessary, but including it suggests several other directions to experiment with (one of which we will see later in part 5). When we eventually incorporate our model into a tool (part 6), the instruction format will also allow us to swap in other off-the-shelf models that support instruction following.\n",
    "\n",
    "\n",
    "Let's convert all of the extracted examples into (prompt, completion) data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Examples/Mathlib\n",
      "num_train_decls\t59587\n",
      "num_dev_decls\t1568\n",
      "num_test_decls\t1568\n",
      "num_dev_file_split_decls\t0\n",
      "num_test_file_split_decls\t0\n",
      "num_train\t291262\n",
      "num_dev\t7735\n",
      "num_test\t8016\n",
      "num_file_split_dev\t0\n",
      "num_file_split_test\t0\n",
      "instructions/state_tactic_mathlib_only\n"
     ]
    }
   ],
   "source": [
    "!cd ../ntp-training-data/ && python scripts/instruction_tuning.py --mathlib-only  # flag excludes e.g. Example0.lean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Upload to huggingface (this will only work with authorized accounts, \n",
    "# but we leave this command here so you can use a similar pattern for your own projects)\n",
    "# \n",
    "# !cd ../ntp-training-data/instructions/ && bash upload.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Data on Huggingface\n",
    "\n",
    "We provide formatted fine-tuning data on HuggingFace:\n",
    "\n",
    "- [`l3lab/ntp-mathlib-instruct-st`](https://huggingface.co/datasets/l3lab/ntp-mathlib-instruct-st) (Mathlib)\n",
    "\n",
    "Alternate formats can be produced using variants of [scripts/instruction_tuning.py](../ntp-training-data/scripts/instruction_tuning.py).\n",
    "\n",
    "*If you use this data or code, we kindly ask that you cite this neural theorem proving tutorial*."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset('l3lab/ntp-mathlib-instruct-st')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "291262\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'task': 'tactic_predition',\n",
       " 'prompt': '/- You are proving a theorem in Lean 4.\\nYou are given the following information:\\n- The current proof state, inside [STATE]...[/STATE]\\n\\nYour task is to generate the next tactic in the proof.\\nPut the next tactic inside [TAC]...[/TAC]\\n-/\\n[STATE]\\nJ : Type v\\ninst✝¹ : SmallCategory J\\ninst✝ : IsFiltered J\\nF : J ⥤ GroupCat\\nx y : (j : J) × ↑(F.obj j)\\nh : Types.FilteredColimit.Rel (F ⋙ forget GroupCat) x y\\n⊢ colimitInvAux F x = colimitInvAux F y\\n[/STATE]\\n[TAC]\\n',\n",
       " 'completion': 'apply G.mk_eq\\n[/TAC]',\n",
       " 'metadata': {'task': 'tactic_prediction',\n",
       "  'project': 'Examples/Mathlib',\n",
       "  'file': 'Examples/Mathlib/TacticPrediction/Mathlib_Algebra_Category_GroupCat_FilteredColimits.jsonl',\n",
       "  'declId': 'Mathlib.Algebra.Category.GroupCat.FilteredColimits.83_0.OlIvs5vXvq1jJtZ',\n",
       "  'target': 'apply G.mk_eq',\n",
       "  'split': 'train'}}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(len(dataset['train']))\n",
    "dataset['train'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/- You are proving a theorem in Lean 4.\n",
      "You are given the following information:\n",
      "- The current proof state, inside [STATE]...[/STATE]\n",
      "\n",
      "Your task is to generate the next tactic in the proof.\n",
      "Put the next tactic inside [TAC]...[/TAC]\n",
      "-/\n",
      "[STATE]\n",
      "J : Type v\n",
      "inst✝¹ : SmallCategory J\n",
      "inst✝ : IsFiltered J\n",
      "F : J ⥤ GroupCat\n",
      "x y : (j : J) × ↑(F.obj j)\n",
      "h : Types.FilteredColimit.Rel (F ⋙ forget GroupCat) x y\n",
      "⊢ colimitInvAux F x = colimitInvAux F y\n",
      "[/STATE]\n",
      "[TAC]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(dataset['train'][0]['prompt'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "apply G.mk_eq\n",
      "[/TAC]\n"
     ]
    }
   ],
   "source": [
    "print(dataset['train'][0]['completion'])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Next steps\n",
    "\n",
    "In part 2, we'll train a next-step generation model on the fine-tuning dataset."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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
}