src/models/gpt_neo/gpt_neo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dee8b62b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ['MASTER_ADDR'] = 'localhost'\n",
    "os.environ['MASTER_PORT'] = '9994'\n",
    "os.environ['RANK'] = \"0\"\n",
    "os.environ['LOCAL_RANK'] = \"0\"\n",
    "os.environ['WORLD_SIZE'] = \"1\"\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "da90efac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import GPT2Tokenizer, GPTNeoForCausalLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b0824339",
   "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": [
    "tokenizer = GPT2Tokenizer.from_pretrained(\"EleutherAI/gpt-neo-125M\",    \n",
    "                            bos_token=\"<|startoftext|>\",\n",
    "                            eos_token=\"<|endoftext|>\",\n",
    "                            pad_token=\"<|pad|>\")\n",
    "# Download the pre-trained GPT-Neo model and transfer it to the GPU\n",
    "model = GPTNeoForCausalLM.from_pretrained(\"EleutherAI/gpt-neo-125M\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c6b811e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "model= model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0a5c2592",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Embedding(50259, 768)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.resize_token_embeddings(len(tokenizer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b52d7da2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/workspace'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3eda9f60",
   "metadata": {},
   "outputs": [],
   "source": [
    "DATA_DIRECTORY = \"workspace/lean_proof_recording/data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "1b1a169e",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"workspace/lean_proof_recording/data/cleaned_training_data/train.src\") as f:\n",
    "    train_src = f.readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1bee10d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"workspace/lean_proof_recording/data/cleaned_training_data/train.tgt\") as f:\n",
    "    train_tgt = f.readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5b68200e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data_and_metadata.csv  test.src     train.names  valid.index  valid.tgt\r\n",
      "test.index             test.tgt     train.src    valid.names\r\n",
      "test.names             train.index  train.tgt    valid.src\r\n"
     ]
    }
   ],
   "source": [
    "ls workspace/lean_proof_recording/data/cleaned_training_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7962a7cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"workspace/lean_proof_recording/data/cleaned_training_data/valid.tgt\") as f:\n",
    "    val_tgt = f.readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "04e68a8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"workspace/lean_proof_recording/data/cleaned_training_data/valid.src\") as f:\n",
    "    val_src = f.readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "fccfc8a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"workspace/lean_proof_recording/data/cleaned_training_data/test.src\") as f:\n",
    "    test_src = f.readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1e061617",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"workspace/lean_proof_recording/data/cleaned_training_data/test.tgt\") as f:\n",
    "    test_tgt = f.readlines()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "1ed37b6e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "7291"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(test_tgt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "25015d8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "b2329888",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.utils.data import Dataset\n",
    "class TacticDataset(Dataset):\n",
    "    def __init__(self, txt_list,tgt_list, tokenizer):\n",
    "        self.input_ids = []\n",
    "        self.attn_masks = []\n",
    "        self.labels = []\n",
    "        for txt,txt2 in zip(txt_list,tgt_list):\n",
    "            # Encode the descriptions using the GPT-Neo tokenizer\n",
    "            encodings_dict = tokenizer('<|startoftext|>'\n",
    "                                        + ' <GOAL> ' + txt +  ' <PROOFSTEP> ' +  txt2+ \n",
    "                                        '<|endoftext|>',\n",
    "                                        truncation=True,\n",
    "                                        padding=\"max_length\")\n",
    "            \n",
    "            input_ids = torch.tensor(encodings_dict['input_ids'])    \n",
    "            self.input_ids.append(input_ids)\n",
    "            mask = torch.tensor(encodings_dict['attention_mask'])\n",
    "            self.attn_masks.append(mask)\n",
    "#             break\n",
    "    def __len__(self):\n",
    "        return len(self.input_ids)\n",
    "    def __getitem__(self, idx):\n",
    "        return self.input_ids[idx], self.attn_masks[idx]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "5d62012c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# del val_src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "856a11b4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "13"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import gc\n",
    "gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "11508122",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "validDataset = TacticDataset(val_src,val_tgt,tokenizer)\n",
    "testDataset = TacticDataset(test_src,test_tgt,tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b0bb272c",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainDataset = TacticDataset(train_src,train_tgt,tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "5e994dbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# x = TacticDataset(test_src,test_tgt,tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "683109ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "# x[0][0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "ae215c09",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "from transformers import Trainer\n",
    "from  transformers import TrainingArguments, Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38b57afa",
   "metadata": {},
   "outputs": [],
   "source": [
    "{\n",
    "  \"train_batch_size\": 15,\n",
    "  \"fp16\": {\n",
    "    \"enabled\": true,\n",
    "    \"min_loss_scale\": 1,\n",
    "    \"opt_level\": \"O2\"\n",
    "  },\n",
    "  \"zero_optimization\": {\n",
    "    \"stage\": 2,\n",
    "    \"offload_param\": {\n",
    "      \"device\": \"cpu\"\n",
    "    },\n",
    "    \"offload_optimizer\": {\n",
    "      \"device\": \"cpu\"\n",
    "    },\n",
    "    \"allgather_partitions\": true,\n",
    "    \"allgather_bucket_size\": 5e8,\n",
    "    \"contiguous_gradients\": true\n",
    "  },\n",
    "  \"optimizer\": {\n",
    "    \"type\": \"AdamW\",\n",
    "    \"params\": {\n",
    "      \"lr\": 5e-05,\n",
    "      \"betas\": [\n",
    "        0.9,\n",
    "        0.999\n",
    "      ],\n",
    "      \"eps\": 1e-08\n",
    "    }\n",
    "  },\n",
    "  \"scheduler\": {\n",
    "    \"type\": \"WarmupLR\",\n",
    "    \"params\": {\n",
    "      \"warmup_min_lr\": 0,\n",
    "      \"warmup_max_lr\": 5e-05,\n",
    "      \"warmup_num_steps\": 50\n",
    "    }\n",
    "  }\n",
    "}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "499308c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, TrainingArguments, Trainer, AutoModelForCausalLM, IntervalStrategy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "a78f9777",
   "metadata": {},
   "outputs": [],
   "source": [
    "training_args = TrainingArguments(output_dir='./results',\n",
    "                                  num_train_epochs=5,\n",
    "                                  logging_steps=5000,\n",
    "                                  save_steps=20000,                                   \n",
    "                                  per_device_train_batch_size=1,\n",
    "                                  per_device_eval_batch_size=1,\n",
    "                                  warmup_steps=100,\n",
    "                                  weight_decay=0.01,  \n",
    "                                  logging_dir='./logs')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "55fb65ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import mpi4py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "4adb63c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(model=model, args=training_args,  \n",
    "                  train_dataset=trainDataset,\n",
    "                  eval_dataset=validDataset, \n",
    "                  # This custom collate function is necessary \n",
    "                  # to built batches of data\n",
    "                  data_collator=lambda data: \n",
    "              {'input_ids': torch.stack([f[0] for f in data]),       \n",
    "               'attention_mask': torch.stack([f[1] for f in data]),\n",
    "               'labels': torch.stack([f[0] for f in data])})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "97f1b9cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "PyTorch: setting up devices\n",
      "The default value for the training argument `--report_to` will change in v5 (from all installed integrations to none). In v5, you will need to use `--report_to all` to get the same behavior as now. You should start updating your code and make this info disappear :-).\n"
     ]
    }
   ],
   "source": [
    "training_args = TrainingArguments(output_dir='./results',\n",
    "                                  num_train_epochs=5,\n",
    "                                  logging_steps=5000,\n",
    "                                  save_steps=20000,                                   \n",
    "                                  per_device_train_batch_size=6,\n",
    "                                  per_device_eval_batch_size=6,\n",
    "                                  warmup_steps=100,\n",
    "                                  weight_decay=0.01,  \n",
    "                                  logging_dir='./logs',fp16= True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "1c465bfc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using amp half precision backend\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(model=model, args=training_args,  \n",
    "                  train_dataset=trainDataset,\n",
    "                  eval_dataset=validDataset, \n",
    "                  # This custom collate function is necessary \n",
    "                  # to built batches of data\n",
    "                  data_collator=lambda data: \n",
    "              {'input_ids': torch.stack([f[0] for f in data]),       \n",
    "               'attention_mask': torch.stack([f[1] for f in data]),\n",
    "               'labels': torch.stack([f[0] for f in data])})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "22f4059a",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "***** Running training *****\n",
      "  Num examples = 168649\n",
      "  Num Epochs = 5\n",
      "  Instantaneous batch size per device = 6\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 6\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 140545\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='19' max='140545' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [    19/140545 00:14 < 34:01:47, 1.15 it/s, Epoch 0.00/5]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fc655a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "1+!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf646b05",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "422ad755",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "902240e7",
   "metadata": {},
   "outputs": [],
   "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.8.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}