01_data_prep.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9f16e594-1ba8-4caf-8127-30d3a51c5790",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install -qqqU sagemaker wandb datasets transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7082c65b-dabb-4b77-9df8-c5687a78753c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "WANDB_PROJECT = \"aws_llm_workshop\"\n",
    "USE_S3 = True\n",
    "\n",
    "RAW_TRAIN_DATASET_ARTIFACT = 'capecape/wandbot/run-m6nz6yrl-wandbot_questions:v0'\n",
    "RAW_EVAL_DATASET_ARTIFACT  = \"wandbot/wandbot-eval/run-kinbxic4-responses:v0\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5380e7bd-66f3-4237-a901-fe935a249982",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "MODEL_NAME = \"codellama/CodeLlama-7b-Instruct-hf\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75f2b7e0-dab7-46aa-b287-9721efb8dae3",
   "metadata": {},
   "source": [
    "# Preparing Data\n",
    "How to prepare our dataset for model Finetune"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25424b4f-e709-4663-beae-5a708e491db0",
   "metadata": {
    "tags": []
   },
   "source": [
    "## SageMaker auth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "ee341747-68ae-4b11-ab7c-2e5b86bf89ec",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /root/.config/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /root/.config/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /root/.config/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /etc/xdg/sagemaker/config.yaml\n",
      "sagemaker.config INFO - Not applying SDK defaults from location: /root/.config/sagemaker/config.yaml\n",
      "sagemaker role arn: arn:aws:iam::372108735839:role/SageMakerExecutionRole\n",
      "sagemaker bucket: sagemaker-us-east-1-372108735839\n",
      "sagemaker session region: us-east-1\n"
     ]
    }
   ],
   "source": [
    "import sagemaker\n",
    "import boto3\n",
    "sess = sagemaker.Session()\n",
    "# sagemaker session bucket -> used for uploading data, models and logs\n",
    "# sagemaker will automatically create this bucket if it not exists\n",
    "sagemaker_session_bucket=None\n",
    "if sagemaker_session_bucket is None and sess is not None:\n",
    "    # set to default bucket if a bucket name is not given\n",
    "    sagemaker_session_bucket = sess.default_bucket()\n",
    "\n",
    "try:\n",
    "    role = sagemaker.get_execution_role()\n",
    "except ValueError:\n",
    "    iam = boto3.client('iam')\n",
    "    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']\n",
    "\n",
    "sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)\n",
    "\n",
    "print(f\"sagemaker role arn: {role}\")\n",
    "print(f\"sagemaker bucket: {sess.default_bucket()}\")\n",
    "print(f\"sagemaker session region: {sess.boto_region_name}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d0a3457-8941-45d5-9ff8-eef4fa936a7b",
   "metadata": {},
   "source": [
    "## Formatting the data for the LLM"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04ca2daa-0c44-4ac1-ba98-1d70f87fd3e8",
   "metadata": {},
   "source": [
    "A big part of training LLMs lives in getting the data formatted correctly!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c6df869f-739b-409d-a50a-495dad432777",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import wandb\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d27c251-56b5-4f64-9268-ebc900720a7f",
   "metadata": {},
   "source": [
    "let's create a run and monitor our work from there"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "569d476a-9eae-4bcd-837d-2cc7f55f62b0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mcapecape\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.12"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/root/sagemaker-llama2-wandb/training/wandb/run-20231017_102757-j8jigch2</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/capecape/aws_llm_workshop/runs/j8jigch2' target=\"_blank\">unique-smoke-1</a></strong> to <a href='https://wandb.ai/capecape/aws_llm_workshop' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/capecape/aws_llm_workshop' target=\"_blank\">https://wandb.ai/capecape/aws_llm_workshop</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/capecape/aws_llm_workshop/runs/j8jigch2' target=\"_blank\">https://wandb.ai/capecape/aws_llm_workshop/runs/j8jigch2</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m:   1 of 1 files downloaded.  \n"
     ]
    }
   ],
   "source": [
    "wandb.init(project=WANDB_PROJECT, job_type=\"text_formatting\")\n",
    "\n",
    "# this way we get tracebility\n",
    "dataset_artifact = wandb.use_artifact(RAW_TRAIN_DATASET_ARTIFACT)\n",
    "table = dataset_artifact.get(\"wandbot_questions\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f902fa8-1b3c-4f19-af81-94372c0a16f7",
   "metadata": {},
   "source": [
    "this is a W&B table, so we can convert it to whatever format we may need"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "372833dc-d535-464e-936e-765731418b46",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>context</th>\n",
       "      <th>question</th>\n",
       "      <th>answer</th>\n",
       "      <th>page_content</th>\n",
       "      <th>metadata</th>\n",
       "      <th>context_len</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>A user has just started using the Weights &amp; Bi...</td>\n",
       "      <td>What is a 'run' in W&amp;B and what can I use it for?</td>\n",
       "      <td>A 'run' in W&amp;B is the fundamental unit that yo...</td>\n",
       "      <td>import Tabs from ‘@theme/Tabs’;  \\n\\nimport Ta...</td>\n",
       "      <td>{'file_type': '.md', 'language': 'en', 'source...</td>\n",
       "      <td>1291.111111</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>A user has just started using W&amp;B and they are...</td>\n",
       "      <td>Hi! I'm new to W&amp;B and I'm a bit stuck. Can yo...</td>\n",
       "      <td>Sure, you can install the W&amp;B library on your ...</td>\n",
       "      <td>import Tabs from ‘@theme/Tabs’;  \\n\\nimport Ta...</td>\n",
       "      <td>{'file_type': '.md', 'language': 'en', 'source...</td>\n",
       "      <td>1291.111111</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>The user is getting started with Weights and B...</td>\n",
       "      <td>I need to track my experiment's hyperparameter...</td>\n",
       "      <td>Certainly! You can pass your hyperparameters t...</td>\n",
       "      <td>import Tabs from ‘@theme/Tabs’;  \\n\\nimport Ta...</td>\n",
       "      <td>{'file_type': '.md', 'language': 'en', 'source...</td>\n",
       "      <td>1291.111111</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>A user is trying to do some mathematical opera...</td>\n",
       "      <td>How do I multiply two numbers using W&amp;B functi...</td>\n",
       "      <td>To multiply two numbers, you can use the `numb...</td>\n",
       "      <td>Value\\n\\n\\nWhether the two values are not equ...</td>\n",
       "      <td>{'file_type': '.md', 'language': 'en', 'source...</td>\n",
       "      <td>2186.111111</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>A user has been working on making sense of a d...</td>\n",
       "      <td>Hi, I have some confusing numbers that represe...</td>\n",
       "      <td>Yes, there is a function in W&amp;B that allows yo...</td>\n",
       "      <td>Value\\n\\n\\nWhether the two values are not equ...</td>\n",
       "      <td>{'file_type': '.md', 'language': 'en', 'source...</td>\n",
       "      <td>2186.111111</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                             context  \\\n",
       "0  A user has just started using the Weights & Bi...   \n",
       "2  A user has just started using W&B and they are...   \n",
       "4  The user is getting started with Weights and B...   \n",
       "5  A user is trying to do some mathematical opera...   \n",
       "6  A user has been working on making sense of a d...   \n",
       "\n",
       "                                            question  \\\n",
       "0  What is a 'run' in W&B and what can I use it for?   \n",
       "2  Hi! I'm new to W&B and I'm a bit stuck. Can yo...   \n",
       "4  I need to track my experiment's hyperparameter...   \n",
       "5  How do I multiply two numbers using W&B functi...   \n",
       "6  Hi, I have some confusing numbers that represe...   \n",
       "\n",
       "                                              answer  \\\n",
       "0  A 'run' in W&B is the fundamental unit that yo...   \n",
       "2  Sure, you can install the W&B library on your ...   \n",
       "4  Certainly! You can pass your hyperparameters t...   \n",
       "5  To multiply two numbers, you can use the `numb...   \n",
       "6  Yes, there is a function in W&B that allows yo...   \n",
       "\n",
       "                                        page_content  \\\n",
       "0  import Tabs from ‘@theme/Tabs’;  \\n\\nimport Ta...   \n",
       "2  import Tabs from ‘@theme/Tabs’;  \\n\\nimport Ta...   \n",
       "4  import Tabs from ‘@theme/Tabs’;  \\n\\nimport Ta...   \n",
       "5   Value\\n\\n\\nWhether the two values are not equ...   \n",
       "6   Value\\n\\n\\nWhether the two values are not equ...   \n",
       "\n",
       "                                            metadata  context_len  \n",
       "0  {'file_type': '.md', 'language': 'en', 'source...  1291.111111  \n",
       "2  {'file_type': '.md', 'language': 'en', 'source...  1291.111111  \n",
       "4  {'file_type': '.md', 'language': 'en', 'source...  1291.111111  \n",
       "5  {'file_type': '.md', 'language': 'en', 'source...  2186.111111  \n",
       "6  {'file_type': '.md', 'language': 'en', 'source...  2186.111111  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.DataFrame(table.data, columns=table.columns)\n",
    "df = df.dropna()\n",
    "df = df.assign(context_len = lambda df: df.page_content.str.len()/3.6)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d61c89c2-831f-4996-9f07-3965ba839858",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2091"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90391542-9e78-46ea-a197-0bc915f54a44",
   "metadata": {},
   "source": [
    "Let's prepare the training dataset now"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dae08b1-5519-43e4-ba28-099950a0bd63",
   "metadata": {},
   "source": [
    "If you use CodeLLama we need to format the instructions accordingly!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "bbb5948d-1b37-46b6-b6e1-5b1c65498bba",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "B_INST, E_INST = \"[INST] \", \" [/INST]\"\n",
    "B_SYS, E_SYS = \"<<SYS>>\\n\", \"\\n<</SYS>>\\n\\n\"\n",
    "EOS = \"</s>\"\n",
    "\n",
    "prompt_format = (\n",
    "    B_INST\n",
    "    + B_SYS\n",
    "    + \"You are an AI assistant designed to assist developers with everyday tasks related to Weight & Biases \"\n",
    "    + \"and provide helpful information. As an expert in the open-source python SDK wandb answer the following \"\n",
    "    # + \"question based on the context below. Answer in formatted Markdown.\\n\"\n",
    "    + \"question below. Answer in formatted Markdown.\\n\"\n",
    "    + \"{page_content}\"\n",
    "    + E_SYS\n",
    "    + \"{question}\"\n",
    "    + E_INST\n",
    "    + \"\\n[W&B]\\n\"\n",
    "    + \"{answer}\"\n",
    "    + \"\\n[/W&B]\"\n",
    "    + EOS\n",
    ")\n",
    "\n",
    "def format_text(row):\n",
    "    return prompt_format.format_map(row)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "20ab80e2-9206-421e-acf9-d1d8d67baab3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INST] <<SYS>>\n",
      "You are an AI assistant designed to assist developers with everyday tasks related to Weight & Biasesand provide helpful information. As an expert in the open-source python SDK wandb answer the followingquestion below. Answer in formatted Markdown.\n",
      "{page_content}\n",
      "<</SYS>>\n",
      "\n",
      "{question} [/INST]\n",
      "[W&B]\n",
      "{answer}\n",
      "[/W&B]</s>\n"
     ]
    }
   ],
   "source": [
    "print(prompt_format)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "da903861-dce9-466d-839e-67d937b228e4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[INST] <<SYS>>\n",
      "You are an AI assistant designed to assist developers with everyday tasks related to Weight & Biasesand provide helpful information. As an expert in the open-source python SDK wandb answer the followingquestion below. Answer in formatted Markdown.\n",
      "import Tabs from ‘@theme/Tabs’;  \n",
      "\n",
      "import TabItem from ‘@theme/TabItem’;\n",
      "\n",
      "\n",
      "# Quickstart\n",
      "\n",
      "\n",
      "Install W&B and start tracking your machine learning experiments in minutes.\n",
      "\n",
      "\n",
      "## 1. Create an account and install W&B\n",
      "\n",
      "\n",
      "Before you get started, make sure you create an account and install W&B:\n",
      "\n",
      "\n",
      "1. Sign up for a free account at <https://wandb.ai/site> and then login to your wandb account.\n",
      "2. Install the wandb library on your machine in a Python 3 environment using `pip`.\n",
      "\n",
      "\n",
      "\n",
      "The following code snippets demonstrate how to install and log into W&B using the W&B CLI and Python Library:\n",
      "\n",
      "\n",
      "\n",
      "Install the CLI and Python library for interacting with the Weights and Biases API:\n",
      "\n",
      "\n",
      "\n",
      "```\n",
      "pip install wandb\n",
      "\n",
      "```\n",
      "\n",
      "\n",
      "Install the CLI and Python library for interacting with the Weights and Biases API:\n",
      "\n",
      "\n",
      "\n",
      "```\n",
      "!pip install wandb\n",
      "\n",
      "```\n",
      "\n",
      "\n",
      "## 2. Log in to W&B\n",
      "\n",
      "\n",
      "\n",
      "Next, log in to W&B:\n",
      "\n",
      "\n",
      "\n",
      "```\n",
      "wandb login\n",
      "\n",
      "```\n",
      "\n",
      "Or if you are using W&B Server:\n",
      "\n",
      "\n",
      "\n",
      "```\n",
      "wandb login --host=http://wandb.your-shared-local-host.com\n",
      "\n",
      "```\n",
      "\n",
      "Provide your API key when prompted.\n",
      "\n",
      "\n",
      "\n",
      "Next, import the W&B Python SDK and log in:\n",
      "\n",
      "\n",
      "\n",
      "```\n",
      "wandb.login()\n",
      "\n",
      "```\n",
      "\n",
      "Provide your API key when prompted.  \n",
      "\n",
      "  \n",
      "\n",
      "\n",
      "\n",
      "\n",
      "## 3. Start a run and track hyperparameters\n",
      "\n",
      "\n",
      "Initialize a W&B Run object in your Python script or notebook with `wandb.init()` and pass a dictionary to the `config` parameter with key-value pairs of hyperparameter names and values:\n",
      "\n",
      "\n",
      "\n",
      "```\n",
      "run = wandb.init(\n",
      "    # Set the project where this run will be logged\n",
      "    project=\"my-awesome-project\",\n",
      "    # Track hyperparameters and run metadata\n",
      "    config={\n",
      "        \"learning_rate\": 0.01,\n",
      "        \"epochs\": 10,\n",
      "    })\n",
      "\n",
      "```\n",
      "\n",
      "\n",
      "A run is the basic building block of W&B. You will use them often to track metrics, create logs, create jobs, and more.\n",
      "\n",
      "\n",
      "\n",
      "\n",
      "## Putting it all together\n",
      "\n",
      "\n",
      "Putting it all together, your training script might look similar to the following code example. The highlighted code shows W&B-specific code.   \n",
      "\n",
      "Note that we added code that mimics machine learning training.\n",
      "\n",
      "\n",
      "\n",
      "```\n",
      "# train.py\n",
      "import wandb\n",
      "import random # for demo script\n",
      "\n",
      "# highlight-next-line\n",
      "wandb.login()\n",
      "\n",
      "epochs=10\n",
      "lr=0.01\n",
      "\n",
      "# highlight-start\n",
      "run = wandb.init(\n",
      "    # Set the project where this run will be logged\n",
      "    project=\"my-awesome-project\",\n",
      "    # Track hyperparameters and run metadata\n",
      "    config={\n",
      "        \"learning_rate\": lr,\n",
      "        \"epochs\": epochs,\n",
      "    })\n",
      "# highlight-end    \n",
      "\n",
      "offset = random.random() / 5\n",
      "print(f\"lr: {lr}\")\n",
      "\n",
      "# simulating a training run\n",
      "for epoch in range(2, epochs):\n",
      "    acc = 1 - 2 ** -epoch - random.random() / epoch - offset\n",
      "    loss = 2 ** -epoch + random.random() / epoch + offset\n",
      "    print(f\"epoch={epoch}, accuracy={acc}, loss={loss}\")\n",
      "    # highlight-next-line\n",
      "    wandb.log({\"accuracy\": acc, \"loss\": loss})\n",
      "\n",
      "# run.log_code()\n",
      "\n",
      "```\n",
      "\n",
      "That’s it! Navigate to the W&B App at <https://wandb.ai/home> to view how the metrics we logged with W&B (accuracy and loss) improved during each training step.\n",
      "\n",
      "\n",
      "\n",
      "The image above (click to expand) shows the loss and accuracy that was tracked from each time we ran the script above. Each run object that was created is show within the **Runs** column. Each run name is randomly generated.\n",
      "\n",
      "\n",
      "## What’s next?\n",
      "\n",
      "\n",
      "\n",
      "Explore the rest of the W&B ecosystem.\n",
      "\n",
      "\n",
      "1. Check out W&B Integrations to learn how to integrate W&B with your ML framework such as PyTorch, ML library such as Hugging Face, or ML service such as SageMaker.\n",
      "2. Organize runs, embed and automate visualizations, describe your findings, and share updates with collaborators with W&B Reports.\n",
      "3. Create W&B Artifacts to track datasets, models, dependencies, and results through each step of your machine learning pipeline.\n",
      "4. Automate hyperparameter search and explore the space of possible models with W&B Sweeps.\n",
      "5. Understand your datasets, visualize model predictions, and share insights in a central dashboard.\n",
      "\n",
      "\n",
      "\n",
      "## Common Questions\n",
      "\n",
      "\n",
      "**Where do I find my API key?**  \n",
      "\n",
      "Once you’ve signed in to www.wandb.ai, the API key will be on the Authorize page.\n",
      "\n",
      "\n",
      "**How do I use W&B in an automated environment?**  \n",
      "\n",
      "If you are training models in an automated environment where it’s inconvenient to run shell commands, such as Google’s CloudML, you should look at our guide to configuration with Environment Variables.\n",
      "\n",
      "\n",
      "**Do you offer local, on-prem installs?**  \n",
      "\n",
      "Yes, you can privately host W&B locally on your own machines or in a private cloud, try this quick tutorial notebook to see how. Note, to login to wandb local server you can set the host flag to the address of the local instance. \n",
      "\n",
      "\n",
      "**How do I turn off wandb logging temporarily?**  \n",
      "\n",
      "If are testing code and want to disable wandb syncing, set the environment variable `WANDB_MODE=offline`.\n",
      "\n",
      "\n",
      "<</SYS>>\n",
      "\n",
      "What is a 'run' in W&B and what can I use it for? [/INST]\n",
      "[W&B]\n",
      "A 'run' in W&B is the fundamental unit that you can use to track various aspects of your machine learning experiments. Once you initialize a W&B Run object in your Python script or notebook using `wandb.init()`, you can pass a dictionary to the `config` parameter where the key-value pairs are your hyperparameters with their corresponding values. You can use a 'run' to track metrics, create logs, jobs and do many more things.\n",
      "[/W&B]</s>\n"
     ]
    }
   ],
   "source": [
    "one_example = format_text(df.iloc[0])\n",
    "print(one_example)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdf19de4-c29a-4803-aad1-8eb0b9c32da5",
   "metadata": {
    "tags": []
   },
   "source": [
    "Let's compute the format over all the dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "369e92cb-d51d-4f97-93e0-bb3c709eb126",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df[\"text\"] = df.apply(format_text, axis=1)\n",
    "\n",
    "# print(df.iloc[200][\"text\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c18f6fe4-6573-4901-ac44-aa51d0400452",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df.to_json(\"wandb_questions_ds.jsonl\", orient='records', lines=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01efd5c0-1fec-4c76-bd6e-1ab182deb7c3",
   "metadata": {},
   "source": [
    "## Saving your work to W&B"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3867767-df18-4866-9c78-ca623ad9b205",
   "metadata": {},
   "source": [
    "We should log this to W&B so we can inspect the dataset interactively using W&B Tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "902581de-ada6-44c4-b399-ef24845948e2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3f35019cd7964ebb882eec4a3fa864c9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='89.800 MB of 89.800 MB uploaded (0.000 MB deduped)\\r'), FloatProgress(value=1.0, m…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "wandb: WARNING Source type is set to 'repo' but some required information is missing from the environment. A job will not be created from this run. See https://docs.wandb.ai/guides/launch/create-job\n"
     ]
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">unique-smoke-1</strong> at: <a href='https://wandb.ai/capecape/aws_llm_workshop/runs/j8jigch2' target=\"_blank\">https://wandb.ai/capecape/aws_llm_workshop/runs/j8jigch2</a><br/>Synced 7 W&B file(s), 1 media file(s), 2 artifact file(s) and 1 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20231017_102757-j8jigch2/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "table = wandb.Table(dataframe=df)\n",
    "wandb.log({\"wandb_questions_ds\": table})\n",
    "\n",
    "# let's also save a the dataset at this stage\n",
    "at = wandb.Artifact(\n",
    "    name=\"wandb_questions_ds\", \n",
    "    type=\"dataset\",\n",
    "    description=\"A wandbot dataset of questions and answers about W&B for training (non tokenized)\",\n",
    "    metadata={\"prompt_format\": prompt_format,\n",
    "              \"length\": len(df),\n",
    "             }\n",
    ")\n",
    "at.add_file(\"wandb_questions_ds.jsonl\")\n",
    "wandb.log_artifact(at)\n",
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdebe295-f927-406f-b790-27bc059abab1",
   "metadata": {},
   "source": [
    "## Tokenizing and saving the preprocessing\n",
    "We can save time during training by pre-processing the dataset and loading directly a tokenized dataset!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "20aaffb3-d181-4c34-a58d-3e9c78a16d14",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/transformers/models/auto/tokenization_auto.py:655: FutureWarning: The `use_auth_token` argument is deprecated and will be removed in v5 of Transformers.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "import wandb\n",
    "from transformers import AutoTokenizer\n",
    "from datasets import load_dataset\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=True)\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c03ed5fc-a178-424d-8bc5-56f452dc3487",
   "metadata": {},
   "source": [
    "we can convert the data to a huggingface parquet-based dataset for fast loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "33ad3e77-c60c-46a1-9fef-432d412bccc3",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "Finishing last run (ID:vei4s8us) before initializing another..."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "wandb: WARNING Source type is set to 'repo' but some required information is missing from the environment. A job will not be created from this run. See https://docs.wandb.ai/guides/launch/create-job\n"
     ]
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">sandy-dream-4</strong> at: <a href='https://wandb.ai/capecape/aws_llm_workshop/runs/vei4s8us' target=\"_blank\">https://wandb.ai/capecape/aws_llm_workshop/runs/vei4s8us</a><br/>Synced 6 W&B file(s), 0 media file(s), 0 artifact file(s) and 1 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20231017_103256-vei4s8us/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Successfully finished last run (ID:vei4s8us). Initializing new run:<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b2724e06813d4d4fb758c0884264e958",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.011112336688852843, max=1.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.12"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/root/sagemaker-llama2-wandb/training/wandb/run-20231017_103308-6fojvwiz</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/capecape/aws_llm_workshop/runs/6fojvwiz' target=\"_blank\">wandering-totem-5</a></strong> to <a href='https://wandb.ai/capecape/aws_llm_workshop' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/capecape/aws_llm_workshop' target=\"_blank\">https://wandb.ai/capecape/aws_llm_workshop</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/capecape/aws_llm_workshop/runs/6fojvwiz' target=\"_blank\">https://wandb.ai/capecape/aws_llm_workshop/runs/6fojvwiz</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m:   1 of 1 files downloaded.  \n"
     ]
    }
   ],
   "source": [
    "wandb.init(project=WANDB_PROJECT, job_type=\"tokenizing\")\n",
    "artifact = wandb.use_artifact('capecape/aws_llm_workshop/wandb_questions_ds:v0', type='dataset')\n",
    "artifact_dir = artifact.download()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "6d894909-a161-4806-832f-0da72a580860",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "de7beeac7fd640b0a3150287b7f2d2e1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "50f99f5a849946178201584a95c8932b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "da27a6e8e2b244e9b55a7504627b6987",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['context', 'question', 'answer', 'page_content', 'metadata', 'context_len', 'text'],\n",
       "    num_rows: 2091\n",
       "})"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_dataset = load_dataset(\n",
    "    path=\".\", \n",
    "    data_files=f\"{artifact_dir}/wandb_questions_ds.jsonl\", \n",
    "    split=\"train\")\n",
    "train_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "676fbde0-d780-4beb-94dd-fff1a5a617ac",
   "metadata": {},
   "source": [
    "one sample looks like this 👇"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "6a4aa0e2-c9b3-4065-8c4a-ca24f8bfd856",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_dataset = train_dataset.select_columns([\"text\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b5df85d-be1f-4726-92c8-73e94146ea5c",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Packing and chunking"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bac52401-ea01-44bb-95d7-b11f8aa34b59",
   "metadata": {},
   "source": [
    "We define some helper functions to pack our samples into sequences of a given length and then tokenize them.tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "f0eb36f2-6dbf-4f13-b8ea-1ba2effc709c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a5d873161f074ec495f879de5025fad3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2091 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'text': '[INST] <<SYS>>\\nYou are an AI assistant designed to assist developers with everyday tasks related to Weight & Biasesand provide helpful information. As an expert in the open-source python SDK wandb answer the followingquestion below. Answer in formatted Markdown.\\n Prompts** is a suite of LLMOps tools built for the development of LLM-powered applications. \\n\\n\\nUse W&B Prompts to visualize and inspect the execution flow of your LLMs, analyze the inputs and outputs of your LLMs, view the intermediate results and securely store and manage your prompts and LLM chain configurations.\\n\\n\\n#### 🪄 View Prompts In Action\\n\\n\\n**In this notebook we will demostrate W&B Prompts:**\\n\\n\\n* Using our 1-line LangChain integration\\n* Using our Trace class when building your own LLM Pipelines\\n\\n\\nSee here for the full W&B Prompts documentation\\n\\n\\n## Installation\\n\\n\\n\\n```\\n!pip install \"wandb>=0.15.4\" -qqq\\n!pip install \"langchain>=0.0.218\" openai -qqq\\n\\n```\\n\\n\\n```\\nimport langchain\\nassert langchain.__version__ >= \"0.0.218\", \"Please ensure you are using LangChain v0.0.188 or higher\"\\n\\n```\\n\\n## Setup\\n\\n\\nThis demo requires that you have an OpenAI key\\n\\n\\n\\n```\\nimport os\\nfrom getpass import getpass\\n\\nif os.getenv(\"OPENAI_API_KEY\") is None:\\n  os.environ[\"OPENAI_API_KEY\"] = getpass(\"Paste your OpenAI key from: https://platform.openai.com/account/api-keys\\\\n\")\\nassert os.getenv(\"OPENAI_API_KEY\", \"\").startswith(\"sk-\"), \"This doesn\\'t look like a valid OpenAI API key\"\\nprint(\"OpenAI API key configured\")\\n\\n```\\n\\n# W&B Prompts\\n\\n\\nW&B Prompts consists of three main components:\\n\\n\\n**Trace table**: Overview of the inputs and outputs of a chain.\\n\\n\\n**Trace timeline**: Displays the execution flow of the chain and is color-coded according to component types.\\n\\n\\n**Model architecture**: View details about the structure of the chain and the parameters used to initialize each component of the chain.\\n\\n\\nAfter running this section, you will see a new panel automatically created in your workspace, showing each execution, the trace, and the model architecture\\n\\n\\n\\n## Maths with LangChain\\n\\n\\nSet the `LANGCHAIN_WANDB_TRACING` environment variable as well as any other relevant W&B environment variables. This could includes a W&B project name, team name, and more. See wandb.init for a full list of arguments.\\n\\n\\n\\n```\\nos.environ[\"LANGCHAIN_WANDB_TRACING\"] = \"true\"\\nos.environ[\"WANDB_PROJECT\"] = \"langchain-testing\"\\n\\n```\\n\\n\\n```\\nfrom langchain.chat_models import ChatOpenAI\\nfrom langchain.agents import load_tools, initialize_agent, AgentType\\n\\n```\\n\\nCreate a standard math Agent using LangChain\\n\\n\\n\\n```\\nllm = ChatOpenAI(temperature=0)\\ntools = load_tools([\"llm-math\"], llm=llm)\\nmath_agent = initialize_agent(tools, \\n                              llm, \\n                              agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION)\\n\\n```\\n\\nUse LangChain as normal by calling your Agent.\\n\\n\\nYou will see a Weights & Biases run start and you will be asked for your Weights & Biases API key. Once your enter your API key, the inputs and outputs of your Agent calls will start to be streamed to the Weights & Biases App.\\n\\n\\n\\n```\\n# some sample maths questions\\nquestions = [\\n  \"Find the square root of 5.4.\",\\n  \"What is 3 divided by 7.34 raised to the power of pi?\",\\n  \"What is the sin of 0.47 radians, divided by the cube root of 27?\"\\n]\\n\\nfor question in questions:\\n  try:\\n    # call your Agent as normal\\n    answer = math_agent.run(question)\\n    print(answer)\\n  except Exception as e:\\n    # any errors will be also logged to Weights & Biases\\n    print(e)\\n    pass\\n\\n```\\n\\nOnce each Agent execution completes, all calls in your LangChain object will be logged to Weights & Biases\\n\\n\\n### LangChain Context Manager\\n\\n\\nDepending on your use case, you might instead prefer to use a context manager to manage your logging to W&B:\\n\\n\\n\\n```\\nfrom langchain.callbacks import wandb_tracing_enabled\\n\\n# unset the environment variable and use a context manager instead\\nif \"LANGCHAIN_WANDB_TRACING\" in os.environ:\\n    del os.environ[\"LANGCHAIN_WANDB_TRACING\"]\\n\\n# enable tracing using a context manager\\nwith wandb_tracing_enabled():\\n    math_agent.run(\"What is 5 raised to .123243 power?\")  # this should be traced\\n\\nmath_agent.run(\"What is 2 raised to .123243 power?\")  # this should not be traced\\n\\n```\\n\\n# Non-Lang Chain Implementation\\n\\n\\nA W&B Trace is created by logging 1 or more “spans”. A root span is expected, which can accept nested child spans, which can in turn accept their own child spans. A Span represents a unit of work, Spans can have type `AGENT`, `TOOL`, `LLM` or `CHAIN`\\n\\n\\nWhen logging with Trace, a single W&B run can have multiple calls to a LLM, Tool, Chain or Agent logged to it, there is no need to start a new W&B run after each generation from your model or pipeline, instead each call will be appended to the Trace Table.\\n\\n\\nIn this quickstart, we will how to log a single call to an OpenAI model to W&B Trace as a single span. Then we will show how to log a more complex series of nested spans.\\n\\n\\n## Logging with W&B Trace\\n\\n\\nCall wandb.init to start a W&B run. Here you can pass a W&B project name as well as an entity name (if logging to a W&B Team), as well as a config and more. See wandb.init for the full list of arguments.\\n\\n\\nYou will see a Weights & Biases run start and be asked for your Weights & Biases API key. Once your enter your API key, the inputs and outputs of your Agent calls will start to be streamed to the Weights & Biases App.\\n\\n\\n**Note:** A W&B run supports logging as many traces you needed to a single run, i.e. you can make multiple calls of `run.log` without the need to create a new run each time\\n\\n\\n\\n```\\nimport wandb\\n\\n# start a wandb run to log to\\nwandb.init(project=\"trace-example\")\\n\\n```\\n\\nYou can also set the entity argument in wandb.init if logging to a W&B Team.\\n\\n\\n### Logging a single Span\\n\\n\\nNow we will query OpenAI times and log the results to a W&B Trace. We will log the inputs and outputs, start and end times, whether the OpenAI call was successful, the token usage, and additional metadata.\\n\\n\\nYou can see the full description of the arguments to the Trace class here.\\n\\n\\n\\n```\\nimport openai\\nimport datetime\\nfrom wandb.sdk.data_types.trace_tree import Trace\\n\\nopenai.api_key = os.environ[\"OPENAI_API_KEY\"]\\n\\n# define your conifg\\nmodel_name = \"gpt-3.5-turbo\"\\ntemperature = 0.7\\nsystem_message = \"You are a helpful assistant that always replies in 3 concise bullet points using markdown.\"\\n\\nqueries_ls = [\\n  \"What is the capital of France?\",\\n  \"How do I boil an egg?\" * 10000,  # deliberately trigger an openai error\\n  \"What to do if the aliens arrive?\" \\n]\\n\\nfor query in queries_ls:\\n    messages=[\\n      {\"role\": \"system\", \"content\": system_message},\\n      {\"role\": \"user\", \"content\": query}\\n    ]\\n\\n    start_time_ms = datetime.datetime.now().timestamp() * 1000\\n    try:\\n      response = openai.ChatCompletion.create(model=model_name,\\n                                              messages=messages,\\n                                              temperature=temperature\\n                                              )   \\n\\n      end_time_ms = round(datetime.datetime.now().timestamp() * 1000)  # logged in milliseconds\\n      status=\"success\"\\n      status_message=None,\\n      response_text = response[\"choices\"][0][\"message\"][\"content\"]\\n      token_usage = response[\"usage\"].to_dict()\\n\\n\\n    except Exception as e:\\n      end_time_ms = round(datetime.datetime.now().timestamp() * 1000)  # logged in milliseconds\\n      status=\"error\"\\n      status_message=str(e)\\n      response_text = \"\"\\n      token_usage = {}\\n\\n    # create a span in wandb\\n    root_span = Trace(\\n          name=\"root_span\",\\n          kind=\"llm\",  # kind can be \"llm\", \"chain\", \"agent\" or \"tool\"\\n          status_code=status,\\n          status_message=status_message,\\n          metadata={\"temperature\": temperature,\\n                    \"token_usage\": token_usage, \\n                    \"model_name\": model_name},\\n          start_time_ms=start_time_ms,\\n          end_time_ms=end_time_ms,\\n          inputs={\"system_prompt\": system_message, \"query\": query},\\n          outputs={\"response\": response_text},\\n          )\\n\\n    # log the span to wandb\\n    root_span.log(name=\"openai_trace\")\\n\\n```\\n\\n### Logging a LLM pipeline using nested Spans\\n\\n\\nIn this example we will simulate an Agent being called, which then calls a LLM Chain, which calls an OpenAI LLM and then the Agent “calls” a Calculator tool.\\n\\n\\nThe inputs, outputs and metadata for each step in the execution of our “Agent” is logged in its own span. Spans can have child\\n\\n\\n\\n```\\nimport time\\n\\nopenai.api_key = os.environ[\"OPENAI_API_KEY\"]\\n\\n# The query our agent has to answer\\nquery = \"How many days until the next US election?\"\\n\\n# part 1 - an Agent is started...\\nstart_time_ms = round(datetime.datetime.now().timestamp() * 1000)\\n\\nroot_span = Trace(\\n      name=\"MyAgent\",\\n      kind=\"agent\",\\n      start_time_ms=start_time_ms,\\n      metadata={\"user\": \"optimus_12\"})\\n\\n\\n# part 2\\n<</SYS>>\\n\\nI saw some documentation mentioning the use of a context manager to manage logging when using LangChain with W&B. Could you explain how to use this in my pipelines? [/INST]\\n[W&B]\\nYes, you can use a context manager with W&B in the LangChain environment to handle logging. Here\\'s how you can do it:\\n\\nFirstly, you need to unset the environment variable for LangChain tracing.\\n\\n```\\nif \\'LANGCHAIN_WANDB_TRACING\\' in os.environ:\\n    del os.environ[\\'LANGCHAIN_WANDB_TRACING\\']\\n```\\n\\nAfter that, you can use the `wandb_tracing_enabled()` function to start the context manager that manages your logging. Any code executed inside this context will be traced and logged to your W&B runs.\\n\\n```python\\nfrom langchain.callbacks import wandb_tracing_enabled\\n\\n# Enable tracing using a context manager\\nwith wandb_tracing_enabled():\\n    math_agent.run(\\'What is 5 raised to .123243 power?\\')  # this should be traced\\n\\nmath_agent.run(\\'What is 2 raised to .123243 power?\\')  # this should not be traced\\n```\\n\\nIn this example, the first call to `math_agent.run()` is inside the context manager and will be traced. The second call, which is made outside the context manager, will not be traced.\\n[/W&B]</s>'}</s>\n"
     ]
    }
   ],
   "source": [
    "from random import randint\n",
    "from itertools import chain\n",
    "from functools import partial\n",
    "\n",
    "\n",
    "# template dataset to add prompt to each sample\n",
    "def template_dataset(sample):\n",
    "    sample[\"text\"] = f\"{sample}{tokenizer.eos_token}\"\n",
    "    return sample\n",
    "\n",
    "\n",
    "# apply prompt template per sample\n",
    "train_dataset = train_dataset.map(template_dataset)\n",
    "# print random sample\n",
    "print(train_dataset[randint(0, len(train_dataset))][\"text\"])\n",
    "\n",
    "# empty list to save remainder from batches to use in next batch\n",
    "remainder = {\"input_ids\": [], \"attention_mask\": [], \"token_type_ids\": []}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "ec8e2766-5b3a-4424-82e6-1cf5ef45da1e",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2649cbea122d440c8677aa51075c8ad7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2091 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c7d49da21fd04c29a1d2bfab71b6413a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/2091 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of samples: 4906\n"
     ]
    }
   ],
   "source": [
    "def chunk(sample, chunk_length=1024):\n",
    "    # define global remainder variable to save remainder from batches to use in next batch\n",
    "    global remainder\n",
    "    # Concatenate all texts and add remainder from previous batch\n",
    "    concatenated_examples = {k: list(chain(*sample[k])) for k in sample.keys()}\n",
    "    concatenated_examples = {k: remainder[k] + concatenated_examples[k] for k in concatenated_examples.keys()}\n",
    "    # get total number of tokens for batch\n",
    "    batch_total_length = len(concatenated_examples[list(sample.keys())[0]])\n",
    "\n",
    "    # get max number of chunks for batch\n",
    "    if batch_total_length >= chunk_length:\n",
    "        batch_chunk_length = (batch_total_length // chunk_length) * chunk_length\n",
    "\n",
    "    # Split by chunks of max_len.\n",
    "    result = {\n",
    "        k: [t[i : i + chunk_length] for i in range(0, batch_chunk_length, chunk_length)]\n",
    "        for k, t in concatenated_examples.items()\n",
    "    }\n",
    "    # add remainder to global variable for next batch\n",
    "    remainder = {k: concatenated_examples[k][batch_chunk_length:] for k in concatenated_examples.keys()}\n",
    "    # prepare labels\n",
    "    result[\"labels\"] = result[\"input_ids\"].copy()\n",
    "    return result\n",
    "\n",
    "\n",
    "# tokenize and chunk dataset\n",
    "lm_dataset = train_dataset.map(\n",
    "    lambda sample: tokenizer(sample[\"text\"]), batched=True, remove_columns=list(train_dataset.features)\n",
    ").map(\n",
    "    partial(chunk, chunk_length=1024),\n",
    "    batched=True,\n",
    ")\n",
    "\n",
    "# Print total number of samples\n",
    "print(f\"Total number of samples: {len(lm_dataset)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cd4726a-5381-4a56-8e5a-59fdfffb8c7c",
   "metadata": {},
   "source": [
    "## Save to a bucket and W&B"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bb0ecae-ed90-4f83-a34d-99e46daa4d1b",
   "metadata": {},
   "source": [
    "We are now going to use W&B Aritfacts integration with S3 buckets, so our dataset is close to the training compute SM instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "322268d9-80a8-4f67-bf55-76cf4773ca95",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a2d1500c46d74a4caaaaa5a4e3710cc5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/4906 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "uploaded data to:\n",
      "training dataset to: s3://sagemaker-us-east-1-372108735839/processed/wandbot/train\n"
     ]
    }
   ],
   "source": [
    "# save train_dataset to s3\n",
    "if USE_S3:\n",
    "    training_input_path = f's3://{sess.default_bucket()}/processed/wandbot/train'\n",
    "else:\n",
    "    training_input_path = \"./wandbot_train_ds\"\n",
    "\n",
    "lm_dataset.save_to_disk(training_input_path)\n",
    "\n",
    "print(\"uploaded data to:\")\n",
    "print(f\"training dataset to: {training_input_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "f92d3c9a-f01d-4b6e-8a20-ab707918a0e3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "at = wandb.Artifact(\n",
    "    name=\"wandbot_dataset_tokenized\", \n",
    "    type=\"dataset\",\n",
    "    description=\"A wandbot dataset of questions and answers about W&B - CodeLLama tokenized\",\n",
    "    metadata={\"model_name\": MODEL_NAME, \"tokenizer\": MODEL_NAME},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "7edcf658-cd05-46de-bc30-45c27a8f4853",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Generating checksum for up to 10000 objects in \"sagemaker-us-east-1-372108735839/processed/wandbot/train\"... Done. 0.1s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Artifact wandbot_dataset_tokenized>"
      ]
     },
     "execution_count": 56,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if USE_S3:\n",
    "    at.add_reference(training_input_path)\n",
    "else:\n",
    "    at.add_dir(training_input_path)\n",
    "wandb.log_artifact(at)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bbd2e08-939d-4269-8d65-8ae3684fb3de",
   "metadata": {},
   "source": [
    "Let's finish this run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "3c3900fa-e7b2-4b59-be6f-d9d2459e237c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "243fbf98439b4f31b66bd29f8bbb3d89",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='0.051 MB of 0.051 MB uploaded (0.000 MB deduped)\\r'), FloatProgress(value=1.0, max…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "wandb: WARNING Source type is set to 'repo' but some required information is missing from the environment. A job will not be created from this run. See https://docs.wandb.ai/guides/launch/create-job\n"
     ]
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">wandering-totem-5</strong> at: <a href='https://wandb.ai/capecape/aws_llm_workshop/runs/6fojvwiz' target=\"_blank\">https://wandb.ai/capecape/aws_llm_workshop/runs/6fojvwiz</a><br/>Synced 7 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20231017_103308-6fojvwiz/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19b63367-19f7-4bb6-b2f8-44e200c74df6",
   "metadata": {},
   "source": [
    "# Eval Dataset\n",
    "We prepared a set of questions from `wandbot` that were gathered and curated by my colleague Ayush T. God's work here!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "8d607bff-cbbd-4554-822f-a0fbabe0c659",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "Finishing last run (ID:9axrufsd) before initializing another..."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "wandb: WARNING Source type is set to 'repo' but some required information is missing from the environment. A job will not be created from this run. See https://docs.wandb.ai/guides/launch/create-job\n"
     ]
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">happy-deluge-21</strong> at: <a href='https://wandb.ai/capecape/aws_llm_workshop/runs/9axrufsd' target=\"_blank\">https://wandb.ai/capecape/aws_llm_workshop/runs/9axrufsd</a><br/>Synced 6 W&B file(s), 0 media file(s), 0 artifact file(s) and 2 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20231026_111723-9axrufsd/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Successfully finished last run (ID:9axrufsd). Initializing new run:<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b4f8d441d89445769826f531e25fe5fd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='Waiting for wandb.init()...\\r'), FloatProgress(value=0.011112445388991747, max=1.0…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.15.12"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/root/sagemaker-llama2-wandb/wandb/run-20231026_111744-qmosgmka</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/capecape/aws_llm_workshop/runs/qmosgmka' target=\"_blank\">whole-forest-22</a></strong> to <a href='https://wandb.ai/capecape/aws_llm_workshop' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/capecape/aws_llm_workshop' target=\"_blank\">https://wandb.ai/capecape/aws_llm_workshop</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/capecape/aws_llm_workshop/runs/qmosgmka' target=\"_blank\">https://wandb.ai/capecape/aws_llm_workshop/runs/qmosgmka</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import json\n",
    "import wandb\n",
    "import pandas as pd\n",
    "\n",
    "wandb.init(project=WANDB_PROJECT, job_type=\"eval_preprocessing\")\n",
    "question_artifacts = wandb.use_artifact(RAW_EVAL_DATASET_ARTIFACT)\n",
    "\n",
    "with open(question_artifacts.file()) as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "columns = data[\"columns\"]\n",
    "data = data[\"data\"]\n",
    "eval_df = pd.DataFrame(data, columns=columns)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6c78e291-ce2d-4061-9720-71f8bdf317ab",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query</th>\n",
       "      <th>retrieved_context</th>\n",
       "      <th>generated_response</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Hey I have a question about using wandb with f...</td>\n",
       "      <td>[`wandb.init()` returns a run object, and you ...</td>\n",
       "      <td>The initialization of `wandb.init()` should be...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Hey with wandb is it possible to link from the...</td>\n",
       "      <td>[### How do I best log models from runs in a s...</td>\n",
       "      <td>Yes, you can link to the best run from a sweep...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>I am training a spacy textcat model. This proc...</td>\n",
       "      <td>[### Model Architecture\\n\\nOur config also def...</td>\n",
       "      <td>To log the best model's metrics instead of the...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Explain how I can version datasets with Weight...</td>\n",
       "      <td>[ # A 3-in-1 Intro to Weights &amp; Biases: Comput...</td>\n",
       "      <td>Weights &amp; Biases provides a feature called Art...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>I'm using   env = SubprocVecEnv(\\n        [mak...</td>\n",
       "      <td>[## WandbCallback 引数\\n\\n| 引数 | 使用法 | | --- | -...</td>\n",
       "      <td>To record a video for a specific subprocess en...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               query  \\\n",
       "0  Hey I have a question about using wandb with f...   \n",
       "1  Hey with wandb is it possible to link from the...   \n",
       "2  I am training a spacy textcat model. This proc...   \n",
       "3  Explain how I can version datasets with Weight...   \n",
       "4  I'm using   env = SubprocVecEnv(\\n        [mak...   \n",
       "\n",
       "                                   retrieved_context  \\\n",
       "0  [`wandb.init()` returns a run object, and you ...   \n",
       "1  [### How do I best log models from runs in a s...   \n",
       "2  [### Model Architecture\\n\\nOur config also def...   \n",
       "3  [ # A 3-in-1 Intro to Weights & Biases: Comput...   \n",
       "4  [## WandbCallback 引数\\n\\n| 引数 | 使用法 | | --- | -...   \n",
       "\n",
       "                                  generated_response  \n",
       "0  The initialization of `wandb.init()` should be...  \n",
       "1  Yes, you can link to the best run from a sweep...  \n",
       "2  To log the best model's metrics instead of the...  \n",
       "3  Weights & Biases provides a feature called Art...  \n",
       "4  To record a video for a specific subprocess en...  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffb58d1b-0b17-49e3-8768-b9baeb3d9046",
   "metadata": {},
   "source": [
    "Let's remove retrieved Japanese text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "646a2941-5ea9-47be-98e2-8731387f6114",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Clean up and prepare (pandas workout)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fee50101-b6a7-4e08-ab2f-e426b5e9840c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def contains_japanese(text):\n",
    "    for char in text:\n",
    "        if '一' <= char <= '龥':\n",
    "            return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "365bff6a-c0bc-4765-8044-854df497d29b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "s = \"## WandbCallback 引数\\n\\n| 引数 | 使用法 |\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "addabb0a-5569-4c26-9b29-17484a20832d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "eval_df[\"retrieved_context_en\"] = [[ctx for ctx in ctxs if not contains_japanese(ctx)] for ctxs in eval_df.retrieved_context.values]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cd18bd4e-aeb5-4ac7-a547-08252188b048",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query</th>\n",
       "      <th>retrieved_context</th>\n",
       "      <th>generated_response</th>\n",
       "      <th>retrieved_context_en</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Hey I have a question about using wandb with f...</td>\n",
       "      <td>[`wandb.init()` returns a run object, and you ...</td>\n",
       "      <td>The initialization of `wandb.init()` should be...</td>\n",
       "      <td>[`wandb.init()` returns a run object, and you ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Hey with wandb is it possible to link from the...</td>\n",
       "      <td>[### How do I best log models from runs in a s...</td>\n",
       "      <td>Yes, you can link to the best run from a sweep...</td>\n",
       "      <td>[### How do I best log models from runs in a s...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>I am training a spacy textcat model. This proc...</td>\n",
       "      <td>[### Model Architecture\\n\\nOur config also def...</td>\n",
       "      <td>To log the best model's metrics instead of the...</td>\n",
       "      <td>[### Model Architecture\\n\\nOur config also def...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Explain how I can version datasets with Weight...</td>\n",
       "      <td>[ # A 3-in-1 Intro to Weights &amp; Biases: Comput...</td>\n",
       "      <td>Weights &amp; Biases provides a feature called Art...</td>\n",
       "      <td>[ # A 3-in-1 Intro to Weights &amp; Biases: Comput...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>I'm using   env = SubprocVecEnv(\\n        [mak...</td>\n",
       "      <td>[## WandbCallback 引数\\n\\n| 引数 | 使用法 | | --- | -...</td>\n",
       "      <td>To record a video for a specific subprocess en...</td>\n",
       "      <td>[## Basic Example\\n\\nThe W&amp;B SB3 integration u...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               query  \\\n",
       "0  Hey I have a question about using wandb with f...   \n",
       "1  Hey with wandb is it possible to link from the...   \n",
       "2  I am training a spacy textcat model. This proc...   \n",
       "3  Explain how I can version datasets with Weight...   \n",
       "4  I'm using   env = SubprocVecEnv(\\n        [mak...   \n",
       "\n",
       "                                   retrieved_context  \\\n",
       "0  [`wandb.init()` returns a run object, and you ...   \n",
       "1  [### How do I best log models from runs in a s...   \n",
       "2  [### Model Architecture\\n\\nOur config also def...   \n",
       "3  [ # A 3-in-1 Intro to Weights & Biases: Comput...   \n",
       "4  [## WandbCallback 引数\\n\\n| 引数 | 使用法 | | --- | -...   \n",
       "\n",
       "                                  generated_response  \\\n",
       "0  The initialization of `wandb.init()` should be...   \n",
       "1  Yes, you can link to the best run from a sweep...   \n",
       "2  To log the best model's metrics instead of the...   \n",
       "3  Weights & Biases provides a feature called Art...   \n",
       "4  To record a video for a specific subprocess en...   \n",
       "\n",
       "                                retrieved_context_en  \n",
       "0  [`wandb.init()` returns a run object, and you ...  \n",
       "1  [### How do I best log models from runs in a s...  \n",
       "2  [### Model Architecture\\n\\nOur config also def...  \n",
       "3  [ # A 3-in-1 Intro to Weights & Biases: Comput...  \n",
       "4  [## Basic Example\\n\\nThe W&B SB3 integration u...  "
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0157fa00-a477-4394-9517-0568a58d392c",
   "metadata": {},
   "source": [
    "Just keep the first page results, to save memory..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "12b8732a-e42e-4f19-b8ee-baceb673f32d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "eval_df[\"retrieved_context_first\"] = [ctxs[0] for ctxs in eval_df.retrieved_context_en.values]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "14459a92-98fe-4674-9e8a-7afdaf7c1095",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query</th>\n",
       "      <th>retrieved_context</th>\n",
       "      <th>generated_response</th>\n",
       "      <th>retrieved_context_en</th>\n",
       "      <th>retrieved_context_first</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Hey I have a question about using wandb with f...</td>\n",
       "      <td>[`wandb.init()` returns a run object, and you ...</td>\n",
       "      <td>The initialization of `wandb.init()` should be...</td>\n",
       "      <td>[`wandb.init()` returns a run object, and you ...</td>\n",
       "      <td>`wandb.init()` returns a run object, and you c...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Hey with wandb is it possible to link from the...</td>\n",
       "      <td>[### How do I best log models from runs in a s...</td>\n",
       "      <td>Yes, you can link to the best run from a sweep...</td>\n",
       "      <td>[### How do I best log models from runs in a s...</td>\n",
       "      <td>### How do I best log models from runs in a sw...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>I am training a spacy textcat model. This proc...</td>\n",
       "      <td>[### Model Architecture\\n\\nOur config also def...</td>\n",
       "      <td>To log the best model's metrics instead of the...</td>\n",
       "      <td>[### Model Architecture\\n\\nOur config also def...</td>\n",
       "      <td>### Model Architecture\\n\\nOur config also defi...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Explain how I can version datasets with Weight...</td>\n",
       "      <td>[ # A 3-in-1 Intro to Weights &amp; Biases: Comput...</td>\n",
       "      <td>Weights &amp; Biases provides a feature called Art...</td>\n",
       "      <td>[ # A 3-in-1 Intro to Weights &amp; Biases: Comput...</td>\n",
       "      <td># A 3-in-1 Intro to Weights &amp; Biases: Compute...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>I'm using   env = SubprocVecEnv(\\n        [mak...</td>\n",
       "      <td>[## WandbCallback 引数\\n\\n| 引数 | 使用法 | | --- | -...</td>\n",
       "      <td>To record a video for a specific subprocess en...</td>\n",
       "      <td>[## Basic Example\\n\\nThe W&amp;B SB3 integration u...</td>\n",
       "      <td>## Basic Example\\n\\nThe W&amp;B SB3 integration us...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               query  \\\n",
       "0  Hey I have a question about using wandb with f...   \n",
       "1  Hey with wandb is it possible to link from the...   \n",
       "2  I am training a spacy textcat model. This proc...   \n",
       "3  Explain how I can version datasets with Weight...   \n",
       "4  I'm using   env = SubprocVecEnv(\\n        [mak...   \n",
       "\n",
       "                                   retrieved_context  \\\n",
       "0  [`wandb.init()` returns a run object, and you ...   \n",
       "1  [### How do I best log models from runs in a s...   \n",
       "2  [### Model Architecture\\n\\nOur config also def...   \n",
       "3  [ # A 3-in-1 Intro to Weights & Biases: Comput...   \n",
       "4  [## WandbCallback 引数\\n\\n| 引数 | 使用法 | | --- | -...   \n",
       "\n",
       "                                  generated_response  \\\n",
       "0  The initialization of `wandb.init()` should be...   \n",
       "1  Yes, you can link to the best run from a sweep...   \n",
       "2  To log the best model's metrics instead of the...   \n",
       "3  Weights & Biases provides a feature called Art...   \n",
       "4  To record a video for a specific subprocess en...   \n",
       "\n",
       "                                retrieved_context_en  \\\n",
       "0  [`wandb.init()` returns a run object, and you ...   \n",
       "1  [### How do I best log models from runs in a s...   \n",
       "2  [### Model Architecture\\n\\nOur config also def...   \n",
       "3  [ # A 3-in-1 Intro to Weights & Biases: Comput...   \n",
       "4  [## Basic Example\\n\\nThe W&B SB3 integration u...   \n",
       "\n",
       "                             retrieved_context_first  \n",
       "0  `wandb.init()` returns a run object, and you c...  \n",
       "1  ### How do I best log models from runs in a sw...  \n",
       "2  ### Model Architecture\\n\\nOur config also defi...  \n",
       "3   # A 3-in-1 Intro to Weights & Biases: Compute...  \n",
       "4  ## Basic Example\\n\\nThe W&B SB3 integration us...  "
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "61dd549f-1923-4e81-ba84-a5fa5dcbb297",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "eval_df = eval_df.assign(tokens = eval_df['retrieved_context_first'].str.len())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "a37a312e-08cf-4406-99d5-87342709dcc7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>query</th>\n",
       "      <th>retrieved_context</th>\n",
       "      <th>generated_response</th>\n",
       "      <th>retrieved_context_en</th>\n",
       "      <th>retrieved_context_first</th>\n",
       "      <th>tokens</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Hey I have a question about using wandb with f...</td>\n",
       "      <td>[`wandb.init()` returns a run object, and you ...</td>\n",
       "      <td>The initialization of `wandb.init()` should be...</td>\n",
       "      <td>[`wandb.init()` returns a run object, and you ...</td>\n",
       "      <td>`wandb.init()` returns a run object, and you c...</td>\n",
       "      <td>505</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Hey with wandb is it possible to link from the...</td>\n",
       "      <td>[### How do I best log models from runs in a s...</td>\n",
       "      <td>Yes, you can link to the best run from a sweep...</td>\n",
       "      <td>[### How do I best log models from runs in a s...</td>\n",
       "      <td>### How do I best log models from runs in a sw...</td>\n",
       "      <td>1116</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>I am training a spacy textcat model. This proc...</td>\n",
       "      <td>[### Model Architecture\\n\\nOur config also def...</td>\n",
       "      <td>To log the best model's metrics instead of the...</td>\n",
       "      <td>[### Model Architecture\\n\\nOur config also def...</td>\n",
       "      <td>### Model Architecture\\n\\nOur config also defi...</td>\n",
       "      <td>1183</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Explain how I can version datasets with Weight...</td>\n",
       "      <td>[ # A 3-in-1 Intro to Weights &amp; Biases: Comput...</td>\n",
       "      <td>Weights &amp; Biases provides a feature called Art...</td>\n",
       "      <td>[ # A 3-in-1 Intro to Weights &amp; Biases: Comput...</td>\n",
       "      <td># A 3-in-1 Intro to Weights &amp; Biases: Compute...</td>\n",
       "      <td>576</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>I'm using   env = SubprocVecEnv(\\n        [mak...</td>\n",
       "      <td>[## WandbCallback 引数\\n\\n| 引数 | 使用法 | | --- | -...</td>\n",
       "      <td>To record a video for a specific subprocess en...</td>\n",
       "      <td>[## Basic Example\\n\\nThe W&amp;B SB3 integration u...</td>\n",
       "      <td>## Basic Example\\n\\nThe W&amp;B SB3 integration us...</td>\n",
       "      <td>1265</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               query  \\\n",
       "0  Hey I have a question about using wandb with f...   \n",
       "1  Hey with wandb is it possible to link from the...   \n",
       "2  I am training a spacy textcat model. This proc...   \n",
       "3  Explain how I can version datasets with Weight...   \n",
       "4  I'm using   env = SubprocVecEnv(\\n        [mak...   \n",
       "\n",
       "                                   retrieved_context  \\\n",
       "0  [`wandb.init()` returns a run object, and you ...   \n",
       "1  [### How do I best log models from runs in a s...   \n",
       "2  [### Model Architecture\\n\\nOur config also def...   \n",
       "3  [ # A 3-in-1 Intro to Weights & Biases: Comput...   \n",
       "4  [## WandbCallback 引数\\n\\n| 引数 | 使用法 | | --- | -...   \n",
       "\n",
       "                                  generated_response  \\\n",
       "0  The initialization of `wandb.init()` should be...   \n",
       "1  Yes, you can link to the best run from a sweep...   \n",
       "2  To log the best model's metrics instead of the...   \n",
       "3  Weights & Biases provides a feature called Art...   \n",
       "4  To record a video for a specific subprocess en...   \n",
       "\n",
       "                                retrieved_context_en  \\\n",
       "0  [`wandb.init()` returns a run object, and you ...   \n",
       "1  [### How do I best log models from runs in a s...   \n",
       "2  [### Model Architecture\\n\\nOur config also def...   \n",
       "3  [ # A 3-in-1 Intro to Weights & Biases: Comput...   \n",
       "4  [## Basic Example\\n\\nThe W&B SB3 integration u...   \n",
       "\n",
       "                             retrieved_context_first  tokens  \n",
       "0  `wandb.init()` returns a run object, and you c...     505  \n",
       "1  ### How do I best log models from runs in a sw...    1116  \n",
       "2  ### Model Architecture\\n\\nOur config also defi...    1183  \n",
       "3   # A 3-in-1 Intro to Weights & Biases: Compute...     576  \n",
       "4  ## Basic Example\\n\\nThe W&B SB3 integration us...    1265  "
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "24f78893-9828-4c35-b5e6-a84b5a32134b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "eval_df = eval_df[[\"query\", \"generated_response\", \"retrieved_context_first\", \"tokens\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "70e8f2c5-1c1a-4cc7-9419-bf2e8749cb9c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "eval_df.columns = [\"question\", \"answer\", \"retrieved_context\", \"char_len\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "50e7ff7d-27fe-4e9d-b04a-98a9494387c0",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>question</th>\n",
       "      <th>answer</th>\n",
       "      <th>page_content</th>\n",
       "      <th>char_len</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Hey I have a question about using wandb with f...</td>\n",
       "      <td>The initialization of `wandb.init()` should be...</td>\n",
       "      <td>`wandb.init()` returns a run object, and you c...</td>\n",
       "      <td>505</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Hey with wandb is it possible to link from the...</td>\n",
       "      <td>Yes, you can link to the best run from a sweep...</td>\n",
       "      <td>### How do I best log models from runs in a sw...</td>\n",
       "      <td>1116</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>I am training a spacy textcat model. This proc...</td>\n",
       "      <td>To log the best model's metrics instead of the...</td>\n",
       "      <td>### Model Architecture\\n\\nOur config also defi...</td>\n",
       "      <td>1183</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Explain how I can version datasets with Weight...</td>\n",
       "      <td>Weights &amp; Biases provides a feature called Art...</td>\n",
       "      <td># A 3-in-1 Intro to Weights &amp; Biases: Compute...</td>\n",
       "      <td>576</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>I'm using   env = SubprocVecEnv(\\n        [mak...</td>\n",
       "      <td>To record a video for a specific subprocess en...</td>\n",
       "      <td>## Basic Example\\n\\nThe W&amp;B SB3 integration us...</td>\n",
       "      <td>1265</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            question  \\\n",
       "0  Hey I have a question about using wandb with f...   \n",
       "1  Hey with wandb is it possible to link from the...   \n",
       "2  I am training a spacy textcat model. This proc...   \n",
       "3  Explain how I can version datasets with Weight...   \n",
       "4  I'm using   env = SubprocVecEnv(\\n        [mak...   \n",
       "\n",
       "                                              answer  \\\n",
       "0  The initialization of `wandb.init()` should be...   \n",
       "1  Yes, you can link to the best run from a sweep...   \n",
       "2  To log the best model's metrics instead of the...   \n",
       "3  Weights & Biases provides a feature called Art...   \n",
       "4  To record a video for a specific subprocess en...   \n",
       "\n",
       "                                        page_content  char_len  \n",
       "0  `wandb.init()` returns a run object, and you c...       505  \n",
       "1  ### How do I best log models from runs in a sw...      1116  \n",
       "2  ### Model Architecture\\n\\nOur config also defi...      1183  \n",
       "3   # A 3-in-1 Intro to Weights & Biases: Compute...       576  \n",
       "4  ## Basic Example\\n\\nThe W&B SB3 integration us...      1265  "
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_df = eval_df.rename({\"retrieved_context\": \"page_content\"}, axis=1)\n",
    "eval_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3f8de07-59c4-40f0-810b-7a1805ab5a8a",
   "metadata": {},
   "source": [
    "## Save to W&B"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "108ce808-876b-404d-bc98-eba585672252",
   "metadata": {
    "tags": []
   },
   "source": [
    "Let's format the dataset in the same way we created the training dataset, we have to be consisten with naming\n",
    "- We remove the answer, but we are going to keep it on the dataset to evaluate the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "664f5165-0d3d-4e19-b237-2d9c198c5c1f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "B_INST, E_INST = \"[INST] \", \" [/INST]\"\n",
    "B_SYS, E_SYS = \"<<SYS>>\\n\", \"\\n<</SYS>>\\n\\n\"\n",
    "EOS = \"</s>\"\n",
    "\n",
    "eval_prompt_format = (\n",
    "    B_INST\n",
    "    + B_SYS\n",
    "    + \"You are an AI assistant designed to assist developers with everyday tasks related to Weight & Biases\"\n",
    "    + \"and provide helpful information. As an expert in the open-source python SDK wandb answer the following\"\n",
    "    # + \"question based on the context below. Answer in formatted Markdown.\\n\"\n",
    "    + \"question below. Answer in formatted Markdown.\\n\"\n",
    "    + \"{page_content}\"\n",
    "    + E_SYS\n",
    "    + \"{question}\"\n",
    "    + E_INST\n",
    "    + \"\\n[W&B]\\n\"\n",
    "    # + \"{answer}\"\n",
    "    # + \"\\n[/W&B]\"\n",
    "    # + EOS\n",
    ")\n",
    "\n",
    "def eval_format_text(row):\n",
    "    return eval_prompt_format.format_map(row)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "bb560592-dbef-4da7-a9cb-a3463341910c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "eval_df[\"text\"] = eval_df.apply(eval_format_text, axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a29f448b-07a4-4dc9-8621-98d50a9edcd0",
   "metadata": {},
   "source": [
    "Save to disk and create HF dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "6808da60-9b6e-43d7-a893-df1c677087b0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "2b6c968a-892d-4d76-a921-63a473bc71dd",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Repo card metadata block was not found. Setting CardData to empty.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "92e52e13f8694eb1905d3bac4fa6006d",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ca944389540f4b87b7e091cadca46ba4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Extracting data files:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a2d4a90a35974e77be605959cf2b79f7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['question', 'answer', 'page_content', 'char_len', 'text'],\n",
       "        num_rows: 132\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_df.to_json(\"wandbot_eval.jsonl\", orient='records', lines=True)\n",
    "eval_dataset = load_dataset(\".\", data_files=\"wandbot_eval.jsonl\")\n",
    "eval_dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "3dd86384-2cae-41f8-83a5-6597242244db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save train_dataset to s3\n",
    "if USE_S3:\n",
    "    eval_input_path = f's3://{sess.default_bucket()}/processed/wandbot/eval'\n",
    "else:\n",
    "    eval_input_path = \"./wandbot_eval_ds\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "7a4c5caf-c577-4360-881d-ef36b5bb39d2",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/fsspec/registry.py:272: UserWarning: Your installed version of s3fs is very old and known to cause\n",
      "severe performance issues, see also https://github.com/dask/dask/issues/10276\n",
      "\n",
      "To fix, you should specify a lower version bound on s3fs, or\n",
      "update the current installation.\n",
      "\n",
      "  warnings.warn(s3_msg)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "54777fd7aea84e838c6af6084d5d3771",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Saving the dataset (0/1 shards):   0%|          | 0/132 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "eval_dataset.save_to_disk(eval_input_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "cd9a87e0-d691-4076-8894-c5f4157bdf7f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "table = wandb.Table(dataframe=eval_df)\n",
    "wandb.log({\"wandbot_eval_dataset\": table})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "4d615bb7-f385-4771-a593-c2845561d0fd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "at = wandb.Artifact(\n",
    "    name=\"wandbot_eval_dataset\", \n",
    "    type=\"dataset\",\n",
    "    description=\"A wandbot dataset of questions and answers about W&B for evaluation\",\n",
    "    metadata={\"prompt_format\": eval_prompt_format,\n",
    "              \"length\": len(eval_dataset),\n",
    "             })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "2b87fb8c-c12e-4f63-b9d2-6ef7d014db1f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Generating checksum for up to 10000 objects in \"sagemaker-us-east-1-372108735839/processed/wandbot/eval\"... Done. 0.1s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<Artifact wandbot_eval_dataset>"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if USE_S3:\n",
    "    at.add_reference(eval_input_path)\n",
    "else:\n",
    "    at.add_dir(eval_input_path)\n",
    "wandb.log_artifact(at)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "ab877b6f-28e2-4d2d-bf28-86cba3736dcd",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "Waiting for W&B process to finish... <strong style=\"color:green\">(success).</strong>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "wandb: WARNING Source type is set to 'repo' but some required information is missing from the environment. A job will not be created from this run. See https://docs.wandb.ai/guides/launch/create-job\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "878cb2c9d7414f20a59dd72c1b2f7573",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='1.528 MB of 1.528 MB uploaded (0.000 MB deduped)\\r'), FloatProgress(value=1.0, max…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">whole-forest-22</strong> at: <a href='https://wandb.ai/capecape/aws_llm_workshop/runs/qmosgmka' target=\"_blank\">https://wandb.ai/capecape/aws_llm_workshop/runs/qmosgmka</a><br/>Synced 7 W&B file(s), 1 media file(s), 1 artifact file(s) and 1 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20231026_111744-qmosgmka/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36ffd90e-d069-42bf-a3de-bbcfa74f7e23",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "availableInstances": [
   {
    "_defaultOrder": 0,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.t3.medium",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 1,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.t3.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 2,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.t3.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 3,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.t3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 4,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 5,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 6,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 7,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 8,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 9,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 10,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 11,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 12,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5d.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 13,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5d.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 14,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5d.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 15,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5d.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 16,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5d.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 17,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5d.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 18,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5d.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 19,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 20,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": true,
    "memoryGiB": 0,
    "name": "ml.geospatial.interactive",
    "supportedImageNames": [
     "sagemaker-geospatial-v1-0"
    ],
    "vcpuNum": 0
   },
   {
    "_defaultOrder": 21,
    "_isFastLaunch": true,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.c5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 22,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.c5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 23,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.c5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 24,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.c5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 25,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 72,
    "name": "ml.c5.9xlarge",
    "vcpuNum": 36
   },
   {
    "_defaultOrder": 26,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 96,
    "name": "ml.c5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 27,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 144,
    "name": "ml.c5.18xlarge",
    "vcpuNum": 72
   },
   {
    "_defaultOrder": 28,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.c5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 29,
    "_isFastLaunch": true,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g4dn.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 30,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g4dn.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 31,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g4dn.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 32,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g4dn.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 33,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g4dn.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 34,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g4dn.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 35,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 61,
    "name": "ml.p3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 36,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 244,
    "name": "ml.p3.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 37,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 488,
    "name": "ml.p3.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 38,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.p3dn.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 39,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.r5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 40,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.r5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 41,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.r5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 42,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.r5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 43,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.r5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 44,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.r5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 45,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.r5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 46,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.r5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 47,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 48,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 49,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 50,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 51,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 52,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 53,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.g5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 54,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.g5.48xlarge",
    "vcpuNum": 192
   },
   {
    "_defaultOrder": 55,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 56,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4de.24xlarge",
    "vcpuNum": 96
   }
  ],
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science 3.0)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/sagemaker-data-science-310-v1"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}