falcon-7b-qlora-remote-decorator_qa.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-tune Falcon-7B with QLoRA and SageMaker remote decorator\n",
    "\n",
    "## Question & Answering"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "In this demo notebook, we demonstrate how to fine-tune the Falcon-7B model using QLoRA, Hugging Face PEFT, and bitsandbytes.\n",
    "\n",
    "We are using SageMaker remote decorator for runinng the fine-tuning job on Amazon SageMaker Training job\n",
    "---\n",
    "\n",
    "SageMaker Distribution image: SageMaker Distribution 2.0\n",
    "\n",
    "PyTorch version: 3.11\n",
    "\n",
    "JupyterLab Instance Type: ml.t3.medium\n",
    "\n",
    "Fine-Tuning:\n",
    "* Instance Type: ml.g5.8xlarge"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Install the required libriaries, including the Hugging Face libraries, and restart the kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-07-20T12:38:06.851473Z",
     "start_time": "2023-07-20T12:38:04.440644Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Setup Configuration file path\n",
    "\n",
    "We are setting the directory in which the config.yaml file resides so that remote decorator can make use of the settings through [SageMaker Defaults](https://sagemaker.readthedocs.io/en/stable/overview.html#configuring-and-using-defaults-with-the-sagemaker-python-sdk).\n",
    "\n",
    "This notebook is using the Hugging Face container for the `us-east-1` region. Make sure you are using the right image for your AWS region, otherwise edit [config.yaml](./config.yaml). Container Images are available [here](https://github.com/aws/deep-learning-containers/blob/master/available_images.md)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T18:18:46.946401Z",
     "start_time": "2023-09-02T18:18:46.887254Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Set path to config file\n",
    "os.environ[\"SAGEMAKER_USER_CONFIG_OVERRIDE\"] = os.getcwd()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualize and upload the dataset\n",
    "\n",
    "Read train dataset in a Pandas dataframe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T17:16:04.934328Z",
     "start_time": "2023-09-02T17:16:04.591322Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "df = pd.read_csv('train.csv.gz', compression='gzip', sep=';')\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T17:16:05.500607Z",
     "start_time": "2023-09-02T17:16:05.115256Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "train, test = train_test_split(df, test_size=0.3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a prompt template and load the dataset with a random sample to try summarization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T17:16:06.169362Z",
     "start_time": "2023-09-02T17:16:06.162912Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from random import randint\n",
    "\n",
    "# custom instruct prompt start\n",
    "prompt_template = f\"{{question}}\\n---\\nAnswer:\\n{{answer}}\"\n",
    "\n",
    "# template dataset to add prompt to each sample\n",
    "def template_dataset(sample):\n",
    "    sample[\"text\"] = prompt_template.format(question=sample[\"question\"],\n",
    "                                            answer=sample[\"answers\"])\n",
    "    return sample"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Use the Hugging Face Trainer class to fine-tune the model. Define the hyperparameters we want to use. We also create a DataCollator that will take care of padding our inputs and labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T17:16:08.239370Z",
     "start_time": "2023-09-02T17:16:07.804881Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import Dataset, DatasetDict\n",
    "\n",
    "train_dataset = Dataset.from_pandas(train)\n",
    "test_dataset = Dataset.from_pandas(test)\n",
    "\n",
    "dataset = DatasetDict({\"train\": train_dataset, \"test\": test_dataset})\n",
    "\n",
    "train_dataset = dataset[\"train\"].map(template_dataset, remove_columns=list(dataset[\"train\"].features))\n",
    "\n",
    "print(train_dataset[randint(0, len(dataset))][\"text\"])\n",
    "\n",
    "test_dataset = dataset[\"test\"].map(template_dataset, remove_columns=list(dataset[\"test\"].features))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n",
    "To train our model, we need to convert our inputs (text) to token IDs. This is done by a Hugging Face Transformers Tokenizer. In addition to QLoRA, we will use bitsanbytes 4-bit precision to quantize out frozen LLM to 4-bit and attach LoRA adapters on it.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T17:16:05.736006Z",
     "start_time": "2023-09-02T17:16:05.730949Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def print_trainable_parameters(model):\n",
    "    \"\"\"\n",
    "    Prints the number of trainable parameters in the model.\n",
    "    \"\"\"\n",
    "    trainable_params = 0\n",
    "    all_param = 0\n",
    "    for _, param in model.named_parameters():\n",
    "        all_param += param.numel()\n",
    "        if param.requires_grad:\n",
    "            trainable_params += param.numel()\n",
    "    print(\n",
    "        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Utility method for finding the target modules and update the necessary matrices. Visit [this](https://github.com/artidoro/qlora/blob/main/qlora.py) link for additional info."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T17:16:09.725306Z",
     "start_time": "2023-09-02T17:16:08.992201Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import bitsandbytes as bnb\n",
    "\n",
    "def find_all_linear_names(hf_model):\n",
    "    lora_module_names = set()\n",
    "    for name, module in hf_model.named_modules():\n",
    "        if isinstance(module, bnb.nn.Linear4bit):\n",
    "            names = name.split(\".\")\n",
    "            lora_module_names.add(names[0] if len(names) == 1 else names[-1])\n",
    "\n",
    "    if \"lm_head\" in lora_module_names:  # needed for 16-bit\n",
    "        lora_module_names.remove(\"lm_head\")\n",
    "    return list(lora_module_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Define the train function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = \"tiiuae/falcon-7b\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T17:16:10.561652Z",
     "start_time": "2023-09-02T17:16:10.245128Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from accelerate import Accelerator\n",
    "from huggingface_hub import login\n",
    "from peft import AutoPeftModelForCausalLM, LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
    "from sagemaker.remote_function import remote\n",
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, set_seed\n",
    "import transformers\n",
    "\n",
    "# Start training\n",
    "@remote(volume_size=100, job_name_prefix=f\"train-{model_id.split('/')[-1].replace('.', '-')}\")\n",
    "def train_fn(\n",
    "        model_name,\n",
    "        train_ds,\n",
    "        test_ds=None,\n",
    "        lora_r=8,\n",
    "        lora_alpha=16,\n",
    "        lora_dropout=0.1,\n",
    "        per_device_train_batch_size=8,\n",
    "        per_device_eval_batch_size=8,\n",
    "        gradient_accumulation_steps=1,\n",
    "        learning_rate=2e-4,\n",
    "        num_train_epochs=1,\n",
    "        gradient_checkpointing=False,\n",
    "        merge_weights=False,\n",
    "        seed=42,\n",
    "        token=None\n",
    "):\n",
    "    set_seed(seed)\n",
    "\n",
    "    accelerator = Accelerator()\n",
    "\n",
    "    if token is not None:\n",
    "        login(token=token)\n",
    "\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "    # Set Tokenizer pad Token\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "    with accelerator.main_process_first():\n",
    "        # tokenize and chunk dataset\n",
    "        lm_train_dataset = train_ds.map(\n",
    "            lambda sample: tokenizer(sample[\"text\"]), remove_columns=list(train_ds.features)\n",
    "        )\n",
    "\n",
    "        print(f\"Total number of train samples: {len(lm_train_dataset)}\")\n",
    "\n",
    "        if test_ds is not None:\n",
    "            lm_test_dataset = test_ds.map(\n",
    "                lambda sample: tokenizer(sample[\"text\"]), remove_columns=list(test_ds.features)\n",
    "            )\n",
    "\n",
    "            print(f\"Total number of test samples: {len(lm_test_dataset)}\")\n",
    "        else:\n",
    "            lm_test_dataset = None\n",
    "\n",
    "    bnb_config = BitsAndBytesConfig(\n",
    "        load_in_4bit=True,\n",
    "        bnb_4bit_use_double_quant=True,\n",
    "        bnb_4bit_quant_type=\"nf4\",\n",
    "        bnb_4bit_compute_dtype=torch.bfloat16\n",
    "    )\n",
    "    \n",
    "    model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_name,\n",
    "        trust_remote_code=True,\n",
    "        quantization_config=bnb_config,\n",
    "        attn_implementation=\"flash_attention_2\",\n",
    "        use_cache=False if gradient_checkpointing else True,\n",
    "        cache_dir=\"/tmp/.cache\"\n",
    "    )\n",
    "\n",
    "    model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=gradient_checkpointing)\n",
    "\n",
    "    if gradient_checkpointing:\n",
    "        model.gradient_checkpointing_enable()\n",
    "\n",
    "    # get lora target modules\n",
    "    modules = find_all_linear_names(model)\n",
    "    print(f\"Found {len(modules)} modules to quantize: {modules}\")\n",
    "    \n",
    "    config = LoraConfig(\n",
    "        r=lora_r,\n",
    "        lora_alpha=lora_alpha,\n",
    "        target_modules=modules,\n",
    "        lora_dropout=lora_dropout,\n",
    "        bias=\"none\",\n",
    "        task_type=\"CAUSAL_LM\"\n",
    "    )\n",
    "\n",
    "    model = get_peft_model(model, config)\n",
    "    print_trainable_parameters(model)\n",
    "    \n",
    "    trainer = transformers.Trainer(\n",
    "        model=model,\n",
    "        train_dataset=lm_train_dataset,\n",
    "        eval_dataset=lm_test_dataset if lm_test_dataset is not None else None,\n",
    "        args=transformers.TrainingArguments(\n",
    "            per_device_train_batch_size=per_device_train_batch_size,\n",
    "            per_device_eval_batch_size=per_device_eval_batch_size,\n",
    "            gradient_accumulation_steps=gradient_accumulation_steps,\n",
    "            gradient_checkpointing=gradient_checkpointing,\n",
    "            logging_strategy=\"steps\",\n",
    "            logging_steps=1,\n",
    "            log_on_each_node=False,\n",
    "            num_train_epochs=num_train_epochs,\n",
    "            learning_rate=learning_rate,\n",
    "            bf16=True,\n",
    "            ddp_find_unused_parameters=False,\n",
    "            save_strategy=\"no\",\n",
    "            output_dir=\"outputs\"\n",
    "        ),\n",
    "        data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
    "    )\n",
    "\n",
    "    trainer.train()\n",
    "\n",
    "    if merge_weights:\n",
    "        output_dir = \"/tmp/model\"\n",
    "\n",
    "        # merge adapter weights with base model and save\n",
    "        # save int 4 model\n",
    "        trainer.model.save_pretrained(output_dir, safe_serialization=False)\n",
    "        \n",
    "        if accelerator.is_main_process:\n",
    "            # clear memory\n",
    "            del model\n",
    "            del trainer\n",
    "    \n",
    "            torch.cuda.empty_cache()\n",
    "    \n",
    "            # load PEFT model\n",
    "            model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "                output_dir,\n",
    "                torch_dtype=torch.float32,\n",
    "                low_cpu_mem_usage=True,\n",
    "                trust_remote_code=True,\n",
    "            )\n",
    "    \n",
    "            # Merge LoRA and base model and save\n",
    "            model = model.merge_and_unload()\n",
    "            model.save_pretrained(\n",
    "                \"/opt/ml/model\", safe_serialization=True, max_shard_size=\"2GB\"\n",
    "            )\n",
    "    else:\n",
    "        trainer.model.save_pretrained(\"/opt/ml/model\", safe_serialization=True)\n",
    "\n",
    "    if accelerator.is_main_process:\n",
    "        tokenizer.save_pretrained(\"/opt/ml/model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "Define LoRA parameters for fine-tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T17:16:11.825267Z",
     "start_time": "2023-09-02T17:16:11.814193Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "learning_rate=2e-4\n",
    "num_train_epochs=1\n",
    "per_device_train_batch_size=8\n",
    "per_device_eval_batch_size=8\n",
    "\n",
    "lora_r=8\n",
    "lora_alpha=32\n",
    "lora_dropout=0.05"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T17:51:05.194404Z",
     "start_time": "2023-09-02T17:16:13.190756Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_fn(\n",
    "    model_id,\n",
    "    train_ds=train_dataset,\n",
    "    test_ds=test_dataset,\n",
    "    lora_r=lora_r,\n",
    "    lora_alpha=lora_alpha,\n",
    "    lora_dropout=lora_dropout,\n",
    "    per_device_train_batch_size=per_device_train_batch_size,\n",
    "    per_device_eval_batch_size=per_device_eval_batch_size,\n",
    "    learning_rate=learning_rate,\n",
    "    num_train_epochs=num_train_epochs,\n",
    "    merge_weights=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Load Fine-Tuned model\n",
    "\n",
    "Note: Run `train_fn` with `merge_weights=False`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Download model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T18:18:47.299354Z",
     "start_time": "2023-09-02T18:18:47.059439Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "s3_client = boto3.client(\"s3\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T18:18:47.629480Z",
     "start_time": "2023-09-02T18:18:47.609236Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_id = \"tiiuae/falcon-7b\"\n",
    "\n",
    "bucket_name = \"<S3_BUCKET>\"\n",
    "job_prefix = f\"train-{model_id.split('/')[-1].replace('.', '-')}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_last_job_name(job_name_prefix):\n",
    "    import boto3\n",
    "    sagemaker_client = boto3.client('sagemaker')\n",
    "    \n",
    "    search_response = sagemaker_client.search(\n",
    "        Resource='TrainingJob',\n",
    "        SearchExpression={\n",
    "            'Filters': [\n",
    "                {\n",
    "                    'Name': 'TrainingJobName',\n",
    "                    'Operator': 'Contains',\n",
    "                    'Value': job_name_prefix\n",
    "                },\n",
    "                {\n",
    "                    'Name': 'TrainingJobStatus',\n",
    "                    'Operator': 'Equals',\n",
    "                    'Value': \"Completed\"\n",
    "                }\n",
    "            ]\n",
    "        },\n",
    "        SortBy='CreationTime',\n",
    "        SortOrder='Descending',\n",
    "        MaxResults=1)\n",
    "    \n",
    "    return search_response['Results'][0]['TrainingJob']['TrainingJobName']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job_name = get_last_job_name(job_prefix)\n",
    "\n",
    "job_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Donwload fine-tuned Peft model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T18:18:58.270521Z",
     "start_time": "2023-09-02T18:18:48.067777Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "s3_client.download_file(bucket_name, f\"{job_name}/{job_name}/output/model.tar.gz\", \"model.tar.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T18:18:59.409318Z",
     "start_time": "2023-09-02T18:18:58.272830Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "! rm -rf ./model && mkdir -p ./model && tar -xf model.tar.gz -C ./model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we are loading the PEFT weights trained"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T18:19:00.689996Z",
     "start_time": "2023-09-02T18:18:59.411744Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "model_id = \"tiiuae/falcon-7b\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "# Set the Falcon tokenizer\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T18:23:57.101169Z",
     "start_time": "2023-09-02T18:19:00.692057Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from peft import PeftModel, PeftConfig\n",
    "import torch\n",
    "from transformers import AutoModelForCausalLM\n",
    "\n",
    "device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'\n",
    "\n",
    "config = PeftConfig.from_pretrained(\"./model\")\n",
    "model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, trust_remote_code=True)\n",
    "model = PeftModel.from_pretrained(model, \"./model\")\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load a test dataset and try a random sample for Q&A."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T18:24:16.707887Z",
     "start_time": "2023-09-02T18:24:15.358844Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "df = pd.read_csv('train.csv.gz', compression='gzip', sep=';')\n",
    "\n",
    "sample = df.sample()\n",
    "\n",
    "# format sample\n",
    "prompt_template = f\"{{question}}\\n---\\nAnswer:\\n\"\n",
    "\n",
    "test_sample = prompt_template.format(question=sample.iloc()[0][\"question\"])\n",
    "\n",
    "print(test_sample)\n",
    "\n",
    "print(\"Original answer:\\n\", sample.iloc()[0][\"answers\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T18:24:21.678317Z",
     "start_time": "2023-09-02T18:24:21.642328Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_ids = tokenizer(test_sample, return_tensors=\"pt\").input_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-02T23:56:54.056915Z",
     "start_time": "2023-09-02T18:24:29.329156Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "#set the tokens for the summary evaluation\n",
    "tokens_for_answer = 100\n",
    "output_tokens = input_ids.shape[1] + tokens_for_answer\n",
    "\n",
    "outputs = model.generate(inputs=input_ids.to(device), do_sample=True, max_length=output_tokens)\n",
    "gen_text = tokenizer.batch_decode(outputs)[0]\n",
    "\n",
    "print(gen_text)"
   ]
  }
 ],
 "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.g5.12xlarge",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}