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 neural next-step predictor $p_\\theta(y_t|x_t)$ on the dataset that we collected in the previous notebook.\n",
    "\n",
    "To do so, we will fine-tune a pretrained language model on the dataset $\\mathcal{D}=\\{(x_t,y_t)\\}$ using the standard supervised fine-tuning approach:\n",
    "\n",
    "$$\n",
    "\\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 next-step $y_t$ given the context $x_t$. \\\n",
    "This corresponds to minimizing a cross-entropy loss at each position of the next-step, $\\sum_{\\ell=1}^{{|y_t|}}-\\log p_\\theta(y_t^\\ell|y_t^{<\\ell})$."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation\n",
    "\n",
    "The implementation consists of two steps:\n",
    "\n",
    "1. **Data formatting** ([data.py](../ntp_python/data.py)): formatting the examples.\n",
    "2. **Tuning**  ([tune.py](../ntp_python/tune.py)): using a standard language model fine-tuning script.\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1. Data formatting\n",
    "\n",
    "We format each (tactic-state, next-step) pair $(x_t, y_t)$ as:\n",
    "\n",
    "        [GOAL]tacticstate[PROOFSTEP]next-step<|endoftext|>\n",
    "\n",
    "Here, `[GOAL]...[PROOFSTEP]` is the input and `next-step<|endoftext|>` is the output.\n",
    "\n",
    "This format comes from [Han et al ICLR 2022]: \\\n",
    "[Proof Artifact Co-training for Theorem Proving with Language Models](https://arxiv.org/pdf/2102.06203.pdf).\n",
    "\n",
    "<!-- *Exercise:* can you think of other auxiliary tasks that might be useful? -->\n",
    "\n",
    "<!-- *Exercise:* can you think of alternative formats, e.g. which provide additional context? -->"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving split to disk...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:01<00:00,  2.87it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train\t169530\n",
      "val\t4053\n",
      "test\t3606\n"
     ]
    }
   ],
   "source": [
    "import sys\n",
    "sys.path.append('../ntp_python')\n",
    "import data\n",
    "\n",
    "datasets = data.proofstep(\n",
    "    data_dir='../data'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Input:\n",
      "[GOAL]ι : Type u_1\n",
      "I✝ J✝ : Box ι\n",
      "x y : ι → ℝ\n",
      "I J : WithBot (Box ι)\n",
      "⊢ ↑I = ↑J ↔ I = J[PROOFSTEP]\n",
      "\n",
      "Output:\n",
      "simp only [Subset.antisymm_iff, ← le_antisymm_iff, withBotCoe_subset_iff]<|endoftext|>\n"
     ]
    }
   ],
   "source": [
    "example = datasets['train'][0]\n",
    "print(\"Input:\", example['input'], '', sep='\\n')\n",
    "print(\"Output:\", example['output'], sep='\\n')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4. Tuning\n",
    "\n",
    "We minimally adapt a standard language-model fine-tuning script from [Stanford Alpaca](https://github.com/tatsu-lab/stanford_alpaca/blob/main/train.py). \n",
    "\n",
    "You can check out the full script at [partI_nextstep/ntp_python/tune.py](../ntp_python/tune.py). \\\n",
    "See [partI_nextstep/scripts/tune_proofstep.sh](../scripts/tune_proofstep.sh) for a command that trains on 8 GPUs with deepspeed."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here's an example command for training a 1.4b model on 1 GPU (and you can adjust the model size to be smaller to fit your compute constraints):"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```bash\n",
    "    REPO_DIR=\"..\"\n",
    "    TRAIN_FILE=${REPO_DIR}/data/leandojo_benchmark_4/processed/proofstep-train.jsonl\n",
    "    VALID_FILE=${REPO_DIR}/data/leandojo_benchmark_4/processed/proofstep-val.jsonl\n",
    "    MODEL=EleutherAI/pythia-1.4b-deduped\n",
    "\n",
    "    OUTDIR=${REPO_DIR}/model/${MODEL}\n",
    "\n",
    "    python ../ntp_python/tune.py \\\n",
    "        --model_name_or_path ${MODEL} \\\n",
    "        --train_data_path ${TRAIN_FILE} \\\n",
    "        --valid_data_path ${VALID_FILE} \\\n",
    "        --fp16 \\\n",
    "        --output_dir ${OUTDIR} \\\n",
    "        --num_train_epochs 10 \\\n",
    "        --per_device_train_batch_size 4 \\\n",
    "        --per_device_eval_batch_size 4 \\\n",
    "        --gradient_accumulation_steps 16 \\\n",
    "        --evaluation_strategy \"steps\" \\\n",
    "        --eval_steps 500 \\\n",
    "        --save_strategy \"steps\" \\\n",
    "        --save_steps 500 \\\n",
    "        --save_total_limit 1 \\\n",
    "        --learning_rate 1e-5 \\\n",
    "        --load_best_model_at_end 1 \\\n",
    "        --weight_decay 0. \\\n",
    "        --warmup_ratio 0.03 \\\n",
    "        --lr_scheduler_type \"cosine\" \\\n",
    "        --logging_steps 10 \\\n",
    "        --logging_dir \"$OUTDIR\" \\\n",
    "        --report_to=\"tensorboard\"\n",
    "\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### After training\n",
    "\n",
    "If everything went well, you should have a model in `../model/{MODEL_NAME}/checkpoint-{BEST_STEP}`.\n",
    "\n",
    "We have fine-tuned an `EleutherAI/pythia-2.8b-deduped` model that can be accessed through HuggingFace ([link](https://huggingface.co/wellecks/llmstep-mathlib4-pythia2.8b)):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import transformers\n",
    "\n",
    "MODEL = 'wellecks/llmstep-mathlib4-pythia2.8b'\n",
    "model = transformers.GPTNeoXForCausalLM.from_pretrained(MODEL)\n",
    "tokenizer = transformers.GPTNeoXTokenizerFast.from_pretrained(MODEL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can use your own model by setting `MODEL = \"../model/{MODEL_NAME}/checkpoint-{BEST_STEP}\"` \\\n",
    "(e.g., `../model/EleutherAI/pythia-2.8b-deduped/checkpoint-5000`)."
   ]
  },
  {
   "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": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rw [← h.gcd_eq_one]\n"
     ]
    }
   ],
   "source": [
    "prompt = \"\"\"[GOAL]m n : ℕ\n",
    "  h : Nat.coprime m n\n",
    "  ⊢ Nat.gcd m n = 1[PROOFSTEP]\"\"\"\n",
    "\n",
    "input_ids = tokenizer.encode(prompt, return_tensors='pt')\n",
    "out = model.generate(\n",
    "    input_ids,\n",
    "    max_new_tokens=256,\n",
    "    pad_token_id=tokenizer.eos_token_id\n",
    ")\n",
    "text = tokenizer.decode(out[0][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 build a VSCode 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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}