structTuningStarCoder2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8d068150",
   "metadata": {},
   "source": [
    "This notebook is calling [StarCoder2-3b-NEAR](https://huggingface.co/bigcode/starcoder2-3b) from the huggingface platform, a code LLM pre-trained on a code dataset called [The Stack v2 smol](https://huggingface.co/datasets/bigcode/the-stack-v2-train-smol-ids) that includes 17 programming languages (C, C#, C++, Go, Java, JavaScript, Kotlin, Lua, PHP, Python, R, Ruby, Rust, SQL, Shell, Swift, TypeScript) and incorporates 7 code documentation languages (including HTML), and further trained in a continued pre-training process on the NEAR Protocol blockchain from the [preTrainingNEAR](https://huggingface.co/datasets/jcarbonnell/preTrainingNEAR) dataset.\n",
    "\n",
    "The Structure-Aware fine-tuning step consists in instructing the model with NEAR dApps trees and their corresponding description in natural language extracted from readme files transformed into user prompts. This structure-aware approach aim to give the pre-trained model a good knowledge of the whole dApp logic, and a corresponding user prompt, so that when a user will ask the model to create an app on the NEAR Protocol blockchain, the LLM will answer with a proposition of an app tree augmented with a extensive description of the functionality of each tree branches. \n",
    "\n",
    "What we are targetting here is to overpass the limitation of the 'next-token prediction' logic, that can spin the model in 'dumb loops' while iterating over coding complex code completion. It will also be of great support for code understanding on both the user side and the model side, since the model and the user will be aware of the big picture at first. \n",
    "\n",
    "More details on the model original StarCoder2 is available in the official paper [StarCoder 2 and The Stack v2: The Next Generation](https://arxiv.org/abs/2402.19173) (2024)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "65eab062",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from datasets import Dataset, DatasetDict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "20ced6ee",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'load_dataset' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[55], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# Load the dataset from Hugging Face\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m dataset \u001b[38;5;241m=\u001b[39m load_dataset(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mjcarbonnell/structTuningNEAR\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m      4\u001b[0m \u001b[38;5;66;03m# Verify the loaded dataset\u001b[39;00m\n\u001b[1;32m      5\u001b[0m \u001b[38;5;28mprint\u001b[39m(dataset)\n",
      "\u001b[0;31mNameError\u001b[0m: name 'load_dataset' is not defined"
     ]
    }
   ],
   "source": [
    "# Load the dataset from Hugging Face\n",
    "dataset = load_dataset('jcarbonnell/structTuningNEAR')\n",
    "\n",
    "# Verify the loaded dataset\n",
    "print(dataset)\n",
    "\n",
    "# Print the first example from the training dataset\n",
    "print(dataset['train'][0])  # Print the first example from the train dataset\n",
    "print(dataset['val'][0])  # Print the first example from the validation dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "6571149a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/anaconda3/lib/python3.11/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9b7b690297fa442185dacf3a0008251a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/1022 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "74881d189db34b1fb50a17d17f647ac0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/114 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Initialize the tokenizer and the model\n",
    "model_name = \"bigcode/starcoder2-3b\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_name)\n",
    "\n",
    "# Set the pad_token to eos_token if not already set\n",
    "if tokenizer.pad_token is None:\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# Tokenize the dataset into chunks of 64 tokens max\n",
    "def tokenize_function(examples):\n",
    "    tokenized_inputs = tokenizer(examples['text'], truncation=True, padding=\"longest\", max_length=64, return_special_tokens_mask=True)\n",
    "    tokenized_inputs['labels'] = tokenized_inputs['input_ids'].copy()\n",
    "    return tokenized_inputs\n",
    "\n",
    "# Apply the tokenization\n",
    "tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=[\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34a8fa62",
   "metadata": {},
   "source": [
    "#### Chunk the tokenized dataset into 64 tokens to avoid crashing\n",
    "def group_texts(examples):\n",
    "    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}\n",
    "    total_length = len(concatenated_examples[list(examples.keys())[0]])\n",
    "    if total_length >= 64:\n",
    "        total_length = (total_length // 64) * 64\n",
    "    result = {\n",
    "        k: [t[i:i + 64] for i in range(0, total_length, 64)]\n",
    "        for k, t in concatenated_examples.items()\n",
    "    }\n",
    "    result[\"labels\"] = result[\"input_ids\"].copy()\n",
    "    return result\n",
    "\n",
    "#### Apply chunking\n",
    "tokenized_datasets = tokenized_datasets.map(group_texts, batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "a5141fbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# setup training arguments\n",
    "training_args = TrainingArguments(\n",
    "    output_dir='./results',\n",
    "    overwrite_output_dir=True,\n",
    "    num_train_epochs=5,\n",
    "    per_device_train_batch_size=1,  # reduced batch size 4 to 1 to avoid crashing\n",
    "    per_device_eval_batch_size=1,\n",
    "    evaluation_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    learning_rate=2e-5,\n",
    "    logging_dir='./logs',\n",
    "    logging_steps=10,\n",
    "    gradient_accumulation_steps=8,  # Simulate larger batch size\n",
    "    #fp16=True,  # mixed precision training not supported on mps\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "de1d227b",
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "MPS backend out of memory (MPS allocated: 18.12 GB, other allocations: 704.00 KB, max allowed: 18.13 GB). Tried to allocate 36.00 MB on private pool. Use PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0 to disable upper limit for memory allocations (may cause system failure).",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[35], line 2\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# create the trainer and start training\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m trainer \u001b[38;5;241m=\u001b[39m Trainer(\n\u001b[1;32m      3\u001b[0m     model\u001b[38;5;241m=\u001b[39mmodel,\n\u001b[1;32m      4\u001b[0m     args\u001b[38;5;241m=\u001b[39mtraining_args,\n\u001b[1;32m      5\u001b[0m     train_dataset\u001b[38;5;241m=\u001b[39mtokenized_datasets[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mtrain\u001b[39m\u001b[38;5;124m'\u001b[39m],\n\u001b[1;32m      6\u001b[0m     eval_dataset\u001b[38;5;241m=\u001b[39mtokenized_datasets[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124mval\u001b[39m\u001b[38;5;124m'\u001b[39m],\n\u001b[1;32m      7\u001b[0m )\n\u001b[1;32m      9\u001b[0m trainer\u001b[38;5;241m.\u001b[39mtrain()\n",
      "File \u001b[0;32m/opt/anaconda3/lib/python3.11/site-packages/transformers/trainer.py:514\u001b[0m, in \u001b[0;36mTrainer.__init__\u001b[0;34m(self, model, args, data_collator, train_dataset, eval_dataset, tokenizer, model_init, compute_metrics, callbacks, optimizers, preprocess_logits_for_metrics)\u001b[0m\n\u001b[1;32m    509\u001b[0m \u001b[38;5;66;03m# Bnb Quantized models doesn't support `.to` operation.\u001b[39;00m\n\u001b[1;32m    510\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[1;32m    511\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mplace_model_on_device\n\u001b[1;32m    512\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mgetattr\u001b[39m(model, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mquantization_method\u001b[39m\u001b[38;5;124m\"\u001b[39m, \u001b[38;5;28;01mNone\u001b[39;00m) \u001b[38;5;241m==\u001b[39m QuantizationMethod\u001b[38;5;241m.\u001b[39mBITS_AND_BYTES\n\u001b[1;32m    513\u001b[0m ):\n\u001b[0;32m--> 514\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_move_model_to_device(model, args\u001b[38;5;241m.\u001b[39mdevice)\n\u001b[1;32m    516\u001b[0m \u001b[38;5;66;03m# Force n_gpu to 1 to avoid DataParallel as MP will manage the GPUs\u001b[39;00m\n\u001b[1;32m    517\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mis_model_parallel:\n",
      "File \u001b[0;32m/opt/anaconda3/lib/python3.11/site-packages/transformers/trainer.py:757\u001b[0m, in \u001b[0;36mTrainer._move_model_to_device\u001b[0;34m(self, model, device)\u001b[0m\n\u001b[1;32m    756\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_move_model_to_device\u001b[39m(\u001b[38;5;28mself\u001b[39m, model, device):\n\u001b[0;32m--> 757\u001b[0m     model \u001b[38;5;241m=\u001b[39m model\u001b[38;5;241m.\u001b[39mto(device)\n\u001b[1;32m    758\u001b[0m     \u001b[38;5;66;03m# Moving a model to an XLA device disconnects the tied weights, so we have to retie them.\u001b[39;00m\n\u001b[1;32m    759\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39margs\u001b[38;5;241m.\u001b[39mparallel_mode \u001b[38;5;241m==\u001b[39m ParallelMode\u001b[38;5;241m.\u001b[39mTPU \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mhasattr\u001b[39m(model, \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtie_weights\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n",
      "File \u001b[0;32m/opt/anaconda3/lib/python3.11/site-packages/transformers/modeling_utils.py:2692\u001b[0m, in \u001b[0;36mPreTrainedModel.to\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   2687\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m dtype_present_in_args:\n\u001b[1;32m   2688\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m   2689\u001b[0m             \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mYou cannot cast a GPTQ model in a new `dtype`. Make sure to load the model using `from_pretrained` using the desired\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   2690\u001b[0m             \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m `dtype` by passing the correct `torch_dtype` argument.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m   2691\u001b[0m         )\n\u001b[0;32m-> 2692\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28msuper\u001b[39m()\u001b[38;5;241m.\u001b[39mto(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[0;32m/opt/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py:1173\u001b[0m, in \u001b[0;36mModule.to\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1170\u001b[0m         \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m   1171\u001b[0m             \u001b[38;5;28;01mraise\u001b[39;00m\n\u001b[0;32m-> 1173\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_apply(convert)\n",
      "File \u001b[0;32m/opt/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py:779\u001b[0m, in \u001b[0;36mModule._apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    777\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m recurse:\n\u001b[1;32m    778\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchildren():\n\u001b[0;32m--> 779\u001b[0m         module\u001b[38;5;241m.\u001b[39m_apply(fn)\n\u001b[1;32m    781\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcompute_should_use_set_data\u001b[39m(tensor, tensor_applied):\n\u001b[1;32m    782\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_has_compatible_shallow_copy_type(tensor, tensor_applied):\n\u001b[1;32m    783\u001b[0m         \u001b[38;5;66;03m# If the new tensor has compatible tensor type as the existing tensor,\u001b[39;00m\n\u001b[1;32m    784\u001b[0m         \u001b[38;5;66;03m# the current behavior is to change the tensor in-place using `.data =`,\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    789\u001b[0m         \u001b[38;5;66;03m# global flag to let the user control whether they want the future\u001b[39;00m\n\u001b[1;32m    790\u001b[0m         \u001b[38;5;66;03m# behavior of overwriting the existing tensor or not.\u001b[39;00m\n",
      "File \u001b[0;32m/opt/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py:779\u001b[0m, in \u001b[0;36mModule._apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    777\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m recurse:\n\u001b[1;32m    778\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchildren():\n\u001b[0;32m--> 779\u001b[0m         module\u001b[38;5;241m.\u001b[39m_apply(fn)\n\u001b[1;32m    781\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcompute_should_use_set_data\u001b[39m(tensor, tensor_applied):\n\u001b[1;32m    782\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_has_compatible_shallow_copy_type(tensor, tensor_applied):\n\u001b[1;32m    783\u001b[0m         \u001b[38;5;66;03m# If the new tensor has compatible tensor type as the existing tensor,\u001b[39;00m\n\u001b[1;32m    784\u001b[0m         \u001b[38;5;66;03m# the current behavior is to change the tensor in-place using `.data =`,\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    789\u001b[0m         \u001b[38;5;66;03m# global flag to let the user control whether they want the future\u001b[39;00m\n\u001b[1;32m    790\u001b[0m         \u001b[38;5;66;03m# behavior of overwriting the existing tensor or not.\u001b[39;00m\n",
      "    \u001b[0;31m[... skipping similar frames: Module._apply at line 779 (2 times)]\u001b[0m\n",
      "File \u001b[0;32m/opt/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py:779\u001b[0m, in \u001b[0;36mModule._apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    777\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m recurse:\n\u001b[1;32m    778\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m module \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mchildren():\n\u001b[0;32m--> 779\u001b[0m         module\u001b[38;5;241m.\u001b[39m_apply(fn)\n\u001b[1;32m    781\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mcompute_should_use_set_data\u001b[39m(tensor, tensor_applied):\n\u001b[1;32m    782\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m torch\u001b[38;5;241m.\u001b[39m_has_compatible_shallow_copy_type(tensor, tensor_applied):\n\u001b[1;32m    783\u001b[0m         \u001b[38;5;66;03m# If the new tensor has compatible tensor type as the existing tensor,\u001b[39;00m\n\u001b[1;32m    784\u001b[0m         \u001b[38;5;66;03m# the current behavior is to change the tensor in-place using `.data =`,\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    789\u001b[0m         \u001b[38;5;66;03m# global flag to let the user control whether they want the future\u001b[39;00m\n\u001b[1;32m    790\u001b[0m         \u001b[38;5;66;03m# behavior of overwriting the existing tensor or not.\u001b[39;00m\n",
      "File \u001b[0;32m/opt/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py:804\u001b[0m, in \u001b[0;36mModule._apply\u001b[0;34m(self, fn, recurse)\u001b[0m\n\u001b[1;32m    800\u001b[0m \u001b[38;5;66;03m# Tensors stored in modules are graph leaves, and we don't want to\u001b[39;00m\n\u001b[1;32m    801\u001b[0m \u001b[38;5;66;03m# track autograd history of `param_applied`, so we have to use\u001b[39;00m\n\u001b[1;32m    802\u001b[0m \u001b[38;5;66;03m# `with torch.no_grad():`\u001b[39;00m\n\u001b[1;32m    803\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m torch\u001b[38;5;241m.\u001b[39mno_grad():\n\u001b[0;32m--> 804\u001b[0m     param_applied \u001b[38;5;241m=\u001b[39m fn(param)\n\u001b[1;32m    805\u001b[0m p_should_use_set_data \u001b[38;5;241m=\u001b[39m compute_should_use_set_data(param, param_applied)\n\u001b[1;32m    807\u001b[0m \u001b[38;5;66;03m# subclasses may have multiple child tensors so we need to use swap_tensors\u001b[39;00m\n",
      "File \u001b[0;32m/opt/anaconda3/lib/python3.11/site-packages/torch/nn/modules/module.py:1159\u001b[0m, in \u001b[0;36mModule.to.<locals>.convert\u001b[0;34m(t)\u001b[0m\n\u001b[1;32m   1152\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m convert_to_format \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;129;01mand\u001b[39;00m t\u001b[38;5;241m.\u001b[39mdim() \u001b[38;5;129;01min\u001b[39;00m (\u001b[38;5;241m4\u001b[39m, \u001b[38;5;241m5\u001b[39m):\n\u001b[1;32m   1153\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m t\u001b[38;5;241m.\u001b[39mto(\n\u001b[1;32m   1154\u001b[0m             device,\n\u001b[1;32m   1155\u001b[0m             dtype \u001b[38;5;28;01mif\u001b[39;00m t\u001b[38;5;241m.\u001b[39mis_floating_point() \u001b[38;5;129;01mor\u001b[39;00m t\u001b[38;5;241m.\u001b[39mis_complex() \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m   1156\u001b[0m             non_blocking,\n\u001b[1;32m   1157\u001b[0m             memory_format\u001b[38;5;241m=\u001b[39mconvert_to_format,\n\u001b[1;32m   1158\u001b[0m         )\n\u001b[0;32m-> 1159\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m t\u001b[38;5;241m.\u001b[39mto(\n\u001b[1;32m   1160\u001b[0m         device,\n\u001b[1;32m   1161\u001b[0m         dtype \u001b[38;5;28;01mif\u001b[39;00m t\u001b[38;5;241m.\u001b[39mis_floating_point() \u001b[38;5;129;01mor\u001b[39;00m t\u001b[38;5;241m.\u001b[39mis_complex() \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m   1162\u001b[0m         non_blocking,\n\u001b[1;32m   1163\u001b[0m     )\n\u001b[1;32m   1164\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mNotImplementedError\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[1;32m   1165\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mstr\u001b[39m(e) \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCannot copy out of meta tensor; no data!\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n",
      "\u001b[0;31mRuntimeError\u001b[0m: MPS backend out of memory (MPS allocated: 18.12 GB, other allocations: 704.00 KB, max allowed: 18.13 GB). Tried to allocate 36.00 MB on private pool. Use PYTORCH_MPS_HIGH_WATERMARK_RATIO=0.0 to disable upper limit for memory allocations (may cause system failure)."
     ]
    }
   ],
   "source": [
    "# create the trainer and start training\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_datasets['train'],\n",
    "    eval_dataset=tokenized_datasets['val'],\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c2d90ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# model evaluation\n",
    "import math\n",
    "\n",
    "eval_results = trainer.evaluate()\n",
    "perplexity = math.exp(eval_results['eval_loss'])\n",
    "print(f\"Perplexity: {perplexity}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc022617",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Assuming you have access to the training logs\n",
    "training_loss = [log['loss'] for log in trainer.state.log_history if 'loss' in log]\n",
    "validation_loss = [log['eval_loss'] for log in trainer.state.log_history if 'eval_loss' in log]\n",
    "epochs = range(len(training_loss))\n",
    "\n",
    "plt.plot(epochs, training_loss, label='Training Loss')\n",
    "plt.plot(epochs, validation_loss, label='Validation Loss')\n",
    "plt.xlabel('Epochs')\n",
    "plt.ylabel('Loss')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ac07371",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"Once upon a time\"\n",
    "inputs = tokenizer(prompt, return_tensors='pt')\n",
    "output = model.generate(inputs['input_ids'], max_length=50)\n",
    "generated_text = tokenizer.decode(output[0], skip_special_tokens=True)\n",
    "print(generated_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "280f9190",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}