llama32_python_code.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-02-18T08:08:47.378057Z",
     "iopub.status.busy": "2025-02-18T08:08:47.377596Z",
     "iopub.status.idle": "2025-02-18T08:09:07.854047Z",
     "shell.execute_reply": "2025-02-18T08:09:07.853541Z",
     "shell.execute_reply.started": "2025-02-18T08:08:47.378033Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: datasets in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 1)) (2.14.5)\n",
      "Collecting trl (from -r requirements.txt (line 2))\n",
      "  Downloading trl-0.15.0-py3-none-any.whl.metadata (11 kB)\n",
      "Requirement already satisfied: peft in /usr/local/lib/python3.11/dist-packages (from -r requirements.txt (line 3)) (0.6.2)\n",
      "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.11/dist-packages (from datasets->-r requirements.txt (line 1)) (1.26.3)\n",
      "Requirement already satisfied: pyarrow>=8.0.0 in /usr/local/lib/python3.11/dist-packages (from datasets->-r requirements.txt (line 1)) (15.0.0)\n",
      "Requirement already satisfied: dill<0.3.8,>=0.3.0 in /usr/local/lib/python3.11/dist-packages (from datasets->-r requirements.txt (line 1)) (0.3.7)\n",
      "Requirement already satisfied: pandas in /usr/local/lib/python3.11/dist-packages (from datasets->-r requirements.txt (line 1)) (2.2.0)\n",
      "Requirement already satisfied: requests>=2.19.0 in /usr/local/lib/python3.11/dist-packages (from datasets->-r requirements.txt (line 1)) (2.31.0)\n",
      "Requirement already satisfied: tqdm>=4.62.1 in /usr/local/lib/python3.11/dist-packages (from datasets->-r requirements.txt (line 1)) (4.66.1)\n",
      "Requirement already satisfied: xxhash in /usr/local/lib/python3.11/dist-packages (from datasets->-r requirements.txt (line 1)) (3.4.1)\n",
      "Requirement already satisfied: multiprocess in /usr/local/lib/python3.11/dist-packages (from datasets->-r requirements.txt (line 1)) (0.70.15)\n",
      "Requirement already satisfied: fsspec<2023.9.0,>=2023.1.0 in /usr/local/lib/python3.11/dist-packages (from fsspec[http]<2023.9.0,>=2023.1.0->datasets->-r requirements.txt (line 1)) (2023.6.0)\n",
      "Requirement already satisfied: aiohttp in /usr/local/lib/python3.11/dist-packages (from datasets->-r requirements.txt (line 1)) (3.9.1)\n",
      "Requirement already satisfied: huggingface-hub<1.0.0,>=0.14.0 in /usr/local/lib/python3.11/dist-packages (from datasets->-r requirements.txt (line 1)) (0.20.3)\n",
      "Requirement already satisfied: packaging in /usr/local/lib/python3.11/dist-packages (from datasets->-r requirements.txt (line 1)) (23.2)\n",
      "Requirement already satisfied: pyyaml>=5.1 in /usr/lib/python3/dist-packages (from datasets->-r requirements.txt (line 1)) (5.4.1)\n",
      "Collecting accelerate>=0.34.0 (from trl->-r requirements.txt (line 2))\n",
      "  Downloading accelerate-1.4.0-py3-none-any.whl.metadata (19 kB)\n",
      "Collecting datasets (from -r requirements.txt (line 1))\n",
      "  Downloading datasets-3.3.1-py3-none-any.whl.metadata (19 kB)\n",
      "Collecting rich (from trl->-r requirements.txt (line 2))\n",
      "  Downloading rich-13.9.4-py3-none-any.whl.metadata (18 kB)\n",
      "Collecting transformers>=4.46.0 (from trl->-r requirements.txt (line 2))\n",
      "  Downloading transformers-4.49.0-py3-none-any.whl.metadata (44 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m44.0/44.0 kB\u001b[0m \u001b[31m5.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from datasets->-r requirements.txt (line 1)) (3.13.1)\n",
      "Collecting requests>=2.32.2 (from datasets->-r requirements.txt (line 1))\n",
      "  Downloading requests-2.32.3-py3-none-any.whl.metadata (4.6 kB)\n",
      "Collecting tqdm>=4.66.3 (from datasets->-r requirements.txt (line 1))\n",
      "  Downloading tqdm-4.67.1-py3-none-any.whl.metadata (57 kB)\n",
      "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m57.7/57.7 kB\u001b[0m \u001b[31m14.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hCollecting huggingface-hub>=0.24.0 (from datasets->-r requirements.txt (line 1))\n",
      "  Downloading huggingface_hub-0.28.1-py3-none-any.whl.metadata (13 kB)\n",
      "Requirement already satisfied: psutil in /usr/local/lib/python3.11/dist-packages (from peft->-r requirements.txt (line 3)) (5.9.8)\n",
      "Requirement already satisfied: torch>=1.13.0 in /usr/local/lib/python3.11/dist-packages (from peft->-r requirements.txt (line 3)) (2.1.1+cu121)\n",
      "Requirement already satisfied: safetensors in /usr/local/lib/python3.11/dist-packages (from peft->-r requirements.txt (line 3)) (0.4.0)\n",
      "Collecting safetensors (from peft->-r requirements.txt (line 3))\n",
      "  Downloading safetensors-0.5.2-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (3.8 kB)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets->-r requirements.txt (line 1)) (23.1.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets->-r requirements.txt (line 1)) (6.0.4)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets->-r requirements.txt (line 1)) (1.9.4)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets->-r requirements.txt (line 1)) (1.4.1)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.11/dist-packages (from aiohttp->datasets->-r requirements.txt (line 1)) (1.3.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.11/dist-packages (from huggingface-hub>=0.24.0->datasets->-r requirements.txt (line 1)) (4.9.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.11/dist-packages (from requests>=2.32.2->datasets->-r requirements.txt (line 1)) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /usr/lib/python3/dist-packages (from requests>=2.32.2->datasets->-r requirements.txt (line 1)) (3.3)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.11/dist-packages (from requests>=2.32.2->datasets->-r requirements.txt (line 1)) (2.0.7)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /usr/lib/python3/dist-packages (from requests>=2.32.2->datasets->-r requirements.txt (line 1)) (2020.6.20)\n",
      "Requirement already satisfied: sympy in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.0->peft->-r requirements.txt (line 3)) (1.12)\n",
      "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.0->peft->-r requirements.txt (line 3)) (3.2.1)\n",
      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.0->peft->-r requirements.txt (line 3)) (3.1.3)\n",
      "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.11/dist-packages (from torch>=1.13.0->peft->-r requirements.txt (line 3)) (2.1.0)\n",
      "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.11/dist-packages (from transformers>=4.46.0->trl->-r requirements.txt (line 2)) (2023.12.25)\n",
      "Collecting tokenizers<0.22,>=0.21 (from transformers>=4.46.0->trl->-r requirements.txt (line 2))\n",
      "  Downloading tokenizers-0.21.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (6.7 kB)\n",
      "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.11/dist-packages (from pandas->datasets->-r requirements.txt (line 1)) (2.8.2)\n",
      "Requirement already satisfied: pytz>=2020.1 in /usr/lib/python3/dist-packages (from pandas->datasets->-r requirements.txt (line 1)) (2022.1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /usr/local/lib/python3.11/dist-packages (from pandas->datasets->-r requirements.txt (line 1)) (2023.4)\n",
      "Collecting markdown-it-py>=2.2.0 (from rich->trl->-r requirements.txt (line 2))\n",
      "  Downloading markdown_it_py-3.0.0-py3-none-any.whl.metadata (6.9 kB)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.11/dist-packages (from rich->trl->-r requirements.txt (line 2)) (2.17.2)\n",
      "Collecting mdurl~=0.1 (from markdown-it-py>=2.2.0->rich->trl->-r requirements.txt (line 2))\n",
      "  Downloading mdurl-0.1.2-py3-none-any.whl.metadata (1.6 kB)\n",
      "Requirement already satisfied: six>=1.5 in /usr/lib/python3/dist-packages (from python-dateutil>=2.8.2->pandas->datasets->-r requirements.txt (line 1)) (1.16.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2->torch>=1.13.0->peft->-r requirements.txt (line 3)) (2.1.4)\n",
      "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.11/dist-packages (from sympy->torch>=1.13.0->peft->-r requirements.txt (line 3)) (1.3.0)\n",
      "Downloading trl-0.15.0-py3-none-any.whl (318 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m318.3/318.3 kB\u001b[0m \u001b[31m56.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading datasets-3.3.1-py3-none-any.whl (484 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m484.9/484.9 kB\u001b[0m \u001b[31m80.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading accelerate-1.4.0-py3-none-any.whl (342 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m342.1/342.1 kB\u001b[0m \u001b[31m70.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading huggingface_hub-0.28.1-py3-none-any.whl (464 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m464.1/464.1 kB\u001b[0m \u001b[31m78.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading requests-2.32.3-py3-none-any.whl (64 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m64.9/64.9 kB\u001b[0m \u001b[31m19.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading safetensors-0.5.2-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (461 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m462.0/462.0 kB\u001b[0m \u001b[31m76.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading tqdm-4.67.1-py3-none-any.whl (78 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m78.5/78.5 kB\u001b[0m \u001b[31m19.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading transformers-4.49.0-py3-none-any.whl (10.0 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m10.0/10.0 MB\u001b[0m \u001b[31m153.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hDownloading rich-13.9.4-py3-none-any.whl (242 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m242.4/242.4 kB\u001b[0m \u001b[31m47.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading markdown_it_py-3.0.0-py3-none-any.whl (87 kB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m87.5/87.5 kB\u001b[0m \u001b[31m23.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading tokenizers-0.21.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.0 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.0/3.0 MB\u001b[0m \u001b[31m121.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hDownloading mdurl-0.1.2-py3-none-any.whl (10.0 kB)\n",
      "Installing collected packages: tqdm, safetensors, requests, mdurl, markdown-it-py, huggingface-hub, tokenizers, rich, accelerate, transformers, datasets, trl\n",
      "  Attempting uninstall: tqdm\n",
      "    Found existing installation: tqdm 4.66.1\n",
      "    Uninstalling tqdm-4.66.1:\n",
      "      Successfully uninstalled tqdm-4.66.1\n",
      "  Attempting uninstall: safetensors\n",
      "    Found existing installation: safetensors 0.4.0\n",
      "    Uninstalling safetensors-0.4.0:\n",
      "      Successfully uninstalled safetensors-0.4.0\n",
      "  Attempting uninstall: requests\n",
      "    Found existing installation: requests 2.31.0\n",
      "    Uninstalling requests-2.31.0:\n",
      "      Successfully uninstalled requests-2.31.0\n",
      "  Attempting uninstall: huggingface-hub\n",
      "    Found existing installation: huggingface-hub 0.20.3\n",
      "    Uninstalling huggingface-hub-0.20.3:\n",
      "      Successfully uninstalled huggingface-hub-0.20.3\n",
      "  Attempting uninstall: tokenizers\n",
      "    Found existing installation: tokenizers 0.15.1\n",
      "    Uninstalling tokenizers-0.15.1:\n",
      "      Successfully uninstalled tokenizers-0.15.1\n",
      "  Attempting uninstall: accelerate\n",
      "    Found existing installation: accelerate 0.24.1\n",
      "    Uninstalling accelerate-0.24.1:\n",
      "      Successfully uninstalled accelerate-0.24.1\n",
      "  Attempting uninstall: transformers\n",
      "    Found existing installation: transformers 4.35.2\n",
      "    Uninstalling transformers-4.35.2:\n",
      "      Successfully uninstalled transformers-4.35.2\n",
      "  Attempting uninstall: datasets\n",
      "    Found existing installation: datasets 2.14.5\n",
      "    Uninstalling datasets-2.14.5:\n",
      "      Successfully uninstalled datasets-2.14.5\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "gradient 2.0.6 requires attrs<=19, but you have attrs 23.1.0 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0mSuccessfully installed accelerate-1.4.0 datasets-3.3.1 huggingface-hub-0.28.1 markdown-it-py-3.0.0 mdurl-0.1.2 requests-2.32.3 rich-13.9.4 safetensors-0.5.2 tokenizers-0.21.0 tqdm-4.67.1 transformers-4.49.0 trl-0.15.0\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0mRequirement already satisfied: bitsandbytes in /usr/local/lib/python3.11/dist-packages (0.41.2)\n",
      "Collecting bitsandbytes\n",
      "  Downloading bitsandbytes-0.45.2-py3-none-manylinux_2_24_x86_64.whl.metadata (5.8 kB)\n",
      "Requirement already satisfied: torch<3,>=2.0 in /usr/local/lib/python3.11/dist-packages (from bitsandbytes) (2.1.1+cu121)\n",
      "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.11/dist-packages (from bitsandbytes) (1.26.3)\n",
      "Requirement already satisfied: filelock in /usr/local/lib/python3.11/dist-packages (from torch<3,>=2.0->bitsandbytes) (3.13.1)\n",
      "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.11/dist-packages (from torch<3,>=2.0->bitsandbytes) (4.9.0)\n",
      "Requirement already satisfied: sympy in /usr/local/lib/python3.11/dist-packages (from torch<3,>=2.0->bitsandbytes) (1.12)\n",
      "Requirement already satisfied: networkx in /usr/local/lib/python3.11/dist-packages (from torch<3,>=2.0->bitsandbytes) (3.2.1)\n",
      "Requirement already satisfied: jinja2 in /usr/local/lib/python3.11/dist-packages (from torch<3,>=2.0->bitsandbytes) (3.1.3)\n",
      "Requirement already satisfied: fsspec in /usr/local/lib/python3.11/dist-packages (from torch<3,>=2.0->bitsandbytes) (2023.6.0)\n",
      "Requirement already satisfied: triton==2.1.0 in /usr/local/lib/python3.11/dist-packages (from torch<3,>=2.0->bitsandbytes) (2.1.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.11/dist-packages (from jinja2->torch<3,>=2.0->bitsandbytes) (2.1.4)\n",
      "Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.11/dist-packages (from sympy->torch<3,>=2.0->bitsandbytes) (1.3.0)\n",
      "Downloading bitsandbytes-0.45.2-py3-none-manylinux_2_24_x86_64.whl (69.7 MB)\n",
      "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m69.7/69.7 MB\u001b[0m \u001b[31m47.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m:00:01\u001b[0m00:01\u001b[0m\n",
      "\u001b[?25hInstalling collected packages: bitsandbytes\n",
      "  Attempting uninstall: bitsandbytes\n",
      "    Found existing installation: bitsandbytes 0.41.2\n",
      "    Uninstalling bitsandbytes-0.41.2:\n",
      "      Successfully uninstalled bitsandbytes-0.41.2\n",
      "Successfully installed bitsandbytes-0.45.2\n",
      "\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "!pip install -r requirements.txt\n",
    "!pip install --upgrade bitsandbytes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2025-02-18T08:09:07.855191Z",
     "iopub.status.busy": "2025-02-18T08:09:07.854983Z",
     "iopub.status.idle": "2025-02-18T08:09:08.149020Z",
     "shell.execute_reply": "2025-02-18T08:09:08.148460Z",
     "shell.execute_reply.started": "2025-02-18T08:09:07.855174Z"
    },
    "id": "W2yGmrJ9LZgl",
    "outputId": "c577413f-968b-4af1-c918-2793b1ad9dc5"
   },
   "outputs": [],
   "source": [
    "from huggingface_hub import login\n",
    "import json\n",
    "\n",
    "with open(\"config.json\", \"r\") as config_file:\n",
    "    config = json.load(config_file)\n",
    "    access_token = config[\"HF_ACCESS_TOKEN\"]\n",
    "\n",
    "login(token=access_token)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wUDTL3n8MLLv"
   },
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-02-18T08:09:08.150048Z",
     "iopub.status.busy": "2025-02-18T08:09:08.149849Z",
     "iopub.status.idle": "2025-02-18T08:09:15.792339Z",
     "shell.execute_reply": "2025-02-18T08:09:15.791825Z",
     "shell.execute_reply.started": "2025-02-18T08:09:08.150033Z"
    },
    "id": "ooBAfXQxL00m"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c143281f64c44d04ad60df300ad8fb10",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "28ca23499b774926bd82f6ad2336d156",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "README.md:   0%|          | 0.00/905 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9a5149f7d91544de9088588cc3a9408e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "from datasets.arrow_dataset import Dataset\n",
    "\n",
    "def format_sample(sample):\n",
    "    \"\"\" Helper function to format a single input sample\"\"\"\n",
    "    instruction=sample['instruction']\n",
    "    input_text=sample['input']\n",
    "    output_text=sample['output']\n",
    "\n",
    "    if input_text is None or input_text==\"\":\n",
    "        formatted_prompt=(\n",
    "            f\"<|start_header_id|>user<|end_header_id|>\\n\\n\"\n",
    "            f\"Below is an instruction that describes a task. Write a response that appropriately completes the request.\\n\\n\"\n",
    "            f\"### Instruction:\\n{instruction}\\n\\n\"\n",
    "            f\"### Response:\\n<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
    "            f\"{output_text}<|eot_id|>\"\n",
    "        )\n",
    "    else:\n",
    "        formatted_prompt=(\n",
    "            f\"<|start_header_id|>user<|end_header_id|>\\n\\n\"\n",
    "            f\"Below is an instruction that describes a task. Write a response that appropriately completes the request.\\n\\n\"\n",
    "            f\"### Instruction:\\n{instruction}\\n\\n### Input:\\n{input_text}\\n\\n\"\n",
    "            f\"### Response:\\n<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
    "            f\"{output_text}<|eot_id|>\"\n",
    "        )\n",
    "    formatted_prompt=\"\".join(formatted_prompt) # exclude trailing white spaces\n",
    "    return formatted_prompt                    # stream text into the dataloader, one by one\n",
    "\n",
    "\n",
    "\n",
    "def gen_train_input():\n",
    "    \"\"\" Format all data input in alpaca style\n",
    "        Return:\n",
    "            A generator on train data \"train_gen\"\n",
    "    \"\"\"\n",
    "    # load data\n",
    "    ds=load_dataset(\"iamtarun/python_code_instructions_18k_alpaca\",streaming=True, split=\"train\")\n",
    "    # datata set has 18.6k samples, we use 16.8k (90%) for training + 1.8k for validation\n",
    "    num_samples=16800\n",
    "    counter=0\n",
    "    for sample in iter(ds):\n",
    "        if counter>=num_samples:\n",
    "            break\n",
    "        formatted_prompt=format_sample(sample)\n",
    "        yield {'text': formatted_prompt}\n",
    "        counter+=1\n",
    "\n",
    "\n",
    "def gen_val_input():\n",
    "    \"\"\" Format all data input in alpaca style\n",
    "        Return:\n",
    "            A generator on val data \"val_gen\"\n",
    "    \"\"\"\n",
    "    # load data\n",
    "    ds=load_dataset(\"iamtarun/python_code_instructions_18k_alpaca\",streaming=True, split=\"train\")\n",
    "    # datata set has 18.6k samples, we use 16.8k (90%) for training + 1.8k for validation\n",
    "    num_samples=16800\n",
    "    counter=0\n",
    "    for sample in iter(ds):\n",
    "        if counter<num_samples:\n",
    "            counter+=1\n",
    "            continue\n",
    "\n",
    "        formatted_prompt=format_sample(sample)\n",
    "        yield {'text': formatted_prompt}\n",
    "        counter+=1\n",
    "\n",
    "dataset_train = Dataset.from_generator(gen_train_input)\n",
    "dataset_val=Dataset.from_generator(gen_val_input)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "execution": {
     "iopub.execute_input": "2025-02-18T08:12:57.243357Z",
     "iopub.status.busy": "2025-02-18T08:12:57.242929Z",
     "iopub.status.idle": "2025-02-18T08:12:57.247997Z",
     "shell.execute_reply": "2025-02-18T08:12:57.247490Z",
     "shell.execute_reply.started": "2025-02-18T08:12:57.243338Z"
    },
    "id": "GwPH7uvML02s",
    "outputId": "2dd352d8-a949-41f8-bb76-68196b253c59"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train dataset size: 16800\n",
      "Validation dataset size: 1812\n",
      "Sample train:\n",
      "{'text': '<|start_header_id|>user<|end_header_id|>\\n\\nBelow is an instruction that describes a task. Write a response that appropriately completes the request.\\n\\n### Instruction:\\nCreate a function to calculate the sum of a sequence of integers.\\n\\n### Input:\\n[1, 2, 3, 4, 5]\\n\\n### Response:\\n<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\n# Python code\\ndef sum_sequence(sequence):\\n  sum = 0\\n  for num in sequence:\\n    sum += num\\n  return sum<|eot_id|>'}\n"
     ]
    }
   ],
   "source": [
    "print(f\"Train dataset size: {len(dataset_train)}\")\n",
    "print(f\"Validation dataset size: {len(dataset_val)}\")\n",
    "\n",
    "print(f\"Sample train:\\n{dataset_train[0]}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hyqrYyOvSYFO"
   },
   "source": [
    "## Model and tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-02-18T08:19:22.919097Z",
     "iopub.status.busy": "2025-02-18T08:19:22.918650Z",
     "iopub.status.idle": "2025-02-18T08:19:25.357067Z",
     "shell.execute_reply": "2025-02-18T08:19:25.356436Z",
     "shell.execute_reply.started": "2025-02-18T08:19:22.919078Z"
    },
    "id": "6ZIveBRhL04g"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from peft import LoraConfig, AutoPeftModelForCausalLM\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, BitsAndBytesConfig\n",
    "from trl import SFTTrainer\n",
    "\n",
    "model_name=\"meta-llama/Llama-3.2-1B-Instruct\"\n",
    "def create_and_prepare_model():\n",
    "    bnb_config = BitsAndBytesConfig(\n",
    "        load_in_4bit=True,\n",
    "        bnb_4bit_quant_type=\"nf4\",\n",
    "        bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "        bnb_4bit_use_double_quant=True,\n",
    "    )\n",
    "\n",
    "    model=AutoModelForCausalLM.from_pretrained(\n",
    "        model_name,\n",
    "        quantization_config=bnb_config,\n",
    "        device_map=\"auto\",\n",
    "        token=access_token\n",
    "    )\n",
    "\n",
    "    peft_config=LoraConfig(\n",
    "        lora_alpha=16,\n",
    "        r=8,\n",
    "        lora_dropout=0.05,\n",
    "        bias=\"none\",\n",
    "        task_type=\"CAUSAL_LM\",\n",
    "        target_modules=['q_proj', 'k_proj', 'v_proj'],\n",
    "    )\n",
    "    \n",
    "    tokenizer=AutoTokenizer.from_pretrained(model_name,token=access_token)\n",
    "    tokenizer.pad_token = \"<|finetune_right_pad_id|>\"\n",
    "    tokenizer.padding_side = \"right\"\n",
    "\n",
    "    return model,peft_config,tokenizer\n",
    "\n",
    "model,peft_config,tokenizer=create_and_prepare_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4-q42NQzSbK_"
   },
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-02-18T08:19:31.464023Z",
     "iopub.status.busy": "2025-02-18T08:19:31.463410Z",
     "iopub.status.idle": "2025-02-18T09:36:42.257159Z",
     "shell.execute_reply": "2025-02-18T09:36:42.256689Z",
     "shell.execute_reply.started": "2025-02-18T08:19:31.463996Z"
    },
    "id": "Vma9uK6yaXLS"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.11/dist-packages/transformers/training_args.py:1594: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead\n",
      "  warnings.warn(\n",
      "/tmp/ipykernel_87/3941464550.py:27: FutureWarning: `tokenizer` is deprecated and removed starting from version 0.16.0 for `SFTTrainer.__init__`. Use `processing_class` instead.\n",
      "  trainer=SFTTrainer(\n",
      "/usr/local/lib/python3.11/dist-packages/transformers/training_args.py:1594: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0bd65a90bdc747a999dd5a43a2b3fa26",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Applying chat template to train dataset:   0%|          | 0/16800 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "48fc41c06bcf446fbe369fbdea461090",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Tokenizing train dataset:   0%|          | 0/16800 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "be1b5f7918d24cf19fa143b4e0747d6e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Tokenizing train dataset:   0%|          | 0/16800 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "179d7896ea914017a0990fc2ca14932f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Applying chat template to eval dataset:   0%|          | 0/1812 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "85d29f1b88df4787bfa05b01f8cd7002",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Tokenizing eval dataset:   0%|          | 0/1812 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "473ff8434df949f6a0e8feecb0d30e60",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Tokenizing eval dataset:   0%|          | 0/1812 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No label_names provided for model class `PeftModelForCausalLM`. Since `PeftModel` hides base models input arguments, if label_names is not given, label_names can't be set automatically within `Trainer`. Note that empty label_names list will be used instead.\n",
      "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='393' max='393' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [393/393 1:16:42, Epoch 2/3]\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",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>50</td>\n",
       "      <td>1.290900</td>\n",
       "      <td>0.875727</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>0.870100</td>\n",
       "      <td>0.839814</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>150</td>\n",
       "      <td>0.845300</td>\n",
       "      <td>0.829108</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>0.837300</td>\n",
       "      <td>0.823679</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>250</td>\n",
       "      <td>0.834600</td>\n",
       "      <td>0.819919</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>300</td>\n",
       "      <td>0.848100</td>\n",
       "      <td>0.818405</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>350</td>\n",
       "      <td>0.823700</td>\n",
       "      <td>0.817808</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=393, training_loss=0.8984256191108063, metrics={'train_runtime': 4617.4541, 'train_samples_per_second': 10.915, 'train_steps_per_second': 0.085, 'total_flos': 1.5350980937613312e+17, 'train_loss': 0.8984256191108063})"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from trl import SFTConfig, SFTTrainer\n",
    "\n",
    "args=SFTConfig(\n",
    "    output_dir=\"./llama32-python\",\n",
    "    num_train_epochs=3,\n",
    "    per_device_train_batch_size=16,\n",
    "    per_device_eval_batch_size=16,\n",
    "    gradient_accumulation_steps=8,\n",
    "    gradient_checkpointing=True, # to save memmory\n",
    "    optim=\"adamw_torch_fused\",\n",
    "    logging_steps=50,            # log train loss\n",
    "    # Add evaluation strategy to compute validation loss during training\n",
    "    evaluation_strategy=\"steps\",  # Evaluate at the end of each epoch\n",
    "    eval_steps=50,  # Evaluate every 50 steps    \n",
    "    save_strategy=\"epoch\",\n",
    "    learning_rate=2e-4,\n",
    "    bf16=True, \n",
    "    tf32=True, # enable true for faster speed (supported in higher-end gpu)\n",
    "    max_grad_norm=0.3,\n",
    "    warmup_ratio=0.03, # follow QLoRA paper\n",
    "    lr_scheduler_type=\"cosine\",\n",
    "    report_to=\"tensorboard\",\n",
    "    gradient_checkpointing_kwargs={\"use_reentrant\": False},\n",
    "    dataset_text_field=\"text\",\n",
    ")\n",
    "\n",
    "trainer=SFTTrainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=dataset_train,\n",
    "    eval_dataset=dataset_val,\n",
    "    peft_config=peft_config,\n",
    "    tokenizer=tokenizer\n",
    ")\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-02-18T09:41:48.341018Z",
     "iopub.status.busy": "2025-02-18T09:41:48.340517Z",
     "iopub.status.idle": "2025-02-18T09:41:48.364800Z",
     "shell.execute_reply": "2025-02-18T09:41:48.364217Z",
     "shell.execute_reply.started": "2025-02-18T09:41:48.341000Z"
    }
   },
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'model' 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[17], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# save model state_dict\u001b[39;00m\n\u001b[1;32m      2\u001b[0m model_file_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mLLAMA32_ft_python_code.pth\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m----> 3\u001b[0m torch\u001b[38;5;241m.\u001b[39msave(\u001b[43mmodel\u001b[49m\u001b[38;5;241m.\u001b[39mstate_dict(), model_file_name)\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mModel saved as \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mmodel_file_name\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      6\u001b[0m \u001b[38;5;66;03m# free the memory \u001b[39;00m\n",
      "\u001b[0;31mNameError\u001b[0m: name 'model' is not defined"
     ]
    }
   ],
   "source": [
    "# save model state_dict\n",
    "model_file_name=\"LLAMA32_ft_python_code.pth\"\n",
    "torch.save(model.state_dict(), model_file_name)\n",
    "print(f\"Model saved as {model_file_name}\")\n",
    "\n",
    "# free the memory \n",
    "del model\n",
    "del trainer\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load finetune model and run inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-02-18T09:45:29.269012Z",
     "iopub.status.busy": "2025-02-18T09:45:29.268551Z",
     "iopub.status.idle": "2025-02-18T09:45:35.149545Z",
     "shell.execute_reply": "2025-02-18T09:45:35.149112Z",
     "shell.execute_reply.started": "2025-02-18T09:45:29.268993Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading fine-tuned model ...\n",
      "Total parameters: 1,236,994,048\n",
      "Trainable parameters: 1,179,648\n"
     ]
    }
   ],
   "source": [
    "from peft import PeftModel, LoraConfig\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import torch\n",
    "\n",
    "def load_fine_tune_model(base_model_id, saved_weights):\n",
    "    # Load tokenizer and base model\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "    tokenizer = AutoTokenizer.from_pretrained(base_model_id)\n",
    "    base_model = AutoModelForCausalLM.from_pretrained(base_model_id)\n",
    "    base_model.to(device)\n",
    "    \n",
    "    # Create LoRA config - make sure these parameters match your training configuration\n",
    "    peft_config = LoraConfig(\n",
    "        lora_alpha=16,\n",
    "        r=8,\n",
    "        lora_dropout=0.05,\n",
    "        bias=\"none\",\n",
    "        task_type=\"CAUSAL_LM\",\n",
    "        target_modules=['q_proj','k_proj','v_proj'],\n",
    "    )\n",
    "    \n",
    "    # Initialize PeftModel\n",
    "    lora_model = PeftModel(base_model, peft_config)\n",
    "    \n",
    "    # Load the saved weights\n",
    "    state_dict = torch.load(saved_weights,map_location=device)\n",
    "        \n",
    "    # Create new state dict with correct prefixes and structure\n",
    "    new_state_dict = {}\n",
    "    for key, value in state_dict.items():\n",
    "        # key start with \"model\" -> add \"base_\" to the new key for base_model\n",
    "        new_key = f\"base_{key}\"        \n",
    "        new_state_dict[new_key] = value\n",
    "    \n",
    "    # Load the weights with strict=False to allow partial loading\n",
    "    lora_model.load_state_dict(new_state_dict, strict=False)\n",
    "    \n",
    "    # Set to evaluation mode\n",
    "    lora_model = lora_model.eval()\n",
    "    \n",
    "    return lora_model, tokenizer\n",
    "\n",
    "# Original model and saved weight\n",
    "base_model_id = \"meta-llama/Llama-3.2-1B-Instruct\"\n",
    "lora_weights = \"LLAMA32_ft_python_code.pth\"\n",
    "\n",
    "# Load model\n",
    "print(\"Loading fine-tuned model ...\")\n",
    "model_ft, tokenizer = load_fine_tune_model(base_model_id, lora_weights)\n",
    "total_params=sum(p.numel() for p in model_ft.parameters())\n",
    "trainable_params=sum(p.numel() for p in model_ft.parameters() if p.requires_grad)\n",
    "\n",
    "print(f\"Total parameters: {total_params:,}\")\n",
    "print(f\"Trainable parameters: {trainable_params:,}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-02-18T09:52:30.282414Z",
     "iopub.status.busy": "2025-02-18T09:52:30.282183Z",
     "iopub.status.idle": "2025-02-18T09:52:30.288577Z",
     "shell.execute_reply": "2025-02-18T09:52:30.287942Z",
     "shell.execute_reply.started": "2025-02-18T09:52:30.282398Z"
    }
   },
   "outputs": [],
   "source": [
    "def generate(model, prompt, tokenizer, max_new_tokens, context_size=512, temperature=0.0, top_k=1, eos_id=[128001, 128009]):\n",
    "    \"\"\"\n",
    "    Generate text using a language model with proper dtype handling and improved sampling.\n",
    "    \"\"\"\n",
    "    # Get model's expected dtype\n",
    "    model_dtype = next(model.parameters()).dtype\n",
    "    model_device = next(model.parameters()).device\n",
    "    \n",
    "    formatted_prompt = (\n",
    "        f\"<|start_header_id|>user<|end_header_id|>\\n\\n\"\n",
    "        f\"Below is an instruction that describes a task. Write a response that appropriately completes the request.\\n\\n\"\n",
    "        f\"### Instruction:\\n{prompt}\"\n",
    "        f\"### Response:\\n<|eot_id|><|start_header_id|>assistant<|end_header_id|>\\n\\n\"\n",
    "    )\n",
    "    \n",
    "    # Encode and prepare input\n",
    "    idx = tokenizer.encode(formatted_prompt)\n",
    "    idx = torch.tensor(idx, dtype=torch.long, device=model_device).unsqueeze(0)\n",
    "    num_tokens = idx.shape[1]\n",
    "    \n",
    "    # Generation loop\n",
    "    for _ in range(max_new_tokens):\n",
    "        # Get conditioning context\n",
    "        idx_cond = idx[:, -context_size:]\n",
    "        \n",
    "        # Generate logits\n",
    "        with torch.no_grad():            \n",
    "            # Forward pass - get output directly from model\n",
    "            outputs = model(\n",
    "                input_ids=idx_cond,\n",
    "                use_cache=False\n",
    "            )\n",
    "            \n",
    "            # Get logits directly from the output\n",
    "            logits = outputs.logits\n",
    "        \n",
    "        # Focus on last time step\n",
    "        logits = logits[:, -1, :]\n",
    "        \n",
    "        # Apply top-k filtering\n",
    "        if top_k is not None and top_k > 0:\n",
    "            top_logits, _ = torch.topk(logits, top_k)\n",
    "            min_val = top_logits[:, [-1]]\n",
    "            logits = torch.where(\n",
    "                logits < min_val,\n",
    "                torch.tensor(float('-inf'), device=model_device, dtype=model_dtype),\n",
    "                logits\n",
    "            )\n",
    "        \n",
    "        # Apply temperature and sample\n",
    "        if temperature > 0.0:\n",
    "            logits = logits / temperature\n",
    "            probs = torch.softmax(logits, dim=-1)\n",
    "            idx_next = torch.multinomial(probs, num_samples=1)\n",
    "        else:\n",
    "            idx_next = torch.argmax(logits, dim=-1, keepdim=True)\n",
    "        \n",
    "        # Check for EOS\n",
    "        if idx_next.item() in eos_id:\n",
    "            break\n",
    "            \n",
    "        # Append new token\n",
    "        idx = torch.cat((idx, idx_next), dim=1)\n",
    "    \n",
    "    # Decode generated text\n",
    "    generated_ids = idx.squeeze(0)[num_tokens:]\n",
    "    generated_text = tokenizer.decode(generated_ids)\n",
    "    \n",
    "    return generated_text\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-02-18T09:52:46.711291Z",
     "iopub.status.busy": "2025-02-18T09:52:46.710849Z",
     "iopub.status.idle": "2025-02-18T09:52:49.836817Z",
     "shell.execute_reply": "2025-02-18T09:52:49.835986Z",
     "shell.execute_reply.started": "2025-02-18T09:52:46.711272Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "```python\n",
      "def fibonacci(n):\n",
      "    \"\"\"\n",
      "    This function computes the nth Fibonacci number.\n",
      "\n",
      "    Args:\n",
      "        n (int): The position of the Fibonacci number to be computed.\n",
      "\n",
      "    Returns:\n",
      "        int: The nth Fibonacci number.\n",
      "    \"\"\"\n",
      "\n",
      "    # Base cases for the Fibonacci sequence\n",
      "    if n <= 0:\n",
      "        return \"Input should be a positive integer.\"\n",
      "    elif n == 1:\n",
      "        return 0\n",
      "    elif n == 2:\n",
      "        return 1\n",
      "\n",
      "    # Initialize the first two Fibonacci numbers\n",
      "    a, b = 0, 1\n",
      "\n",
      "    # Compute the nth Fibonacci number\n",
      "    for _ in range(2, n):\n",
      "        # Update a and b to be the sum of the previous two\n",
      "        a, b = b, a + b\n",
      "\n",
      "    # Return the nth Fibonacci number\n",
      "    return b\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "prompt = (\"Write a function that computes fibonacci numbers.\")\n",
    "print(generate(model_ft, prompt, tokenizer, max_new_tokens=512))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-02-18T09:53:56.685333Z",
     "iopub.status.busy": "2025-02-18T09:53:56.684829Z",
     "iopub.status.idle": "2025-02-18T09:54:05.269256Z",
     "shell.execute_reply": "2025-02-18T09:54:05.268650Z",
     "shell.execute_reply.started": "2025-02-18T09:53:56.685314Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### Solution\n",
      "\n",
      "```python\n",
      "def findNumbers(nums):\n",
      "    \"\"\"\n",
      "    This function takes an array of integers as input and returns the count of numbers \n",
      "    that contain an even number of digits.\n",
      "\n",
      "    Args:\n",
      "        nums (list): A list of integers.\n",
      "\n",
      "    Returns:\n",
      "        int: The count of numbers with an even number of digits.\n",
      "    \"\"\"\n",
      "    count = 0  # Initialize a counter variable to store the count of numbers with even digits\n",
      "    for num in nums:  # Iterate over each number in the input list\n",
      "        digits = len(str(abs(num)))  # Convert the number to a string and count the number of digits\n",
      "        if digits % 2 == 0:  # Check if the number of digits is even\n",
      "            count += 1  # If the number of digits is even, increment the counter\n",
      "    return count  # Return the total count of numbers with an even number of digits\n",
      "```\n",
      "\n",
      "### Example Use Cases\n",
      "\n",
      "```python\n",
      "# Test the function with an array of numbers\n",
      "print(findNumbers([1, 2, 3, 4, 5]))  # Output: 2\n",
      "\n",
      "# Test the function with an array of numbers with an even number of digits\n",
      "print(findNumbers([10, 20, 30, 40, 50]))  # Output: 5\n",
      "\n",
      "# Test the function with an array of numbers with an odd number of digits\n",
      "print(findNumbers([100, 200, 300, 400, 500]))  # Output: 0\n",
      "```\n",
      "\n",
      "This solution works by iterating over each number in the input list, converting it to a string to count the number of digits, and checking if the number of digits is even. If it is, the counter is incremented. Finally, the total count of numbers with an even number of digits is returned.\n"
     ]
    }
   ],
   "source": [
    "prompt = \"Given an array nums of integers, return how many of them contain an even number of digits.\"\n",
    "print(generate(model_ft, prompt, tokenizer, max_new_tokens=512))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "GNN_M1",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}