NLP_workloads/Text_generation/LLM_finetuning_and_batch_inference.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Large Language Model Fine-Tuning and Batch Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Generic/ray_logo.png\" width=\"20%\" loading=\"lazy\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/NLP_workloads/Text_generation/nlp_pipeline_full.png\" width=\"100%\" loading=\"lazy\">\n",
    "\n",
    "Welcome to this tutorial notebook, where you'll explore how to leverage [Ray AI Runtime (AIR)](https://docs.ray.io/en/latest/ray-air/getting-started.html) to perform distributed data preprocessing, fine-tuning, hyperparameter tuning, and batch inference using the [FLAN-T5](https://huggingface.co/docs/transformers/model_doc/flan-t5) model applied to the [Alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca) dataset.\n",
    "\n",
    "[FLAN-T5](https://arxiv.org/pdf/2210.11416.pdf) is transformer-based language model based on [Google's T5](https://ai.googleblog.com/2020/02/exploring-transfer-learning-with-t5.html) architecture and fine-tuned on instruction data. You will be further training this model on [Alpaca](https://crfm.stanford.edu/2023/03/13/alpaca.html), a set of 52k instructions and demonstrations. Through Ray AIR's integration with the Hugging Face hub, these components are easily accessible, and this example can be adapted for use with other similar models.\n",
    "\n",
    "By the end of this tutorial, you'll have a comprehensive understanding of how to harness Ray AIR to efficiently distribute complex machine learning tasks, allowing you to scale your projects easily."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting started"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up imports and utilities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import torch\n",
    "import transformers\n",
    "import warnings\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from IPython.display import display, HTML\n",
    "from typing import Any, Dict, List, Optional\n",
    "\n",
    "transformers.set_seed(42)\n",
    "warnings.simplefilter(\"ignore\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize Ray runtime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ray.init()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By calling `ray.init()`, you will initialize a Ray cluster. Follow the link outputted above to open the Ray Dashboard——a vital observability tool for understanding your infrastructure and application."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data ingest"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load the dataset\n",
    "\n",
    "You will be fine-tuning the model on [Alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca) to hopefully further refine the question answering and text generation ability of the original model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "from utils import get_random_elements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hf_dataset = load_dataset(\"tatsu-lab/alpaca\", split=\"train\").train_test_split(\n",
    "    test_size=0.2, seed=57\n",
    ")\n",
    "hf_dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Display sample data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = get_random_elements(dataset=hf_dataset[\"train\"], num_examples=3)\n",
    "display(HTML(df.to_html()))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that there are four feature columns in the dataset:\n",
    "\n",
    "* `instruction` - The original prompt or query such as \"How do we reduce air pollution?\"\n",
    "* `input` - Any additional context that wasn't provided by the instruction.\n",
    "* `output` - A sample generated response as generated by [Open AI's](https://platform.openai.com/docs/models/gpt-3-5) `text-davinci-003`.\n",
    "* `text` - The instruction, input, output, along with an [instructional prefix](https://github.com/tatsu-lab/stanford_alpaca#data-release)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert to Ray Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ray_dataset = ray.data.from_huggingface(hf_dataset)\n",
    "ray_dataset"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Ray Datasets](https://docs.ray.io/en/master/data/dataset.html#datasets) are the standard method for loading and exchanging data in Ray AIR libraries. They are specifically designed for easy distributed batch preprocessing, and you can easily convert from a Hugging Face dataset to Ray by using [`ray.data.from_huggingface()`](https://docs.ray.io/en/master/data/api/doc/ray.data.from_huggingface.html#ray.data.from_huggingface)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up train and validation Ray datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SMALL_DATA = True\n",
    "\n",
    "if SMALL_DATA:\n",
    "    train_dataset = ray_dataset[\"train\"].limit(100)\n",
    "    validation_dataset = ray_dataset[\"test\"].limit(100)\n",
    "else:\n",
    "    train_dataset = ray_dataset[\"train\"]\n",
    "    validation_dataset = ray_dataset[\"test\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note the `SMALL_DATA` flag which, when `True`, limits the number of samples used for downstream steps. This is to reduce training time for demonstration purposes. However, if you have more time, it is advised to set this flag to `False` to utilize the full dataset."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed preprocessing\n",
    "\n",
    "<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/NLP_workloads/Text_generation/nlp_pipeline_data.png\" width=\"100%\" loading=\"lazy\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implement preprocessing function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.data.preprocessors import BatchMapper\n",
    "from transformers import T5Tokenizer, T5ForConditionalGeneration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def preprocess_function(batch: Dict[str, Any]) -> Dict[str, Any]:\n",
    "    \"\"\"\n",
    "    Tokenizes the input and instruction pairs in a batch using the T5 tokenizer\n",
    "    from the Google/flan-t5-base model, and returns a dictionary containing the\n",
    "    encoded inputs and labels.\n",
    "\n",
    "    Args:\n",
    "        batch: A dictionary containing at least two keys, \"instruction\" and\n",
    "        \"input\", whose values are lists of strings.\n",
    "\n",
    "    Returns:\n",
    "        A dictionary containing the encoded inputs and labels, as returned by\n",
    "        the T5 tokenizer.\n",
    "    \"\"\"\n",
    "    model_name = \"google/flan-t5-base\"\n",
    "    tokenizer = T5Tokenizer.from_pretrained(model_name)\n",
    "\n",
    "    encoded_inputs = tokenizer(\n",
    "        list(batch[\"instruction\"]),\n",
    "        list(batch[\"input\"]),\n",
    "        padding=\"max_length\",\n",
    "        truncation=True,\n",
    "        return_tensors=\"np\",\n",
    "    )\n",
    "\n",
    "    encoded_inputs[\"labels\"] = encoded_inputs[\"input_ids\"].copy()\n",
    "\n",
    "    return dict(encoded_inputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_preprocessor = BatchMapper(preprocess_function, batch_format=\"pandas\", batch_size=4096)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You need to define a preprocessing function to convert a batch of data from Alpaca to a format that the FLAN-T5 model can accept. [Ray AIR's `BatchMapper`](https://docs.ray.io/en/latest/ray-air/api/doc/ray.data.preprocessors.BatchMapper.html#ray-data-preprocessors-batchmapper) will then map this function onto each incoming batch during the fine-tuning step.\n",
    "\n",
    "Unpacking this function a bit, the most important component is the [tokenizer](https://huggingface.co/docs/transformers/main_classes/tokenizer), which is a Hugging Face component associated with the FLAN-T5 model that turns natural language into formatted tokens with the right padding and truncation necessary for training."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed finetuning\n",
    "\n",
    "Now you have the dataset prepared, and a batch preprocessor initialized, it is time to configure [Ray AIR's `HuggingFaceTrainer`](https://docs.ray.io/en/master/train/api/doc/ray.train.huggingface.HuggingFaceTrainer.html#ray.train.huggingface.HuggingFaceTrainer) to distribute FLAN-T5 fine-tuning on Alpaca.\n",
    "\n",
    "<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/NLP_workloads/Text_generation/nlp_pipeline_finetune.png\" width=\"100%\" loading=\"lazy\">\n",
    "\n",
    "### Ray AIR Distributed Fine-Tuning Flow\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/NLP_workloads/Text_generation/nlp_train.png\" width=\"100%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Each worker node houses a preprocessor copy to process partitioned batches of the Ray Dataset, and then individual model copies train on these batches. PyTorch DDP synchronizes their weights, resulting in an integrated, fine-tuned model.|"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initialize training logic for each worker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments, Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 2\n",
    "use_gpu = True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before getting started, set the batch size (use a small number here since training requires a large amount of compute) and specify use of GPUs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def trainer_init_per_worker(\n",
    "    train_dataset: ray.data.Dataset,\n",
    "    eval_dataset: Optional[ray.data.Dataset] = None,\n",
    "    **config,\n",
    ") -> Trainer:\n",
    "    \"\"\"\n",
    "    Initializes a Hugging Face Trainer for training a T5 text generation model.\n",
    "\n",
    "    Args:\n",
    "        train_dataset (ray.data.Dataset): The dataset for training the model.\n",
    "        eval_dataset (ray.data.Dataset, optional): The dataset for evaluating\n",
    "        the model.\n",
    "            Defaults to None.\n",
    "        config: Additional arguments to configure the Trainer.\n",
    "\n",
    "    Returns:\n",
    "        Trainer: A Hugging Face Trainer for training the T5 model.\n",
    "    \"\"\"\n",
    "    device = torch.device(\"cuda\" if use_gpu and torch.cuda.is_available() else \"cpu\")\n",
    "    print(f\"Using device: {device}\")\n",
    "\n",
    "    model_name = \"google/flan-t5-base\"\n",
    "\n",
    "    tokenizer = T5Tokenizer.from_pretrained(model_name)\n",
    "    model = T5ForConditionalGeneration.from_pretrained(model_name)\n",
    "\n",
    "    training_args = TrainingArguments(\n",
    "        \"flan-t5-base-finetuned-alpaca\",\n",
    "        evaluation_strategy=\"epoch\",\n",
    "        save_strategy=\"epoch\",\n",
    "        logging_strategy=\"epoch\",\n",
    "        learning_rate=config.get(\"learning_rate\", 2e-5),\n",
    "        per_device_train_batch_size=batch_size,\n",
    "        per_device_eval_batch_size=batch_size,\n",
    "        num_train_epochs=config.get(\"epochs\", 4),\n",
    "        weight_decay=config.get(\"weight_decay\", 0.01),\n",
    "        push_to_hub=False,\n",
    "        disable_tqdm=True,\n",
    "    )\n",
    "\n",
    "    hf_trainer = Trainer(\n",
    "        model=model,\n",
    "        args=training_args,\n",
    "        train_dataset=train_dataset,\n",
    "        eval_dataset=eval_dataset,\n",
    "        tokenizer=tokenizer,\n",
    "    )\n",
    "\n",
    "    print(\"Starting training...\")\n",
    "    return hf_trainer"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `trainer_init_per_worker` function creates a Hugging Face Transformers Trainer that will be distributed by Ray using Distributed Data Parallelism (using PyTorch Distributed backend internally). This means that each worker will have its own copy of the model, but operate on different data. At the end of each step, all the workers will sync gradients.\n",
    "\n",
    "Note: The Hugging Face hub offers different versions of [FLAN-T5](https://huggingface.co/docs/transformers/model_doc/flan-t5) with increasing size. Here, the model and associated tokenizer are [\"flan_t5_base\"](https://huggingface.co/google/flan-t5-base), the smallest variant, in order to expedite fine-tuning for demonstration purposes. You can try this notebook with larger models, and you might find [this related tutorial](https://docs.ray.io/en/master/ray-air/examples/gptj_deepspeed_fine_tuning.html#train) helpful if the model does not fit on a single GPU."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.air.config import RunConfig, ScalingConfig, CheckpointConfig\n",
    "from ray.train.huggingface import HuggingFaceTrainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_workers = 2"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since you have access to two GPUs, set the number of workers to match in order to utilize the full cluster for fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = HuggingFaceTrainer(\n",
    "    trainer_init_per_worker=trainer_init_per_worker,\n",
    "    scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),\n",
    "    datasets={\n",
    "        \"train\": train_dataset,\n",
    "        \"evaluation\": validation_dataset,\n",
    "    },\n",
    "    run_config=RunConfig(\n",
    "        checkpoint_config=CheckpointConfig(\n",
    "            num_to_keep=1,\n",
    "            checkpoint_score_attribute=\"eval_loss\",\n",
    "            checkpoint_score_order=\"min\",\n",
    "        ),\n",
    "    ),\n",
    "    preprocessor=batch_preprocessor,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Ray AIR's HuggingFaceTrainer](https://docs.ray.io/en/latest/train/api/doc/ray.train.huggingface.HuggingFaceTrainer.html?highlight=ray%20air%20hugging%20face%20trainer) integrates with the Hugging Face Transformers library to scale training and fine-tuning across multiple workers, each with its own copy of the Hugging Face `transformers.Trainer` set up in the previous step.\n",
    "\n",
    "Here, you specify the following:\n",
    "* `trainer_init_per_worker` - Training logic copied onto each worker node.\n",
    "* `scaling_config` - Specify how to scale and the hardware to run on.\n",
    "* `datasets` - Which datasets to run training and evaluation on.\n",
    "* `run_config` - Specify checkpointing behavior (how many times to save the model and how to compare between saved models).\n",
    "* `preprocessor` - The same [Ray AIR preprocessor](https://docs.ray.io/en/latest/ray-air/preprocessors.html) defined above used to transform raw data into tokenized batches."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run finetuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result = trainer.fit()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Try the finetuned model\n",
    "\n",
    "Now that you have a fine-tuned model stored in a Checkpoint, you can retrieve it and test out your own instructions. In a later section, you will implement inference at scale."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"google/flan-t5-base\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = T5Tokenizer.from_pretrained(model_name)\n",
    "model = T5ForConditionalGeneration.from_pretrained(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = result.checkpoint\n",
    "finetuned_model = checkpoint.get_model(model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note: You are fetching the fine-tuned FLAN-T5 from the saved [checkpoint object](https://docs.ray.io/en/latest/ray-air/api/doc/ray.air.checkpoint.Checkpoint.html#ray.air.checkpoint.Checkpoint), which requires passing in what kind of model you expect to receive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "instruction = \"How many bees do I have?\"  # Enter your own instruction here.\n",
    "input_query = (\n",
    "    \"I don't have enough bees.\"  # Write additional context for the model here.\n",
    ")\n",
    "\n",
    "inputs = tokenizer(instruction, input_query, return_tensors=\"pt\")\n",
    "outputs = finetuned_model.generate(**inputs)\n",
    "\n",
    "print(tokenizer.batch_decode(outputs, skip_special_tokens=True))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## [Optional] Distributed hyperparameter tuning\n",
    "\n",
    "If you would like to tune hyperparameters in pursuit of a better performing model, you can pass the previous `HuggingFaceTrainer` into a [Ray AIR `Tuner`](https://docs.ray.io/en/latest/ray-air/tuner.html) and define the parameter search space to conduct experiments.\n",
    "\n",
    "<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/NLP_workloads/Text_generation/nlp_pipeline_tune.png\" width=\"100%\" loading=\"lazy\">\n",
    "\n",
    "### Ray AIR Distributed Hyperparameter Tuning Flow\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/NLP_workloads/Text_generation/nlp_tune.png\" width=\"100%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|To achieve the best configuration for the fine-tuned model, define a Tuner object with a customized search space and behavioral settings for scheduling, scaling, and checkpointing. Running multiple trial experiments using this approach can help converge on the optimal configuration.|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray import tune\n",
    "from ray.tune import Tuner\n",
    "from ray.tune.schedulers.async_hyperband import ASHAScheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "total_num_trials = 4\n",
    "max_tune_epochs = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "num_workers = 1\n",
    "use_gpu = True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set the number of workers to 1 for each `Trainer` so that hyperparameter tuning can run in parallel rather than sequentially with each trial utilizing all resources per experiment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = HuggingFaceTrainer(\n",
    "    trainer_init_per_worker=trainer_init_per_worker,\n",
    "    scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),\n",
    "    datasets={\n",
    "        \"train\": train_dataset,\n",
    "        \"evaluation\": validation_dataset,\n",
    "    },\n",
    "    run_config=RunConfig(\n",
    "        checkpoint_config=CheckpointConfig(\n",
    "            num_to_keep=1,\n",
    "            checkpoint_score_attribute=\"eval_loss\",\n",
    "            checkpoint_score_order=\"min\",\n",
    "        ),\n",
    "    ),\n",
    "    preprocessor=batch_preprocessor,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the same `HuggingFaceTrainer` created previously, just with a different number of workers for fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tuner = Tuner(\n",
    "    trainer,\n",
    "    param_space={\n",
    "        \"trainer_init_config\": {\n",
    "            \"learning_rate\": tune.choice([2e-5, 2e-4, 2e-3, 2e-2]),\n",
    "            \"epochs\": tune.choice([2, 4, 8, max_tune_epochs]),\n",
    "            \"weight_decay\": tune.choice([0.01, 0.1, 1.0, 10.0]),\n",
    "        }\n",
    "    },\n",
    "    tune_config=tune.TuneConfig(\n",
    "        metric=\"eval_loss\",\n",
    "        mode=\"min\",\n",
    "        num_samples=total_num_trials,\n",
    "        scheduler=ASHAScheduler(\n",
    "            max_t=max_tune_epochs,\n",
    "        ),\n",
    "    ),\n",
    "    run_config=RunConfig(\n",
    "        checkpoint_config=CheckpointConfig(\n",
    "            num_to_keep=1,\n",
    "            checkpoint_score_attribute=\"eval_loss\",\n",
    "            checkpoint_score_order=\"min\",\n",
    "        )\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are four major components passed into the Tuner:\n",
    "1. `trainer` - The `HuggingFaceTrainer` with scaling, preprocessing, and fine-tuning logic from earlier.\n",
    "2. `param_space` - The [possibilities of hyperparameters](https://docs.ray.io/en/latest/ray-air/tuner.html#how-to-configure-a-search-space) to tune and search for any given trial.\n",
    "3. `tune_config` - Specify how to compare different experiments, the number of trials, as well as any advanced [search algorithms](https://docs.ray.io/en/latest/tune/key-concepts.html#search-alg-ref) and [schedulers](https://docs.ray.io/en/latest/tune/key-concepts.html#schedulers-ref) like [ASHA](https://openreview.net/forum?id=S1Y7OOlRZ).\n",
    "4. `run_config` - Used to specify checkpointing behavior, custom callbacks, failure/retry configurations, [and more.](https://docs.ray.io/en/latest/ray-air/api/doc/ray.air.RunConfig.html#ray.air.RunConfig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_grid = tuner.fit()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Distributed batch inference\n",
    "\n",
    "Once you have a fine-tuned model, you can apply it to batches of inputs to generate predictions at scale, which is exactly what [Ray AIR's `BatchPredictor`](https://docs.ray.io/en/latest/ray-air/predictors.html#batch-prediction) is designed to facilitate.\n",
    "\n",
    "<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/NLP_workloads/Text_generation/nlp_pipeline_inference.png\" width=\"100%\" loading=\"lazy\">\n",
    "\n",
    "### Ray AIR Distributed Batch Inference Flow\n",
    "\n",
    "|<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/NLP_workloads/Text_generation/nlp_batchpredict.png\" width=\"100%\" loading=\"lazy\">|\n",
    "|:--|\n",
    "|Using the best fine-tuned model stored in a Checkpoint object, apply BatchPredictor to new batches of data to generate predictions.|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.train.predictor import Predictor\n",
    "from ray.train.batch_predictor import BatchPredictor\n",
    "from transformers import AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HuggingFaceModelPredictor(Predictor):\n",
    "    \"\"\"\n",
    "    A Ray Predictor for Hugging Face models that generates text given input data.\n",
    "\n",
    "    Args:\n",
    "        model (transformers.PreTrainedModel): A trained Hugging Face model.\n",
    "        tokenizer (Optional[transformers.PreTrainedTokenizerBase]): A tokenizer\n",
    "        that can tokenize input text.\n",
    "        preprocessor (Optional[Callable]): A function that takes raw input data\n",
    "        and returns tokenized input data.\n",
    "        use_gpu (bool): Whether to use a GPU or CPU for prediction.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        model: Any,\n",
    "        tokenizer: Optional[Any] = None,\n",
    "        preprocessor: Optional[Any] = None,\n",
    "        use_gpu: bool = False,\n",
    "    ) -> None:\n",
    "        super().__init__(preprocessor)\n",
    "        self.model = model\n",
    "        self.use_gpu = use_gpu\n",
    "        self.tokenizer = tokenizer\n",
    "\n",
    "    @classmethod\n",
    "    def from_checkpoint(\n",
    "        cls,\n",
    "        checkpoint: Any,\n",
    "        model_cls: Any,\n",
    "        *,\n",
    "        tokenizer: Optional[Any] = None,\n",
    "        use_gpu: bool = False,\n",
    "        **get_model_kwargs: Any,\n",
    "    ) -> \"HuggingFaceModelPredictor\":\n",
    "        \"\"\"\n",
    "        Create a HuggingFaceModelPredictor from a checkpoint.\n",
    "\n",
    "        Args:\n",
    "            checkpoint (Any): A checkpoint containing a trained Hugging Face model.\n",
    "            model_cls (Any): The type of Hugging Face model to load from the checkpoint.\n",
    "            tokenizer (Optional[Any]): A tokenizer that can tokenize input text.\n",
    "            use_gpu (bool): Whether to use a GPU or CPU for prediction.\n",
    "            **get_model_kwargs (Any): Additional keyword arguments for loading\n",
    "            the Hugging Face model.\n",
    "\n",
    "        Returns:\n",
    "            HuggingFaceModelPredictor: A Ray Predictor for the Hugging Face model.\n",
    "        \"\"\"\n",
    "        if not tokenizer:\n",
    "            tokenizer = AutoTokenizer\n",
    "        if isinstance(tokenizer, type):\n",
    "            tokenizer = checkpoint.get_tokenizer(tokenizer)\n",
    "        return cls(\n",
    "            checkpoint.get_model(model_cls, **get_model_kwargs),\n",
    "            tokenizer=tokenizer,\n",
    "            preprocessor=checkpoint.get_preprocessor(),\n",
    "            use_gpu=use_gpu,\n",
    "        )\n",
    "\n",
    "    def _predict_numpy(\n",
    "        self,\n",
    "        data: Dict[str, Any],\n",
    "        feature_columns: Optional[List[str]] = None,\n",
    "        **generate_kwargs: Any,\n",
    "    ) -> pd.DataFrame:\n",
    "        \"\"\"\n",
    "        Generates text given input data.\n",
    "\n",
    "        Args:\n",
    "            data (Dict[str, Any]): A dictionary of input data.\n",
    "            feature_columns (Optional[List[str]]): A list of feature column names\n",
    "            to use for prediction.\n",
    "            **generate_kwargs (Any): Additional keyword arguments for generating text.\n",
    "\n",
    "        Returns:\n",
    "            pd.DataFrame: A Pandas DataFrame with a single column \"generated_output\"\n",
    "            containing the generated text.\n",
    "        \"\"\"\n",
    "        # we get already tokenized text here because we have the tokenizer as an AIR preprocessor\n",
    "        if feature_columns:\n",
    "            data = {k: v for k, v in data.items() if k in feature_columns}\n",
    "\n",
    "        data = {\n",
    "            k: torch.from_numpy(v).to(device=self.model.device) for k, v in data.items()\n",
    "        }\n",
    "        generate_kwargs = {**data, **generate_kwargs}\n",
    "\n",
    "        outputs = self.model.generate(**generate_kwargs)\n",
    "        return pd.DataFrame(\n",
    "            self.tokenizer.batch_decode(outputs, skip_special_tokens=True),\n",
    "            columns=[\"generated_output\"],\n",
    "        )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Establish a custom class for prediction, `HugginFaceModelPredictor`, which extends the base Ray AIR [`Predictor`](https://docs.ray.io/en/latest/ray-air/api/doc/ray.train.predictor.Predictor.html?highlight=ray%20air%20predictor) to generate text responses to input instructions:\n",
    "\n",
    "* The predictor takes a trained Hugging Face model, a tokenizer, and a preprocessor (which can be a function that takes raw input data and returns tokenized input data). \n",
    "\n",
    "* `from_checkpoint` creates a `HuggingFaceModelPredictor` from a checkpoint containing a trained Hugging Face model. \n",
    "\n",
    "* `_predict_numpy` generates text given input data in the form of a dictionary, and returns a Pandas DataFrame with a single column \"generated_output\" containing the generated text. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = BatchPredictor.from_checkpoint(\n",
    "    checkpoint=result.checkpoint,\n",
    "    predictor_cls=HuggingFaceModelPredictor,\n",
    "    model_cls=T5ForConditionalGeneration,\n",
    "    tokenizer=T5Tokenizer,\n",
    "    use_gpu=use_gpu,\n",
    "    device_map=\"auto\",\n",
    "    torch_dtype=torch.float16,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a Ray AIR `BatchPredictor` from a [Checkpoint](https://docs.ray.io/en/latest/ray-air/api/checkpoint.html?highlight=checkpoint) and specify the custom predictor, model class, tokenizer, as well as any additional arguments."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run batch inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction = predictor.predict(\n",
    "    validation_dataset,\n",
    "    num_gpus_per_worker=int(use_gpu),\n",
    "    batch_size=256,\n",
    "    max_new_tokens=128,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inspect predictions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Display inputs and generated outputs side by side.\n",
    "input_data_pd = validation_dataset.to_pandas()\n",
    "prediction_pd = prediction.to_pandas()\n",
    "\n",
    "input_data_pd.join(prediction_pd, how='inner')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Connect with the Ray community\n",
    "\n",
    "You can learn and get more involved with the Ray community of developers and researchers:\n",
    "\n",
    "* [**Ray documentation**](https://docs.ray.io/en/latest)\n",
    "\n",
    "* [**Official Ray site**](https://www.ray.io/)  \n",
    "Browse the ecosystem and use this site as a hub to get the information that you need to get going and building with Ray.\n",
    "\n",
    "* [**Join the community on Slack**](https://forms.gle/9TSdDYUgxYs8SA9e8)  \n",
    "Find friends to discuss your new learnings in our Slack space.\n",
    "\n",
    "* [**Use the discussion board**](https://discuss.ray.io/)  \n",
    "Ask questions, follow topics, and view announcements on this community forum.\n",
    "\n",
    "* [**Join a meetup group**](https://www.meetup.com/Bay-Area-Ray-Meetup/)  \n",
    "Tune in on meet-ups to listen to compelling talks, get to know other users, and meet the team behind Ray.\n",
    "\n",
    "* [**Open an issue**](https://github.com/ray-project/ray/issues/new/choose)  \n",
    "Ray is constantly evolving to improve developer experience. Submit feature requests, bug-reports, and get help via GitHub issues.\n",
    "\n",
    "* [**Become a Ray contributor**](https://docs.ray.io/en/latest/ray-contribute/getting-involved.html)  \n",
    "We welcome community contributions to improve our documentation and Ray framework."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "<img src=\"https://technical-training-assets.s3.us-west-2.amazonaws.com/Generic/ray_logo.png\" width=\"20%\" loading=\"lazy\">"
   ]
  }
 ],
 "metadata": {
  "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "567405a8058597909526349386224fe35dd047505a91307e44ed44be00113429"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}