partI_nextstep/notebooks/I_nextstep_lean__part2_learn.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Neural next-step prediction | part 2: learning\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 next-step generator $p_\\theta(y_t|x_t)$ with the dataset that we collected in the previous notebook.\n",
    "\n",
    "To do so, we will fine-tune a pretrained language model with the dataset $\\mathcal{D}=\\{(x_t,y_t)\\}$ using standard supervised fine-tuning:\n",
    "\n",
    "$$\n",
    "\\arg\\max_\\theta \\sum_{(x_t,y_t)\\in \\mathcal{D}}-\\log p_\\theta(y_t|x_t).\n",
    "$$\n",
    "\n",
    "That is, we maximize the conditional likelihood of a completion $y_t$ (which contains a next-tactic) given the prompt $x_t$ (which contains a proof state). \n",
    "\n",
    "This corresponds to minimizing a cross-entropy loss at each position of the completion, $\\sum_{\\ell=1}^{{|y_t|}}-\\log p_\\theta(y_t^\\ell|y_t^{<\\ell})$."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation\n",
    "\n",
    "In the previous notebook, we saw how to use [instruction_tuning.py](../ntp-training-data/scripts/instruction_tuning.py) to format the extracted Mathlib data into (prompt, completion) examples.\n",
    "\n",
    "We provide formatted fine-tuning data on HuggingFace:\n",
    "\n",
    "- [`l3lab/ntp-mathlib-instruct-st`](l3lab/ntp-mathlib-instruct-st)\n",
    "\n",
    "*If you use this data or code, we kindly ask that you cite this neural theorem proving tutorial*.\n",
    "\n",
    "\n",
    "We will finetune on Mathlib (`ntp-lean-mathlib-tactic-instruct`). First, we download the training and validation set:\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd ../ntp-tune && bash prepare_data.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!head -n 2 ../ntp-tune/data/state_tactic_mathlib_only_train.jsonl"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4. Fine-tuning\n",
    "\n",
    "We can now use an off-the-shelf fine-tuning script. We minimally adapt a standard language-model fine-tuning script from [open-instruct](https://github.com/allenai/open-instruct). \n",
    "\n",
    "You can check out the full script at [ntp-tune/finetune.py](../ntp-tune/finetune.py). \\\n",
    "See [ntp-tune/finetune.sh](../ntp-tune/finetune.sh) for a command that finetunes a `deepseek-coder-1.3b-base` model on 4 GPUs. \n",
    "\n",
    "Please see the [ntp-tune](../ntp-tune) directory for setup instructions."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### After training\n",
    "\n",
    "We have fine-tuned a `deepseek-coder-1.3b-base` model that can be accessed through HuggingFace:\n",
    "- [`l3lab/ntp-mathlib-st-deepseek-coder-1.3b`](https://huggingface.co/l3lab/ntp-mathlib-st-deepseek-coder-1.3b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "import transformers\n",
    "\n",
    "MODEL = 'l3lab/ntp-mathlib-st-deepseek-coder-1.3b'\n",
    "model = transformers.AutoModelForCausalLM.from_pretrained(MODEL)\n",
    "tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use your own model by setting `MODEL = \"/path/to/checkpoint-{BEST_STEP}\"`"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's generate a next-step suggestion for the proof state from our original example:\n",
    "\n",
    "```lean\n",
    "    theorem test_thm (m n : Nat) (h : m.Coprime n) : m.gcd n = 1\n",
    "```\n",
    "Recal from the previous notebook that the initial proof state $x_0$ is:\n",
    "\n",
    "    m n : ℕ\n",
    "    h : Nat.Coprime m n\n",
    "    ⊢ Nat.gcd m n = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rw [← Nat.gcd_comm, Nat.gcd_eq_one_iff_coprime]\n",
      "[/TAC]\n"
     ]
    }
   ],
   "source": [
    "state = \"\"\"m n : ℕ\n",
    "h : Nat.Coprime m n\n",
    "⊢ Nat.gcd m n = 1\"\"\"\n",
    "\n",
    "prompt = \"\"\"/- 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",
    "%s\n",
    "[/STATE]\n",
    "[TAC]\n",
    "\"\"\" % state\n",
    "\n",
    "model.eval()\n",
    "input_ids = tokenizer(prompt, return_tensors='pt')\n",
    "out = model.generate(\n",
    "    input_ids['input_ids'],\n",
    "    attention_mask=input_ids['attention_mask'],\n",
    "    max_new_tokens=256,\n",
    "    pad_token_id=tokenizer.eos_token_id,\n",
    ")\n",
    "for item in out:\n",
    "    text = tokenizer.decode(item[input_ids['input_ids'].shape[1]:], skip_special_tokens=True)\n",
    "    print(text)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Next steps\n",
    "\n",
    "In the next notebook, we will prove theorems with the trained model by interacting with the Lean proof assistant.\n",
    "\n",
    "This will let us automatically check whether a generated proof (e.g., one containing the step above) is correct.\n",
    "\n",
    "Later on, we will train a language model that uses additional file context, then we will build a VS Code plugin that returns next-step suggestions from the language model."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "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
}