From d25250a89735187ed124308eb83ee596e3f4d96e Mon Sep 17 00:00:00 2001
From: Balaji Veeramani <balaji@anyscale.com>
Date: Mon, 10 Jul 2023 15:14:00 -0500
Subject: [PATCH] [AIR][Docs] Remove `BatchPredictor` from examples (#37178)

The following examples already use updated APIs:
* Stable Diffusion Batch Prediction with Ray AIR
* GPT-J-6B Batch Prediction with Ray AIR (LLM)

The following examples have been updated to use updated APIs:
* Training a model with distributed XGBoost
* Training a model with distributed LightGBM

I've removed batch prediction sections from the other examples, and, where appropriate, linked to the batch inference user guide.

Signed-off-by: Balaji Veeramani <balaji@anyscale.com>
---
 doc/source/_toc.yml                           |   1 -
 .../ray-air/doc_code/pytorch_starter.py       |  14 -
 doc/source/ray-air/doc_code/tf_starter.py     |  21 -
 doc/source/ray-air/examples/BUILD             |  13 +-
 ...ert_existing_pytorch_code_to_ray_air.ipynb | 386 +-------
 .../convert_existing_tf_code_to_ray_air.ipynb | 260 +----
 .../examples/gptj_deepspeed_fine_tuning.ipynb |   2 +-
 .../huggingface_text_classification.ipynb     | 127 +--
 doc/source/ray-air/examples/index.rst         |   1 -
 .../ray-air/examples/lightgbm_example.ipynb   | 450 +++++----
 .../opt_deepspeed_batch_inference.ipynb       | 911 ------------------
 .../pytorch_tabular_batch_prediction.py       |  46 -
 .../examples/pytorch_tabular_starter.py       |  19 -
 .../ray-air/examples/sklearn_example.ipynb    | 106 +-
 .../examples/tf_tabular_batch_prediction.py   |  49 -
 .../ray-air/examples/tf_tabular_starter.py    |  21 -
 .../ray-air/examples/torch_detection.ipynb    | 149 +--
 .../examples/torch_image_batch_pretrained.py  |  30 -
 .../examples/torch_incremental_learning.ipynb | 416 +-------
 .../examples/xgboost_batch_prediction.py      |  42 -
 .../ray-air/examples/xgboost_example.ipynb    | 433 ++++-----
 .../ray-air/examples/xgboost_starter.py       |  18 -
 doc/source/ray-air/getting-started.rst        |  28 +-
 doc/source/ray-air/predictors.rst             |  45 -
 doc/source/ray-overview/examples.rst          |   7 -
 doc/source/ray-overview/getting-started.md    |   9 +-
 .../lightning/lightning_cola_advanced.ipynb   | 126 +--
 ...ingface_basic_language_modeling_example.py |  20 +-
 .../pytorch/torch_regression_example.py       |  21 +-
 .../tf/tensorflow_autoencoder_example.py      |  63 +-
 .../tf/tensorflow_regression_example.py       |  29 +-
 release/release_tests.yaml                    |  57 +-
 32 files changed, 584 insertions(+), 3336 deletions(-)
 delete mode 100644 doc/source/ray-air/examples/opt_deepspeed_batch_inference.ipynb
 delete mode 100644 doc/source/ray-air/examples/pytorch_tabular_batch_prediction.py
 delete mode 100644 doc/source/ray-air/examples/tf_tabular_batch_prediction.py
 delete mode 100644 doc/source/ray-air/examples/torch_image_batch_pretrained.py
 delete mode 100644 doc/source/ray-air/examples/xgboost_batch_prediction.py

diff --git a/doc/source/_toc.yml b/doc/source/_toc.yml
index 99ff9dc9ee4c..4f4cfd850e9c 100644
--- a/doc/source/_toc.yml
+++ b/doc/source/_toc.yml
@@ -57,7 +57,6 @@ parts:
               - file: ray-air/deployment
           - file: ray-air/examples/index
             sections:
-              - file: ray-air/examples/opt_deepspeed_batch_inference
               - file: ray-air/examples/torch_image_example
               - file: ray-air/examples/torch_detection
               - file: ray-air/examples/convert_existing_pytorch_code_to_ray_air
diff --git a/doc/source/ray-air/doc_code/pytorch_starter.py b/doc/source/ray-air/doc_code/pytorch_starter.py
index edd711923515..f6f93e3459e6 100644
--- a/doc/source/ray-air/doc_code/pytorch_starter.py
+++ b/doc/source/ray-air/doc_code/pytorch_starter.py
@@ -129,17 +129,3 @@ def train_func(config):
 result = trainer.fit()
 print(f"Last result: {result.metrics}")
 # __air_pytorch_train_end__
-
-
-# # __air_pytorch_batchpred_start__
-# import random
-# from ray.train.batch_predictor import BatchPredictor
-# from ray.train.torch import TorchPredictor
-
-# batch_predictor = BatchPredictor.from_checkpoint(result.checkpoint, TorchPredictor)
-
-# items = [{"x": random.uniform(0, 1) for _ in range(10)}]
-# prediction_dataset = ray.data.from_items(items)
-
-# predictions = batch_predictor.predict(prediction_dataset, dtype=torch.float)
-# # __air_pytorch_batchpred_end__
diff --git a/doc/source/ray-air/doc_code/tf_starter.py b/doc/source/ray-air/doc_code/tf_starter.py
index 5ea80d19ad90..acd78dc49288 100644
--- a/doc/source/ray-air/doc_code/tf_starter.py
+++ b/doc/source/ray-air/doc_code/tf_starter.py
@@ -75,24 +75,3 @@ def train_func(config: dict):
 result = trainer.fit()
 print(result.metrics)
 # __air_tf_train_end__
-
-# __air_tf_batchpred_start__
-import numpy as np
-
-from ray.train.batch_predictor import BatchPredictor
-from ray.train.tensorflow import TensorflowPredictor
-
-
-batch_predictor = BatchPredictor.from_checkpoint(
-    result.checkpoint, TensorflowPredictor, model_definition=build_model
-)
-
-items = [{"x": np.random.uniform(0, 1)} for _ in range(10)]
-prediction_dataset = ray.data.from_items(items)
-
-predictions = batch_predictor.predict(prediction_dataset, dtype=tf.float32)
-
-print("PREDICTIONS")
-predictions.show()
-
-# __air_tf_batchpred_end__
diff --git a/doc/source/ray-air/examples/BUILD b/doc/source/ray-air/examples/BUILD
index ecbfa5b1a41d..9e0c1905a587 100644
--- a/doc/source/ray-air/examples/BUILD
+++ b/doc/source/ray-air/examples/BUILD
@@ -15,22 +15,12 @@ filegroup(
 py_test_run_all_subdirectory(
     size = "medium",
     include = ["*.py"],
-    exclude = ["torch_image_batch_pretrained.py"],
-    extra_srcs = [],
-    data = ["//doc/source/ray-air/examples:air_examples"],
-    tags = ["exclusive", "team:ml", "ray_air"],
-)
-
-py_test_run_all_subdirectory(
-    size = "medium",
-    include = ["torch_image_batch_pretrained.py"],
     exclude = [],
     extra_srcs = [],
     data = ["//doc/source/ray-air/examples:air_examples"],
-    tags = ["exclusive", "team:ml", "ray_air", "gpu"],
+    tags = ["exclusive", "team:ml", "ray_air"],
 )
 
-
 # --------------------------------------------------------------------
 # Test all doc/source/ray-air/examples notebooks.
 # --------------------------------------------------------------------
@@ -50,7 +40,6 @@ py_test_run_all_notebooks(
         "gptj_serving.ipynb",  # Requires GPUs
         "stablediffusion_batch_prediction.ipynb",  # Requires GPUs
         "gptj_deepspeed_fine_tuning.ipynb",  # Requires release test
-        "opt_deepspeed_batch_inference.ipynb", # Requires release test
         "dolly_lightning_fsdp_finetuning.ipynb", # Requires release test
     ],
     data = ["//doc/source/ray-air/examples:air_examples"],
diff --git a/doc/source/ray-air/examples/convert_existing_pytorch_code_to_ray_air.ipynb b/doc/source/ray-air/examples/convert_existing_pytorch_code_to_ray_air.ipynb
index d56c70a745cb..08cffc200169 100644
--- a/doc/source/ray-air/examples/convert_existing_pytorch_code_to_ray_air.ipynb
+++ b/doc/source/ray-air/examples/convert_existing_pytorch_code_to_ray_air.ipynb
@@ -18,13 +18,13 @@
     "- Automatic checkpointing/fault tolerance and result tracking\n",
     "- Parallel data preprocessing\n",
     "- Seamless integration with hyperparameter tuning\n",
-    "- Scalable batch prediction\n",
     "- Scalable model serving\n",
     "\n",
-    "This tutorial will show you how to start with Ray AIR from your existing PyTorch training code. We will learn how to **distribute your training** and do **scalable batch prediction**.\n"
+    "This tutorial will show you how to start with Ray AIR from your existing PyTorch training code and learn how to **distribute your training**.\n"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "9a4855cf",
    "metadata": {},
@@ -37,6 +37,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "a42faedb",
    "metadata": {},
@@ -97,6 +98,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "9795c146",
    "metadata": {},
@@ -119,6 +121,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "ae11399e",
    "metadata": {},
@@ -178,6 +181,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "b692d06a",
    "metadata": {},
@@ -197,6 +201,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "681d5798",
    "metadata": {},
@@ -232,6 +237,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "6621cffa",
    "metadata": {},
@@ -264,6 +270,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "d915d788",
    "metadata": {},
@@ -389,6 +396,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "6655d903",
    "metadata": {},
@@ -397,6 +405,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "d0b98b1c",
    "metadata": {},
@@ -442,6 +451,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "60f7341a",
    "metadata": {},
@@ -554,6 +564,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "b3df2581",
    "metadata": {},
@@ -672,6 +683,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "6e260f44",
    "metadata": {},
@@ -704,6 +716,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "cf280e6a",
    "metadata": {},
@@ -750,6 +763,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "0fc52cc7",
    "metadata": {},
@@ -778,6 +792,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "341f4fd8",
    "metadata": {},
@@ -788,6 +803,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "3bbe06f3",
    "metadata": {},
@@ -875,6 +891,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "d2af219d",
    "metadata": {},
@@ -898,6 +915,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "534ed4df",
    "metadata": {},
@@ -926,367 +944,7 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "id": "b6b15d88",
-   "metadata": {},
-   "source": [
-    "## Loading the model for prediction\n",
-    "You may have noticed that we skipped one part of the original tutorial - loading the model and using it for inference. The original code looks like this (we've wrapped it in a function):"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 19,
-   "id": "68e664ff",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def predict_from_model(model):\n",
-    "    classes = [\n",
-    "        \"T-shirt/top\",\n",
-    "        \"Trouser\",\n",
-    "        \"Pullover\",\n",
-    "        \"Dress\",\n",
-    "        \"Coat\",\n",
-    "        \"Sandal\",\n",
-    "        \"Shirt\",\n",
-    "        \"Sneaker\",\n",
-    "        \"Bag\",\n",
-    "        \"Ankle boot\",\n",
-    "    ]\n",
-    "\n",
-    "    model.eval()\n",
-    "    x, y = test_data[0][0], test_data[0][1]\n",
-    "    with torch.no_grad():\n",
-    "        pred = model(x)\n",
-    "        predicted, actual = classes[pred[0].argmax(0)], classes[y]\n",
-    "        print(f'Predicted: \"{predicted}\", Actual: \"{actual}\"')\n"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "1abf022a",
-   "metadata": {},
-   "source": [
-    "We can use our saved model with the existing code to do prediction:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 20,
-   "id": "0c135a17",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Predicted: \"Ankle boot\", Actual: \"Ankle boot\"\n"
-     ]
-    }
-   ],
-   "source": [
-    "from ray.train.torch import TorchCheckpoint\n",
-    "\n",
-    "model = TorchCheckpoint.from_checkpoint(result.checkpoint).get_model(NeuralNetwork())\n",
-    "\n",
-    "predict_from_model(model)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "f6fc1441",
-   "metadata": {},
-   "source": [
-    "To predict more than one example, we can use a loop:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 21,
-   "id": "17652fa1",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "classes = [\n",
-    "    \"T-shirt/top\",\n",
-    "    \"Trouser\",\n",
-    "    \"Pullover\",\n",
-    "    \"Dress\",\n",
-    "    \"Coat\",\n",
-    "    \"Sandal\",\n",
-    "    \"Shirt\",\n",
-    "    \"Sneaker\",\n",
-    "    \"Bag\",\n",
-    "    \"Ankle boot\",\n",
-    "]\n",
-    "\n",
-    "def predict_from_model(model, data):\n",
-    "    model.eval()\n",
-    "    with torch.no_grad():\n",
-    "        for x, y in data:\n",
-    "            pred = model(x)\n",
-    "            predicted, actual = classes[pred[0].argmax(0)], classes[y]\n",
-    "            print(f'Predicted: \"{predicted}\", Actual: \"{actual}\"')\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 22,
-   "id": "3bc14ed6",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Predicted: \"Ankle boot\", Actual: \"Ankle boot\"\n",
-      "Predicted: \"Pullover\", Actual: \"Pullover\"\n",
-      "Predicted: \"Trouser\", Actual: \"Trouser\"\n",
-      "Predicted: \"Trouser\", Actual: \"Trouser\"\n",
-      "Predicted: \"Pullover\", Actual: \"Shirt\"\n",
-      "Predicted: \"Trouser\", Actual: \"Trouser\"\n",
-      "Predicted: \"Coat\", Actual: \"Coat\"\n",
-      "Predicted: \"Pullover\", Actual: \"Shirt\"\n",
-      "Predicted: \"Sneaker\", Actual: \"Sandal\"\n",
-      "Predicted: \"Sneaker\", Actual: \"Sneaker\"\n"
-     ]
-    }
-   ],
-   "source": [
-    "predict_from_model(model, [test_data[i] for i in range(10)])"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "a0ce0733",
-   "metadata": {},
-   "source": [
-    "## Using Ray AIR for scalable batch prediction\n",
-    "However, we can also use Ray AIRs `BatchPredictor` class to do scalable prediction."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 24,
-   "id": "4d8b0f50",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from ray.train.batch_predictor import BatchPredictor\n",
-    "from ray.train.torch import TorchPredictor\n",
-    "\n",
-    "batch_predictor = BatchPredictor.from_checkpoint(result.checkpoint, TorchPredictor, model=NeuralNetwork())"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "ad556eeb",
-   "metadata": {},
-   "source": [
-    "Batch predictors work with Ray Data. Here we convert our test dataset into a Dataset - note that this is not very efficient, and you can look at our {ref}`other tutorials <air-examples-ref>` to see more efficient ways to generate a Dataset."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 25,
-   "id": "8cb0556f",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import ray.data\n",
-    "\n",
-    "ds = ray.data.from_items([x.numpy() for x, y in test_data], parallelism=8)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "264dd2e4",
-   "metadata": {},
-   "source": [
-    "We can then trigger prediction with two workers:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 26,
-   "id": "8a823f7a",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Map Progress (2 actors 1 pending): 100%|██████████| 8/8 [00:02<00:00, 70.01it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "results = batch_predictor.predict(ds, batch_size=32, min_scoring_workers=2)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "41094a55",
-   "metadata": {},
-   "source": [
-    "`results` is another Dataset. We can use `results.show()` to see our prediction results:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 27,
-   "id": "d3dce40d",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "{'predictions': array([-1.6813023 , -1.80252   , -0.7062941 , -1.311813  , -0.73570144,\n",
-      "        1.5710734 , -0.7933277 ,  2.0013504 ,  1.3654878 ,  2.3410547 ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 0.7655406 , -2.3314183 ,  2.7599745 , -0.9481916 ,  2.381936  ,\n",
-      "       -1.7827132 ,  1.9278868 , -3.1977224 ,  0.99582016, -1.4932251 ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 1.3619348 ,  3.6063552 , -0.31104898,  2.543014  ,  0.35176522,\n",
-      "       -2.2156405 ,  0.33978355, -2.346588  , -1.7794112 , -2.3220763 ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 1.0049653 ,  2.828181  , -0.29254377,  2.0342605 ,  0.12778719,\n",
-      "       -1.6141529 ,  0.17694427, -1.7565594 , -1.4074212 , -1.6818824 ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 0.7404187 , -1.0129585 ,  1.0854365 , -0.20976087,  1.0174558 ,\n",
-      "       -0.9567458 ,  1.0075954 , -1.7656276 ,  0.42417505, -0.82513285],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 1.4985809 ,  2.499547  ,  0.12339873,  1.9594493 ,  0.717446  ,\n",
-      "       -2.0457497 ,  0.6526047 , -2.4334526 , -1.4454234 , -2.2310004 ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 0.45234615, -0.23714153,  0.63517165,  0.04347774,  0.6996659 ,\n",
-      "       -0.5516397 ,  0.64028525, -1.0785    ,  0.10881007, -0.9026278 ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 0.38904738, -0.80522966,  1.1767559 , -0.21403429,  1.1468315 ,\n",
-      "       -0.84129035,  0.95365965, -1.6148682 ,  0.27161083, -0.96888554],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([-0.54510164, -0.31364274, -0.22182664, -0.25785953, -0.25741974,\n",
-      "        0.48500216, -0.2174497 ,  0.7817588 ,  0.34047806,  0.24852225],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([-1.2857382 , -0.9965143 , -0.64847904, -0.7487341 , -0.60564923,\n",
-      "        1.1155919 , -0.59477496,  2.0135763 ,  0.88436544,  1.067797  ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 0.43335694, -0.8999133 ,  1.7488041 , -0.31407052,  1.6201458 ,\n",
-      "       -1.2921515 ,  1.2184532 , -2.068122  ,  0.17047453, -1.2746251 ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([-1.267686  , -1.2830508 , -0.4776874 , -0.94430155, -0.51243144,\n",
-      "        1.167536  , -0.48850274,  1.4446495 ,  1.00295   ,  1.4936616 ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([-1.313108  , -1.2630323 , -0.4338272 , -0.9408438 , -0.42691046,\n",
-      "        1.0805027 , -0.47953707,  1.6175348 ,  1.2289674 ,  0.99234164],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 0.8392371 ,  2.005179  , -0.51027215,  2.2383528 ,  0.11543664,\n",
-      "       -1.418318  ,  0.10795547, -1.5231588 , -0.9388958 , -1.2481594 ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 0.9623028 , -1.5077128 ,  1.9832315 , -0.06346714,  2.3645868 ,\n",
-      "       -2.1186042 ,  1.7628006 , -3.423348  ,  0.84258574, -1.9048262 ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 1.1201253 ,  2.6718287 , -0.22753508,  2.1176536 ,  0.23477581,\n",
-      "       -1.691438  ,  0.2711372 , -1.9383426 , -1.3917452 , -1.7475704 ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 0.70388055, -0.6840804 ,  1.1767206 , -0.21303988,  0.96372414,\n",
-      "       -0.94062155,  0.92242914, -1.689395  ,  0.23195787, -1.00324   ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 0.78447473, -1.2020342 ,  1.51774   , -0.36963996,  1.368768  ,\n",
-      "       -1.3143553 ,  1.2229909 , -2.284686  ,  0.6896354 , -1.0750523 ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([-1.5986964 , -2.6742263 , -0.04187664, -1.7070676 , -0.00644506,\n",
-      "        1.1022365 , -0.31155828,  1.3389733 ,  2.226508  ,  1.72136   ],\n",
-      "      dtype=float32)}\n",
-      "{'predictions': array([ 2.9001627 ,  0.791762  ,  1.148489  ,  1.6756771 ,  1.5494249 ,\n",
-      "       -2.8295102 ,  1.7419    , -4.0650196 , -0.98189455, -2.9981184 ],\n",
-      "      dtype=float32)}\n"
-     ]
-    }
-   ],
-   "source": [
-    "results.show()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "427b68e8",
-   "metadata": {},
-   "source": [
-    "If we want to convert these predictions into class names (as in the original example), we can use a `map` function to do this:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 29,
-   "id": "f17b5c10",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Map_Batches: 100%|██████████| 8/8 [00:02<00:00, 80.05it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "predicted_classes = results.map_batches(\n",
-    "    lambda batch: {\"pred\": [classes[pred.argmax(0)] for pred in batch[\"predictions\"]]}, \n",
-    "    batch_size=32,\n",
-    "    batch_format=\"pandas\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "cb7040db",
-   "metadata": {},
-   "source": [
-    "To see how well our prediction did, let's zip the predicted labels together with some of the actual labels to compare them:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 30,
-   "id": "207e13b9",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "('Ankle boot', 'Ankle boot')\n",
-      "('Pullover', 'Pullover')\n",
-      "('Trouser', 'Trouser')\n",
-      "('Trouser', 'Trouser')\n",
-      "('Pullover', 'Shirt')\n",
-      "('Trouser', 'Trouser')\n",
-      "('Coat', 'Coat')\n",
-      "('Pullover', 'Shirt')\n",
-      "('Sneaker', 'Sandal')\n",
-      "('Sneaker', 'Sneaker')\n",
-      "('Pullover', 'Coat')\n",
-      "('Ankle boot', 'Sandal')\n",
-      "('Sneaker', 'Sneaker')\n",
-      "('Dress', 'Dress')\n",
-      "('Coat', 'Coat')\n",
-      "('Trouser', 'Trouser')\n",
-      "('Pullover', 'Pullover')\n",
-      "('Pullover', 'Coat')\n",
-      "('Bag', 'Bag')\n",
-      "('T-shirt/top', 'T-shirt/top')\n"
-     ]
-    }
-   ],
-   "source": [
-    "real_classes = [classes[y] for x, y in test_data]\n",
-    "for predicted, real in zip(predicted_classes.take_batch()[\"pred\"], real_classes):\n",
-    "    print((predicted, real))"
-   ]
-  },
-  {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "2963e1f7",
    "metadata": {},
@@ -1300,7 +958,7 @@
     "- save and retrieve model checkpoints via Ray AIR\n",
     "- load a model for batch prediction\n",
     "\n",
-    "In our {ref}`other examples <air-examples-ref>` you can learn how to do more things with the Ray AIR API, such as **serving your model with Ray Serve** or **tune your hyperparameters with Ray Tune.** You can also learn how to **construct Ray Data** to leverage Ray AIR's **preprocessing** API.\n",
+    "In our {ref}`other examples <air-examples-ref>` you can learn how to do more things with the Ray AIR API, such as **serving your model with Ray Serve** or **tune your hyperparameters with Ray Tune.** You can also learn how to perform {ref}`offline batch inference <batch_inference_home>` with Ray Data.\n",
     "\n",
     "We hope this tutorial gave you a good starting point to leverage Ray AIR. If you have any questions, suggestions, or run into any problems pelase reach out on [Discuss](https://discuss.ray.io/) or [GitHub](https://github.com/ray-project/ray)!"
    ]
diff --git a/doc/source/ray-air/examples/convert_existing_tf_code_to_ray_air.ipynb b/doc/source/ray-air/examples/convert_existing_tf_code_to_ray_air.ipynb
index 400e4cfa090a..c4a16461e0d6 100644
--- a/doc/source/ray-air/examples/convert_existing_tf_code_to_ray_air.ipynb
+++ b/doc/source/ray-air/examples/convert_existing_tf_code_to_ray_air.ipynb
@@ -18,13 +18,13 @@
     "- Automatic checkpointing/fault tolerance and result tracking\n",
     "- Parallel data preprocessing\n",
     "- Seamless integration with hyperparameter tuning\n",
-    "- Scalable batch prediction\n",
     "- Scalable model serving\n",
     "\n",
-    "This tutorial will show you how to start with Ray AIR from your existing Tensorflow training code. We will learn how to perform **distributed data-parallel training** and do **scalable batch prediction**."
+    "This tutorial will show you how to start with Ray AIR from your existing Tensorflow training code. We will learn how to perform **distributed data-parallel training**."
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "2b4c7109",
    "metadata": {},
@@ -37,6 +37,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "de8dfe91",
    "metadata": {},
@@ -66,6 +67,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "42445269",
    "metadata": {},
@@ -105,6 +107,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "f898c4b1",
    "metadata": {},
@@ -132,6 +135,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "58ceff67",
    "metadata": {},
@@ -160,6 +164,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "310c5f98",
    "metadata": {},
@@ -195,6 +200,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "0e70160c",
    "metadata": {},
@@ -229,6 +235,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "2d741835",
    "metadata": {},
@@ -281,6 +288,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "d5ce3827",
    "metadata": {},
@@ -293,6 +301,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "4b122ade",
    "metadata": {},
@@ -448,6 +457,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "1d8630f9-ed38-4b36-8dd5-1e1acc63c66e",
    "metadata": {},
@@ -458,6 +468,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "c2949ebd",
    "metadata": {},
@@ -490,6 +501,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "6774b3dd",
    "metadata": {},
@@ -498,6 +510,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "cfdda09f",
    "metadata": {},
@@ -539,6 +552,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "dbf8d1ac",
    "metadata": {},
@@ -547,6 +561,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "8b461680",
    "metadata": {},
@@ -586,236 +601,7 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "id": "68aa1557",
-   "metadata": {},
-   "source": [
-    "## Using the trained model for prediction\n",
-    "\n",
-    "The original Tensorflow quickstart tutorial doesn't actually go over loading the model for prediction, but let's see how you would complete the machine learning lifecycle with scalable batch prediction with Ray AIR!"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "e8841935",
-   "metadata": {},
-   "source": [
-    "### Loading the model for prediction\n",
-    "\n",
-    "Here's how to load an [AIR Predictor](air-predictors) `TensorflowPredictor` from a checkpoint and perform model inference on some sample data."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 20,
-   "id": "7ec884ee",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from ray.train.tensorflow import TensorflowCheckpoint, TensorflowPredictor\n",
-    "\n",
-    "tf_checkpoint: TensorflowCheckpoint = result.checkpoint\n",
-    "predictor = TensorflowPredictor.from_checkpoint(\n",
-    "    tf_checkpoint, model_definition=build_model\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 27,
-   "id": "c3890dfe",
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAOQAAADcCAYAAABgfg1NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/av/WaAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAaoUlEQVR4nO3de1RTV74H8G8CSXgHQUhCBcTWV6XilApSH4NKiWJ9t9bazkLrVFuDz5nrUtdUq3WVXtur+ECnzmphpgPioFdZUmtHEfFR1ApYLz4YxotKiwRs5SEgr+z7h5fYeA6YQCAb+H3WOmvJL/uc7IN8s8/ZJyeRMMYYCCFckNq6A4SQxyiQhHCEAkkIRyiQhHCEAkkIRyiQhHCEAkkIRyiQhHCEAkkIR7pdIPv374/58+cbfz516hQkEglOnTplteeQSCT48MMPrbY9Ik4ikSAmJsZq27t16xYkEgkSExOtts2uZlEgExMTIZFIjIuDgwMGDRqEmJgY6PX6zupjpzh69GiPDt38+fNN/q+eXH766SeLt9ny4nfgwIFO6HH39GQmnlySkpIs2p59ezqxadMmBAQE4OHDhzh79iz27NmDo0ePIj8/H05OTu3ZZLuNGzcOdXV1kMvlFq139OhRxMfHi4ayrq4O9vbt+tVwY/HixYiIiDCpMcbw3nvvoX///njmmWds1LOeZdy4cfjqq68E9W3btuGHH37AxIkTLdpeu/7qJk+ejJdeegkA8Pvf/x6enp7YunUr0tLS8Oabb4quU1NTA2dn5/Y8XZukUikcHBysuk1rb88WwsLCEBYWZlI7e/Ysamtr8dZbb9moVz3PgAEDMGDAAJNaXV0dlixZggkTJkCtVlu0PaucQ06YMAEAUFRUBODR4ZKLiwtu3ryJqKgouLq6Gv8IDAYD4uLiMGzYMDg4OEClUmHx4sW4f/++yTYZY9i8eTP69esHJycnjB8/HlevXhU8d2vnkBcuXEBUVBT69OkDZ2dnDB8+HNu3bzf2Lz4+HgBMDi9aiJ1D5uXlYfLkyXBzc4OLiwsmTpyI8+fPm7RpOXw5d+4cVq1aBS8vLzg7O2PmzJkoLy+38LdqfcnJyZBIJJg3b16nPs9nn32Gl19+GZ6ennB0dERwcHCbh7lJSUkYPHgwHBwcEBwcjNOnTwva/PTTT3jnnXegUqmgUCgwbNgwfPnll525G+125MgRVFdXt+uFzyrHZTdv3gQAeHp6GmtNTU3QarUYM2YMPvvsM+Oh7OLFi5GYmIgFCxZg2bJlKCoqwq5du5CXl4dz585BJpMBANavX4/NmzcjKioKUVFRyM3NRWRkJBoaGp7an+PHj+PVV1+FRqPB8uXLoVarcf36daSnp2P58uVYvHgxSkpKcPz4cdHDjSddvXoVY8eOhZubG1avXg2ZTIbPP/8c4eHhyMrKQmhoqEn7pUuXok+fPtiwYQNu3bqFuLg4xMTEYP/+/W0+T319Paqrq5/aHwDo27evWe1aNDY24h//+Adefvll9O/f36J1LbV9+3ZMmzYNb731FhoaGpCSkoLXX38d6enpmDJliknbrKws7N+/H8uWLYNCocDu3bsxadIkXLx4EYGBgQAAvV6PUaNGGSeBvLy88M0332DhwoWoqqrCihUrLOqfwWDAL7/8YlZbpVJp/Js0V1JSEhwdHTFr1iyL1gMAMAskJCQwAOzEiROsvLycFRcXs5SUFObp6ckcHR3Zjz/+yBhjLDo6mgFga9asMVn/zJkzDABLSkoyqR87dsykXlZWxuRyOZsyZQozGAzGduvWrWMAWHR0tLGWmZnJALDMzEzGGGNNTU0sICCA+fv7s/v375s8z6+3pdPpWGu7D4Bt2LDB+POMGTOYXC5nN2/eNNZKSkqYq6srGzdunOD3ExERYfJcK1euZHZ2dqyiokL0+Z5c35zFUkeOHGEA2O7duy1et0XL7zo1NbXNdrW1tSY/NzQ0sMDAQDZhwgSTesu+XLp0yVi7ffs2c3BwYDNnzjTWFi5cyDQaDbt3757J+nPnzmVKpdL4fEVFRQwAS0hIaLN/Le3MWVr+rsz1888/M7lczubMmWPRei3aNUI+OVng7++PpKQkwUTB+++/b/JzamoqlEolXnnlFdy7d89YDw4OhouLCzIzMzFv3jycOHECDQ0NWLp0qcmh5IoVK/Dxxx+32be8vDwUFRVh27ZtcHd3N3ns19syV3NzM/75z39ixowZJucKGo0G8+bNw1/+8hdUVVXBzc3N+NiiRYtMnmvs2LHYtm0bbt++jeHDh7f6XFqtFsePH7e4j+ZITk6GTCbDnDlzOmX7v+bo6Gj89/3799Hc3IyxY8di3759grZhYWEIDg42/uzn54fp06fjyJEjaG5uhlQqxcGDBzFnzhwwxkz+brRaLVJSUpCbm4vRo0eb3T+1Wm327zkoKMjs7QLAgQMH0NDQ0O7z9HYFMj4+HoMGDYK9vT1UKhUGDx4MqdT0dNTe3h79+vUzqRUWFqKyshLe3t6i2y0rKwMA3L59GwAwcOBAk8e9vLzQp0+fNvvWcvjccrjTUeXl5aitrcXgwYMFjw0dOhQGgwHFxcUYNmyYse7n52fSrqXPT54nP0mj0UCj0Vih16YePHiAtLQ0aLVak9OKzpKeno7Nmzfj8uXLqK+vN9bFXhCf/D8GgEGDBqG2thbl5eWQSqWoqKjA3r17sXfvXtHna/m7MZeDg4NgULGWpKQkeHh4YPLkye1av12BDAkJMc6ytkahUAhCajAY4O3t3eq1GS8vr/Z0hzt2dnaidfaUT0upq6tDZWWlWc9hyezd4cOHu2x29cyZM5g2bRrGjRuH3bt3Q6PRQCaTISEhAcnJyRZvz2AwAADefvttREdHi7Zp66hDTHNzs9mTbB4eHmZfUrtz5w7OnDmDRYsWWXze2aJLL7Y9++yzOHHiBEaPHm1yWPMkf39/AI9G1F8fJpaXlz91lHn22WcBAPn5+W2+Cpp7+Orl5QUnJycUFBQIHrtx4wakUil8fX3N2tbT7N+/HwsWLDCr7dPC/WtJSUlwcXHBtGnT2ts1sx08eBAODg749ttvoVAojPWEhATR9oWFhYLav/71Lzg5ORlfoF1dXdHc3Gy1Ua24uBgBAQFmtc3MzER4eLhZbfft2wfGWIde+Lo0kHPmzMHu3bvx0UcfCc4Fm5qa8ODBA7i7uyMiIgIymQw7d+5EZGSkMTxxcXFPfY4XX3wRAQEBiIuLw/z5803OIxljxm21XBOtqKgQnGv+mp2dHSIjI5GWloZbt24ZZyj1ej2Sk5MxZswYk/PHjuiMc8jy8nKcOHECb775Zpe8acPOzg4SiQTNzc3G2q1bt3D48GHR9tnZ2cjNzcWLL74I4FFY0tLSMGnSJOORxuzZs5GcnIz8/HzBqUh5ebnFR1addQ6ZnJwMPz8/jBkzxqL+/FqXBvK3v/0tFi9ejNjYWFy+fBmRkZGQyWQoLCxEamoqtm/fjtdeew1eXl744x//iNjYWLz66quIiopCXl4evvnmm6dO90ulUuzZswdTp07FiBEjsGDBAmg0Gty4cQNXr17Ft99+CwDGiYRly5ZBq9XCzs4Oc+fOFd3m5s2bcfz4cYwZMwZLliyBvb09Pv/8c9TX12PLli1W+/10xjnk/v370dTU1Oar9ocffoiNGzeaPRocPHgQN27cENSjo6MxZcoUbN26FZMmTcK8efNQVlaG+Ph4PPfcc7hy5YpgncDAQGi1WpPLHgCwceNGY5tPPvkEmZmZCA0Nxbvvvovnn38ev/zyC3Jzc3HixAmzL2G06IxzyPz8fFy5cgVr1qxp1+ShkSVTsi3T8t9//32b7aKjo5mzs3Orj+/du5cFBwczR0dH5urqyl544QW2evVqVlJSYmzT3NzMNm7cyDQaDXN0dGTh4eEsPz+f+fv7t3nZo8XZs2fZK6+8wlxdXZmzszMbPnw427lzp/HxpqYmtnTpUubl5cUkEonJpQQ8cdmDMcZyc3OZVqtlLi4uzMnJiY0fP5599913Zv1+WutjVxg1ahTz9vZmTU1Nrbb5wx/+wCQSCbt+/Xqb22rZj9aWM2fOMMYY++KLL9jAgQOZQqFgQ4YMYQkJCWzDhg2CyzUAmE6nY3//+9+N7X/zm9+I/p70ej3T6XTM19eXyWQyplar2cSJE9nevXuNbcy97NEZ1qxZwwCwK1eudGg7Esboc1l7u5CQEPj7+yM1NdXWXen1KJC9XFVVFby8vHD58mUMHTrU1t3p9SiQhHCk292gTEhPRoEkhCMUSEI4QoEkhCOd9saA+Ph4fPrppygtLUVQUBB27tyJkJCQp65nMBhQUlICV1fXjl1gJaQNjDFUV1fDx8dH8J5rm+roBVExKSkpTC6Xsy+//JJdvXqVvfvuu8zd3Z3p9fqnrltcXGz2vWq00NLRpbi4uDMi0G6dctkjNDQUI0eOxK5duwA8GvV8fX2xdOlSrFmzps11Kysr4e7ujjGIgj3a9455Qp6mCY04i6OoqKiAUqm0dXeMrH7I2tDQgJycHKxdu9ZYk0qliIiIQHZ2tqB9fX29yT1zLR9hYQ8Z7CUUSNJJ/n8Y4u20yOoHz/fu3UNzczNUKpVJXaVSobS0VNA+NjYWSqXSuFjrViZCuiObn82uXbsWlZWVxqW4uNjWXSLEZqx+yNq3b1/Y2dkJPslcr9eL3uWuUChMbmQlpDez+ggpl8sRHByMjIwMY81gMCAjI0Pwwb2EEFOdch1y1apViI6OxksvvYSQkBDExcWhpqbG7I+nIKS36pRAvvHGGygvL8f69etRWlqKESNG4NixY4KJHkKIKe5uv6qqqoJSqUQ4ptNlD9JpmlgjTiENlZWVVvtMJGuw+SwrIeQxCiQhHKFAEsIRCiQhHKFAEsIRCiQhHKFAEsIRCiQhHKFAEsIRCiQhHKFAEsIRCiQhHKFAEsIRCiQhHKFAEsIRCiQhHKFAEsIRCiQhHKFAEsIRCiQhHKFAEsIRCiQhHKFAEsIRCiQhHKFAEsIRCiQhHKFAEsIRCiQhHOmUb7/qCX5+V/hdln6/+7do2xtlwm/1aqgX/6KgZ/YJ604/PhBta7h8ra0ukh6IRkhCOEKBJIQjFEhCOEKBJIQjFEhCOEKzrK1Y/R/Jgtps5/vijZ+1YMPhwtKtplrRptvLx1uwYdu7WOYvqDn/l1K0rX1GTmd3p1uiEZIQjlAgCeEIBZIQjlAgCeEITeq0Yse6uYLa+uHir199rjNB7f5QiWhb+fAKQW1L4H+Ltt2muSCofV3rItp2ipP42+/MVccaROsX6p0FtXCHRvGNiPT3uTcWizYdlGF+33oTGiEJ4QgFkhCOUCAJ4QgFkhCOWBzI06dPY+rUqfDx8YFEIsHhw4dNHmeMYf369dBoNHB0dERERAQKCwut1V9CejSLZ1lramoQFBSEd955B7NmzRI8vmXLFuzYsQN//etfERAQgA8++ABarRbXrl2Dg4ODVTrdFZwPCGcMnQ+Yv76bBc+1Ux0uWt88ur9wu1niN0lvCX/OgmcUsq8ziNadr9wV1DxPHxRt+4Jc5ObrW+I3ahNxFgdy8uTJmDx5suhjjDHExcXhT3/6E6ZPnw4A+Nvf/gaVSoXDhw9j7lzhpQRCyGNWPYcsKipCaWkpIiIijDWlUonQ0FBkZ2eLrlNfX4+qqiqThZDeyqqBLC0tBQCoVKafMaNSqYyPPSk2NhZKpdK4+Pr6WrNLhHQrNp9lXbt2LSorK41LcXGxrbtEiM1Y9a1zarUaAKDX66HRaIx1vV6PESNGiK6jUCigUCis2Y1up6lUL1p3PiisN7eyDecDP1uxR4/pfy/89L1hcvE/m89+GSyo9U/4X9G2TR3rVo9l1REyICAAarUaGRmP36hYVVWFCxcuICxM+B9LCDFl8Qj54MED/Pvfj6fei4qKcPnyZXh4eMDPzw8rVqzA5s2bMXDgQONlDx8fH8yYMcOa/SakR7I4kJcuXcL48Y8/WmLVqlUAgOjoaCQmJmL16tWoqanBokWLUFFRgTFjxuDYsWPd6hokIbZicSDDw8PBmPB2oxYSiQSbNm3Cpk2bOtQxQnojm8+yEkIeoxuUCQDA3l/8+u+udbsENZnETrRt6vYIQc3zrvgbQog4GiEJ4QgFkhCOUCAJ4QgFkhCO0KQOAQDcWPmMaH2kQvjpeVcb6kTbelwT/0oEYj4aIQnhCAWSEI5QIAnhCAWSEI5QIAnhCM2y9kL1U0YKarmvbWultfDm8feXLxdt6fjdxY50i4BGSEK4QoEkhCMUSEI4QoEkhCM0qdML3ZksfB12kYh/8t+bRa8Iak7HfhBt2/rnSBBz0QhJCEcokIRwhAJJCEcokIRwhAJJCEdolrUHk7q6itZ/N/asoFZleCjatuzjAYKaov77jnWMtIpGSEI4QoEkhCMUSEI4QoEkhCM0qdODFX44TLSe3ne3oDa9cLZoW8VRmsDpSjRCEsIRCiQhHKFAEsIRCiQhHKFAEsIRmmXtISrfHiWoXXljh2jbm02NgtqD/+wn2laBux3rGLEIjZCEcIQCSQhHKJCEcIQCSQhHaFKnm7F/xke0vuKD/YKaQiL+3zv3h98Jal7f0FvkeEAjJCEcoUASwhEKJCEcoUASwhGLAhkbG4uRI0fC1dUV3t7emDFjBgoKCkzaPHz4EDqdDp6ennBxccHs2bOh1+ut2mlCeiqLZlmzsrKg0+kwcuRINDU1Yd26dYiMjMS1a9fg7OwMAFi5ciW+/vprpKamQqlUIiYmBrNmzcK5c+c6ZQd6Mom98L8nKP1H0bavu/wsqCVVe4u2VX0gfB02WNg30jksCuSxY8dMfk5MTIS3tzdycnIwbtw4VFZW4osvvkBycjImTJgAAEhISMDQoUNx/vx5jBolfL8lIeSxDp1DVlZWAgA8PDwAADk5OWhsbERERISxzZAhQ+Dn54fs7GzRbdTX16OqqspkIaS3ancgDQYDVqxYgdGjRyMwMBAAUFpaCrlcDnd3d5O2KpUKpaWlotuJjY2FUqk0Lr6+vu3tEiHdXrsDqdPpkJ+fj5SUlA51YO3ataisrDQuxcXFHdoeId1Zu946FxMTg/T0dJw+fRr9+j2+j06tVqOhoQEVFRUmo6Rer4darRbdlkKhgEIh/mWhvV7QYEHpI++vzF49/uPXRevuP4ifPhDbs2iEZIwhJiYGhw4dwsmTJxEQEGDyeHBwMGQyGTIyMoy1goIC3LlzB2FhYdbpMSE9mEUjpE6nQ3JyMtLS0uDq6mo8L1QqlXB0dIRSqcTChQuxatUqeHh4wM3NDUuXLkVYWBjNsBJiBosCuWfPHgBAeHi4ST0hIQHz588HAGzbtg1SqRSzZ89GfX09tFotdu8WfjAvIUTIokAyxp7axsHBAfHx8YiPj293pwjprei9rIRwhG5Q5oDd84NE64tS0szexvNf6gS1/l+db3efiG3QCEkIRyiQhHCEAkkIRyiQhHCEJnU4cGNJH9H6VCfz73zpd6pBWDTjMhXhC42QhHCEAkkIRyiQhHCEAkkIRyiQhHCEZlm72MOpIYJaxtT/aqW1U+d2hnCHRkhCOEKBJIQjFEhCOEKBJIQjNKnTxUpG2wlqfvbmT9609vUAsirhW+fojXPdD42QhHCEAkkIRyiQhHCEAkkIRyiQhHCEZlk5Fvvz84Jatra/aFt29386uTekK9AISQhHKJCEcIQCSQhHKJCEcIQmdbrYgDXCL0uNWvOiBVsQ/2p40jPQCEkIRyiQhHCEAkkIRyiQhHCEu0mdlm9pbkIj3dBHOk0TGgGY963gXYm7QFZXVwMAzuKojXtCeoPq6moolUpbd8NIwjh7iTAYDCgpKYGrqyuqq6vh6+uL4uJiuLm52bprVlVVVUX7ZkOMMVRXV8PHxwdSKT9nbtyNkFKpFP369QMASCQSAICbmxu3/7EdRftmOzyNjC34eWkghFAgCeEJ14FUKBTYsGEDFAqFrbtidbRvRAx3kzqE9GZcj5CE9DYUSEI4QoEkhCMUSEI4wnUg4+Pj0b9/fzg4OCA0NBQXL160dZcsdvr0aUydOhU+Pj6QSCQ4fPiwyeOMMaxfvx4ajQaOjo6IiIhAYWGhbTprgdjYWIwcORKurq7w9vbGjBkzUFBQYNLm4cOH0Ol08PT0hIuLC2bPng29Xm+jHncP3AZy//79WLVqFTZs2IDc3FwEBQVBq9WirKzM1l2zSE1NDYKCghAfHy/6+JYtW7Bjxw78+c9/xoULF+Ds7AytVouHDx92cU8tk5WVBZ1Oh/Pnz+P48eNobGxEZGQkampqjG1WrlyJI0eOIDU1FVlZWSgpKcGsWbNs2OtugHEqJCSE6XQ648/Nzc3Mx8eHxcbG2rBXHQOAHTp0yPizwWBgarWaffrpp8ZaRUUFUygUbN++fTboYfuVlZUxACwrK4sx9mg/ZDIZS01NNba5fv06A8Cys7Nt1U3ucTlCNjQ0ICcnBxEREcaaVCpFREQEsrOFn0nTXRUVFaG0tNRkP5VKJUJDQ7vdflZWVgIAPDw8AAA5OTlobGw02bchQ4bAz8+v2+1bV+IykPfu3UNzczNUKpVJXaVSobS053zIU8u+dPf9NBgMWLFiBUaPHo3AwEAAj/ZNLpfD3d3dpG1327euxt3dHqT70el0yM/Px9mzZ23dlW6PyxGyb9++sLOzE8zI6fV6qNVqG/XK+lr2pTvvZ0xMDNLT05GZmWm8bQ54tG8NDQ2oqKgwad+d9s0WuAykXC5HcHAwMjIyjDWDwYCMjAyEhYXZsGfWFRAQALVabbKfVVVVuHDhAvf7yRhDTEwMDh06hJMnTyIgIMDk8eDgYMhkMpN9KygowJ07d7jfN5uy9axSa1JSUphCoWCJiYns2rVrbNGiRczd3Z2VlpbaumsWqa6uZnl5eSwvL48BYFu3bmV5eXns9u3bjDHGPvnkE+bu7s7S0tLYlStX2PTp01lAQACrq6uzcc/b9v777zOlUslOnTrF7t69a1xqa2uNbd577z3m5+fHTp48yS5dusTCwsJYWFiYDXvNP24DyRhjO3fuZH5+fkwul7OQkBB2/vx5W3fJYpmZmQyPPq7LZImOjmaMPbr08cEHHzCVSsUUCgWbOHEiKygosG2nzSC2TwBYQkKCsU1dXR1bsmQJ69OnD3NycmIzZ85kd+/etV2nuwG6/YoQjnB5DklIb0WBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQj/wdXjCvjyw2SbwAAAABJRU5ErkJggg==",
-      "text/plain": [
-       "<Figure size 200x200 with 1 Axes>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAOQAAADcCAYAAABgfg1NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/av/WaAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAbl0lEQVR4nO3de1gTZ74H8G+4JCCEgAoJVG7WUisWaKlavF8QHlFa1K26dbfg9qy0BizSXVfcVqt1i7ejtIi221NxdbW62PVaF4uIWileivS03pBaULaYqF0JFyVc8p4/PKQdM0gCwbzo7/M8eR7zyzsz7xC/mZl3ZhIJY4yBEMIFO1t3gBDyMwokIRyhQBLCEQokIRyhQBLCEQokIRyhQBLCEQokIRyhQBLCkW4XyICAACQkJBifHzlyBBKJBEeOHLHaMiQSCd555x2rzY+ICwgIwKRJk6w6z+7+3lkUyE2bNkEikRgfTk5OCAoKQlJSErRabVf1sUscOHCgW79x7Tl9+jSSkpIQHBwMFxcX+Pn5Ydq0abh06VKH51lRUQGJRILVq1dbsafd208//YRVq1Zh5MiR8PT0hLu7O55//nns2LGjQ/Nz6MhES5cuRWBgIBoaGnD8+HFs2LABBw4cwNmzZ9GjR48OdaSjRo4ciTt37kAqlVo03YEDB5CVlSUayjt37sDBoUN/Gm6sWLEChYWFeOmllxASEgKNRoN169bh2WefxYkTJzBw4EBbd/GhUFRUhD//+c+IiYnBW2+9BQcHB3z22WeYMWMGzp8/jyVLllg2Q2aB7OxsBoCdPn1aUE9NTWUA2LZt29qctq6uzpJFtcnf35/Fx8d3ej5qtZpZuPrdSmFhIdPr9YLapUuXmEwmYzNnzuzQPMvLyxkAtmrVKmt0kfn7+7OJEydaZV6tALDFixdbdZ7388MPP7CKigpBzWAwsLFjxzKZTGbx/3urHEOOHTsWAFBeXg4ASEhIgKurKy5fvoyYmBjI5XLMnDkTAGAwGJCRkYHg4GA4OTlBqVQiMTERt27duveDAsuWLUOfPn3Qo0cPjBkzBufOnTNZdlvHkCdPnkRMTAw8PDzg4uKCkJAQvP/++8b+ZWVlAYBgF7yV2HFISUkJJkyYADc3N7i6umLcuHE4ceKEoE3rLn1hYSFSU1Ph6ekJFxcXTJ48GTdu3LDwr9o5Q4cONdlreOKJJxAcHIwLFy506bKzs7MxduxYeHl5QSaTYcCAAdiwYUOb7b/44guEhYXByckJAwYMwD//+U+TNtXV1UhJSYGvry9kMhn69euHFStWwGAwdOWqtCswMBD+/v6CmkQiQVxcHPR6PX744QeL5meV/bLLly8DAHr16mWsNTc3Izo6GsOHD8fq1auNu7KJiYnYtGkTZs2ahblz56K8vBzr1q1DSUkJCgsL4ejoCABYtGgRli1bhpiYGMTExODMmTOIiopCY2Nju/3Jy8vDpEmT4O3tjTfeeAMqlQoXLlzA/v378cYbbyAxMRFVVVXIy8vDli1b2p3fuXPnMGLECLi5uWH+/PlwdHTERx99hNGjR+Po0aMYMmSIoH1ycjI8PDywePFiVFRUICMjA0lJSe0eV+j1etTW1rbbHwDo3bu3We1+iTEGrVaL4OBgi6e1xIYNGxAcHIwXXngBDg4O2LdvH+bMmQODwQC1Wi1oW1ZWhunTp+O1115DfHw8srOz8dJLLyE3Nxfjx48HANy+fRujRo3Cjz/+iMTERPj5+eGrr75CWloarl27hoyMDIv7ePPmTbPayeVyyGQyi+ev0WgAdOB9smRz2rrLeujQIXbjxg1WWVnJtm/fznr16sWcnZ3Zv//9b8YYY/Hx8QwAW7BggWD6L7/8kgFgW7duFdRzc3MF9evXrzOpVMomTpzIDAaDsd3ChQsZAMEua0FBAQPACgoKGGOMNTc3s8DAQObv789u3bolWM4v53W/XVbcs9sTFxfHpFIpu3z5srFWVVXF5HI5GzlypMnfJzIyUrCsefPmMXt7e1ZdXS26vHunN+fREVu2bGEA2CeffNKh6c3dZb19+7ZJLTo6mvXt21dQ8/f3ZwDYZ599ZqzpdDrm7e3NnnnmGWPt3XffZS4uLuzSpUuC6RcsWMDs7e3Z1atXjbV737u2mPt3zs7Obnde9/rpp5+Yl5cXGzFihMXTdmgLGRkZKXju7++PrVu34rHHHhPUX3/9dcHznJwcKBQKjB8/XvAJFR4eDldXVxQUFODll1/GoUOH0NjYiOTkZMGuZEpKCt5777379q2kpATl5eVYu3Yt3N3dBa/9cl7mamlpwRdffIG4uDj07dvXWPf29sbLL7+Mjz/+GDU1NXBzczO+Nnv2bMGyRowYgbVr1+LKlSsICQlpc1nR0dHIy8uzuI/muHjxItRqNSIiIhAfH98ly2jl7Oxs/LdOp0NTUxNGjRqFgwcPQqfTQaFQGF/38fHB5MmTjc/d3NzwyiuvYMWKFdBoNFCpVMjJycGIESPg4eEh+H8TGRmJ5cuX49ixY8ZDInOZ+3e2dG/CYDBg5syZqK6uRmZmpkXTAh3cZc3KykJQUBAcHBygVCrx5JNPws5OeDjq4OCAPn36CGplZWXQ6XTw8vISne/169cBAFeuXAFw95jnlzw9PeHh4XHfvrXuPltrFPHGjRu4ffs2nnzySZPXnnrqKRgMBlRWVgreOD8/P0G71j7fe5x8L29vb3h7e1uh10IajQYTJ06EQqHAzp07YW9vb/Vl/FJhYSEWL16MoqIi3L59W/DavYHs16+fyQdlUFAQgLunWVQqFcrKyvDtt9/C09NTdHmt/28sce9GxVqSk5ORm5uLzZs3IzQ01OLpOxTIwYMH47nnnrtvG5lMZhJSg8EALy8vbN26VXSatv7g3U1b/+FZO9+WcufOHeh0OrOWoVKpzGqn0+kwYcIEVFdX48svv4SPj49Z03XU5cuXMW7cOPTv3x9r1qyBr68vpFIpDhw4gLVr13ZoEMZgMGD8+PGYP3++6OutAbZE6zFeexQKhWCLfz9LlizB+vXrsXz5cvz2t7+1uE+AlQZ1zPX444/j0KFDGDZs2H1XsnXUqqysTLCbeOPGjXa3Mo8//jgA4OzZs/f9FDR399XT0xM9evRAaWmpyWsXL16EnZ0dfH19zZpXe3bs2IFZs2aZ1ba9cANAQ0MDYmNjcenSJRw6dAgDBgzobBfbtW/fPuj1euzdu1ewp1BQUCDa/vvvvwdjTPB+tF68EBAQAODue1pXV2fVrZq5eyLZ2dmCK8Pa0npOOyUlBX/605863K8HGshp06Zh/fr1ePfdd02OBZubm1FXVwd3d3dERkbC0dERmZmZiIqKMr5Z5oymPfvsswgMDERGRgYSEhIEx5G/fONdXFwA3B1Ov/dY85fs7e0RFRWFPXv2oKKiwvifRKvVYtu2bRg+fLjg+LEzrHkM2dLSgunTp6OoqAh79uxBRESEVebbnta9g19+YOh0OmRnZ4u2r6qqwq5duzBlyhQAQE1NDTZv3oywsDDjXsC0adPwzjvv4ODBg4iOjhZMX11dDVdXV4sv5LDmMeSOHTswd+5czJw5E2vWrLGoH/d6oIEcNWoUEhMTkZ6ejm+++QZRUVFwdHREWVkZcnJy8P777+NXv/oVPD098Yc//AHp6emYNGkSYmJiUFJSgn/961/tDiPb2dlhw4YNiI2NRVhYGGbNmgVvb29cvHgR586dw8GDBwHcHUgCgLlz5yI6Ohr29vaYMWOG6DyXLVuGvLw8DB8+HHPmzIGDgwM++ugj6PV6rFy50mp/H2seQ7755pvYu3cvYmNj8Z///Ad///vfBa//5je/Mf679TSUuVuD/Px8NDQ0mNTj4uIQFRUFqVSK2NhYJCYmoq6uDh9//DG8vLxw7do1k2mCgoLw6quv4vTp01Aqldi4cSO0Wq0gwH/84x+xd+9eTJo0CQkJCQgPD0d9fT2+++477Ny5ExUVFRafXrDW1vbUqVN45ZVX0KtXL4wbN87kcGzo0KGCvbx2WTIk29aVOveKj49nLi4ubb7+17/+lYWHhzNnZ2cml8vZ008/zebPn8+qqqqMbVpaWtiSJUuYt7c3c3Z2ZqNHj2Znz541uVLn3tMerY4fP87Gjx/P5HI5c3FxYSEhISwzM9P4enNzM0tOTmaenp5MIpEITiVAZOj8zJkzLDo6mrm6urIePXqwMWPGsK+++sqsv09bfexKo0aNMvu0SWZmJgPAcnNz7zvP1tMebT22bNnCGGNs7969LCQkhDk5ObGAgAC2YsUKtnHjRgaAlZeXG+fXeqXOwYMHWUhICJPJZKx///4sJyfHZNm1tbUsLS2N9evXj0mlUta7d282dOhQtnr1atbY2GhsJ/bedaX2TlVZetpEwhh9L+ujbtq0aaioqMCpU6ds3ZVHXve+gpp0GmMMR44cMdmlJbZBW0hCONLtblAm5GFGgSSEIxRIQjhCgSSEI102ypqVlYVVq1ZBo9EgNDQUmZmZGDx4cLvTGQwGVFVVQS6Xd+juDELMwRhDbW0tfHx8TK65tilrnyhljLHt27czqVTKNm7cyM6dO8d+//vfM3d3d6bVatudtrKy0ux71ehBj84+KisruyICHdYlpz2GDBmCQYMGYd26dQDubvV8fX2RnJyMBQsW3HdanU4Hd3d3DEcMHOBo7a4RAgBoRhOO4wCqq6sFt4PZmtV3WRsbG1FcXIy0tDRjzc7ODpGRkSgqKjJpr9frodfrjc9bv8LCAY5wkFAgSRf5/80Qb4dFVt95vnnzJlpaWqBUKgV1pVIpeg9aeno6FAqF8WGtW5kI6Y5sfjSblpYGnU5nfFRWVtq6S4TYjNV3WXv37g17e3uTbzLXarWid7nLZLIOfasXIQ8jq28hpVIpwsPDkZ+fb6wZDAbk5+c/sJtkCemuuuQ8ZGpqKuLj4/Hcc89h8ODByMjIQH19vdlfT0HIo6pLAjl9+nTcuHEDixYtgkajQVhYGHJzc00GegghQtzdflVTUwOFQoHReJFOe5Au08yacAR7oNPprPadSNZg81FWQsjPKJCEcIQCSQhHKJCEcIQCSQhHKJCEcIS+BvIBq1hmerVSi5P4mSfPYNNfXS4K/czsZT1+WPxCDPkp099VUX7wldnzJV2HtpCEcIQCSQhHKJCEcIQCSQhHKJCEcIRGWbvIrc+fEK2fDVvXqfk2WXArwMUx/yNa3/qc6W9Q/iNvlGjblgtl5i+QdBptIQnhCAWSEI5QIAnhCAWSEI7QoI4ViA3gFIZt7/R8P6zua1JbUzRetG2Av+lldl8M+Kdo25nyaya1vyT0Fm3b9080qPMg0RaSEI5QIAnhCAWSEI5QIAnhCAWSEI7QKKsFmseFi9YPh2aJVMW/UzbjVpBJrWD6c+ILrLpuUgq69bVoUzsnJ5PaeyefFm27sPd3JrVmj2bxPpAHiraQhHCEAkkIRyiQhHCEAkkIR2hQxwJ1j0lF63Yin2tigzcAcOQF04GWlh9KO9cxAN8vecaktq3nf7fR2vQHcvvk0mczD+hdIIQjFEhCOEKBJIQjFEhCOEKBJIQjNMpqAffNRaL1X339G5Oa5FaNaNvmaxXW7JLRf8UcMqm52pmOphK+0RaSEI5QIAnhCAWSEI5QIAnhCA3qWEHL+UsPbFkVfzH9wVcAeNV9tUjV9B5JAHjz2vMmNfmhC6JtW8zuGbEG2kISwhEKJCEcoUASwhEKJCEcsTiQx44dQ2xsLHx8fCCRSLB7927B64wxLFq0CN7e3nB2dkZkZCTKyujr6Akxh8WjrPX19QgNDcXvfvc7TJkyxeT1lStX4oMPPsDf/vY3BAYG4u2330Z0dDTOnz8PJ5FvRiNtq/6t6Yhq4Stio6mAws70b1uktxdt+80y05uZnWtOWdg70hUsDuSECRMwYcIE0dcYY8jIyMBbb72FF198EQCwefNmKJVK7N69GzNmzOhcbwl5yFn1GLK8vBwajQaRkZHGmkKhwJAhQ1BUJH5htl6vR01NjeBByKPKqoHUaDQAAKVSKagrlUrja/dKT0+HQqEwPnx9fa3ZJUK6FZuPsqalpUGn0xkflZWVtu4SITZj1UvnVCoVAECr1cLb29tY12q1CAsLE51GJpNBJqP79sTcfJaZ1MQGb9oSf+S/ROtBu2kAh1dW3UIGBgZCpVIhPz/fWKupqcHJkycRESF+DSYh5GcWbyHr6urw/fffG5+Xl5fjm2++Qc+ePeHn54eUlBQsW7YMTzzxhPG0h4+PD+Li4qzZb0IeShYH8uuvv8aYMWOMz1NTUwEA8fHx2LRpE+bPn4/6+nrMnj0b1dXVGD58OHJzc+kcJCFmsDiQo0ePBmOmxzatJBIJli5diqVLl3aqY4Q8imw+ykoI+RndoMyBxjx/0XpRf7Hf5hDf9Q8tijepPfXmZdG2dNMxv2gLSQhHKJCEcIQCSQhHKJCEcIQGdR4wh74BJrV3++WItvUQuUyuWC8+X/93TYdqWm7dsqhvxPZoC0kIRyiQhHCEAkkIRyiQhHCEAkkIR2iU9QF7/B8/mtSekZr/ufjr/NdE60H/e7rDfSL8oC0kIRyhQBLCEQokIRyhQBLCERrU6SK34sW/1GuJUuweR/Fv3YuviDSpPTX/e5GWdI/jw4K2kIRwhAJJCEcokIRwhAJJCEcokIRwhEZZrcDhMR+T2oi5J0XbutqZ/zsmRef7mdSCbtElcg8z2kISwhEKJCEcoUASwhEKJCEcoUEdK7iw0PRn2Her9pk9/ZjvXhKti10mR5fIPdxoC0kIRyiQhHCEAkkIRyiQhHCEAkkIR2iU1QqKX1grUjX/EjnFHINovZl+m+ORQ1tIQjhCgSSEIxRIQjhCgSSEIzSow4EmpUK07tj4WJcsr+XGTZMa04v/EqxEZjo4Ze/Z2/xlebqL1svelJo9DzGsRSJa758scrlhTU2nlvUg0RaSEI5QIAnhCAWSEI5QIAnhiEWBTE9Px6BBgyCXy+Hl5YW4uDiUlpYK2jQ0NECtVqNXr15wdXXF1KlTodVqrdppQh5WFo2yHj16FGq1GoMGDUJzczMWLlyIqKgonD9/Hi4uLgCAefPm4fPPP0dOTg4UCgWSkpIwZcoUFBYWdskKPAw+37nxgS5vaMmvTWo3tW6ibT08a01qJ8O3Wb1P1jLgrSSTWt/5RTboScdYFMjc3FzB802bNsHLywvFxcUYOXIkdDodPvnkE2zbtg1jx44FAGRnZ+Opp57CiRMn8Pzzz1uv54Q8hDp1DKnT6QAAPXv2BAAUFxejqakJkZE//2pT//794efnh6Ii8U8pvV6PmpoawYOQR1WHA2kwGJCSkoJhw4Zh4MCBAACNRgOpVAp3d3dBW6VSCY1GIzqf9PR0KBQK48PX1/T7aQh5VHQ4kGq1GmfPnsX27ds71YG0tDTodDrjo7KyslPzI6Q769Clc0lJSdi/fz+OHTuGPn36GOsqlQqNjY2orq4WbCW1Wi1UKpXovGQyGWQil2d1Jy+en2lSyx+40wY9Mc9Xz3zaJfO9zRpNak1M/F5PMTHfJojWdd+Yf6neY8ebzW7LI4u2kIwxJCUlYdeuXTh8+DACAwMFr4eHh8PR0RH5+fnGWmlpKa5evYqICPFfFCaE/MyiLaRarca2bduwZ88eyOVy43GhQqGAs7MzFAoFXn31VaSmpqJnz55wc3NDcnIyIiIiaISVEDNYFMgNGzYAAEaPHi2oZ2dnIyEhAQCwdu1a2NnZYerUqdDr9YiOjsb69eut0llCHnYWBZIx1m4bJycnZGVlISsrq8OdIuRRRdeyEsIRukHZCpyjy01qwe+ZXsIFAKyTf3F5//+I1jt7OVvwl7NE6+yqi9nz6LuzzrR46juzp/dAmUX1hxFtIQnhCAWSEI5QIAnhCAWSEI7QoE4XCVz4YO/Bm4TwTk0fiG+t1BPSGbSFJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOMLdTwm0/kpzM5qA9n+wmZAOaUYTAPN+FfxB4i6QtbW1AIDjOGDjnpBHQW1tLRQKha27YSRhnH1EGAwGVFVVQS6Xo7a2Fr6+vqisrISbm5utu2ZVNTU1tG42xBhDbW0tfHx8YGfHz5Ebd1tIOzs79OnTBwAgkUgAAG5ubty+sZ1F62Y7PG0ZW/Hz0UAIoUASwhOuAymTybB48WLIZDJbd8XqaN2IGO4GdQh5lHG9hSTkUUOBJIQjFEhCOEKBJIQjXAcyKysLAQEBcHJywpAhQ3Dq1Clbd8lix44dQ2xsLHx8fCCRSLB7927B64wxLFq0CN7e3nB2dkZkZCTKysps01kLpKenY9CgQZDL5fDy8kJcXBxKS0sFbRoaGqBWq9GrVy+4urpi6tSp0Gq1Nupx98BtIHfs2IHU1FQsXrwYZ86cQWhoKKKjo3H9+nVbd80i9fX1CA0NRVZWlujrK1euxAcffIAPP/wQJ0+ehIuLC6Kjo9HQ0PCAe2qZo0ePQq1W48SJE8jLy0NTUxOioqJQX19vbDNv3jzs27cPOTk5OHr0KKqqqjBlyhQb9robYJwaPHgwU6vVxuctLS3Mx8eHpaen27BXnQOA7dq1y/jcYDAwlUrFVq1aZaxVV1czmUzGPv30Uxv0sOOuX7/OALCjR48yxu6uh6OjI8vJyTG2uXDhAgPAioqKbNVN7nG5hWxsbERxcTEiIyONNTs7O0RGRqKoqMiGPbOu8vJyaDQawXoqFAoMGTKk262nTqcDAPTs2RMAUFxcjKamJsG69e/fH35+ft1u3R4kLgN58+ZNtLS0QKlUCupKpRIajcZGvbK+1nXp7utpMBiQkpKCYcOGYeDAgQDurptUKoW7u7ugbXdbtweNu7s9SPejVqtx9uxZHD9+3NZd6fa43EL27t0b9vb2JiNyWq0WKpXKRr2yvtZ16c7rmZSUhP3796OgoMB42xxwd90aGxtRXV0taN+d1s0WuAykVCpFeHg48vPzjTWDwYD8/HxERETYsGfWFRgYCJVKJVjPmpoanDx5kvv1ZIwhKSkJu3btwuHDhxEYGCh4PTw8HI6OjoJ1Ky0txdWrV7lfN5uy9ahSW7Zv385kMhnbtGkTO3/+PJs9ezZzd3dnGo3G1l2zSG1tLSspKWElJSUMAFuzZg0rKSlhV65cYYwxtnz5cubu7s727NnDvv32W/biiy+ywMBAdufOHRv3/P5ef/11plAo2JEjR9i1a9eMj9u3bxvbvPbaa8zPz48dPnyYff311ywiIoJFRETYsNf84zaQjDGWmZnJ/Pz8mFQqZYMHD2YnTpywdZcsVlBQwHD367oEj/j4eMbY3VMfb7/9NlMqlUwmk7Fx48ax0tJS23baDGLrBIBlZ2cb29y5c4fNmTOHeXh4sB49erDJkyeza9eu2a7T3QDdfkUIR7g8hiTkUUWBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQjFEhCOEKBJIQj/wceQqurc2vaTQAAAABJRU5ErkJggg==",
-      "text/plain": [
-       "<Figure size 200x200 with 1 Axes>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAOQAAADcCAYAAABgfg1NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/av/WaAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAZJklEQVR4nO3de1SUdf4H8PdwmeE+CMLMkICUt9SQXUJjvesIq2mh7lrb7ll026QESz27/qSzSWueyMtRSonaTsJmmh5qvWaYIkqamiKdVi0iDyqGM+r+YkDkPt/fH/6YbZxRGRicL/h+nfOc43yf7zzzeWZ8832e7zMXhRBCgIik4ObqAojovxhIIokwkEQSYSCJJMJAEkmEgSSSCANJJBEGkkgiDCSRRLpdIPv27YvZs2dbbh88eBAKhQIHDx502mMoFAq8+uqrTtse2de3b19MnTrVqdvs7q+dQ4HMy8uDQqGwLF5eXhgwYADS0tJgNBq7qsYusWfPnm79wt3N9evXkZGRgV//+tcICgqCQqFAXl5ep7Z5/vx5KBQKrF692jlF9hBbt27FH/7wB/Tv3x8KhQLjxo3r8LY8OnKnZcuWISoqCg0NDTh8+DBycnKwZ88enD59Gj4+Ph0upiPGjBmD+vp6KJVKh+63Z88eZGdn2w1lfX09PDw69NRI49q1a1i2bBkiIiIwbNgwpx5BkLWcnByUlJQgLi4O//nPfzq1rQ79r5s8eTIeffRRAMCf//xnBAcHY82aNdixYwd+97vf2b1PXV0dfH19O17pbbi5ucHLy8up23T29lxBp9Ph8uXL0Gq1OHnyJOLi4lxdUo+1ceNGPPDAA3Bzc8PQoUM7tS2nnENOmDABAFBRUQEAmD17Nvz8/HDu3DlMmTIF/v7++P3vfw8AMJvNyMrKwpAhQ+Dl5QWNRoOUlBT89NNPVtsUQmD58uXo06cPfHx8MH78eJw5c8bmsW93Dnn8+HFMmTIFvXr1gq+vL6Kjo/Hmm29a6svOzgYAq0PwNvbOQ0pLSzF58mQEBATAz88PEydOxLFjx6z6tB3SHzlyBIsWLUJISAh8fX0xffp0XL161cFntXNUKhW0Wu09fcw2ubm5mDBhAkJDQ6FSqTB48GDk5OTctv/nn3+OmJgYeHl5YfDgwfjXv/5l06e6uhoLFixAeHg4VCoV+vXrhxUrVsBsNnflrrRLeHg43NycMx3jlOOyc+fOAQCCg4MtbS0tLUhMTMSoUaOwevVqy6FsSkoK8vLyMGfOHLz44ouoqKjA+vXrUVpaiiNHjsDT0xMAsHTpUixfvhxTpkzBlClTcOrUKSQkJKCpqemu9ezbtw9Tp06FTqfDSy+9BK1Wi2+//Ra7d+/GSy+9hJSUFFRVVWHfvn3YuHHjXbd35swZjB49GgEBAVi8eDE8PT3x7rvvYty4cTh06BBGjBhh1X/+/Pno1asXMjIycP78eWRlZSEtLQ1bt2694+M0Njaitrb2rvUAQO/evdvVzxVycnIwZMgQPPHEE/Dw8MCuXbswb948mM1mpKamWvUtLy/HU089heeffx7JycnIzc3Fb3/7WxQUFGDSpEkAgBs3bmDs2LH48ccfkZKSgoiICHz55ZdIT0/H5cuXkZWV5XCN165da1c/f39/qFQqh7ffYcIBubm5AoDYv3+/uHr1qqisrBRbtmwRwcHBwtvbW1y6dEkIIURycrIAIJYsWWJ1/y+++EIAEJs2bbJqLygosGq/cuWKUCqV4vHHHxdms9nS7+WXXxYARHJysqWtqKhIABBFRUVCCCFaWlpEVFSUiIyMFD/99JPV4/x8W6mpqeJ2uw9AZGRkWG4nJSUJpVIpzp07Z2mrqqoS/v7+YsyYMTbPj16vt3qshQsXCnd3d1FdXW338W69f3sWR5w4cUIAELm5uQ7d71YVFRUCgFi1atUd+924ccOmLTExUTz44INWbZGRkQKA+OSTTyxtJpNJ6HQ68Ytf/MLS9tprrwlfX1/x/fffW91/yZIlwt3dXVy8eNHSdutrdzvtfZ4dfc6GDBkixo4d69B9fq5DI6Rer7e6HRkZiU2bNuGBBx6wan/hhResbufn50OtVmPSpElWf6FiY2Ph5+eHoqIiPPPMM9i/fz+ampowf/58q0PJBQsW4PXXX79jbaWlpaioqMDatWsRGBhote7n22qv1tZWfP7550hKSsKDDz5oadfpdHjmmWfw3nvvoaamBgEBAZZ1c+fOtXqs0aNHY+3atbhw4QKio6Nv+1iJiYnYt2+fwzXKxtvb2/Jvk8mE5uZmjB07Fnv37oXJZIJarbasDwsLw/Tp0y23AwIC8Mc//hErVqyAwWCAVqtFfn4+Ro8ejV69eln9v9Hr9XjjjTdQXFxsOSVqr/Y+z0OGDHFou53VoUBmZ2djwIAB8PDwgEajwcCBA22OoT08PNCnTx+rtvLycphMJoSGhtrd7pUrVwAAFy5cAAD079/fan1ISAh69ep1x9raDp87e3Ld5urVq7hx4wYGDhxos+7hhx+G2WxGZWWl1QsXERFh1a+t5lvPk2+l0+mg0+mcULVrHTlyBBkZGTh69Chu3Lhhte7WQPbr18/mD+WAAQMA3LzMotVqUV5ejm+++QYhISF2H6/t/40jbh1UZNGhQA4fPtwyy3o7KpXKJqRmsxmhoaHYtGmT3fvc7gnvbtzd3e22i7t8W0p9fT1MJlO7HsNVEzZ3c+7cOUycOBGDBg3CmjVrEB4eDqVSiT179mDt2rUdmoQxm82YNGkSFi9ebHd9W4AdYTAY2tVPrVZbjfhd7Z5ebHvooYewf/9+jBw58o47GRkZCeDmiPrzw8SrV6/edZR56KGHAACnT5++41/B9h6+hoSEwMfHB2VlZTbrvvvuO7i5uSE8PLxd27qbrVu3Ys6cOe3qe7dwu8quXbvQ2NiInTt3Wh0pFBUV2e3/ww8/QAhh9Xp8//33AG6+kwe4+Zpev37dqaNae49EcnNzrd4Z1tXuaSBnzZqFt99+G6+99prNuWBLSwuuX7+OwMBA6PV6eHp6Yt26dUhISLC8WO2ZTfvlL3+JqKgoZGVlYfbs2VbnkT9/4duuiVZXV9uca/6cu7s7EhISsGPHDpw/f97yn8RoNGLz5s0YNWqU1fljZ/SEc8i2o4Of/8EwmUzIzc2127+qqgrbtm3DjBkzAAA1NTX44IMPEBMTYzkKmDVrFl599VXs3bsXiYmJVvevrq6Gn5+fw2/k6FHnkB01duxYpKSkIDMzE19//TUSEhLg6emJ8vJy5Ofn480338RvfvMbhISE4C9/+QsyMzMxdepUTJkyBaWlpfjss8/uOt3v5uaGnJwcTJs2DTExMZgzZw50Oh2+++47nDlzBnv37gVwcyIJAF588UUkJibC3d0dTz/9tN1tLl++HPv27cOoUaMwb948eHh44N1330VjYyNWrlzptOfH2eeQ69evR3V1NaqqqgDcHL0uXboE4OalmbZzubbLUO0dDQoLC9HQ0GDTnpSUhISEBCiVSkybNg0pKSm4fv063nvvPYSGhuLy5cs29xkwYACeffZZnDhxAhqNBhs2bIDRaLQK8F//+lfs3LkTU6dOxezZsxEbG4u6ujr8+9//xscff4zz5887fBnImaNtcXExiouLAdw8iqurq8Py5csB3Hwn2ZgxY9q/MUemZNum5U+cOHHHfsnJycLX1/e26//xj3+I2NhY4e3tLfz9/cUjjzwiFi9eLKqqqix9Wltbxd///neh0+mEt7e3GDdunDh9+rSIjIy842WPNocPHxaTJk0S/v7+wtfXV0RHR4t169ZZ1re0tIj58+eLkJAQoVAorC4lwM7U+alTp0RiYqLw8/MTPj4+Yvz48eLLL79s1/Nzuxq7WttlBXtLRUWFpd+6desEAFFQUHDH7bVd9rjdsnHjRiGEEDt37hTR0dHCy8tL9O3bV6xYsUJs2LDB5nEjIyPF448/Lvbu3Suio6OFSqUSgwYNEvn5+TaPXVtbK9LT00W/fv2EUqkUvXv3Fr/61a/E6tWrRVNTk6Wfvdeuq2VkZNz2OXG0FoUQkp6M0D0za9YsnD9/Hl999ZWrS7nvde93UFOnCSFw8OBBfPjhh64uhQBwhCSSSLf7gDJRT8ZAEkmEgSSSCANJJJEum2XNzs7GqlWrYDAYMGzYMKxbtw7Dhw+/6/3MZjOqqqrg7+/foU9nELWHEAK1tbUICwtz2oeLncLZF0mFEGLLli1CqVSKDRs2iDNnzojnnntOBAYGCqPReNf7VlZWtvuzaly4dHaprKzsigh0WJdc9hgxYgTi4uKwfv16ADdHvfDwcMyfPx9Lliy5431NJhMCAwMxClPgAU9nl0YEAGhBMw5jD6qrq60+DuZqTj9kbWpqQklJCdLT0y1tbm5u0Ov1OHr0qE3/xsZGNDY2Wm63fYWFBzzhoWAgqYv8/zAk22mR0w+er127htbWVmg0Gqt2jUZj9zNomZmZUKvVlsVZH2Ui6o5cfjabnp4Ok8lkWSorK11dEpHLOP2QtXfv3nB3d7f5JnOj0Wj3U+4qlerefqsXkcScPkIqlUrExsaisLDQ0mY2m1FYWIj4+HhnPxxRj9Il1yEXLVqE5ORkPProoxg+fDiysrJQV1fX7q+nILpfdUkgn3rqKVy9ehVLly6FwWBATEwMCgoKbCZ6iMiadB+/qqmpgVqtxjg8ycse1GVaRDMOYgdMJpPTvhPJGVw+y0pE/8VAEkmEgSSSCANJJBEGkkgiDCSRRBhIIokwkEQSYSCJJMJAEkmEgSSSCANJJBH+2M59SBFr+yOkn+7caLfvI++k2bSFv/al02uimzhCEkmEgSSSCANJJBEGkkginNS5D12Js/2EfAta7fb1qZLqCyV6PI6QRBJhIIkkwkASSYSBJJIIA0kkEc6y3od+iradUb3U0minJxD8vu1PCFLX4QhJJBEGkkgiDCSRRBhIIolwUqcHEyNj7LZ/MXWNTdvY4vl2+/ZDqTNLorvgCEkkEQaSSCIMJJFEGEgiiTCQRBLhLGsP9r+Dve2269x9bNoe+Jg/Hy8DjpBEEmEgiSTCQBJJhIEkkggndXqwifPsf5Zxe12gTZvfwTK7fe1/Fx11FY6QRBJhIIkkwkASSYSBJJKIw4EsLi7GtGnTEBYWBoVCge3bt1utF0Jg6dKl0Ol08Pb2hl6vR3l5ubPqJerRHJ5lraurw7Bhw/CnP/0JM2bMsFm/cuVKvPXWW/jnP/+JqKgovPLKK0hMTMTZs2fh5eXllKLJlvuQgTZtr4d+ZLfv+zV9bNpaq01Or4kc53AgJ0+ejMmTJ9tdJ4RAVlYW/va3v+HJJ58EAHzwwQfQaDTYvn07nn766c5VS9TDOfUcsqKiAgaDAXq93tKmVqsxYsQIHD1q/5pYY2MjampqrBai+5VTA2kwGAAAGo3Gql2j0VjW3SozMxNqtdqyhIeHO7Mkom7F5bOs6enpMJlMlqWystLVJRG5jFPfOqfVagEARqMROp3O0m40GhETE2P3PiqVCiqVypll3Jd+nBTc7r4ltZF2WuudVwx1mFNHyKioKGi1WhQWFlraampqcPz4ccTHxzvzoYh6JIdHyOvXr+OHH36w3K6oqMDXX3+NoKAgREREYMGCBVi+fDn69+9vuewRFhaGpKQkZ9ZN1CM5HMiTJ09i/PjxltuLFi0CACQnJyMvLw+LFy9GXV0d5s6di+rqaowaNQoFBQW8BknUDg4Hcty4cRBC3Ha9QqHAsmXLsGzZsk4VRnQ/cvksKxH9Fz+g3EPUDG5ud9+v18fYtAWCP8wqA46QRBJhIIkkwkASSYSBJJIIJ3W6mcbJcXbbdySss2lbdi3Wbt+gT76xaTN3rixyEo6QRBJhIIkkwkASSYSBJJIIA0kkEc6ydjOXJth/yaKVtp+mST7/iN2+oXXfObUmch6OkEQSYSCJJMJAEkmEgSSSCCd1upmQoVfstrcK2ze/eezo1dXlkJNxhCSSCANJJBEGkkgiDCSRRBhIIolwllViHlG2v8GxemC+3b7vmWx/NSxoA79JrrvhCEkkEQaSSCIMJJFEGEgiiXBSR2LlKWE2bY/d5rdtnzs13qYtHKedXRJ1MY6QRBJhIIkkwkASSYSBJJIIA0kkEc6ySswc3tDuvvXVtt86R90PR0giiTCQRBJhIIkkwkASSYSTOhJ7e8SH7e77wGfuXVgJ3SscIYkkwkASSYSBJJIIA0kkEYcCmZmZibi4OPj7+yM0NBRJSUkoKyuz6tPQ0IDU1FQEBwfDz88PM2fOhNFodGrRRD2VQ7Oshw4dQmpqKuLi4tDS0oKXX34ZCQkJOHv2LHx9fQEACxcuxKeffor8/Hyo1WqkpaVhxowZOHLkSJfsQE/QMG243fZRXl/ZaeXEeE/m0KtbUFBgdTsvLw+hoaEoKSnBmDFjYDKZ8P7772Pz5s2YMGECACA3NxcPP/wwjh07hscee8x5lRP1QJ06hzSZTACAoKAgAEBJSQmam5uh1+stfQYNGoSIiAgcPWr/O0IbGxtRU1NjtRDdrzocSLPZjAULFmDkyJEYOnQoAMBgMECpVCIwMNCqr0ajgcFgsLudzMxMqNVqyxIebvuFv0T3iw4HMjU1FadPn8aWLVs6VUB6ejpMJpNlqays7NT2iLqzDs0QpKWlYffu3SguLkafPn0s7VqtFk1NTaiurrYaJY1GI7Rard1tqVQqqFS3+Sq1+8TFJ4TddpXC9uVZdu0Ru339dpTYtNnfKsnMoRFSCIG0tDRs27YNBw4cQFRUlNX62NhYeHp6orCw0NJWVlaGixcvIj4+3jkVE/VgDo2Qqamp2Lx5M3bs2AF/f3/LeaFarYa3tzfUajWeffZZLFq0CEFBQQgICMD8+fMRHx/PGVaidnAokDk5OQCAcePGWbXn5uZi9uzZAIC1a9fCzc0NM2fORGNjIxITE/H22287pViins6hQApx97MSLy8vZGdnIzs7u8NFEd2v+F5WIonwfVj3mHtAgE3b/4zc0+77b/5sjN32B1v446w9AUdIIokwkEQSYSCJJMJAEkmEkzr3mLmx0abt7A3bH2YFAP2Pj9q09X/9jN2+rZ0riyTBEZJIIgwkkUQYSCKJMJBEEmEgiSTCWdZ7TNiZZS2znUwFAChxwaaNs6k9G0dIIokwkEQSYSCJJMJAEkmEgSSSCANJJBEGkkgiDCSRRBhIIokwkEQSYSCJJMJAEkmEgSSSCANJJBEGkkgiDCSRRBhIIokwkEQSYSCJJMJAEklEui+5avuV5hY0A3f/wWaiDmlBM4D2/Sr4vSRdIGtrawEAh9H+HzEl6qja2lqo1WpXl2GhEJL9iTCbzaiqqoK/vz9qa2sRHh6OyspKBNj55eHurKamhvvmQkII1NbWIiwsDG5u8py5STdCurm5oU+fPgAAhUIBAAgICJD2he0s7pvryDQytpHnTwMRMZBEMpE6kCqVChkZGVCpVK4uxem4b2SPdJM6RPczqUdIovsNA0kkEQaSSCIMJJFEpA5kdnY2+vbtCy8vL4wYMQJfffWVq0tyWHFxMaZNm4awsDAoFAps377dar0QAkuXLoVOp4O3tzf0ej3Ky8tdU6wDMjMzERcXB39/f4SGhiIpKQllZWVWfRoaGpCamorg4GD4+flh5syZMBqNLqq4e5A2kFu3bsWiRYuQkZGBU6dOYdiwYUhMTMSVK1dcXZpD6urqMGzYMGRnZ9tdv3LlSrz11lt45513cPz4cfj6+iIxMRENDQ33uFLHHDp0CKmpqTh27Bj27duH5uZmJCQkoK6uztJn4cKF2LVrF/Lz83Ho0CFUVVVhxowZLqy6GxCSGj58uEhNTbXcbm1tFWFhYSIzM9OFVXUOALFt2zbLbbPZLLRarVi1apWlrbq6WqhUKvHRRx+5oMKOu3LligAgDh06JIS4uR+enp4iPz/f0ufbb78VAMTRo0ddVab0pBwhm5qaUFJSAr1eb2lzc3ODXq/H0aNHXViZc1VUVMBgMFjtp1qtxogRI7rdfppMJgBAUFAQAKCkpATNzc1W+zZo0CBERER0u327l6QM5LVr19Da2gqNRmPVrtFoYDAYXFSV87XtS3ffT7PZjAULFmDkyJEYOnQogJv7plQqERgYaNW3u+3bvSbdpz2o+0lNTcXp06dx+PBhV5fS7Uk5Qvbu3Rvu7u42M3JGoxFardZFVTlf27505/1MS0vD7t27UVRUZPnYHHBz35qamlBdXW3VvzvtmytIGUilUonY2FgUFhZa2sxmMwoLCxEfH+/CypwrKioKWq3Waj9rampw/Phx6fdTCIG0tDRs27YNBw4cQFRUlNX62NhYeHp6Wu1bWVkZLl68KP2+uZSrZ5VuZ8uWLUKlUom8vDxx9uxZMXfuXBEYGCgMBoOrS3NIbW2tKC0tFaWlpQKAWLNmjSgtLRUXLlwQQgjxxhtviMDAQLFjxw7xzTffiCeffFJERUWJ+vp6F1d+Zy+88IJQq9Xi4MGD4vLly5blxo0blj7PP/+8iIiIEAcOHBAnT54U8fHxIj4+3oVVy0/aQAohxLp160RERIRQKpVi+PDh4tixY64uyWFFRUUCN7+uy2pJTk4WQty89PHKK68IjUYjVCqVmDhxoigrK3Nt0e1gb58AiNzcXEuf+vp6MW/ePNGrVy/h4+Mjpk+fLi5fvuy6orsBfvyKSCJSnkMS3a8YSCKJMJBEEmEgiSTCQBJJhIEkkggDSSQRBpJIIgwkkUQYSCKJMJBEEmEgiSTyf+peplweLaBHAAAAAElFTkSuQmCC",
-      "text/plain": [
-       "<Figure size 200x200 with 1 Axes>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "import matplotlib.pyplot as plt\n",
-    "\n",
-    "sample_images = x_test[:3]\n",
-    "sample_labels = y_test[:3]\n",
-    "preds = predictor.predict(sample_images)[\"predictions\"].argmax(1)\n",
-    "for image, pred, label in zip(sample_images, preds, sample_labels):\n",
-    "    plt.figure(figsize=(2, 2))\n",
-    "    plt.title(f\"Prediction = {pred}, Label = {label}\")\n",
-    "    plt.imshow(image.reshape((28, 28)))\n",
-    "    plt.show()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "2a9e7dd5",
-   "metadata": {},
-   "source": [
-    "### Using Ray AIR for scalable batch prediction\n",
-    "\n",
-    "Although what we did above works for a small amount of test data, we can use Ray AIR's {class}`BatchPredictor <ray.train.batch_predictor.BatchPredictor>` to do scalable prediction on a much larger dataset.\n",
-    "\n",
-    "We can create a `BatchPredictor` from a checkpoint. We pass in the predictor class `TensorflowPredictor` and the `model_definition` of the checkpointed model."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 22,
-   "id": "94400a99",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import ray\n",
-    "from ray.train.batch_predictor import BatchPredictor\n",
-    "from ray.train.tensorflow import TensorflowPredictor\n",
-    "\n",
-    "checkpoint: TensorflowCheckpoint = result.checkpoint\n",
-    "\n",
-    "batch_predictor = BatchPredictor.from_checkpoint(\n",
-    "    checkpoint,\n",
-    "    TensorflowPredictor,\n",
-    "    # A function that returns the model architecture\n",
-    "    model_definition=build_model,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "fd72830b",
-   "metadata": {},
-   "source": [
-    "Batch predictors work with [Ray Data](data). Here, we create a {class}`Dataset <ray.data.Dataset>` of images from our test set."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 23,
-   "id": "58bf6e2a",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "test_images_ds = ray.data.from_items(x_test)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "6ab1b08a",
-   "metadata": {},
-   "source": [
-    "Let's run {meth}`BatchPredictor.predict <ray.train.batch_predictor.BatchPredictor.predict>` on our Dataset. This will distribute the prediction across a specified number of workers!"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "c715c713",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "predict_results = batch_predictor.predict(test_images_ds)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "9ccadf89",
-   "metadata": {},
-   "source": [
-    "`predict_results` is also a Dataset, and we can take a look at the predictions inside:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "f423e247",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "predict_results.show()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "e31bb5b6",
-   "metadata": {},
-   "source": [
-    "```\n",
-    "{'predictions': array([ -2.5820212,  -5.676518 ,   2.0977738,   3.0335152,  -6.936249 ,\n",
-    "        -2.4465032, -12.235596 ,   9.626698 ,  -2.7761698,  -1.363163 ],\n",
-    "      dtype=float32)}\n",
-    "{'predictions': array([ -2.6007364,   6.3847456,  11.046391 ,   4.3116417, -12.745696 ,\n",
-    "         1.6003606,  -2.794311 ,  -6.191238 ,   2.4532976,  -9.940281 ],\n",
-    "      dtype=float32)}\n",
-    "...\n",
-    "```\n",
-    "\n",
-    "Our model outputs logits, but we want the actual predicted labels. We can convert the logits to labels by taking the `argmax` of each model output in `predict_results` using {meth}`map_batches <ray.data.Dataset.map_batches>`. Then, we can compute the accuracy by comparing to the test set labels!"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 26,
-   "id": "7aea49e1",
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Map_Batches: 100%|██████████| 3/3 [00:01<00:00,  1.56it/s]\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Prediction Accuracy = 0.976\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "(BlockWorker pid=15393) 2022-10-17 16:44:40.531349: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcuda.so.1'; dlerror: libcuda.so.1: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64\n",
-      "(BlockWorker pid=15393) 2022-10-17 16:44:40.531385: W tensorflow/stream_executor/cuda/cuda_driver.cc:269] failed call to cuInit: UNKNOWN ERROR (303)\n"
-     ]
-    }
-   ],
-   "source": [
-    "predicted_classes = predict_results.map_batches(\n",
-    "    lambda batch: {\"pred\": [pred.argmax(0) for pred in batch[\"predictions\"]]}, \n",
-    "    batch_format=\"pandas\"\n",
-    ")\n",
-    "predicted_classes_np = predicted_classes.take_batch(float(\"inf\"))[\"pred\"]\n",
-    "\n",
-    "pred_accuracy = (predicted_classes_np == y_test).astype(int).sum() / len(predicted_classes_np)\n",
-    "print(\"Prediction Accuracy =\", pred_accuracy)"
-   ]
-  },
-  {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "c02abb39",
    "metadata": {},
@@ -829,20 +615,12 @@
     "- save and retrieve model checkpoints via Ray AIR\n",
     "- load a model for batch prediction\n",
     "\n",
-    "In our [other examples](air-examples-ref) you can learn how to do more things with the Ray AIR API, such as **serving your model with Ray Serve** or **tune your hyperparameters with Ray Tune**. You can also learn how to **construct Ray Data** to leverage Ray AIR’s **preprocessing** API.\n",
+    "In our [other examples](air-examples-ref) you can learn how to do more things with the Ray AIR API, such as **serving your model with Ray Serve** or **tune your hyperparameters with Ray Tune**. You can also learn how to perform {ref}`offline batch inference <batch_inference_home>` with Ray Data.\n",
     "\n",
     "See [this table](train-framework-catalog) for a full catalog of frameworks that AIR supports out of the box.\n",
     "\n",
     "We hope this tutorial gave you a good starting point to leverage Ray AIR. If you have any questions, suggestions, or run into any problems pelase reach out on [Discuss](https://discuss.ray.io/), [GitHub](https://github.com/ray-project/ray) or the [Ray Slack](https://forms.gle/9TSdDYUgxYs8SA9e8)!"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "fdc44eb3",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {
diff --git a/doc/source/ray-air/examples/gptj_deepspeed_fine_tuning.ipynb b/doc/source/ray-air/examples/gptj_deepspeed_fine_tuning.ipynb
index 9a63a6c748f3..001160db8262 100644
--- a/doc/source/ray-air/examples/gptj_deepspeed_fine_tuning.ipynb
+++ b/doc/source/ray-air/examples/gptj_deepspeed_fine_tuning.ipynb
@@ -1017,7 +1017,7 @@
     "We can use the {class}`~ray.train.huggingface.huggingface_predictor.TransformersPredictor` to generate predictions from our fine-tuned model.\n",
     "\n",
     "```{tip}\n",
-    "For large scale batch inference, consider configuring cloud checkpointing and then pass the cloud-backed {class}`~ray.air.checkpoint.Checkpoint` to {class}`~ray.train.batch_predictor.BatchPredictor`. More information [here](air-predictors).\n",
+    "For large scale batch inference, see {ref}`End-to-end: Offline Batch Inference <batch_inference_home>`.\n",
     "```\n",
     "\n",
     "Because the {class}`~ray.train.huggingface.huggingface_predictor.TransformersPredictor` uses a 🤗 Transformers [`pipeline`](https://huggingface.co/docs/transformers/en/main_classes/pipelines) under the hood, we disable the tokenizer AIR Preprocessor we have used for training and let the `pipeline` to tokenize the data itself."
diff --git a/doc/source/ray-air/examples/huggingface_text_classification.ipynb b/doc/source/ray-air/examples/huggingface_text_classification.ipynb
index d9a0cb5c3e97..fbac21e8b873 100644
--- a/doc/source/ray-air/examples/huggingface_text_classification.ipynb
+++ b/doc/source/ray-air/examples/huggingface_text_classification.ipynb
@@ -1,6 +1,7 @@
 {
  "cells": [
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -8,6 +9,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "VaFMt6AIhYbK"
@@ -20,11 +22,11 @@
     "2. [Load the dataset](#load)\n",
     "3. [Preprocess the dataset with Ray AIR](#preprocess)\n",
     "4. [Run the training with Ray AIR](#train)\n",
-    "5. [Predict on test data with Ray AIR](#predict)\n",
-    "6. [Optionally, share the model with the community](#share)"
+    "5. [Optionally, share the model with the community](#share)"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "sQbdfyWQhYbO"
@@ -45,6 +47,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "pvSRaEHChYbP"
@@ -54,6 +57,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "LRdL3kWBhYbQ"
@@ -146,6 +150,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "oJiSdWy2hYbR"
@@ -187,6 +192,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "uS6oeJELhYbS"
@@ -212,6 +218,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "rEJBSTyZIrIb"
@@ -221,6 +228,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "kTCFado4IrIc"
@@ -243,6 +251,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "4RRkXuteIrIh"
@@ -265,6 +274,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "whPRbBNbIrIl"
@@ -274,6 +284,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "W7QYTpxXIrIl"
@@ -306,6 +317,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "RzfPtOMoIrIu"
@@ -315,6 +327,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "_TOee7nohYbW"
@@ -338,6 +351,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "lnjDIuQ3IrI-"
@@ -347,6 +361,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "n9qywopnIrJH"
@@ -356,6 +371,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "YVx71GdAIrJH"
@@ -388,6 +404,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "Vl6IidfdIrJK"
@@ -397,6 +414,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "qo_0B1M2IrJM"
@@ -428,6 +446,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "256fOuzjhYbY"
@@ -462,6 +481,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "2C0hcmp9IrJQ"
@@ -502,6 +522,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "545PP3o8IrJV"
@@ -511,6 +532,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "FBiW8UpKIrJW"
@@ -591,6 +613,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "CdzABDVcIrJg"
@@ -633,6 +656,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "XvS136zKhYba"
@@ -1007,6 +1031,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "4cnWqUWmhYba"
@@ -1042,6 +1067,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -1049,6 +1075,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -1816,6 +1843,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -2032,85 +2060,7 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Predict on test data with Ray AIR <a name=\"predict\"></a>"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "Tfoyu1q7hYbb"
-   },
-   "source": [
-    "You can now use the checkpoint to run prediction with `TransformersPredictor`, which wraps around [🤗 Pipelines](https://huggingface.co/docs/transformers/main_classes/pipelines). In order to distribute prediction, we use `BatchPredictor`. While this is not necessary for the very small example we are using (you could use `TransformersPredictor` directly), it will scale well to a large dataset."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 21,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/",
-     "height": 262
-    },
-    "id": "UOUcBkX8IrJi",
-    "outputId": "4dc16812-1400-482d-8c3f-85991ce4b081"
-   },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Map_Batches: 100%|██████████| 1/1 [00:00<00:00, 12.41it/s]\n",
-      "Map_Batches: 100%|██████████| 1/1 [00:00<00:00,  7.46it/s]\n",
-      "Map Progress (1 actors 1 pending): 100%|██████████| 1/1 [00:18<00:00, 18.46s/it]\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "{'label': 'LABEL_1', 'score': 0.6822417974472046}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822402477264404}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822407841682434}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822386980056763}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822428107261658}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822453737258911}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822437047958374}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822428703308105}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822431683540344}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822426915168762}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822447776794434}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822456121444702}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822471022605896}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822477579116821}\n",
-      "{'label': 'LABEL_1', 'score': 0.682244598865509}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822422742843628}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822470426559448}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822417378425598}\n",
-      "{'label': 'LABEL_1', 'score': 0.6822449564933777}\n",
-      "{'label': 'LABEL_1', 'score': 0.682239294052124}\n"
-     ]
-    }
-   ],
-   "source": [
-    "from ray.train.huggingface import TransformersPredictor\n",
-    "from ray.train.batch_predictor import BatchPredictor\n",
-    "import pandas as pd\n",
-    "\n",
-    "predictor = BatchPredictor.from_checkpoint(\n",
-    "    checkpoint=best_result.checkpoint,\n",
-    "    predictor_cls=TransformersPredictor,\n",
-    "    task=\"text-classification\",\n",
-    "    device=0 if use_gpu else -1,  # -1 is CPU, otherwise device index\n",
-    ")\n",
-    "prediction = predictor.predict(ray_datasets[\"test\"].map_batches(lambda x: x[[\"sentence\"]], batch_format=\"pandas\"), num_gpus_per_worker=int(use_gpu))\n",
-    "prediction.show()"
-   ]
-  },
-  {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -2118,6 +2068,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "mS8PId_NhYbb"
@@ -2147,6 +2098,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "SybKUDryhYbb"
@@ -2170,6 +2122,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "5fr6E0e8hYbb"
@@ -2196,6 +2149,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "tgV2xKfFhYbc"
@@ -2219,6 +2173,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "UL-Boc4dhYbc"
@@ -2233,6 +2188,16 @@
     "```"
    ]
   },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Next steps\n",
+    "\n",
+    "- {ref}`End-to-end: Offline Batch Inference <batch_inference_home>`"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
diff --git a/doc/source/ray-air/examples/index.rst b/doc/source/ray-air/examples/index.rst
index ede76789908a..8e41c9215e77 100644
--- a/doc/source/ray-air/examples/index.rst
+++ b/doc/source/ray-air/examples/index.rst
@@ -28,7 +28,6 @@ Text/NLP
 - :doc:`/ray-air/examples/gptj_batch_prediction`: How to use Ray AIR to do batch prediction with the Hugging Face Transformers GPT-J model.
 - :doc:`/ray-air/examples/gptj_serving`: How to use Ray AIR to do online serving with the Hugging Face Transformers GPT-J model.
 - :doc:`/ray-air/examples/dreambooth_finetuning`: How to fine-tune a DreamBooth text-to-image model with your own images.
-- :doc:`/ray-air/examples/opt_deepspeed_batch_inference`: How to run batch inference on a dataset of texts with a 30B OPT model.
 - :doc:`/ray-air/examples/dolly_lightning_fsdp_finetuning`: How to fine-tune a dolly-v2-7b model with Ray AIR LightningTrainer and FSDP.
 
 Image/CV
diff --git a/doc/source/ray-air/examples/lightgbm_example.ipynb b/doc/source/ray-air/examples/lightgbm_example.ipynb
index 2798c05babfc..d0cf8ea7b3e3 100644
--- a/doc/source/ray-air/examples/lightgbm_example.ipynb
+++ b/doc/source/ray-air/examples/lightgbm_example.ipynb
@@ -13,6 +13,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "07d92cee",
    "metadata": {},
@@ -25,12 +26,23 @@
    "execution_count": 1,
    "id": "86131abe",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1.2\u001b[0m\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
+     ]
+    }
+   ],
    "source": [
     "!pip install -qU \"ray[tune]\" lightgbm_ray"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "135fc884",
    "metadata": {},
@@ -40,15 +52,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 2,
    "id": "102ef1ac",
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/balaji/Documents/GitHub/ray/.venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+      "  from .autonotebook import tqdm as notebook_tqdm\n",
+      "2023-07-07 14:34:14,951\tINFO util.py:159 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.\n",
+      "2023-07-07 14:34:15,892\tINFO util.py:159 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.\n"
+     ]
+    }
+   ],
    "source": [
     "from typing import Tuple\n",
     "\n",
     "import ray\n",
-    "from ray.train.batch_predictor import BatchPredictor\n",
     "from ray.train.lightgbm import LightGBMPredictor\n",
     "from ray.data.preprocessors.chain import Chain\n",
     "from ray.data.preprocessors.encoder import Categorizer\n",
@@ -60,6 +82,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "c7d102bd",
    "metadata": {},
@@ -69,7 +92,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 3,
    "id": "f1f35cd7",
    "metadata": {},
    "outputs": [],
@@ -82,6 +105,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "8f7afbce",
    "metadata": {},
@@ -91,7 +115,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 4,
    "id": "fefcbc8a",
    "metadata": {},
    "outputs": [],
@@ -127,6 +151,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "04d278ae",
    "metadata": {},
@@ -136,30 +161,42 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 5,
    "id": "3f1d0c19",
    "metadata": {},
    "outputs": [],
    "source": [
+    "import pandas as pd\n",
+    "from ray.air import Checkpoint\n",
+    "from ray.data import ActorPoolStrategy\n",
+    "\n",
+    "\n",
+    "class Predict:\n",
+    "\n",
+    "    def __init__(self, checkpoint: Checkpoint):\n",
+    "        self.predictor = LightGBMPredictor.from_checkpoint(checkpoint)\n",
+    "\n",
+    "    def __call__(self, batch: pd.DataFrame) -> pd.DataFrame:\n",
+    "        return self.predictor.predict(batch)\n",
+    "\n",
+    "\n",
     "def predict_lightgbm(result: Result):\n",
     "    _, _, test_dataset = prepare_data()\n",
-    "    batch_predictor = BatchPredictor.from_checkpoint(\n",
-    "        result.checkpoint, LightGBMPredictor\n",
-    "    )\n",
     "\n",
-    "    predicted_labels = (\n",
-    "        batch_predictor.predict(test_dataset)\n",
-    "        .map_batches(lambda df: (df > 0.5).astype(int), batch_format=\"pandas\")\n",
+    "    scores = test_dataset.map_batches(\n",
+    "        Predict, \n",
+    "        fn_constructor_args=[result.checkpoint], \n",
+    "        compute=ActorPoolStrategy(), \n",
+    "        batch_format=\"pandas\"\n",
     "    )\n",
+    "    \n",
+    "    predicted_labels = scores.map_batches(lambda df: (df > 0.5).astype(int), batch_format=\"pandas\")\n",
     "    print(f\"PREDICTED LABELS\")\n",
-    "    predicted_labels.show()\n",
-    "\n",
-    "    shap_values = batch_predictor.predict(test_dataset, pred_contrib=True)\n",
-    "    print(f\"SHAP VALUES\")\n",
-    "    shap_values.show()"
+    "    predicted_labels.show()"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "2bb0e5df",
    "metadata": {},
@@ -169,29 +206,75 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 6,
    "id": "8244ff3c",
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2022-06-22 17:26:41,346\tWARNING read_api.py:260 -- The number of blocks in this dataset (1) limits its parallelism to 1 concurrent tasks. This is much less than the number of available CPU slots in the cluster. Use `.repartition(n)` to increase the number of dataset blocks.\n",
-      "Map_Batches: 100%|██████████| 1/1 [00:00<00:00, 46.26it/s]\n"
-     ]
-    },
     {
      "data": {
       "text/html": [
-       "== Status ==<br>Current time: 2022-06-22 17:26:56 (running for 00:00:14.07)<br>Memory usage on this node: 10.0/31.0 GiB<br>Using FIFO scheduling algorithm.<br>Resources requested: 0/8 CPUs, 0/0 GPUs, 0.0/13.32 GiB heap, 0.0/6.66 GiB objects<br>Result logdir: /home/ubuntu/ray_results/LightGBMTrainer_2022-06-22_17-26-41<br>Number of trials: 1/1 (1 TERMINATED)<br><table>\n",
+       "<div class=\"tuneStatus\">\n",
+       "  <div style=\"display: flex;flex-direction: row\">\n",
+       "    <div style=\"display: flex;flex-direction: column;\">\n",
+       "      <h3>Tune Status</h3>\n",
+       "      <table>\n",
+       "<tbody>\n",
+       "<tr><td>Current time:</td><td>2023-07-07 14:34:34</td></tr>\n",
+       "<tr><td>Running for: </td><td>00:00:06.06        </td></tr>\n",
+       "<tr><td>Memory:      </td><td>12.2/64.0 GiB      </td></tr>\n",
+       "</tbody>\n",
+       "</table>\n",
+       "    </div>\n",
+       "    <div class=\"vDivider\"></div>\n",
+       "    <div class=\"systemInfo\">\n",
+       "      <h3>System Info</h3>\n",
+       "      Using FIFO scheduling algorithm.<br>Logical resource usage: 4.0/10 CPUs, 0/0 GPUs\n",
+       "    </div>\n",
+       "    \n",
+       "  </div>\n",
+       "  <div class=\"hDivider\"></div>\n",
+       "  <div class=\"trialStatus\">\n",
+       "    <h3>Trial Status</h3>\n",
+       "    <table>\n",
        "<thead>\n",
-       "<tr><th>Trial name                 </th><th>status    </th><th>loc                  </th><th style=\"text-align: right;\">  iter</th><th style=\"text-align: right;\">  total time (s)</th><th style=\"text-align: right;\">  train-binary_logloss</th><th style=\"text-align: right;\">  train-binary_error</th><th style=\"text-align: right;\">  valid-binary_logloss</th></tr>\n",
+       "<tr><th>Trial name                 </th><th>status    </th><th>loc            </th><th style=\"text-align: right;\">  iter</th><th style=\"text-align: right;\">  total time (s)</th><th style=\"text-align: right;\">  train-binary_logloss</th><th style=\"text-align: right;\">  train-binary_error</th><th style=\"text-align: right;\">  valid-binary_logloss</th></tr>\n",
        "</thead>\n",
        "<tbody>\n",
-       "<tr><td>LightGBMTrainer_7b049_00000</td><td>TERMINATED</td><td>172.31.43.110:1491578</td><td style=\"text-align: right;\">   100</td><td style=\"text-align: right;\">         10.9726</td><td style=\"text-align: right;\">           0.000574522</td><td style=\"text-align: right;\">                   0</td><td style=\"text-align: right;\">              0.171898</td></tr>\n",
+       "<tr><td>LightGBMTrainer_0c5ae_00000</td><td>TERMINATED</td><td>127.0.0.1:10027</td><td style=\"text-align: right;\">   101</td><td style=\"text-align: right;\">          4.5829</td><td style=\"text-align: right;\">           0.000202293</td><td style=\"text-align: right;\">                   0</td><td style=\"text-align: right;\">              0.130232</td></tr>\n",
        "</tbody>\n",
-       "</table><br><br>"
+       "</table>\n",
+       "  </div>\n",
+       "</div>\n",
+       "<style>\n",
+       ".tuneStatus {\n",
+       "  color: var(--jp-ui-font-color1);\n",
+       "}\n",
+       ".tuneStatus .systemInfo {\n",
+       "  display: flex;\n",
+       "  flex-direction: column;\n",
+       "}\n",
+       ".tuneStatus td {\n",
+       "  white-space: nowrap;\n",
+       "}\n",
+       ".tuneStatus .trialStatus {\n",
+       "  display: flex;\n",
+       "  flex-direction: column;\n",
+       "}\n",
+       ".tuneStatus h3 {\n",
+       "  font-weight: bold;\n",
+       "}\n",
+       ".tuneStatus .hDivider {\n",
+       "  border-bottom-width: var(--jp-border-width);\n",
+       "  border-bottom-color: var(--jp-border-color0);\n",
+       "  border-bottom-style: solid;\n",
+       "}\n",
+       ".tuneStatus .vDivider {\n",
+       "  border-left-width: var(--jp-border-width);\n",
+       "  border-left-color: var(--jp-border-color0);\n",
+       "  border-left-style: solid;\n",
+       "  margin: 0.5em 1em 0.5em 1em;\n",
+       "}\n",
+       "</style>\n"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"
@@ -204,155 +287,100 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "UserWarning: cpus_per_actor is set to less than 2. Distributed LightGBM needs at least 2 CPUs per actor to train efficiently. This may lead to a degradation of performance during training.\n",
-      "\u001b[2m\u001b[36m(pid=1491578)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491578)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1491578)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491578)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1491578)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491578)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1491578)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491578)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(LightGBMTrainer pid=1491578)\u001b[0m UserWarning: Dataset 'train' has 1 blocks, which is less than the `num_workers` 2. This dataset will be automatically repartitioned to 2 blocks.\n",
-      "\u001b[2m\u001b[36m(LightGBMTrainer pid=1491578)\u001b[0m UserWarning: Dataset 'valid' has 1 blocks, which is less than the `num_workers` 2. This dataset will be automatically repartitioned to 2 blocks.\n",
-      "\u001b[2m\u001b[36m(LightGBMTrainer pid=1491578)\u001b[0m UserWarning: cpus_per_actor is set to less than 2. Distributed LightGBM needs at least 2 CPUs per actor to train efficiently. This may lead to a degradation of performance during training.\n",
-      "\u001b[2m\u001b[36m(pid=1491651)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491651)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(pid=1491651)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491651)\u001b[0m FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491651)\u001b[0m FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491653)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491653)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1491653)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491653)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1491653)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491653)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1491653)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491653)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(pid=1491652)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491652)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1491652)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491652)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1491652)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491652)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1491652)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491652)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491653)\u001b[0m 2022-06-22 17:26:50,509\tWARNING __init__.py:190 -- DeprecationWarning: `ray.worker.get_resource_ids` is a private attribute and access will be removed in a future Ray version.\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491652)\u001b[0m 2022-06-22 17:26:50,658\tWARNING __init__.py:190 -- DeprecationWarning: `ray.worker.get_resource_ids` is a private attribute and access will be removed in a future Ray version.\n"
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m The `preprocessor` arg to Trainer is deprecated. Apply preprocessor transformations ahead of time by calling `preprocessor.transform(ds)`. Support for the preprocessor arg will be dropped in a future release.\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Executing DAG InputDataBuffer[Input] -> TaskPoolMapOperator[MapBatches(get_pd_value_counts)]\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Execution config: ExecutionOptions(resource_limits=ExecutionResources(cpu=None, gpu=None, object_store_memory=None), locality_with_output=False, preserve_order=False, actor_locality_enabled=True, verbose_progress=False)\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Tip: For detailed progress reporting, run `ray.data.DataContext.get_current().execution_options.verbose_progress = True`\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Tip: Use `take_batch()` instead of `take() / show()` to return records in pandas or numpy batch format.\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Executing DAG InputDataBuffer[Input] -> TaskPoolMapOperator[MapBatches(Categorizer._transform_pandas)] -> AllToAllOperator[Aggregate]\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Execution config: ExecutionOptions(resource_limits=ExecutionResources(cpu=None, gpu=None, object_store_memory=None), locality_with_output=False, preserve_order=False, actor_locality_enabled=True, verbose_progress=False)\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Tip: For detailed progress reporting, run `ray.data.DataContext.get_current().execution_options.verbose_progress = True`\n",
+      "                                                                                                                                     \n",
+      "\u001b[A\n",
+      "\u001b[A\n",
+      "\n",
+      "\u001b[A\u001b[A\n",
+      "\n",
+      "(pid=10027) Running: 0.0/10.0 CPU, 0.0/0.0 GPU, 0.0 MiB/512.0 MiB object_store_memory:   0%|          | 0/14 [00:00<?, ?it/s] \n",
+      "\u001b[A                                                                                                                    \n",
+      "\n",
+      "\u001b[A\u001b[A                                                              \n",
+      "\n",
+      "\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Warning: The Ray cluster currently does not have any available CPUs. The Dataset job will hang unless more CPUs are freed up. A common reason is that cluster resources are used by Actors or Tune trials; see the following link for more details: https://docs.ray.io/en/master/data/dataset-internals.html#datasets-and-tune\n",
+      "\n",
+      "\u001b[A\n",
+      "\n",
+      "(pid=10027) Running: 0.0/10.0 CPU, 0.0/0.0 GPU, 0.0 MiB/512.0 MiB object_store_memory:   7%|▋         | 1/14 [00:00<00:01,  9.53it/s]\n",
+      "\u001b[A                                                                                                                   \n",
+      "\n",
+      "\u001b[A\u001b[A                                                                        \n",
+      "\n",
+      "\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Executing DAG InputDataBuffer[Input] -> TaskPoolMapOperator[MapBatches(Categorizer._transform_pandas)->MapBatches(StandardScaler._transform_pandas)]\n",
+      "\n",
+      "\u001b[A\n",
+      "\n",
+      "(pid=10027) Running: 0.0/10.0 CPU, 0.0/0.0 GPU, 0.0 MiB/512.0 MiB object_store_memory:   7%|▋         | 1/14 [00:00<00:01,  7.59it/s]\n",
+      "\u001b[A                                                                                                                   \n",
+      "\n",
+      "\u001b[A\u001b[A                                                                        \n",
+      "\n",
+      "\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Execution config: ExecutionOptions(resource_limits=ExecutionResources(cpu=None, gpu=None, object_store_memory=None), locality_with_output=False, preserve_order=False, actor_locality_enabled=True, verbose_progress=False)\n",
+      "\n",
+      "\u001b[A\n",
+      "\n",
+      "(pid=10027) Running: 0.0/10.0 CPU, 0.0/0.0 GPU, 0.0 MiB/512.0 MiB object_store_memory:   7%|▋         | 1/14 [00:00<00:01,  6.59it/s]\n",
+      "\u001b[A                                                                                                                   \n",
+      "\n",
+      "\u001b[A\u001b[A                                                                       \n",
+      "\n",
+      "\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Tip: For detailed progress reporting, run `ray.data.DataContext.get_current().execution_options.verbose_progress = True`\n",
+      "\n",
+      "\u001b[A\n",
+      "\n",
+      "                                                                                                                                     \n",
+      "\u001b[A\n",
+      "\n",
+      "\u001b[A\u001b[A\n",
+      "\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Executing DAG InputDataBuffer[Input] -> TaskPoolMapOperator[MapBatches(Categorizer._transform_pandas)->MapBatches(StandardScaler._transform_pandas)]\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Execution config: ExecutionOptions(resource_limits=ExecutionResources(cpu=None, gpu=None, object_store_memory=None), locality_with_output=False, preserve_order=False, actor_locality_enabled=True, verbose_progress=False)\n",
+      "\u001b[2m\u001b[36m(LightGBMTrainer pid=10027)\u001b[0m Tip: For detailed progress reporting, run `ray.data.DataContext.get_current().execution_options.verbose_progress = True`\n",
+      "                                                                                                                             \r"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491653)\u001b[0m [LightGBM] [Info] Trying to bind port 59039...\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491653)\u001b[0m [LightGBM] [Info] Binding port 59039 succeeded\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491653)\u001b[0m [LightGBM] [Info] Listening...\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491652)\u001b[0m [LightGBM] [Info] Trying to bind port 46955...\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491652)\u001b[0m [LightGBM] [Info] Binding port 46955 succeeded\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491652)\u001b[0m [LightGBM] [Info] Listening...\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491652)\u001b[0m [LightGBM] [Warning] Connecting to rank 1 failed, waiting for 200 milliseconds\n"
+      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=10063)\u001b[0m [LightGBM] [Info] Trying to bind port 51134...\n",
+      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=10063)\u001b[0m [LightGBM] [Info] Binding port 51134 succeeded\n",
+      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=10063)\u001b[0m [LightGBM] [Info] Listening...\n",
+      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=10062)\u001b[0m [LightGBM] [Warning] Connecting to rank 1 failed, waiting for 200 milliseconds\n",
+      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=10063)\u001b[0m [LightGBM] [Info] Connected to rank 0\n",
+      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=10063)\u001b[0m [LightGBM] [Info] Local rank: 1, total number of machines: 2\n",
+      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=10063)\u001b[0m [LightGBM] [Warning] num_threads is set=2, n_jobs=-1 will be ignored. Current value: num_threads=2\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491653)\u001b[0m UserWarning: Overriding the parameters from Reference Dataset.\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491653)\u001b[0m UserWarning: categorical_column in param dict is overridden.\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491652)\u001b[0m UserWarning: Overriding the parameters from Reference Dataset.\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491652)\u001b[0m UserWarning: categorical_column in param dict is overridden.\n"
+      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=10062)\u001b[0m /Users/balaji/Documents/GitHub/ray/.venv/lib/python3.11/site-packages/lightgbm/basic.py:1780: UserWarning: Overriding the parameters from Reference Dataset.\n",
+      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=10062)\u001b[0m   _log_warning('Overriding the parameters from Reference Dataset.')\n",
+      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=10062)\u001b[0m /Users/balaji/Documents/GitHub/ray/.venv/lib/python3.11/site-packages/lightgbm/basic.py:1513: UserWarning: categorical_column in param dict is overridden.\n",
+      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=10062)\u001b[0m   _log_warning(f'{cat_alias} in param dict is overridden.')\n",
+      "2023-07-07 14:34:34,087\tINFO tune.py:1148 -- Total run time: 7.18 seconds (6.05 seconds for the tuning loop).\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491653)\u001b[0m [LightGBM] [Info] Connected to rank 0\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491653)\u001b[0m [LightGBM] [Info] Local rank: 1, total number of machines: 2\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491653)\u001b[0m [LightGBM] [Warning] num_threads is set=1, n_jobs=-1 will be ignored. Current value: num_threads=1\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491652)\u001b[0m [LightGBM] [Info] Connected to rank 1\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491652)\u001b[0m [LightGBM] [Info] Local rank: 0, total number of machines: 2\n",
-      "\u001b[2m\u001b[36m(_RemoteRayLightGBMActor pid=1491652)\u001b[0m [LightGBM] [Warning] num_threads is set=1, n_jobs=-1 will be ignored. Current value: num_threads=1\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[2m\u001b[36m(_QueueActor pid=1491650)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(_QueueActor pid=1491650)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(_QueueActor pid=1491650)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(_QueueActor pid=1491650)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(_QueueActor pid=1491650)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(_QueueActor pid=1491650)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(_QueueActor pid=1491650)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(_QueueActor pid=1491650)\u001b[0m   from pandas import MultiIndex, Int64Index\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Result for LightGBMTrainer_7b049_00000:\n",
-      "  date: 2022-06-22_17-26-53\n",
-      "  done: false\n",
-      "  experiment_id: b4a87c26a7604a43baf895755d4f16b3\n",
-      "  hostname: ip-172-31-43-110\n",
-      "  iterations_since_restore: 1\n",
-      "  node_ip: 172.31.43.110\n",
-      "  pid: 1491578\n",
-      "  should_checkpoint: true\n",
-      "  time_since_restore: 8.369545459747314\n",
-      "  time_this_iter_s: 8.369545459747314\n",
-      "  time_total_s: 8.369545459747314\n",
-      "  timestamp: 1655918813\n",
-      "  timesteps_since_restore: 0\n",
-      "  train-binary_error: 0.5175879396984925\n",
-      "  train-binary_logloss: 0.6302848981539763\n",
-      "  training_iteration: 1\n",
-      "  trial_id: 7b049_00000\n",
-      "  valid-binary_error: 0.2\n",
-      "  valid-binary_logloss: 0.558752017793943\n",
-      "  warmup_time: 0.008721590042114258\n",
-      "  \n",
-      "Result for LightGBMTrainer_7b049_00000:\n",
-      "  date: 2022-06-22_17-26-56\n",
-      "  done: true\n",
-      "  experiment_id: b4a87c26a7604a43baf895755d4f16b3\n",
-      "  experiment_tag: '0'\n",
-      "  hostname: ip-172-31-43-110\n",
-      "  iterations_since_restore: 100\n",
-      "  node_ip: 172.31.43.110\n",
-      "  pid: 1491578\n",
-      "  should_checkpoint: true\n",
-      "  time_since_restore: 10.972588300704956\n",
-      "  time_this_iter_s: 0.027977466583251953\n",
-      "  time_total_s: 10.972588300704956\n",
-      "  timestamp: 1655918816\n",
-      "  timesteps_since_restore: 0\n",
-      "  train-binary_error: 0.0\n",
-      "  train-binary_logloss: 0.0005745220956391456\n",
-      "  training_iteration: 100\n",
-      "  trial_id: 7b049_00000\n",
-      "  valid-binary_error: 0.058823529411764705\n",
-      "  valid-binary_logloss: 0.17189847605331432\n",
-      "  warmup_time: 0.008721590042114258\n",
-      "  \n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2022-06-22 17:26:56,406\tINFO tune.py:734 -- Total run time: 14.73 seconds (14.06 seconds for the tuning loop).\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "{'train-binary_logloss': 0.0005745220956391456, 'train-binary_error': 0.0, 'valid-binary_logloss': 0.17189847605331432, 'valid-binary_error': 0.058823529411764705, 'time_this_iter_s': 0.027977466583251953, 'should_checkpoint': True, 'done': True, 'timesteps_total': None, 'episodes_total': None, 'training_iteration': 100, 'trial_id': '7b049_00000', 'experiment_id': 'b4a87c26a7604a43baf895755d4f16b3', 'date': '2022-06-22_17-26-56', 'timestamp': 1655918816, 'time_total_s': 10.972588300704956, 'pid': 1491578, 'hostname': 'ip-172-31-43-110', 'node_ip': '172.31.43.110', 'config': {}, 'time_since_restore': 10.972588300704956, 'timesteps_since_restore': 0, 'iterations_since_restore': 100, 'warmup_time': 0.008721590042114258, 'experiment_tag': '0'}\n"
+      "{'train-binary_logloss': 0.00020229312743896637, 'train-binary_error': 0.0, 'valid-binary_logloss': 0.13023245107091222, 'valid-binary_error': 0.023529411764705882, 'time_this_iter_s': 0.021785974502563477, 'should_checkpoint': True, 'done': True, 'training_iteration': 101, 'trial_id': '0c5ae_00000', 'date': '2023-07-07_14-34-34', 'timestamp': 1688765674, 'time_total_s': 4.582904100418091, 'pid': 10027, 'hostname': 'Balajis-MacBook-Pro-16', 'node_ip': '127.0.0.1', 'config': {}, 'time_since_restore': 4.582904100418091, 'iterations_since_restore': 101, 'experiment_tag': '0'}\n"
      ]
     }
    ],
@@ -361,6 +389,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "d7155d9b",
    "metadata": {},
@@ -370,7 +399,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 7,
    "id": "871c9be6",
    "metadata": {},
    "outputs": [
@@ -378,25 +407,33 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2022-06-22 17:26:57,517\tWARNING read_api.py:260 -- The number of blocks in this dataset (1) limits its parallelism to 1 concurrent tasks. This is much less than the number of available CPU slots in the cluster. Use `.repartition(n)` to increase the number of dataset blocks.\n",
-      "Map_Batches: 100%|██████████| 1/1 [00:00<00:00, 50.96it/s]\n",
-      "Map_Batches:   0%|          | 0/1 [00:00<?, ?it/s]\u001b[2m\u001b[36m(pid=1491998)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491998)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1491998)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491998)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1491998)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491998)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1491998)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1491998)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "Map Progress (1 actors 1 pending): 100%|██████████| 1/1 [00:02<00:00,  2.05s/it]\n",
-      "Map_Batches: 100%|██████████| 1/1 [00:00<00:00, 75.07it/s]\n"
+      "2023-07-07 14:34:36,769\tINFO read_api.py:374 -- To satisfy the requested parallelism of 20, each read task output will be split into 20 smaller blocks.\n",
+      "2023-07-07 14:34:38,655\tWARNING plan.py:567 -- Warning: The Ray cluster currently does not have any available CPUs. The Dataset job will hang unless more CPUs are freed up. A common reason is that cluster resources are used by Actors or Tune trials; see the following link for more details: https://docs.ray.io/en/master/data/dataset-internals.html#datasets-and-tune\n",
+      "2023-07-07 14:34:38,668\tINFO dataset.py:2180 -- Tip: Use `take_batch()` instead of `take() / show()` to return records in pandas or numpy batch format.\n",
+      "2023-07-07 14:34:38,674\tINFO streaming_executor.py:92 -- Executing DAG InputDataBuffer[Input] -> ActorPoolMapOperator[MapBatches(<lambda>)->MapBatches(Predict)] -> TaskPoolMapOperator[MapBatches(<lambda>)]\n",
+      "2023-07-07 14:34:38,674\tINFO streaming_executor.py:93 -- Execution config: ExecutionOptions(resource_limits=ExecutionResources(cpu=None, gpu=None, object_store_memory=None), locality_with_output=False, preserve_order=False, actor_locality_enabled=True, verbose_progress=False)\n",
+      "2023-07-07 14:34:38,676\tINFO streaming_executor.py:95 -- Tip: For detailed progress reporting, run `ray.data.DataContext.get_current().execution_options.verbose_progress = True`\n",
+      "2023-07-07 14:34:38,701\tINFO actor_pool_map_operator.py:117 -- MapBatches(<lambda>)->MapBatches(Predict): Waiting for 1 pool actors to start...\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "PREDICTED LABELS\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                                                                        "
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "PREDICTED LABELS\n",
       "{'predictions': 1}\n",
       "{'predictions': 1}\n",
       "{'predictions': 0}\n",
@@ -423,64 +460,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Map_Batches:   0%|          | 0/1 [00:00<?, ?it/s]\u001b[2m\u001b[36m(pid=1492031)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1492031)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1492031)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1492031)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1492031)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1492031)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1492031)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1492031)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(pid=1492033)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1492033)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1492033)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1492033)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1492033)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1492033)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1492033)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1492033)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "Map Progress (1 actors 1 pending): 100%|██████████| 1/1 [00:02<00:00,  2.09s/it]\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "SHAP VALUES\n",
-      "{'predictions_0': 0.006121974664714535, 'predictions_1': 0.8940294162424869, 'predictions_2': -0.013623909529011522, 'predictions_3': -0.26580572803883, 'predictions_4': 0.2897686828261492, 'predictions_5': -0.03784232120648852, 'predictions_6': 0.021865334852359534, 'predictions_7': 1.1753326094382734, 'predictions_8': -0.02525466292349231, 'predictions_9': 0.0733463992354119, 'predictions_10': 0.09191922035401615, 'predictions_11': -0.0035196096494634313, 'predictions_12': 0.20211476104388482, 'predictions_13': 0.7813488658944929, 'predictions_14': 0.10000464816891827, 'predictions_15': 0.11543593649642907, 'predictions_16': -0.009732477634862284, 'predictions_17': 0.19117650484758314, 'predictions_18': -0.17600075102817322, 'predictions_19': 0.5829434737180024, 'predictions_20': 1.4220773445509465, 'predictions_21': 0.6086211783805069, 'predictions_22': 2.0031654232526925, 'predictions_23': 0.3090376110779834, 'predictions_24': -0.21156467772251453, 'predictions_25': 0.14122943819731193, 'predictions_26': -0.1324700025487787, 'predictions_27': 0.8280650504246968, 'predictions_28': 0.03147457104755769, 'predictions_29': 0.00029604737237433516, 'predictions_30': 0.024336487839325866, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.00565090216762466, 'predictions_1': 0.7173247917145018, 'predictions_2': -0.01352989648419376, 'predictions_3': -0.204508963539279, 'predictions_4': -0.11703564338083555, 'predictions_5': 0.059858710083059874, 'predictions_6': 0.06974454296095976, 'predictions_7': 1.5952991804773315, 'predictions_8': 0.30494490847895245, 'predictions_9': 0.03770331660034111, 'predictions_10': 0.08779844216179675, 'predictions_11': 0.0001818669974550241, 'predictions_12': -0.10871732001356472, 'predictions_13': 0.49872871949407244, 'predictions_14': 0.16083030838859202, 'predictions_15': 0.4071487385487001, 'predictions_16': -0.00920287075428388, 'predictions_17': 0.21519060265555054, 'predictions_18': -0.24141319659570365, 'predictions_19': -0.19394859165532527, 'predictions_20': 1.2358452648954865, 'predictions_21': 0.16127531717942642, 'predictions_22': 1.3397755121893355, 'predictions_23': 0.24271016133964965, 'predictions_24': -0.11296858156987878, 'predictions_25': 0.21775788278030012, 'predictions_26': 0.8594002204044787, 'predictions_27': 1.0571631081079365, 'predictions_28': 0.06338809094380635, 'predictions_29': 0.14952090064808415, 'predictions_30': -0.020191656254497082, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.0011410972769028797, 'predictions_1': -0.023112580054428615, 'predictions_2': 0.0015007474035067395, 'predictions_3': -0.3960490192373774, 'predictions_4': -0.30646108596137317, 'predictions_5': -0.015606280874156383, 'predictions_6': -0.10875176916234583, 'predictions_7': -2.3253286264519457, 'predictions_8': 0.2758843675860649, 'predictions_9': 0.029091310311824298, 'predictions_10': -0.057950348636255644, 'predictions_11': -0.00017555498393944432, 'predictions_12': -0.4136454204716676, 'predictions_13': -0.3629735922139978, 'predictions_14': 0.04232756741319012, 'predictions_15': 0.06936920198392876, 'predictions_16': 0.010307144611165166, 'predictions_17': -0.4063116213440989, 'predictions_18': -0.07826460708005233, 'predictions_19': 0.28668914680505037, 'predictions_20': -2.0034181076720015, 'predictions_21': -0.4289092806234529, 'predictions_22': -2.059807308089095, 'predictions_23': -0.2625534917898286, 'predictions_24': -1.0607560950436483, 'predictions_25': -0.13241418825219023, 'predictions_26': -0.46713657128877134, 'predictions_27': -2.0707325110237127, 'predictions_28': -0.0212343580603297, 'predictions_29': -0.11761200100287779, 'predictions_30': 0.03805635018946682, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.003501920356918757, 'predictions_1': 0.9649889446638613, 'predictions_2': -0.011627077584939034, 'predictions_3': -0.33201537640627937, 'predictions_4': 0.2626117060870051, 'predictions_5': -0.0420997182498785, 'predictions_6': 0.05656763216450521, 'predictions_7': 1.076092179662977, 'predictions_8': -0.1396182169782879, 'predictions_9': -0.09872952353947571, 'predictions_10': 0.04378766056466948, 'predictions_11': 0.002478996394296549, 'predictions_12': 0.25042183813566526, 'predictions_13': 0.8751692867530225, 'predictions_14': 0.18679133739736484, 'predictions_15': 0.046846715006450504, 'predictions_16': -0.009211815518670832, 'predictions_17': 0.22485983912144494, 'predictions_18': -0.2861737431801593, 'predictions_19': -0.2533929278067911, 'predictions_20': 1.316951719635302, 'predictions_21': 1.1964971086769494, 'predictions_22': 1.2740098717427248, 'predictions_23': 0.25042580055967084, 'predictions_24': -0.4015257176668039, 'predictions_25': 0.17935395324361414, 'predictions_26': 1.126933988937795, 'predictions_27': 0.8031626612897146, 'predictions_28': 0.0771850514731471, 'predictions_29': 0.03755423306624511, 'predictions_30': -0.016833253240925238, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': -0.0034560551402153003, 'predictions_1': 0.5230708630376469, 'predictions_2': -0.015562114219360572, 'predictions_3': -0.1196402194436373, 'predictions_4': 0.4106482044292619, 'predictions_5': 0.06220233147046589, 'predictions_6': 0.12716114707514065, 'predictions_7': 1.3356455912614509, 'predictions_8': 0.1447514882444872, 'predictions_9': 0.12370386736447751, 'predictions_10': 0.07410456355721864, 'predictions_11': 0.012016763274156357, 'predictions_12': -0.10513441936331262, 'predictions_13': 0.7484191363603289, 'predictions_14': 0.18707788149117566, 'predictions_15': 0.3327881147491029, 'predictions_16': -0.009219336794413353, 'predictions_17': -0.10065740008750416, 'predictions_18': 0.16625881614886867, 'predictions_19': 0.23084551369873454, 'predictions_20': 1.358717098613538, 'predictions_21': 0.19175435277095332, 'predictions_22': 1.3375643842995248, 'predictions_23': 0.2926283902278477, 'predictions_24': 0.1146310032943002, 'predictions_25': 0.23343399483643015, 'predictions_26': 0.6034462734909513, 'predictions_27': 0.9230214841058666, 'predictions_28': 0.029594344165258104, 'predictions_29': 0.04913153000099999, 'predictions_30': 0.02341707352913655, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': -0.005102561841927789, 'predictions_1': 1.0861119112102469, 'predictions_2': -0.0154828846564582, 'predictions_3': -0.3088905099091714, 'predictions_4': 0.05779026036152443, 'predictions_5': 0.047351932324116885, 'predictions_6': 0.0876371219806605, 'predictions_7': 1.1210466016114495, 'predictions_8': -0.1252369777517682, 'predictions_9': 0.04572512843104436, 'predictions_10': 0.09245771221086214, 'predictions_11': 0.007753500238910626, 'predictions_12': 0.2309698163766563, 'predictions_13': 0.9684988783771291, 'predictions_14': 0.024511467599608535, 'predictions_15': 0.18657179919131872, 'predictions_16': -0.009212652411079585, 'predictions_17': -0.13395842318946233, 'predictions_18': 0.152376407447391, 'predictions_19': -0.28554273302892486, 'predictions_20': 1.3994697511973517, 'predictions_21': 0.5784457048607689, 'predictions_22': 1.3325378201278, 'predictions_23': 0.30730022154186687, 'predictions_24': 0.017237427138293876, 'predictions_25': 0.19484371531419448, 'predictions_26': 1.0716212980249242, 'predictions_27': 0.7424857548065319, 'predictions_28': 0.030110335845485465, 'predictions_29': 0.08677604394991238, 'predictions_30': -0.018230914497616164, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.011026642414565658, 'predictions_1': 1.0433621693813095, 'predictions_2': -0.00702393810808943, 'predictions_3': -0.2962479861350653, 'predictions_4': 0.20838486132625483, 'predictions_5': -0.07568934141814487, 'predictions_6': 0.026798049998736986, 'predictions_7': 1.2233970557267948, 'predictions_8': -0.07215770822854156, 'predictions_9': 0.016138237086580777, 'predictions_10': 0.04908427317188252, 'predictions_11': -0.013274124641011575, 'predictions_12': -0.16059386568879297, 'predictions_13': 0.38386374312584454, 'predictions_14': -0.03476748264814593, 'predictions_15': -0.5225211720205649, 'predictions_16': -0.009220168600202043, 'predictions_17': -0.15278574495418593, 'predictions_18': 0.12911665421378546, 'predictions_19': -0.2782951415110554, 'predictions_20': 1.2470508123020512, 'predictions_21': 1.049830317708393, 'predictions_22': 2.102971796648596, 'predictions_23': 0.2851979247349288, 'predictions_24': -0.0006702647871052775, 'predictions_25': -0.11420596882801563, 'predictions_26': 1.0834575497816143, 'predictions_27': 0.8164104508549398, 'predictions_28': 0.06634783513626033, 'predictions_29': 0.10518170393387423, 'predictions_30': 0.05948635171854934, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.008846265655878418, 'predictions_1': 1.071493056050533, 'predictions_2': 0.0241002358002765, 'predictions_3': -0.2914009217752569, 'predictions_4': -0.1844514405858182, 'predictions_5': 0.09688586653158524, 'predictions_6': 0.1189004872794518, 'predictions_7': 0.019812317046639417, 'predictions_8': -0.0841879790447643, 'predictions_9': 0.0689671067023492, 'predictions_10': 0.057123796305462, 'predictions_11': 0.018751811843757425, 'predictions_12': -0.19278774516524225, 'predictions_13': 0.5521382975001031, 'predictions_14': -0.1961614983559944, 'predictions_15': 0.3352816348185536, 'predictions_16': -0.009197695434128215, 'predictions_17': -0.0600167757501572, 'predictions_18': 0.27488314466683056, 'predictions_19': -0.35962747336476697, 'predictions_20': 1.2317107478669351, 'predictions_21': 0.05530975604521487, 'predictions_22': 2.382011011440535, 'predictions_23': 0.33824065775317486, 'predictions_24': 0.3498540690011901, 'predictions_25': 0.1739274660593352, 'predictions_26': 1.160333734158511, 'predictions_27': 1.033879786485623, 'predictions_28': 0.08158573366246898, 'predictions_29': 0.10563970622307337, 'predictions_30': -0.04267793892712356, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.004750001129083867, 'predictions_1': 3.2244208404334374, 'predictions_2': 0.01622715285279811, 'predictions_3': -0.817260302999862, 'predictions_4': -0.09736090983332732, 'predictions_5': 0.07881792915896496, 'predictions_6': 0.24070898834769355, 'predictions_7': 0.05001221074373508, 'predictions_8': -0.2567854774979608, 'predictions_9': 0.03063087506346955, 'predictions_10': 0.05499599036837444, 'predictions_11': -0.015303644634305683, 'predictions_12': -0.14884606737286166, 'predictions_13': 0.8519672928318166, 'predictions_14': 0.09824149785766935, 'predictions_15': 0.26921023748269235, 'predictions_16': -0.010848751281971217, 'predictions_17': -0.11619083730523652, 'predictions_18': -0.17527472428145596, 'predictions_19': -0.5874677933384177, 'predictions_20': -0.3990904299729458, 'predictions_21': 2.2068328291797936, 'predictions_22': -1.932202847332452, 'predictions_23': -0.3152964245377162, 'predictions_24': 0.7834452171983805, 'predictions_25': 0.2512128072560273, 'predictions_26': -0.6206434154152907, 'predictions_27': 0.08708205787374604, 'predictions_28': 0.040648951231987765, 'predictions_29': 0.06879586583909683, 'predictions_30': 0.043515107484221424, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': -0.011027810657660349, 'predictions_1': -0.5337191669392997, 'predictions_2': -0.0026033241771052282, 'predictions_3': -0.2382336633486158, 'predictions_4': 0.8615636351404935, 'predictions_5': 0.059347121609268944, 'predictions_6': 0.14253423982272048, 'predictions_7': 1.4462830393121449, 'predictions_8': -0.06536550111076092, 'predictions_9': 0.12249420022849346, 'predictions_10': -0.040845467674169966, 'predictions_11': 0.03619973926410233, 'predictions_12': -0.14839345664605622, 'predictions_13': -0.38765958181699983, 'predictions_14': 0.45137385893985227, 'predictions_15': 0.4818261473751218, 'predictions_16': 0.005229328958126197, 'predictions_17': -0.14927291449462546, 'predictions_18': 0.12257473692108792, 'predictions_19': 0.5775523654869467, 'predictions_20': 1.4945158847763393, 'predictions_21': -0.11572127634540279, 'predictions_22': 1.2803791500605577, 'predictions_23': 0.2519454034779557, 'predictions_24': 0.12639705427540554, 'predictions_25': 0.20374090734634412, 'predictions_26': 0.9872077234715891, 'predictions_27': 1.1931782325388345, 'predictions_28': 0.07647609206107736, 'predictions_29': 0.017535160650109134, 'predictions_30': -0.011152247353355573, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': -0.009236814888869137, 'predictions_1': -0.08974532751806996, 'predictions_2': -0.005072446447076212, 'predictions_3': -0.49454476931590674, 'predictions_4': -0.14583165960960504, 'predictions_5': 0.037743362980294, 'predictions_6': -0.09071218034159645, 'predictions_7': -2.076157655204495, 'predictions_8': 0.6915530135596496, 'predictions_9': 0.015305309316520455, 'predictions_10': -0.05407297473599998, 'predictions_11': -0.01202689608724274, 'predictions_12': -0.37048240770136764, 'predictions_13': -1.1222567180136822, 'predictions_14': 0.037999849333804875, 'predictions_15': 0.05179781531623324, 'predictions_16': -0.009442169563784072, 'predictions_17': -0.3518926772423797, 'predictions_18': -0.18168464537700557, 'predictions_19': -0.246308669315598, 'predictions_20': -1.8215267653197431, 'predictions_21': -0.16464307910939846, 'predictions_22': -2.294068720859334, 'predictions_23': -0.3304406806357679, 'predictions_24': -0.8059935139116144, 'predictions_25': -0.15473187742974112, 'predictions_26': -0.44492987082868113, 'predictions_27': -1.706574981012038, 'predictions_28': 0.009928350750753007, 'predictions_29': -0.005531569126011125, 'predictions_30': 0.03400893184303606, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.002893918337539625, 'predictions_1': 0.6953965947651528, 'predictions_2': -0.013300368855470382, 'predictions_3': -0.15693491782098012, 'predictions_4': 0.4052561639196121, 'predictions_5': -0.01785344804083238, 'predictions_6': 0.19406598570732034, 'predictions_7': 1.5796202341560919, 'predictions_8': 0.28954821935673325, 'predictions_9': -0.215897797520852, 'predictions_10': 0.05835282036206641, 'predictions_11': 0.03331176153763488, 'predictions_12': -0.10112958834294049, 'predictions_13': 0.3947745629125056, 'predictions_14': 0.22909135741673778, 'predictions_15': 0.473005657256218, 'predictions_16': -0.009633689567643305, 'predictions_17': -0.09362381604913257, 'predictions_18': 0.14969971629912343, 'predictions_19': -0.1688864705396212, 'predictions_20': 1.3001215347067874, 'predictions_21': -0.21918668227485943, 'predictions_22': 1.3437058168797267, 'predictions_23': 0.3124907025891718, 'predictions_24': 0.14131080537131419, 'predictions_25': 0.2243700411172835, 'predictions_26': 0.9296630907535046, 'predictions_27': 0.41471504174869356, 'predictions_28': 0.020173572275052214, 'predictions_29': 0.04820465228613692, 'predictions_30': -0.020545384469295942, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.0019388150699810927, 'predictions_1': 0.9223705632364128, 'predictions_2': -0.012007685883043798, 'predictions_3': -0.31903137131372966, 'predictions_4': 0.628481905867853, 'predictions_5': -0.06149389221971728, 'predictions_6': 0.07061611794203079, 'predictions_7': 1.3823056189088423, 'predictions_8': -0.133343124664483, 'predictions_9': 0.11603949252367812, 'predictions_10': 0.21857476218484376, 'predictions_11': 0.015902798791774055, 'predictions_12': 0.2913666065699202, 'predictions_13': 0.9315294837553827, 'predictions_14': 0.277372510153019, 'predictions_15': -0.5071083100622337, 'predictions_16': -0.009631147961094873, 'predictions_17': 0.23976024706824375, 'predictions_18': -0.20540519019181294, 'predictions_19': -0.2450413530123813, 'predictions_20': 1.1789579806256083, 'predictions_21': -1.5177833757024324, 'predictions_22': 1.4604002248277352, 'predictions_23': 0.27531878725283415, 'predictions_24': -1.352094841156462, 'predictions_25': 0.16870219247146698, 'predictions_26': 1.2263320807717468, 'predictions_27': 0.8656450275023648, 'predictions_28': 0.04415827467622267, 'predictions_29': 0.049210669003044466, 'predictions_30': 0.03485239596130599, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.003483875604996794, 'predictions_1': 0.9677217395705439, 'predictions_2': -0.01375116195045965, 'predictions_3': -0.3148394441913672, 'predictions_4': -0.2574040676795255, 'predictions_5': 0.07782351238517007, 'predictions_6': 0.09223237164727777, 'predictions_7': 1.359163521325679, 'predictions_8': -0.10520478897286097, 'predictions_9': 0.051820926250002466, 'predictions_10': 0.15651454755052202, 'predictions_11': 0.012354841533717503, 'predictions_12': 0.29314938008831337, 'predictions_13': 1.0134451429783053, 'predictions_14': 0.07334166731849916, 'predictions_15': -0.5580245806930221, 'predictions_16': -0.009637538822743917, 'predictions_17': -0.12931888564696647, 'predictions_18': -0.08985648327837921, 'predictions_19': -0.2838831478457971, 'predictions_20': 1.197739882302604, 'predictions_21': -0.14264086768498266, 'predictions_22': 2.4168915798709034, 'predictions_23': 0.35060520926622657, 'predictions_24': -0.243435195670719, 'predictions_25': 0.15680747277488985, 'predictions_26': 1.2012113470638528, 'predictions_27': 0.9897751319349664, 'predictions_28': 0.05573907097988011, 'predictions_29': 0.06252860717834312, 'predictions_30': -0.05792966463337761, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.015122615065552805, 'predictions_1': -1.1167653489947622, 'predictions_2': -0.008012147012472742, 'predictions_3': -0.20874221256644707, 'predictions_4': 0.4252072619730782, 'predictions_5': 0.0038900875799020296, 'predictions_6': 0.1140119630004244, 'predictions_7': 1.1987104625838227, 'predictions_8': -0.0802347059616203, 'predictions_9': 0.14227487864929314, 'predictions_10': 0.061570146412656145, 'predictions_11': -0.0013235117361348366, 'predictions_12': 0.22496427871452854, 'predictions_13': 0.6826705611065566, 'predictions_14': 0.331084179340632, 'predictions_15': 0.2325873510907064, 'predictions_16': -0.005890948415758354, 'predictions_17': 0.23108082656181192, 'predictions_18': 0.08866538976848874, 'predictions_19': 0.5251741787977718, 'predictions_20': 1.5307505517513718, 'predictions_21': -0.07014338016238107, 'predictions_22': 1.8024293796373567, 'predictions_23': 0.3420427562962711, 'predictions_24': -0.3356060347979862, 'predictions_25': 0.08823113765567157, 'predictions_26': 0.9993112240252872, 'predictions_27': 1.1583364010362838, 'predictions_28': 0.05818942683648322, 'predictions_29': -0.010171113593323115, 'predictions_30': 0.017500344327137828, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.00510007523648401, 'predictions_1': -1.8438586845336085, 'predictions_2': 0.07034588660781765, 'predictions_3': -0.6334911581627888, 'predictions_4': 0.9114601232034509, 'predictions_5': 0.015641139153578926, 'predictions_6': 0.3391581513516312, 'predictions_7': 0.015993612058473987, 'predictions_8': 0.4979057135726034, 'predictions_9': 0.14140896753303245, 'predictions_10': 0.03348118561743431, 'predictions_11': 0.018019313387541973, 'predictions_12': -0.1851865976812716, 'predictions_13': 0.18463673035754868, 'predictions_14': 0.3321862529567762, 'predictions_15': 0.4582953091852766, 'predictions_16': -0.023872509230380146, 'predictions_17': -0.05714457269664822, 'predictions_18': 0.1677010761064405, 'predictions_19': 0.6590215547332258, 'predictions_20': -0.4470726570372422, 'predictions_21': -1.2957188152033094, 'predictions_22': -0.49568168502602117, 'predictions_23': -0.5319175224432703, 'predictions_24': 0.8792904089758667, 'predictions_25': -0.16764333932557407, 'predictions_26': -0.5006140094263773, 'predictions_27': -0.559662593948684, 'predictions_28': 0.009575432219475658, 'predictions_29': 0.03620587401831965, 'predictions_30': -0.022617768077518082, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': -0.0035275503865915213, 'predictions_1': -1.7572879984733045, 'predictions_2': 0.02112961345058588, 'predictions_3': -0.5373126141152578, 'predictions_4': 0.8167634172202621, 'predictions_5': -0.04568688093881375, 'predictions_6': 0.25612518616907237, 'predictions_7': 1.6585574657259259, 'predictions_8': 0.5708485569128593, 'predictions_9': -0.2579041111541445, 'predictions_10': -0.02303431468406031, 'predictions_11': 0.0294850340796527, 'predictions_12': -0.20150167754729115, 'predictions_13': -0.3122660955180186, 'predictions_14': 0.08946441512158586, 'predictions_15': 0.19599051996984723, 'predictions_16': -0.007048153187921441, 'predictions_17': -0.10443059840913398, 'predictions_18': 0.14693551914712497, 'predictions_19': -0.2646088763947528, 'predictions_20': -0.14423328554566225, 'predictions_21': -1.096484854939204, 'predictions_22': -0.5409155996124657, 'predictions_23': -0.2125528115678197, 'predictions_24': 0.2028301059851276, 'predictions_25': 0.15076353496237724, 'predictions_26': 1.863813974679114, 'predictions_27': 1.2625204294969739, 'predictions_28': -0.009542569876103744, 'predictions_29': 0.08892200769099384, 'predictions_30': 0.016344768954324324, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.009144654047150176, 'predictions_1': 0.3616068989438007, 'predictions_2': -0.009713127697925978, 'predictions_3': -0.25503380890174077, 'predictions_4': -0.08841376708924753, 'predictions_5': 0.12905387860630704, 'predictions_6': 0.09872550234229752, 'predictions_7': 1.5461380270279617, 'predictions_8': 0.5934142135359506, 'predictions_9': 0.04672843933802434, 'predictions_10': 0.017982350366210965, 'predictions_11': 0.011836524186964618, 'predictions_12': -0.09329132650766998, 'predictions_13': 0.9816979966957412, 'predictions_14': -0.26131805604494435, 'predictions_15': 0.2573728246698596, 'predictions_16': -0.009616853447343936, 'predictions_17': -0.11778440476199589, 'predictions_18': 0.19894108953925974, 'predictions_19': -0.28976560140618507, 'predictions_20': 1.145755494068452, 'predictions_21': 0.19170884942775918, 'predictions_22': 1.751619931359333, 'predictions_23': 0.31591084941785597, 'predictions_24': -0.9883146017669252, 'predictions_25': 0.3832169744602564, 'predictions_26': 1.3459027320296548, 'predictions_27': 1.0895032649194054, 'predictions_28': 0.054326669111151096, 'predictions_29': 0.11224841710144848, 'predictions_30': -0.01934243236389702, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.05414474380600745, 'predictions_1': -1.149372263278166, 'predictions_2': -0.0034542658503694534, 'predictions_3': -0.11226127026369175, 'predictions_4': -0.2765479593192703, 'predictions_5': 0.057605254673602974, 'predictions_6': 0.04807218948118946, 'predictions_7': 1.627632661158546, 'predictions_8': 0.23594239851080898, 'predictions_9': 0.08102266882022441, 'predictions_10': -0.035797597595999694, 'predictions_11': -0.006940512375528646, 'predictions_12': -0.10341465914545066, 'predictions_13': 0.27134162901025793, 'predictions_14': -0.4589675261254597, 'predictions_15': 0.16906923946657362, 'predictions_16': -0.005805030106413082, 'predictions_17': -0.11670739934889805, 'predictions_18': 0.270439579413901, 'predictions_19': 0.2757024597749045, 'predictions_20': 1.2679444783850085, 'predictions_21': -1.2185063190204835, 'predictions_22': 2.6862730600162457, 'predictions_23': 0.45079291995440945, 'predictions_24': -0.8576927701312551, 'predictions_25': 0.1825880636881889, 'predictions_26': 0.9481775337394789, 'predictions_27': 1.3019845783662138, 'predictions_28': 0.03309325718132554, 'predictions_29': 0.037279537345320794, 'predictions_30': 0.030849407280271066, 'predictions_31': 1.5201632854544105}\n",
-      "{'predictions_0': 0.026701912078513444, 'predictions_1': -0.016049183561005216, 'predictions_2': -0.026512557715316794, 'predictions_3': -0.33992007086017256, 'predictions_4': -0.3231034954783173, 'predictions_5': 0.020522588667874812, 'predictions_6': -0.09818245278711138, 'predictions_7': -1.9632581054922957, 'predictions_8': 0.2796715168175009, 'predictions_9': 0.025963248780199805, 'predictions_10': -0.13243884691329014, 'predictions_11': -0.007600341414574132, 'predictions_12': -0.3505614312588073, 'predictions_13': -0.8449241022454159, 'predictions_14': -0.0623541831245574, 'predictions_15': 0.11533014973600747, 'predictions_16': 0.008322220108907262, 'predictions_17': -0.02930862278171467, 'predictions_18': 0.02496960430979726, 'predictions_19': 0.3997160251519232, 'predictions_20': -2.0119476119311948, 'predictions_21': -0.3601922717542553, 'predictions_22': -2.240466883625807, 'predictions_23': -0.24430626245778664, 'predictions_24': -0.732571668183472, 'predictions_25': -0.14435610495492934, 'predictions_26': -0.4186367055351456, 'predictions_27': -1.7801593987201698, 'predictions_28': 0.014498054148804375, 'predictions_29': -0.10768829118597369, 'predictions_30': -0.02172472974992555, 'predictions_31': 1.5201632854544105}\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[2m\u001b[36m(pid=1492090)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1492090)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1492090)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1492090)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1492090)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/dask/dataframe/backends.py:181: FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1492090)\u001b[0m   _numeric_index_types = (pd.Int64Index, pd.Float64Index, pd.UInt64Index)\n",
-      "\u001b[2m\u001b[36m(pid=1492090)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1492090)\u001b[0m   from pandas import MultiIndex, Int64Index\n"
+      "\r"
      ]
     }
    ],
@@ -510,7 +490,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.16"
+   "version": "3.11.2"
   },
   "vscode": {
    "interpreter": {
diff --git a/doc/source/ray-air/examples/opt_deepspeed_batch_inference.ipynb b/doc/source/ray-air/examples/opt_deepspeed_batch_inference.ipynb
deleted file mode 100644
index ca4b5c787fa3..000000000000
--- a/doc/source/ray-air/examples/opt_deepspeed_batch_inference.ipynb
+++ /dev/null
@@ -1,911 +0,0 @@
-{
- "cells": [
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "id": "dfdf1047",
-   "metadata": {},
-   "source": [
-    "# Batch Inference with OPT 30B and Ray Data\n",
-    "\n",
-    "This notebook was tested on a single p3.16xlarge instance with 8 V100 GPUs.\n",
-    "\n",
-    "## Set Up\n",
-    "Initialize Ray and a runtime environment to ensure that all dependent packages are available."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "36bb842b-b6b6-4cbc-a4f9-a3a65ec069ce",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2023-04-22 11:12:15,071\tINFO worker.py:1314 -- Using address localhost:9031 set in the environment variable RAY_ADDRESS\n",
-      "fatal: not a git repository (or any parent up to mount point /home/ray)\n",
-      "Stopping at filesystem boundary (GIT_DISCOVERY_ACROSS_FILESYSTEM not set).\n",
-      "2023-04-22 11:12:15,676\tINFO worker.py:1432 -- Connecting to existing Ray cluster at address: 172.31.244.129:9031...\n",
-      "2023-04-22 11:12:15,724\tINFO worker.py:1607 -- Connected to Ray cluster. View the dashboard at https://console.anyscale.com/api/v2/sessions/ses_jgkdnu2723aleytwqqhebr12vs/services?redirect_to=dashboard \n",
-      "2023-04-22 11:12:15,732\tINFO packaging.py:347 -- Pushing file package 'gcs://_ray_pkg_7ad665e3661cefc8f8037daeb0b5ba6e.zip' (0.03MiB) to Ray cluster...\n",
-      "2023-04-22 11:12:15,733\tINFO packaging.py:360 -- Successfully pushed file package 'gcs://_ray_pkg_7ad665e3661cefc8f8037daeb0b5ba6e.zip'.\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "<div>\n",
-       "    <div style=\"margin-left: 50px;display: flex;flex-direction: row;align-items: center\">\n",
-       "        <h3 style=\"color: var(--jp-ui-font-color0)\">Ray</h3>\n",
-       "        <svg version=\"1.1\" id=\"ray\" width=\"3em\" viewBox=\"0 0 144.5 144.6\" style=\"margin-left: 3em;margin-right: 3em\">\n",
-       "            <g id=\"layer-1\">\n",
-       "                <path fill=\"#00a2e9\" class=\"st0\" d=\"M97.3,77.2c-3.8-1.1-6.2,0.9-8.3,5.1c-3.5,6.8-9.9,9.9-17.4,9.6S58,88.1,54.8,81.2c-1.4-3-3-4-6.3-4.1\n",
-       "                    c-5.6-0.1-9.9,0.1-13.1,6.4c-3.8,7.6-13.6,10.2-21.8,7.6C5.2,88.4-0.4,80.5,0,71.7c0.1-8.4,5.7-15.8,13.8-18.2\n",
-       "                    c8.4-2.6,17.5,0.7,22.3,8c1.3,1.9,1.3,5.2,3.6,5.6c3.9,0.6,8,0.2,12,0.2c1.8,0,1.9-1.6,2.4-2.8c3.5-7.8,9.7-11.8,18-11.9\n",
-       "                    c8.2-0.1,14.4,3.9,17.8,11.4c1.3,2.8,2.9,3.6,5.7,3.3c1-0.1,2,0.1,3,0c2.8-0.5,6.4,1.7,8.1-2.7s-2.3-5.5-4.1-7.5\n",
-       "                    c-5.1-5.7-10.9-10.8-16.1-16.3C84,38,81.9,37.1,78,38.3C66.7,42,56.2,35.7,53,24.1C50.3,14,57.3,2.8,67.7,0.5\n",
-       "                    C78.4-2,89,4.7,91.5,15.3c0.1,0.3,0.1,0.5,0.2,0.8c0.7,3.4,0.7,6.9-0.8,9.8c-1.7,3.2-0.8,5,1.5,7.2c6.7,6.5,13.3,13,19.8,19.7\n",
-       "                    c1.8,1.8,3,2.1,5.5,1.2c9.1-3.4,17.9-0.6,23.4,7c4.8,6.9,4.6,16.1-0.4,22.9c-5.4,7.2-14.2,9.9-23.1,6.5c-2.3-0.9-3.5-0.6-5.1,1.1\n",
-       "                    c-6.7,6.9-13.6,13.7-20.5,20.4c-1.8,1.8-2.5,3.2-1.4,5.9c3.5,8.7,0.3,18.6-7.7,23.6c-7.9,5-18.2,3.8-24.8-2.9\n",
-       "                    c-6.4-6.4-7.4-16.2-2.5-24.3c4.9-7.8,14.5-11,23.1-7.8c3,1.1,4.7,0.5,6.9-1.7C91.7,98.4,98,92.3,104.2,86c1.6-1.6,4.1-2.7,2.6-6.2\n",
-       "                    c-1.4-3.3-3.8-2.5-6.2-2.6C99.8,77.2,98.9,77.2,97.3,77.2z M72.1,29.7c5.5,0.1,9.9-4.3,10-9.8c0-0.1,0-0.2,0-0.3\n",
-       "                    C81.8,14,77,9.8,71.5,10.2c-5,0.3-9,4.2-9.3,9.2c-0.2,5.5,4,10.1,9.5,10.3C71.8,29.7,72,29.7,72.1,29.7z M72.3,62.3\n",
-       "                    c-5.4-0.1-9.9,4.2-10.1,9.7c0,0.2,0,0.3,0,0.5c0.2,5.4,4.5,9.7,9.9,10c5.1,0.1,9.9-4.7,10.1-9.8c0.2-5.5-4-10-9.5-10.3\n",
-       "                    C72.6,62.3,72.4,62.3,72.3,62.3z M115,72.5c0.1,5.4,4.5,9.7,9.8,9.9c5.6-0.2,10-4.8,10-10.4c-0.2-5.4-4.6-9.7-10-9.7\n",
-       "                    c-5.3-0.1-9.8,4.2-9.9,9.5C115,72.1,115,72.3,115,72.5z M19.5,62.3c-5.4,0.1-9.8,4.4-10,9.8c-0.1,5.1,5.2,10.4,10.2,10.3\n",
-       "                    c5.6-0.2,10-4.9,9.8-10.5c-0.1-5.4-4.5-9.7-9.9-9.6C19.6,62.3,19.5,62.3,19.5,62.3z M71.8,134.6c5.9,0.2,10.3-3.9,10.4-9.6\n",
-       "                    c0.5-5.5-3.6-10.4-9.1-10.8c-5.5-0.5-10.4,3.6-10.8,9.1c0,0.5,0,0.9,0,1.4c-0.2,5.3,4,9.8,9.3,10\n",
-       "                    C71.6,134.6,71.7,134.6,71.8,134.6z\"/>\n",
-       "            </g>\n",
-       "        </svg>\n",
-       "        <table>\n",
-       "            <tr>\n",
-       "                <td style=\"text-align: left\"><b>Python version:</b></td>\n",
-       "                <td style=\"text-align: left\"><b>3.9.15</b></td>\n",
-       "            </tr>\n",
-       "            <tr>\n",
-       "                <td style=\"text-align: left\"><b>Ray version:</b></td>\n",
-       "                <td style=\"text-align: left\"><b> 2.6.0</b></td>\n",
-       "            </tr>\n",
-       "            <tr>\n",
-       "    <td style=\"text-align: left\"><b>Dashboard:</b></td>\n",
-       "    <td style=\"text-align: left\"><b><a href=\"http://console.anyscale.com/api/v2/sessions/ses_jgkdnu2723aleytwqqhebr12vs/services?redirect_to=dashboard\" target=\"_blank\">http://console.anyscale.com/api/v2/sessions/ses_jgkdnu2723aleytwqqhebr12vs/services?redirect_to=dashboard</a></b></td>\n",
-       "</tr>\n",
-       "\n",
-       "        </table>\n",
-       "    </div>\n",
-       "</div>\n"
-      ],
-      "text/plain": [
-       "RayContext(dashboard_url='console.anyscale.com/api/v2/sessions/ses_jgkdnu2723aleytwqqhebr12vs/services?redirect_to=dashboard', python_version='3.9.15', ray_version='2.6.0', ray_commit='17df2ef17983406bb178c251044c9dc654b378c0', address_info={'node_ip_address': '172.31.244.129', 'raylet_ip_address': '172.31.244.129', 'redis_address': None, 'object_store_address': '/tmp/ray/session_2023-04-22_11-09-11_790337_150/sockets/plasma_store', 'raylet_socket_name': '/tmp/ray/session_2023-04-22_11-09-11_790337_150/sockets/raylet', 'webui_url': 'console.anyscale.com/api/v2/sessions/ses_jgkdnu2723aleytwqqhebr12vs/services?redirect_to=dashboard', 'session_dir': '/tmp/ray/session_2023-04-22_11-09-11_790337_150', 'metrics_export_port': 61073, 'gcs_address': '172.31.244.129:9031', 'address': '172.31.244.129:9031', 'dashboard_agent_listen_port': 52365, 'node_id': 'e6e9dfeda4469dd816c080bec2cf1cd12abdd978ae74b87e869164eb'})"
-      ]
-     },
-     "execution_count": 1,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "import ray\n",
-    "\n",
-    "ray.init(\n",
-    "    runtime_env={\n",
-    "        \"pip\": [\n",
-    "            \"numpy==1.23\",\n",
-    "            \"protobuf==3.20.0\",\n",
-    "            \"transformers==4.27.2\",\n",
-    "            \"accelerate==0.17.1\",\n",
-    "            \"deepspeed==0.8.3\",\n",
-    "        ],\n",
-    "        \"env_vars\": {\n",
-    "            \"HF_HUB_DISABLE_PROGRESS_BARS\": \"1\",\n",
-    "        }\n",
-    "    }\n",
-    ")"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "id": "b619a878",
-   "metadata": {},
-   "source": [
-    "## Define Hyperparameters\n",
-    "\n",
-    "Define a list of hyperparameters as a global dataclass.\n",
-    "\n",
-    "Refer to https://deepspeed.readthedocs.io/en/stable/inference-init.html#deepspeed.inference.config.DeepSpeedInferenceConfig for more details about the configurations of a DeepSpeed inference job."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "id": "613df744",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "from dataclasses import dataclass\n",
-    "from typing import Optional\n",
-    "\n",
-    "\n",
-    "@dataclass\n",
-    "class Config:\n",
-    "    model_name: str = \"facebook/opt-30b\"\n",
-    "    # Path to HuggingFace cache directory. Default is ~/.cache/huggingface/.\n",
-    "    cache_dir: Optional[str] = None\n",
-    "    # Path to the directory that actually holds model files.\n",
-    "    # e.g., ~/.cache/huggingface/models--facebook--opt-30b/snapshots/xxx/\n",
-    "    # If this path is not None, we skip download models from HuggingFace.\n",
-    "    repo_root: Optional[str] = None\n",
-    "    # This is how many DeepSpeed-inference replicas to run for\n",
-    "    # this batch inference job.\n",
-    "    num_worker_groups: int = 1\n",
-    "    # Number of DeepSpeed workers per group.\n",
-    "    num_workers_per_group: int = 8\n",
-    "\n",
-    "    batch_size: int = 1\n",
-    "    dtype: str = \"float16\"\n",
-    "    # Maximum number of tokens DeepSpeed inference-engine can work with,\n",
-    "    # including the input and output tokens.\n",
-    "    max_tokens: int = 1024\n",
-    "    # Use meta tensors to initialize model.\n",
-    "    use_meta_tensor: bool = True\n",
-    "    # Use cache for generation.\n",
-    "    use_cache: bool = True\n",
-    "    # The path for which we want to save the loaded model with a checkpoint.\n",
-    "    save_mp_checkpoint_path: Optional[str] = None\n",
-    "\n",
-    "\n",
-    "config = Config()"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "id": "28df05bf",
-   "metadata": {},
-   "source": [
-    "## Download and Cache Model\n",
-    "\n",
-    "Next, we will download and cache model files on all instances of the cluster before we run the job.\n",
-    "\n",
-    "Notice that when we download model snapshots from HuggingFace, we skip files that end with safetensors, msgpack, and h5 extensions. These are Tensorflow and JAX weight files. We only need PyTorch weights for this example.\n",
-    "\n",
-    "We execute the ``download_model()`` function on every node of the cluster by using a ``NodeAffinitySchedulingStrategy`` from Ray Core."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "id": "63b8a84d-57a6-4430-8fe8-9811760b8b7c",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Caching model locally ...\n",
-      "Done. Model saved in /home/ray/.cache/huggingface/hub/models--facebook--opt-30b/snapshots/ceea0a90ac0f6fae7c2c34bcb40477438c152546\n"
-     ]
-    }
-   ],
-   "source": [
-    "\n",
-    "from huggingface_hub import snapshot_download\n",
-    "import ray\n",
-    "from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy\n",
-    "\n",
-    "\n",
-    "@ray.remote\n",
-    "def download_model(config: Config):\n",
-    "    # This function downloads the specified HF model into a local directory.\n",
-    "    # This can also download models from cloud storages like S3.\n",
-    "    return snapshot_download(\n",
-    "        repo_id=config.model_name,\n",
-    "        cache_dir=config.cache_dir,\n",
-    "        allow_patterns=[\"*\"],\n",
-    "        # Skip downloading TF and FLAX weight files.\n",
-    "        ignore_patterns=[\"*.safetensors\", \"*.msgpack\", \"*.h5\"],\n",
-    "        revision=None,\n",
-    "    )\n",
-    "\n",
-    "if config.repo_root is None:\n",
-    "    # Download model files to all GPU nodes, and set correct repo_root.\n",
-    "    refs = []\n",
-    "    for node in ray.nodes():\n",
-    "        if node[\"Alive\"] and node[\"Resources\"].get(\"GPU\", None):\n",
-    "            node_id = node[\"NodeID\"]\n",
-    "            scheduling_strategy = NodeAffinitySchedulingStrategy(\n",
-    "                node_id=node_id, soft=False\n",
-    "            )\n",
-    "            options = {\"scheduling_strategy\": scheduling_strategy}\n",
-    "            refs.append(\n",
-    "                download_model.options(scheduling_strategy=scheduling_strategy).remote(config)\n",
-    "            )\n",
-    "\n",
-    "    print(\"Caching model locally ...\")\n",
-    "\n",
-    "    # Wait for models to finish downloading.\n",
-    "    config.repo_root = ray.get(refs)[0]\n",
-    "\n",
-    "    print(f\"Done. Model saved in {config.repo_root}\")\n",
-    "else:\n",
-    "    print(f\"Using existing model saved in {config.repo_root}\")"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "id": "6b14b7d9",
-   "metadata": {},
-   "source": [
-    "## Define DeepSpeed Utility Classes\n",
-    "\n",
-    "Next, we define a few utility classes and functions that are useful for setting up and running the DeepSpeed inference job.\n",
-    "\n",
-    "Note that the Pipeline is modeled after https://github.com/microsoft/DeepSpeedExamples/tree/efacebb3ddbea86bb20c3af30fd060be0fa41ac8/inference/huggingface/text-generation."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "id": "f9aad2a9",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/ray/anaconda3/lib/python3.9/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "  from pandas import MultiIndex, Int64Index\n"
-     ]
-    }
-   ],
-   "source": [
-    "import gc\n",
-    "import io\n",
-    "import json\n",
-    "import math\n",
-    "import os\n",
-    "from pathlib import Path\n",
-    "from typing import List\n",
-    "\n",
-    "import deepspeed\n",
-    "import torch\n",
-    "from deepspeed.runtime.utils import see_memory_usage\n",
-    "from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer\n",
-    "\n",
-    "\n",
-    "class DSPipeline:\n",
-    "    \"\"\"\n",
-    "    Example helper class for comprehending DeepSpeed Meta Tensors, meant to mimic HF pipelines.\n",
-    "    The DSPipeline can run with and without meta tensors.\n",
-    "    \"\"\"\n",
-    "\n",
-    "    def __init__(\n",
-    "        self,\n",
-    "        model_name,\n",
-    "        dtype=torch.float16,\n",
-    "        is_meta=True,\n",
-    "        device=-1,\n",
-    "        repo_root=None,\n",
-    "    ):\n",
-    "        self.model_name = model_name\n",
-    "        self.dtype = dtype\n",
-    "\n",
-    "        if isinstance(device, torch.device):\n",
-    "            self.device = device\n",
-    "        elif isinstance(device, str):\n",
-    "            self.device = torch.device(device)\n",
-    "        elif device < 0:\n",
-    "            self.device = torch.device(\"cpu\")\n",
-    "        else:\n",
-    "            self.device = torch.device(f\"cuda:{device}\")\n",
-    "\n",
-    "        self.tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side=\"right\")\n",
-    "        self.tokenizer.pad_token = self.tokenizer.eos_token\n",
-    "\n",
-    "        if is_meta:\n",
-    "            # When meta tensors enabled, use checkpoints\n",
-    "            self.config = AutoConfig.from_pretrained(self.model_name)\n",
-    "            self.checkpoints_json = self._generate_json(repo_root)\n",
-    "\n",
-    "            with deepspeed.OnDevice(dtype=dtype, device=\"meta\"):\n",
-    "                self.model = AutoModelForCausalLM.from_config(self.config)\n",
-    "        else:\n",
-    "            self.model = AutoModelForCausalLM.from_pretrained(self.model_name)\n",
-    "\n",
-    "        self.model.eval()\n",
-    "\n",
-    "    def __call__(self, inputs, **kwargs):\n",
-    "        input_list = [inputs] if isinstance(inputs, str) else inputs\n",
-    "        outputs = self.generate_outputs(input_list, **kwargs)\n",
-    "        return outputs\n",
-    "\n",
-    "    def _generate_json(self, repo_root):\n",
-    "        if os.path.exists(os.path.join(repo_root, \"ds_inference_config.json\")):\n",
-    "            # Simply use the available inference config.\n",
-    "            return os.path.join(repo_root, \"ds_inference_config.json\")\n",
-    "\n",
-    "        # Write a checkpoints config file in local directory.\n",
-    "        checkpoints_json = \"checkpoints.json\"\n",
-    "\n",
-    "        with io.open(checkpoints_json, \"w\", encoding=\"utf-8\") as f:\n",
-    "            file_list = [\n",
-    "                str(entry).split(\"/\")[-1]\n",
-    "                for entry in Path(repo_root).rglob(\"*.[bp][it][n]\")\n",
-    "                if entry.is_file()\n",
-    "            ]\n",
-    "            data = {\n",
-    "                # Hardcode bloom for now.\n",
-    "                # Possible choices are \"bloom\", \"ds_model\", \"Megatron\".\n",
-    "                \"type\": \"bloom\",\n",
-    "                \"checkpoints\": file_list,\n",
-    "                \"version\": 1.0\n",
-    "            }\n",
-    "            json.dump(data, f)\n",
-    "\n",
-    "        return checkpoints_json\n",
-    "\n",
-    "    def generate_outputs(self, inputs, **generate_kwargs):\n",
-    "        input_tokens = self.tokenizer.batch_encode_plus(\n",
-    "            inputs, return_tensors=\"pt\", padding=True\n",
-    "        )\n",
-    "        for t in input_tokens:\n",
-    "            if torch.is_tensor(input_tokens[t]):\n",
-    "                input_tokens[t] = input_tokens[t].to(self.device)\n",
-    "\n",
-    "        self.model.cuda().to(self.device)\n",
-    "\n",
-    "        outputs = self.model.generate(**input_tokens, **generate_kwargs)\n",
-    "        outputs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)\n",
-    "\n",
-    "        return outputs\n",
-    "\n",
-    "\n",
-    "def _memory_usage(gpu_id: int, msg: str):\n",
-    "    \"\"\"Print memory usage.\"\"\"\n",
-    "    if gpu_id != 0:\n",
-    "        return\n",
-    "    see_memory_usage(msg, True)\n",
-    "\n",
-    "\n",
-    "def init_model(config: Config, world_size: int, gpu_id: int) -> DSPipeline:\n",
-    "    \"\"\"Initialize the deepspeed model.\"\"\"\n",
-    "    data_type = getattr(torch, config.dtype)\n",
-    "\n",
-    "    _memory_usage(gpu_id, \"before init\")\n",
-    "    pipe = DSPipeline(\n",
-    "        model_name=config.model_name,\n",
-    "        dtype=data_type,\n",
-    "        is_meta=config.use_meta_tensor,\n",
-    "        device=gpu_id,\n",
-    "        repo_root=config.repo_root,\n",
-    "    )\n",
-    "    _memory_usage(gpu_id, \"after init\")\n",
-    "\n",
-    "    if config.use_meta_tensor:\n",
-    "        ds_kwargs = dict(\n",
-    "            base_dir=config.repo_root, checkpoint=pipe.checkpoints_json\n",
-    "        )\n",
-    "    else:\n",
-    "        ds_kwargs = dict()\n",
-    "\n",
-    "    gc.collect()\n",
-    "\n",
-    "    pipe.model = deepspeed.init_inference(\n",
-    "        pipe.model,\n",
-    "        dtype=data_type,\n",
-    "        mp_size=world_size,\n",
-    "        replace_with_kernel_inject=True,\n",
-    "        replace_method=True,\n",
-    "        max_tokens=config.max_tokens,\n",
-    "        save_mp_checkpoint_path=config.save_mp_checkpoint_path,\n",
-    "        **ds_kwargs,\n",
-    "    )\n",
-    "    _memory_usage(gpu_id, \"after init_inference\")\n",
-    "\n",
-    "    return pipe\n",
-    "\n",
-    "\n",
-    "def generate(\n",
-    "    input_sentences: List[str], pipe: DSPipeline, batch_size: int, **generate_kwargs\n",
-    ") -> List[str]:\n",
-    "    \"\"\"Generate predictions using a DSPipeline.\"\"\"\n",
-    "    if batch_size > len(input_sentences):\n",
-    "        # Dynamically extend to support larger bs by repetition.\n",
-    "        input_sentences *= math.ceil(batch_size / len(input_sentences))\n",
-    "\n",
-    "    inputs = input_sentences[:batch_size]\n",
-    "    outputs = pipe(inputs, **generate_kwargs)\n",
-    "    return outputs"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "bd20d4d9",
-   "metadata": {},
-   "source": []
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "id": "62eee91d",
-   "metadata": {},
-   "source": [
-    "## Define a DeepSpeed Predictor\n",
-    "\n",
-    "Define an AIR Predictor to be instantiated by the Dataset pipeline below.\n",
-    "\n",
-    "Each DeepSpeedPredictor is a stateful Ray actor that understands how to process the input prompt using a group of DeepSpeed inference workers.\n",
-    "\n",
-    "More specifically, each DeepSpeedPredictor sets up a proper PyTorch DDP process group before spinning up multiple PredictionWorkers. Since the model is loaded using the DeepSpeed inference framework, each PredictionWorker handles a shard of the entire DeepSpeed inference model.\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "id": "516a200d-14e4-4b52-a615-e09778ba4117",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "from typing import List\n",
-    "\n",
-    "import pandas as pd\n",
-    "import ray\n",
-    "import ray.util\n",
-    "from ray.air import Checkpoint, ScalingConfig\n",
-    "from ray.air.util.torch_dist import (\n",
-    "    TorchDistributedWorker,\n",
-    "    init_torch_dist_process_group,\n",
-    "    shutdown_torch_dist_process_group,\n",
-    ")\n",
-    "from ray.train.predictor import Predictor\n",
-    "from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy\n",
-    "\n",
-    "\n",
-    "@ray.remote\n",
-    "class PredictionWorker(TorchDistributedWorker):\n",
-    "    \"\"\"A PredictionWorker is a Ray remote actor that runs a single shard of a DeepSpeed job.\n",
-    "    \n",
-    "    Multiple PredictionWorkers of the same WorkerGroup form a PyTorch DDP process\n",
-    "    group and work together under the orchestration of DeepSpeed.\n",
-    "    \"\"\"\n",
-    "    def __init__(self, config: Config, world_size: int):\n",
-    "        self.config = config\n",
-    "        self.world_size = world_size\n",
-    "\n",
-    "    def init_model(self, local_rank: int):\n",
-    "        \"\"\"Initialize model for inference.\"\"\"\n",
-    "        # Note: We have to provide the local_rank that was used to initiate\n",
-    "        # the DDP process group here. e.g., a PredictionWorker may be the\n",
-    "        # rank 0 worker of a group, but occupies gpu 7.\n",
-    "        self.generator = init_model(self.config, self.world_size, local_rank)\n",
-    "\n",
-    "    def generate(self, data: pd.DataFrame, column: str, **kwargs) -> List[str]:\n",
-    "        return generate(\n",
-    "            list(data[column]), self.generator, self.config.batch_size, **kwargs\n",
-    "        )\n",
-    "\n",
-    "\n",
-    "# TODO: This Predictor should be part of Ray AIR.\n",
-    "class DeepSpeedPredictor(Predictor):\n",
-    "    def __init__(self, checkpoint: Checkpoint, scaling_config: ScalingConfig) -> None:\n",
-    "        self.checkpoint = checkpoint\n",
-    "        self.scaling_config = scaling_config\n",
-    "        self.init_worker_group(scaling_config)\n",
-    "\n",
-    "    def __del__(self):\n",
-    "        shutdown_torch_dist_process_group(self.prediction_workers)\n",
-    "\n",
-    "    def init_worker_group(self, scaling_config: ScalingConfig):\n",
-    "        \"\"\"Create the worker group.\n",
-    "\n",
-    "        Each worker in the group communicates with other workers through the\n",
-    "        torch distributed backend. The worker group is inelastic (a failure of\n",
-    "        one worker destroys the entire group). Each worker in the group\n",
-    "        recieves the same input data and outputs the same generated text.\n",
-    "        \"\"\"\n",
-    "        config = self.checkpoint.to_dict()[\"config\"]\n",
-    "\n",
-    "        # Start a placement group for the workers.\n",
-    "        self.pg = scaling_config.as_placement_group_factory().to_placement_group()\n",
-    "        prediction_worker_cls = PredictionWorker.options(\n",
-    "            num_cpus=scaling_config.num_cpus_per_worker,\n",
-    "            num_gpus=scaling_config.num_gpus_per_worker,\n",
-    "            resources=scaling_config.additional_resources_per_worker,\n",
-    "            scheduling_strategy=PlacementGroupSchedulingStrategy(\n",
-    "                placement_group=self.pg, placement_group_capture_child_tasks=True\n",
-    "            ),\n",
-    "        )\n",
-    "        # Create the prediction workers.\n",
-    "        self.prediction_workers = [\n",
-    "            prediction_worker_cls.remote(config, scaling_config.num_workers)\n",
-    "            for i in range(scaling_config.num_workers)\n",
-    "        ]\n",
-    "\n",
-    "        # Initialize torch distributed process group for the workers.\n",
-    "        local_ranks = init_torch_dist_process_group(self.prediction_workers, backend=\"nccl\")\n",
-    "\n",
-    "        # Initialize the model on each worker.\n",
-    "        ray.get([\n",
-    "            worker.init_model.remote(local_rank)\n",
-    "            for worker, local_rank in zip(self.prediction_workers, local_ranks)\n",
-    "        ])\n",
-    "\n",
-    "    def _predict_pandas(\n",
-    "        self,\n",
-    "        data: pd.DataFrame,\n",
-    "        input_column: str = \"prompt\",\n",
-    "        output_column: str = \"output\",\n",
-    "        **kwargs\n",
-    "    ) -> pd.DataFrame:\n",
-    "        data_ref = ray.put(data)\n",
-    "        prediction = ray.get(\n",
-    "            [\n",
-    "                worker.generate.remote(data_ref, column=input_column, **kwargs)\n",
-    "                for worker in self.prediction_workers\n",
-    "            ]\n",
-    "        )[0]\n",
-    "\n",
-    "        return pd.DataFrame(prediction, columns=[output_column])\n",
-    "\n",
-    "    @classmethod\n",
-    "    def from_checkpoint(cls, checkpoint: Checkpoint, **kwargs) -> \"Predictor\":\n",
-    "        return cls(checkpoint=checkpoint, **kwargs)\n"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "id": "ca57e150",
-   "metadata": {},
-   "source": [
-    "## Create a Dataset Pipeline\n",
-    "\n",
-    "Finally, we connect all these pieces together, and use a BatchPredictor to run multiple copies of the DeepSpeedPredictor actors.\n",
-    "\n",
-    "This step helps parallelize our batch inference job and utilize all available resources in the cluster."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "48bf4a4f-0ac4-4e77-a05a-710d42e0dc4e",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2023-04-22 11:14:12,074\tWARNING dataset.py:4124 -- Deprecation warning: use Dataset.materialize() instead of fully_executed().\n",
-      "2023-04-22 11:14:12,079\tINFO streaming_executor.py:87 -- Executing DAG InputDataBuffer[Input] -> AllToAllOperator[Repartition] -> AllToAllOperator[RandomShuffle]\n",
-      "2023-04-22 11:14:12,081\tINFO streaming_executor.py:88 -- Execution config: ExecutionOptions(resource_limits=ExecutionResources(cpu=None, gpu=None, object_store_memory=None), locality_with_output=False, preserve_order=False, actor_locality_enabled=True, verbose_progress=False)\n",
-      "2023-04-22 11:14:12,082\tINFO streaming_executor.py:90 -- Tip: To enable per-operator progress reporting, set RAY_DATA_VERBOSE_PROGRESS=1.\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "- Repartition 1:   0%|          | 0/16 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "- RandomShuffle 3:   0%|          | 0/16 [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2023-04-22 11:14:12,680\tINFO streaming_executor.py:87 -- Executing DAG InputDataBuffer[Input] -> ActorPoolMapOperator[MapBatches(ScoringWrapper)]\n",
-      "2023-04-22 11:14:12,682\tINFO streaming_executor.py:88 -- Execution config: ExecutionOptions(resource_limits=ExecutionResources(cpu=None, gpu=None, object_store_memory=None), locality_with_output=False, preserve_order=False, actor_locality_enabled=True, verbose_progress=False)\n",
-      "2023-04-22 11:14:12,683\tINFO streaming_executor.py:90 -- Tip: To enable per-operator progress reporting, set RAY_DATA_VERBOSE_PROGRESS=1.\n",
-      "2023-04-22 11:14:12,785\tINFO actor_pool_map_operator.py:114 -- MapBatches(ScoringWrapper): Waiting for 1 pool actors to start...\n",
-      "(_MapWorker pid=7005) The cache for model files in Transformers v4.22.0 has been updated. Migrating your old cache. This is a one-time only operation. You can interrupt this and resume the migration later on by calling `transformers.utils.move_cache()`.\n",
-      "0it [00:00, ?it/s]05) \n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(PredictionWorker pid=10038) [2023-04-22 11:14:30,762] [INFO] [utils.py:829:see_memory_usage] before init\n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:14:30,762] [INFO] [utils.py:830:see_memory_usage] MA 0.0 GB         Max_MA 0.0 GB         CA 0.0 GB         Max_CA 0 GB \n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:14:30,762] [INFO] [utils.py:838:see_memory_usage] CPU Virtual Memory:  used = 11.63 GB, percent = 2.4%\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "(PredictionWorker pid=10040) --------------------------------------------------------------------------\n",
-      "(PredictionWorker pid=10040)                  Aim collects anonymous usage analytics.                 \n",
-      "(PredictionWorker pid=10040)                         Read how to opt-out here:                         \n",
-      "(PredictionWorker pid=10040)     https://aimstack.readthedocs.io/en/latest/community/telemetry.html    \n",
-      "(PredictionWorker pid=10040) --------------------------------------------------------------------------\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(PredictionWorker pid=10045) [2023-04-22 11:14:33,061] [INFO] [logging.py:93:log_dist] [Rank -1] DeepSpeed info: version=0.8.3, git-hash=unknown, git-branch=unknown\n",
-      "(PredictionWorker pid=10045) [2023-04-22 11:14:33,062] [WARNING] [config_utils.py:75:_process_deprecated_field] Config parameter replace_method is deprecated. This parameter is no longer needed, please remove from your call to DeepSpeed-inference\n",
-      "(PredictionWorker pid=10045) [2023-04-22 11:14:33,062] [WARNING] [config_utils.py:75:_process_deprecated_field] Config parameter mp_size is deprecated use tensor_parallel.tp_size instead\n",
-      "(PredictionWorker pid=10045) [2023-04-22 11:14:33,062] [INFO] [logging.py:93:log_dist] [Rank -1] quantize_bits = 8 mlp_extra_grouping = False, quantize_groups = 1\n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:14:33,074] [INFO] [utils.py:829:see_memory_usage] after init\n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:14:33,075] [INFO] [utils.py:830:see_memory_usage] MA 0.0 GB         Max_MA 0.0 GB         CA 0.0 GB         Max_CA 0 GB \n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:14:33,075] [INFO] [utils.py:838:see_memory_usage] CPU Virtual Memory:  used = 12.25 GB, percent = 2.6%\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "(PredictionWorker pid=10040) Using /home/ray/.cache/torch_extensions/py39_cu116 as PyTorch extensions root...\n",
-      "(PredictionWorker pid=10038) Creating extension directory /home/ray/.cache/torch_extensions/py39_cu116/transformer_inference...\n",
-      "(PredictionWorker pid=10038) Detected CUDA files, patching ldflags\n",
-      "(PredictionWorker pid=10038) Emitting ninja build file /home/ray/.cache/torch_extensions/py39_cu116/transformer_inference/build.ninja...\n",
-      "(PredictionWorker pid=10038) Building extension module transformer_inference...\n",
-      "(PredictionWorker pid=10038) Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(PredictionWorker pid=10038) [1/9] /usr/local/cuda/bin/nvcc  -DTORCH_EXTENSION_NAME=transformer_inference -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/includes -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/includes -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/torch/csrc/api/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/TH -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/THC -isystem /usr/local/cuda/include -isystem /home/ray/anaconda3/include/python3.9 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_70,code=compute_70 -gencode=arch=compute_70,code=sm_70 --compiler-options '-fPIC' -O3 --use_fast_math -std=c++14 -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ -U__CUDA_NO_HALF2_OPERATORS__ -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_70,code=compute_70 -c /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/dequantize.cu -o dequantize.cuda.o \n",
-      "(PredictionWorker pid=10038) [2/9] /usr/local/cuda/bin/nvcc  -DTORCH_EXTENSION_NAME=transformer_inference -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/includes -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/includes -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/torch/csrc/api/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/TH -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/THC -isystem /usr/local/cuda/include -isystem /home/ray/anaconda3/include/python3.9 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_70,code=compute_70 -gencode=arch=compute_70,code=sm_70 --compiler-options '-fPIC' -O3 --use_fast_math -std=c++14 -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ -U__CUDA_NO_HALF2_OPERATORS__ -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_70,code=compute_70 -c /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/relu.cu -o relu.cuda.o \n",
-      "(PredictionWorker pid=10038) [3/9] /usr/local/cuda/bin/nvcc  -DTORCH_EXTENSION_NAME=transformer_inference -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/includes -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/includes -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/torch/csrc/api/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/TH -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/THC -isystem /usr/local/cuda/include -isystem /home/ray/anaconda3/include/python3.9 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_70,code=compute_70 -gencode=arch=compute_70,code=sm_70 --compiler-options '-fPIC' -O3 --use_fast_math -std=c++14 -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ -U__CUDA_NO_HALF2_OPERATORS__ -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_70,code=compute_70 -c /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/apply_rotary_pos_emb.cu -o apply_rotary_pos_emb.cuda.o \n",
-      "(PredictionWorker pid=10038) [4/9] /usr/local/cuda/bin/nvcc  -DTORCH_EXTENSION_NAME=transformer_inference -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/includes -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/includes -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/torch/csrc/api/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/TH -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/THC -isystem /usr/local/cuda/include -isystem /home/ray/anaconda3/include/python3.9 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_70,code=compute_70 -gencode=arch=compute_70,code=sm_70 --compiler-options '-fPIC' -O3 --use_fast_math -std=c++14 -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ -U__CUDA_NO_HALF2_OPERATORS__ -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_70,code=compute_70 -c /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/transform.cu -o transform.cuda.o \n",
-      "(PredictionWorker pid=10038) /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/transform.cu(56): warning #177-D: variable \"lane\" was declared but never referenced\n",
-      "(PredictionWorker pid=10038) \n",
-      "(PredictionWorker pid=10038) /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/transform.cu(93): warning #177-D: variable \"half_dim\" was declared but never referenced\n",
-      "(PredictionWorker pid=10038) \n",
-      "(PredictionWorker pid=10038) /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/transform.cu(110): warning #177-D: variable \"vals_half\" was declared but never referenced\n",
-      "(PredictionWorker pid=10038) \n",
-      "(PredictionWorker pid=10038) /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/transform.cu(111): warning #177-D: variable \"output_half\" was declared but never referenced\n",
-      "(PredictionWorker pid=10038) \n",
-      "(PredictionWorker pid=10038) /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/transform.cu(128): warning #177-D: variable \"lane\" was declared but never referenced\n",
-      "(PredictionWorker pid=10038) \n",
-      "(PredictionWorker pid=10038) [5/9] /usr/local/cuda/bin/nvcc  -DTORCH_EXTENSION_NAME=transformer_inference -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/includes -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/includes -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/torch/csrc/api/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/TH -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/THC -isystem /usr/local/cuda/include -isystem /home/ray/anaconda3/include/python3.9 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_70,code=compute_70 -gencode=arch=compute_70,code=sm_70 --compiler-options '-fPIC' -O3 --use_fast_math -std=c++14 -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ -U__CUDA_NO_HALF2_OPERATORS__ -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_70,code=compute_70 -c /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/softmax.cu -o softmax.cuda.o \n",
-      "(PredictionWorker pid=10038) /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/softmax.cu(272): warning #177-D: variable \"alibi_offset\" was declared but never referenced\n",
-      "(PredictionWorker pid=10038) \n",
-      "(PredictionWorker pid=10038) /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/softmax.cu(427): warning #177-D: variable \"warp_num\" was declared but never referenced\n",
-      "(PredictionWorker pid=10038) \n",
-      "(PredictionWorker pid=10038) [6/9] /usr/local/cuda/bin/nvcc  -DTORCH_EXTENSION_NAME=transformer_inference -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/includes -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/includes -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/torch/csrc/api/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/TH -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/THC -isystem /usr/local/cuda/include -isystem /home/ray/anaconda3/include/python3.9 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_70,code=compute_70 -gencode=arch=compute_70,code=sm_70 --compiler-options '-fPIC' -O3 --use_fast_math -std=c++14 -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ -U__CUDA_NO_HALF2_OPERATORS__ -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_70,code=compute_70 -c /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/gelu.cu -o gelu.cuda.o \n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:14:33,250] [INFO] [logging.py:93:log_dist] [Rank -1] DeepSpeed info: version=0.8.3, git-hash=unknown, git-branch=unknown [repeated 7x across cluster] (Ray deduplicates logs by default. Set RAY_DEDUP_LOGS=0 to disable log deduplication, or see https://docs.ray.io/en/master/ray-observability/ray-logging.html#log-deduplication for more options.)\n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:14:33,251] [WARNING] [config_utils.py:75:_process_deprecated_field] Config parameter replace_method is deprecated. This parameter is no longer needed, please remove from your call to DeepSpeed-inference [repeated 7x across cluster]\n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:14:33,251] [WARNING] [config_utils.py:75:_process_deprecated_field] Config parameter mp_size is deprecated use tensor_parallel.tp_size instead [repeated 7x across cluster]\n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:14:33,251] [INFO] [logging.py:93:log_dist] [Rank -1] quantize_bits = 8 mlp_extra_grouping = False, quantize_groups = 1 [repeated 7x across cluster]\n",
-      "(PredictionWorker pid=10038) [7/9] /usr/local/cuda/bin/nvcc  -DTORCH_EXTENSION_NAME=transformer_inference -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/includes -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/includes -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/torch/csrc/api/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/TH -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/THC -isystem /usr/local/cuda/include -isystem /home/ray/anaconda3/include/python3.9 -D_GLIBCXX_USE_CXX11_ABI=0 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_70,code=compute_70 -gencode=arch=compute_70,code=sm_70 --compiler-options '-fPIC' -O3 --use_fast_math -std=c++14 -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ -U__CUDA_NO_HALF2_OPERATORS__ -gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_70,code=compute_70 -c /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/layer_norm.cu -o layer_norm.cuda.o \n",
-      "(PredictionWorker pid=10038) [8/9] c++ -MMD -MF pt_binding.o.d -DTORCH_EXTENSION_NAME=transformer_inference -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\\\"_gcc\\\" -DPYBIND11_STDLIB=\\\"_libstdcpp\\\" -DPYBIND11_BUILD_ABI=\\\"_cxxabi1011\\\" -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/includes -I/home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/includes -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/torch/csrc/api/include -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/TH -isystem /home/ray/anaconda3/lib/python3.9/site-packages/torch/include/THC -isystem /usr/local/cuda/include -isystem /home/ray/anaconda3/include/python3.9 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++14 -O3 -std=c++14 -g -Wno-reorder -c /home/ray/anaconda3/lib/python3.9/site-packages/deepspeed/ops/csrc/transformer/inference/csrc/pt_binding.cpp -o pt_binding.o \n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "(PredictionWorker pid=10038) Loading extension module transformer_inference...\n",
-      "(PredictionWorker pid=10041) -------------------------------------------------------------------------- [repeated 14x across cluster]\n",
-      "(PredictionWorker pid=10041)                  Aim collects anonymous usage analytics.                  [repeated 7x across cluster]\n",
-      "(PredictionWorker pid=10041)                         Read how to opt-out here:                          [repeated 7x across cluster]\n",
-      "(PredictionWorker pid=10041)     https://aimstack.readthedocs.io/en/latest/community/telemetry.html     [repeated 7x across cluster]\n",
-      "(PredictionWorker pid=10041) Using /home/ray/.cache/torch_extensions/py39_cu116 as PyTorch extensions root... [repeated 7x across cluster]\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(PredictionWorker pid=10038) [9/9] c++ pt_binding.o gelu.cuda.o relu.cuda.o layer_norm.cuda.o softmax.cuda.o dequantize.cuda.o apply_rotary_pos_emb.cuda.o transform.cuda.o -shared -lcurand -L/home/ray/anaconda3/lib/python3.9/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda_cu -ltorch_cuda_cpp -ltorch -ltorch_python -L/usr/local/cuda/lib64 -lcudart -o transformer_inference.so\n",
-      "(PredictionWorker pid=10038) Time to load transformer_inference op: 46.834928035736084 seconds\n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:15:21,799] [INFO] [logging.py:93:log_dist] [Rank 0] DeepSpeed-Inference config: {'layer_id': 0, 'hidden_size': 7168, 'intermediate_size': 28672, 'heads': 56, 'num_hidden_layers': -1, 'fp16': True, 'pre_layer_norm': True, 'local_rank': -1, 'stochastic_mode': False, 'epsilon': 1e-12, 'mp_size': 8, 'q_int8': False, 'scale_attention': True, 'triangular_masking': True, 'local_attention': False, 'window_size': 1, 'rotary_dim': -1, 'rotate_half': False, 'rotate_every_two': True, 'return_tuple': True, 'mlp_after_attn': True, 'mlp_act_func_type': <ActivationFuncType.ReLU: 2>, 'specialized_mode': False, 'training_mp_size': 1, 'bigscience_bloom': False, 'max_out_tokens': 1024, 'scale_attn_by_inverse_layer_idx': False, 'enable_qkv_quantization': False, 'use_mup': False, 'return_single_tuple': False}\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "(PredictionWorker pid=10040) No modifications detected for re-loaded extension module transformer_inference, skipping build step...\n",
-      "Loading 7 checkpoint shards:   0%|          | 0/7 [00:00<?, ?it/s]\n",
-      "Loading 7 checkpoint shards:  14%|█▍        | 1/7 [00:39<03:57, 39.57s/it]\n",
-      "(PredictionWorker pid=10041) Loading extension module transformer_inference... [repeated 15x across cluster]\n",
-      "(PredictionWorker pid=10041) Using /home/ray/.cache/torch_extensions/py39_cu116 as PyTorch extensions root... [repeated 8x across cluster]\n",
-      "(PredictionWorker pid=10041) No modifications detected for re-loaded extension module transformer_inference, skipping build step... [repeated 7x across cluster]\n",
-      "Loading 7 checkpoint shards:   0%|          | 0/7 [00:00<?, ?it/s] [repeated 7x across cluster]\n",
-      "Loading 7 checkpoint shards:  29%|██▊       | 2/7 [01:15<03:06, 37.25s/it] [repeated 8x across cluster]\n",
-      "Loading 7 checkpoint shards:  29%|██▊       | 2/7 [01:28<03:42, 44.58s/it] [repeated 7x across cluster]\n",
-      "Loading 7 checkpoint shards:  43%|████▎     | 3/7 [01:51<02:26, 36.73s/it]\n",
-      "Loading 7 checkpoint shards:  43%|████▎     | 3/7 [01:51<02:26, 36.56s/it]\n",
-      "Loading 7 checkpoint shards:  43%|████▎     | 3/7 [01:57<02:34, 38.58s/it] [repeated 5x across cluster]\n",
-      "Loading 7 checkpoint shards:  43%|████▎     | 3/7 [02:03<02:41, 40.32s/it]\n",
-      "Loading 7 checkpoint shards:  57%|█████▋    | 4/7 [02:24<01:45, 35.29s/it]\n",
-      "Loading 7 checkpoint shards:  57%|█████▋    | 4/7 [02:31<01:50, 36.96s/it] [repeated 6x across cluster]\n",
-      "Loading 7 checkpoint shards:  71%|███████▏  | 5/7 [02:59<01:09, 34.92s/it] [repeated 2x across cluster]\n",
-      "Loading 7 checkpoint shards:  86%|████████▌ | 6/7 [03:05<00:24, 24.84s/it] [repeated 10x across cluster]\n",
-      "Loading 7 checkpoint shards:  86%|████████▌ | 6/7 [03:10<00:25, 25.15s/it] [repeated 4x across cluster]\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(PredictionWorker pid=10044) checkpoint loading time at rank 6: 216.07904958724976 sec\n",
-      "(PredictionWorker pid=10040) Time to load transformer_inference op: 0.03857231140136719 seconds [repeated 15x across cluster]\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Loading 7 checkpoint shards: 100%|██████████| 7/7 [03:36<00:00, 30.87s/it]\n",
-      "Loading 7 checkpoint shards: 100%|██████████| 7/7 [03:36<00:00, 30.87s/it]\n",
-      "Loading 7 checkpoint shards: 100%|██████████| 7/7 [03:43<00:00, 31.88s/it] [repeated 6x across cluster]\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "(PredictionWorker pid=10040) checkpoint loading time at rank 1: 223.18208837509155 sec [repeated 6x across cluster]\n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:19:13,839] [INFO] [utils.py:829:see_memory_usage] after init_inference\n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:19:13,840] [INFO] [utils.py:830:see_memory_usage] MA 7.69 GB         Max_MA 7.69 GB         CA 7.83 GB         Max_CA 8 GB \n",
-      "(PredictionWorker pid=10038) [2023-04-22 11:19:13,840] [INFO] [utils.py:838:see_memory_usage] CPU Virtual Memory:  used = 22.22 GB, percent = 4.6%\n",
-      "(PredictionWorker pid=10039) [2023-04-22 11:19:13,840] [INFO] [utils.py:838:see_memory_usage] CPU Virtual Memory:  used = 22.22 GB, percent = 4.6%\n",
-      "(PredictionWorker pid=10038) ------------------------------------------------------\n",
-      "(PredictionWorker pid=10038) Free memory : 6.587830 (GigaBytes)  \n",
-      "(PredictionWorker pid=10038) Total memory: 15.781921 (GigaBytes)  \n",
-      "(PredictionWorker pid=10038) Requested memory: 0.601562 (GigaBytes) \n",
-      "(PredictionWorker pid=10038) Setting maximum total tokens (input + output) to 1024 \n",
-      "(PredictionWorker pid=10038) ------------------------------------------------------\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "(PredictionWorker pid=10040) 2023-04-22 11:19:26.855845: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64\n",
-      "(PredictionWorker pid=10040) 2023-04-22 11:19:26.856002: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /usr/local/nvidia/lib:/usr/local/nvidia/lib64\n",
-      "(PredictionWorker pid=10040) 2023-04-22 11:19:26.856022: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n",
-      "(PredictionWorker pid=10039) 2023-04-22 11:19:26.856022: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "                                               output\n",
-      "0   DeepSpeed is the one to go with. No need for a...\n",
-      "1   Testimonials:\\n\\nG. SACCHIOULAS (TX)\\n\\n\"We bo...\n",
-      "2   Testimonials\\n\\nI received my order today, I'm...\n",
-      "3   Testimonials\\n\\nWhat do our clients say about ...\n",
-      "4   How can you make them that high?\\nI edited the...\n",
-      "..                                                ...\n",
-      "59  Please complete the form below to request more...\n",
-      "60  DeepSpeed is the most popular way of dealing t...\n",
-      "61  How can you not tell that's not a real tweet?\\...\n",
-      "62  Testimonials\\n\\n\"The staff and community of H....\n",
-      "63  DeepSpeed is an independent, privately held co...\n",
-      "\n",
-      "[64 rows x 1 columns]\n",
-      "(autoscaler +12m27s) Tip: use `ray status` to view detailed cluster status. To disable these messages, set RAY_SCHEDULER_EVENTS=0.\n",
-      "(autoscaler +12m27s) Resized to 64 CPUs, 8 GPUs.\n"
-     ]
-    }
-   ],
-   "source": [
-    "import pandas as pd\n",
-    "import ray\n",
-    "from ray.air import Checkpoint, ScalingConfig\n",
-    "from ray.train.batch_predictor import BatchPredictor\n",
-    "\n",
-    "# Disable terminal progress bar for notebook environments.\n",
-    "ray.data.set_progress_bars(False)\n",
-    "\n",
-    "# Prompts.\n",
-    "# For testing purpose, we create 64 prompts in total.\n",
-    "df = pd.DataFrame(\n",
-    "    [\n",
-    "        \"DeepSpeed is\",\n",
-    "        \"Test\",\n",
-    "        \"Please complete\",\n",
-    "        \"How can you\"\n",
-    "    ] * 16,\n",
-    "    columns=[\"prompt\"]\n",
-    ")\n",
-    "ds = (\n",
-    "    ray.data.from_pandas(df)\n",
-    "    # Make sure there are enough blocks for parallelized execution.\n",
-    "    .repartition(config.num_workers_per_group * 2)\n",
-    "    .random_shuffle()\n",
-    "    .fully_executed()\n",
-    ")\n",
-    "\n",
-    "# Scaling config for one worker group.\n",
-    "group_scaling_config = ScalingConfig(\n",
-    "    use_gpu=True,\n",
-    "    num_workers=config.num_workers_per_group,\n",
-    "    # Should not be necessary after we switch to the new API.\n",
-    "    trainer_resources={\"CPU\": 0},\n",
-    ")\n",
-    "batch_predictor = BatchPredictor.from_checkpoint(\n",
-    "    # TODO: Use HugginFaceDeepSpeedCheckpoint when it's available.\n",
-    "    Checkpoint.from_dict({\"config\": config}),\n",
-    "    DeepSpeedPredictor,\n",
-    "    scaling_config=group_scaling_config,\n",
-    ")\n",
-    "\n",
-    "# Batch prediction.\n",
-    "pred = batch_predictor.predict(\n",
-    "    ds,\n",
-    "    batch_size=1,\n",
-    "    num_cpus_per_worker=0,\n",
-    "    min_scoring_workers=config.num_worker_groups,\n",
-    "    max_scoring_workers=config.num_worker_groups,\n",
-    "    # Kwargs passed to model.generate()\n",
-    "    do_sample=True,\n",
-    "    temperature=0.9,\n",
-    "    max_length=100,\n",
-    ")\n",
-    "\n",
-    "# Let's see the genreated texts.\n",
-    "print(pred.to_pandas())"
-   ]
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "codemirror_mode": {
-    "name": "ipython",
-    "version": 3
-   },
-   "file_extension": ".py",
-   "mimetype": "text/x-python",
-   "name": "python",
-   "nbconvert_exporter": "python",
-   "pygments_lexer": "ipython3",
-   "version": "3.11.1"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 5
-}
diff --git a/doc/source/ray-air/examples/pytorch_tabular_batch_prediction.py b/doc/source/ray-air/examples/pytorch_tabular_batch_prediction.py
deleted file mode 100644
index 771e33f9c9e8..000000000000
--- a/doc/source/ray-air/examples/pytorch_tabular_batch_prediction.py
+++ /dev/null
@@ -1,46 +0,0 @@
-import numpy as np
-import torch.nn as nn
-
-import ray
-from ray.data.preprocessors import Concatenator
-from ray.train.torch import TorchCheckpoint, TorchPredictor
-from ray.train.batch_predictor import BatchPredictor
-
-
-def create_model(input_features: int):
-    return nn.Sequential(
-        nn.Linear(in_features=input_features, out_features=16),
-        nn.ReLU(),
-        nn.Linear(16, 16),
-        nn.ReLU(),
-        nn.Linear(16, 1),
-        nn.Sigmoid(),
-    )
-
-
-dataset = ray.data.read_csv("s3://anonymous@air-example-data/breast_cancer.csv")
-
-# All columns are features except the target column.
-num_features = len(dataset.schema().names) - 1
-
-# Specify a preprocessor to concatenate all feature columns.
-prep = Concatenator(
-    output_column_name="concat_features", exclude=["target"], dtype=np.float32
-)
-
-checkpoint = TorchCheckpoint.from_model(
-    model=create_model(num_features), preprocessor=prep
-)
-# You can also fetch a checkpoint from a Trainer
-# checkpoint = best_result.checkpoint
-
-batch_predictor = BatchPredictor.from_checkpoint(checkpoint, TorchPredictor)
-
-# Predict on the features.
-predicted_probabilities = batch_predictor.predict(
-    dataset, feature_columns=["concat_features"]
-)
-# Call show on the output probabilities to trigger execution.
-predicted_probabilities.show()
-# {'predictions': array([1.], dtype=float32)}
-# {'predictions': array([0.], dtype=float32)}
diff --git a/doc/source/ray-air/examples/pytorch_tabular_starter.py b/doc/source/ray-air/examples/pytorch_tabular_starter.py
index 72654ab7d593..d113cbe9002d 100644
--- a/doc/source/ray-air/examples/pytorch_tabular_starter.py
+++ b/doc/source/ray-air/examples/pytorch_tabular_starter.py
@@ -127,22 +127,3 @@ def train_loop_per_worker(config):
 print("Best Result:", best_result)
 # Best Result: Result(metrics={'loss': 0.278409322102863, ...})
 # __air_tune_generic_end__
-
-# __air_pytorch_batchpred_start__
-from ray.train.batch_predictor import BatchPredictor
-from ray.train.torch import TorchPredictor
-
-# You can also create a checkpoint from a trained model using
-# `TorchCheckpoint.from_model`.
-checkpoint = best_result.checkpoint
-
-batch_predictor = BatchPredictor.from_checkpoint(
-    checkpoint, TorchPredictor, model=create_model(num_features)
-)
-
-predicted_probabilities = batch_predictor.predict(test_dataset)
-predicted_probabilities.show()
-# {'predictions': array([1.], dtype=float32)}
-# {'predictions': array([0.], dtype=float32)}
-# ...
-# __air_pytorch_batchpred_end__
diff --git a/doc/source/ray-air/examples/sklearn_example.ipynb b/doc/source/ray-air/examples/sklearn_example.ipynb
index 7d47f4f1390c..c2500b859d43 100644
--- a/doc/source/ray-air/examples/sklearn_example.ipynb
+++ b/doc/source/ray-air/examples/sklearn_example.ipynb
@@ -1,6 +1,7 @@
 {
  "cells": [
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "c3192ac4",
    "metadata": {},
@@ -10,6 +11,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "5a4823bf",
    "metadata": {},
@@ -32,6 +34,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "c049c692",
    "metadata": {},
@@ -51,7 +54,6 @@
     "\n",
     "import ray\n",
     "from ray.data import Dataset\n",
-    "from ray.train.batch_predictor import BatchPredictor\n",
     "from ray.train.sklearn import SklearnPredictor\n",
     "from ray.data.preprocessors import Chain, OrdinalEncoder, StandardScaler\n",
     "from ray.air.result import Result\n",
@@ -67,6 +69,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "52e017f1",
    "metadata": {},
@@ -89,6 +92,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "8d6c6d17",
    "metadata": {},
@@ -137,39 +141,7 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "id": "7a2efb9d",
-   "metadata": {},
-   "source": [
-    "Once we have the result, we can do batch inference on the obtained model. Let's define a utility function for this."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "id": "59eeadd8",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def predict_sklearn(result: Result, use_gpu: bool = False):\n",
-    "    _, _, test_dataset = prepare_data()\n",
-    "\n",
-    "    batch_predictor = BatchPredictor.from_checkpoint(\n",
-    "        result.checkpoint, SklearnPredictor\n",
-    "    )\n",
-    "\n",
-    "    predicted_labels = (\n",
-    "        batch_predictor.predict(\n",
-    "            test_dataset,\n",
-    "            num_gpus_per_worker=int(use_gpu),\n",
-    "        )\n",
-    "        .map_batches(lambda df: (df > 0.5).astype(int), batch_format=\"pandas\")\n",
-    "    )\n",
-    "    print(f\"PREDICTED LABELS\")\n",
-    "    predicted_labels.show()"
-   ]
-  },
-  {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "7d073994",
    "metadata": {},
@@ -339,70 +311,14 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
-   "id": "0daba603",
+   "id": "e11cf27b",
    "metadata": {},
    "source": [
-    "And perform inference on the obtained model:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "id": "24b16ede",
-   "metadata": {
-    "pycharm": {
-     "name": "#%%\n"
-    }
-   },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2022-06-22 17:27:59,658\tWARNING read_api.py:260 -- The number of blocks in this dataset (1) limits its parallelism to 1 concurrent tasks. This is much less than the number of available CPU slots in the cluster. Use `.repartition(n)` to increase the number of dataset blocks.\n",
-      "Map_Batches: 100%|██████████| 1/1 [00:00<00:00, 64.73it/s]\n",
-      "Map Progress (1 actors 1 pending): 100%|██████████| 1/1 [00:01<00:00,  1.60s/it]\n",
-      "Map_Batches: 100%|██████████| 1/1 [00:00<00:00, 71.41it/s]"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "PREDICTED LABELS\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 0}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 0}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 0}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 0}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 1}\n",
-      "{'predictions': 0}\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    }
-   ],
-   "source": [
-    "predict_sklearn(result, use_gpu=False)"
+    "## Next steps\n",
+    "\n",
+    "- {ref}`End-to-end: Offline Batch Inference <batch_inference_home>`"
    ]
   }
  ],
diff --git a/doc/source/ray-air/examples/tf_tabular_batch_prediction.py b/doc/source/ray-air/examples/tf_tabular_batch_prediction.py
deleted file mode 100644
index c4fae8c9a688..000000000000
--- a/doc/source/ray-air/examples/tf_tabular_batch_prediction.py
+++ /dev/null
@@ -1,49 +0,0 @@
-import numpy as np
-
-import ray
-from ray.data.preprocessors import Concatenator
-from ray.train.tensorflow import TensorflowCheckpoint, TensorflowPredictor
-from ray.train.batch_predictor import BatchPredictor
-
-
-def create_model(input_features):
-    from tensorflow import keras  # this is needed for tf<2.9
-    from tensorflow.keras import layers
-
-    return keras.Sequential(
-        [
-            keras.Input(shape=(input_features,)),
-            layers.Dense(16, activation="relu"),
-            layers.Dense(16, activation="relu"),
-            layers.Dense(1, activation="sigmoid"),
-        ]
-    )
-
-
-dataset = ray.data.read_csv("s3://anonymous@air-example-data/breast_cancer.csv")
-
-# All columns are features except the target column.
-num_features = len(dataset.schema().names) - 1
-
-# Specify a preprocessor to concatenate all feature columns.
-prep = Concatenator(
-    output_column_name="concat_features", exclude=["target"], dtype=np.float32
-)
-
-checkpoint = TensorflowCheckpoint.from_model(
-    model=create_model(num_features), preprocessor=prep
-)
-# You can also fetch a checkpoint from a Trainer
-# checkpoint = trainer.fit().checkpoint
-
-batch_predictor = BatchPredictor.from_checkpoint(
-    checkpoint, TensorflowPredictor, model_definition=lambda: create_model(num_features)
-)
-
-predicted_probabilities = batch_predictor.predict(
-    dataset, feature_columns=["concat_features"]
-)
-# Call show on the output probabilities to trigger execution.
-predicted_probabilities.show()
-# {'predictions': array([1.], dtype=float32)}
-# {'predictions': array([0.], dtype=float32)}
diff --git a/doc/source/ray-air/examples/tf_tabular_starter.py b/doc/source/ray-air/examples/tf_tabular_starter.py
index 56e66f2f60b3..3b862092d5ff 100644
--- a/doc/source/ray-air/examples/tf_tabular_starter.py
+++ b/doc/source/ray-air/examples/tf_tabular_starter.py
@@ -131,24 +131,3 @@ def train_loop_per_worker(config):
 print("Best Result:", best_result)
 # Best Result: Result(metrics={'loss': 4.997025489807129, ...)
 # __air_tune_generic_end__
-
-# __air_tf_batchpred_start__
-from ray.train.batch_predictor import BatchPredictor
-from ray.train.tensorflow import TensorflowPredictor
-
-# You can also create a checkpoint from a trained model using `TensorflowCheckpoint`.
-checkpoint = best_result.checkpoint
-
-batch_predictor = BatchPredictor.from_checkpoint(
-    checkpoint,
-    TensorflowPredictor,
-    model_definition=lambda: create_keras_model(num_features),
-)
-
-predicted_probabilities = batch_predictor.predict(test_dataset)
-predicted_probabilities.show()
-# {'predictions': 0.033036969602108}
-# {'predictions': 0.05944341793656349}
-# {'predictions': 0.1657751202583313}
-# ...
-# __air_tf_batchpred_end__
diff --git a/doc/source/ray-air/examples/torch_detection.ipynb b/doc/source/ray-air/examples/torch_detection.ipynb
index 1daaa8631657..5095caa7b77b 100644
--- a/doc/source/ray-air/examples/torch_detection.ipynb
+++ b/doc/source/ray-air/examples/torch_detection.ipynb
@@ -837,155 +837,12 @@
         },
         {
             "cell_type": "markdown",
-            "id": "224a1139",
+            "id": "838101c2",
             "metadata": {},
             "source": [
-                "## Evaluate the model on test data\n",
+                "## Next steps\n",
                 "\n",
-                "Now that you've fine-tuned the model, you'll evaluate it on the test data.\n",
-                "\n",
-                "### Generate predictions on the test data"
-            ]
-        },
-        {
-            "cell_type": "markdown",
-            "id": "1fc9bac2",
-            "metadata": {},
-            "source": [
-                "`Predictors` let you perform scalable [batch prediction](batch-prediction) and\n",
-                "[online inference](air-serving-guide). To evaluate the model, you'll use\n",
-                "`BatchPredictor` to perform inference in a distributed fashion.\n",
-                "\n",
-                "Create a `BatchPredictor` and pass `TorchDetectionPredictor` to the constructor. Then,\n",
-                "call `BatchPredictor.predict` to detect objects in the test data."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 14,
-            "id": "cc3cc662",
-            "metadata": {},
-            "outputs": [
-                {
-                    "name": "stderr",
-                    "output_type": "stream",
-                    "text": [
-                        "2023-03-01 13:08:48,113\tINFO batch_predictor.py:214 -- `num_gpus_per_worker` is set for `BatchPreditor`.Automatically enabling GPU prediction for this predictor. To disable set `use_gpu` to `False` in `BatchPredictor.predict`.\n",
-                        "2023-03-01 13:08:48,945\tINFO bulk_executor.py:41 -- Executing DAG InputDataBuffer[Input] -> TaskPoolMapOperator[TorchVisionPreprocessor] -> ActorPoolMapOperator[MapBatches(ScoringWrapper)]\n",
-                        "TorchVisionPreprocessor: 100%|██████████| 26/26 [00:17<00:00,  1.49it/s]\n",
-                        "MapBatches(ScoringWrapper), 0 actors [26 locality hits, 0 misses]: 100%|██████████| 26/26 [00:32<00:00,  1.25s/it]           \n"
-                    ]
-                },
-                {
-                    "data": {
-                        "application/vnd.jupyter.widget-view+json": {
-                            "model_id": "468b32006b5f440dae152b288d84d5d3",
-                            "version_major": 2,
-                            "version_minor": 0
-                        },
-                        "text/plain": [
-                            "VBox(children=(HTML(value='<h2>Dataset</h2>'), Tab(children=(HTML(value='<div class=\"scrollableTable jp-Render…"
-                        ]
-                    },
-                    "metadata": {},
-                    "output_type": "display_data"
-                }
-            ],
-            "source": [
-                "from ray.train.batch_predictor import BatchPredictor\n",
-                "from ray.train.torch import TorchDetectionPredictor\n",
-                "\n",
-                "\n",
-                "model = models.detection.fasterrcnn_resnet50_fpn(num_classes=3)\n",
-                "predictor = BatchPredictor.from_checkpoint(results.checkpoint, TorchDetectionPredictor, model=model)\n",
-                "\n",
-                "predictions = predictor.predict(\n",
-                "    test_dataset,\n",
-                "    feature_columns=[\"image\"],\n",
-                "    keep_columns=[\"boxes\", \"labels\"],\n",
-                "    batch_size=4,\n",
-                "    num_gpus_per_worker=1,\n",
-                ")\n",
-                "predictions"
-            ]
-        },
-        {
-            "cell_type": "markdown",
-            "id": "4f740d0a",
-            "metadata": {},
-            "source": [
-                "### Evaluate the model"
-            ]
-        },
-        {
-            "cell_type": "markdown",
-            "id": "491d8c12",
-            "metadata": {},
-            "source": [
-                "Once you've created the `predictions` dataset, iterate over the rows of the dataset\n",
-                "and compute the accuracy of the model."
-            ]
-        },
-        {
-            "cell_type": "code",
-            "execution_count": 15,
-            "id": "39ded656",
-            "metadata": {},
-            "outputs": [
-                {
-                    "name": "stderr",
-                    "output_type": "stream",
-                    "text": [
-                        "/tmp/ipykernel_160001/1622602304.py:8: UserWarning: The given NumPy array is not writable, and PyTorch does not support non-writable tensors. This means writing to this tensor will result in undefined behavior. You may want to copy the array to protect its data or make it writable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at ../torch/csrc/utils/tensor_numpy.cpp:199.)\n",
-                        "  \"boxes\": torch.as_tensor(row[\"pred_boxes\"]),\n"
-                    ]
-                },
-                {
-                    "data": {
-                        "text/plain": [
-                            "{'map': tensor(0.0898),\n",
-                            " 'map_50': tensor(0.3034),\n",
-                            " 'map_75': tensor(0.0244),\n",
-                            " 'map_small': tensor(0.),\n",
-                            " 'map_medium': tensor(0.0158),\n",
-                            " 'map_large': tensor(0.0964),\n",
-                            " 'mar_1': tensor(0.1799),\n",
-                            " 'mar_10': tensor(0.3551),\n",
-                            " 'mar_100': tensor(0.3635),\n",
-                            " 'mar_small': tensor(0.),\n",
-                            " 'mar_medium': tensor(0.1063),\n",
-                            " 'mar_large': tensor(0.3818),\n",
-                            " 'map_per_class': tensor(-1.),\n",
-                            " 'mar_100_per_class': tensor(-1.)}"
-                        ]
-                    },
-                    "execution_count": 15,
-                    "metadata": {},
-                    "output_type": "execute_result"
-                }
-            ],
-            "source": [
-                "from torchmetrics.detection.mean_ap import MeanAveragePrecision\n",
-                "\n",
-                "\n",
-                "metric = MeanAveragePrecision()\n",
-                "for row in predictions.iter_rows():\n",
-                "    preds = [\n",
-                "        {\n",
-                "            \"boxes\": torch.as_tensor(row[\"pred_boxes\"]),\n",
-                "            \"scores\": torch.as_tensor(row[\"pred_scores\"]),\n",
-                "            \"labels\": torch.as_tensor(row[\"pred_labels\"]),\n",
-                "        }\n",
-                "    ]\n",
-                "    target = [\n",
-                "        {\n",
-                "            \"boxes\": torch.as_tensor(row[\"boxes\"]),\n",
-                "            \"labels\": torch.as_tensor(row[\"labels\"]),\n",
-                "        }\n",
-                "    ]\n",
-                "    metric.update(preds, target)\n",
-                "\n",
-                "metric.compute()"
+                "- {ref}`End-to-end: Offline Batch Inference <batch_inference_home>`"
             ]
         }
     ],
diff --git a/doc/source/ray-air/examples/torch_image_batch_pretrained.py b/doc/source/ray-air/examples/torch_image_batch_pretrained.py
deleted file mode 100644
index dc0f4ed874ae..000000000000
--- a/doc/source/ray-air/examples/torch_image_batch_pretrained.py
+++ /dev/null
@@ -1,30 +0,0 @@
-from torchvision import transforms
-from torchvision.models import resnet18
-
-import ray
-from ray.train.torch import TorchCheckpoint, TorchPredictor
-from ray.train.batch_predictor import BatchPredictor
-from ray.data.preprocessors import TorchVisionPreprocessor
-
-
-data_url = "s3://anonymous@air-example-data-2/1G-image-data-synthetic-raw"
-print(f"Running GPU batch prediction with 1GB data from {data_url}")
-dataset = ray.data.read_images(data_url, size=(256, 256)).limit(10)
-
-model = resnet18(pretrained=True)
-
-transform = transforms.Compose(
-    [
-        transforms.ToTensor(),
-        transforms.CenterCrop(224),
-        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
-    ]
-)
-preprocessor = TorchVisionPreprocessor(columns=["image"], transform=transform)
-
-ckpt = TorchCheckpoint.from_model(model=model, preprocessor=preprocessor)
-
-predictor = BatchPredictor.from_checkpoint(ckpt, TorchPredictor)
-predictions = predictor.predict(dataset, batch_size=80, num_gpus_per_worker=1)
-# Call show on the output probabilities to trigger execution
-predictions.show()
diff --git a/doc/source/ray-air/examples/torch_incremental_learning.ipynb b/doc/source/ray-air/examples/torch_incremental_learning.ipynb
index 4bf5fd448fd2..b0cf82f99b1e 100644
--- a/doc/source/ray-air/examples/torch_incremental_learning.ipynb
+++ b/doc/source/ray-air/examples/torch_incremental_learning.ipynb
@@ -1,6 +1,7 @@
 {
  "cells": [
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "TsniIjjg2Pym"
@@ -10,6 +11,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "1VsUrzVm1W-h"
@@ -40,6 +42,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "Q3oGiuqYfj9_"
@@ -50,11 +53,11 @@
     "2. Create an `Iterator[ray.data.Dataset]` abstraction to represent a stream of data to train on for incremental training.\n",
     "3. Implement a custom Ray AIR preprocessor to preprocess the dataset.\n",
     "4. Incrementally train a model using data parallel training.\n",
-    "5. Use our trained model to perform batch prediction on test data.\n",
-    "6. Incrementally deploying our trained model with Ray Serve and performing online prediction queries."
+    "5. Incrementally deploying our trained model with Ray Serve and performing online prediction queries."
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "z52Y8O4q1bIk"
@@ -83,6 +86,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "RpD4STX3g1dq"
@@ -171,6 +175,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "AedcxD_FClQL"
@@ -221,6 +226,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "L2N1U22VC_N9"
@@ -237,6 +243,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "3SVSrkqrDJuc"
@@ -299,6 +306,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "vqrfgfl9YnVe"
@@ -394,6 +402,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "HDGHgtb699kd"
@@ -492,6 +501,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "9HUciluylZbX"
@@ -527,65 +537,13 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "Uto3v90Hagni"
-   },
-   "source": [
-    "## 4c: Define logic for Batch/Offline Prediction.\n",
-    "\n",
-    "After training on each task, we want to use our trained model to do batch (i.e. offline) inference on a test dataset. \n",
-    "\n",
-    "To do this, we leverage the built-in `ray.air.BatchPredictor`. We define a `batch_predict` function that will take in a Checkpoint and a Test Dataset and outputs the accuracy our model achieves on the test dataset."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {
-    "id": "DM2lFHzFa6uI"
-   },
-   "outputs": [],
-   "source": [
-    "from ray.train.batch_predictor import BatchPredictor\n",
-    "from ray.train.torch import TorchPredictor\n",
-    "\n",
-    "def batch_predict(checkpoint: ray.air.Checkpoint, test_dataset: ray.data.Dataset) -> float:\n",
-    "  \"\"\"Perform batch prediction on the provided test dataset, and return accuracy results.\"\"\"\n",
-    "\n",
-    "  batch_predictor = BatchPredictor.from_checkpoint(checkpoint, predictor_cls=TorchPredictor, model=SimpleMLP(num_classes=10))\n",
-    "  model_output = batch_predictor.predict(\n",
-    "            data=test_dataset, feature_columns=[\"image\"], keep_columns=[\"label\"]\n",
-    "        )\n",
-    "  \n",
-    "  # Postprocess model outputs.\n",
-    "  # Convert logits outputted from model into actual class predictions.\n",
-    "  def convert_logits_to_classes(df):\n",
-    "     best_class = df[\"predictions\"].map(lambda x: np.array(x).argmax())\n",
-    "     df[\"predictions\"] = best_class\n",
-    "     return df\n",
-    "    \n",
-    "  prediction_results = model_output.map_batches(convert_logits_to_classes, batch_format=\"pandas\")\n",
-    "  \n",
-    "  # Then, for each prediction output, see if it matches with the ground truth\n",
-    "  # label.\n",
-    "  def calculate_prediction_scores(df):\n",
-    "      return pd.DataFrame({\"correct\": df[\"predictions\"] == df[\"label\"]})\n",
-    "\n",
-    "  correct_dataset = prediction_results.map_batches(\n",
-    "      calculate_prediction_scores, batch_format=\"pandas\"\n",
-    "  )\n",
-    "\n",
-    "  return correct_dataset.sum(on=\"correct\") / correct_dataset.count()"
-   ]
-  },
-  {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "GWiTtsmVbIZP"
    },
    "source": [
-    "## 4d: Define logic for Deploying and Querying our model\n",
+    "## 4c: Define logic for Deploying and Querying our model\n",
     "\n",
     "In addition to batch inference, we also want to deploy our model so that we can submit live queries to it for online inference. We use Ray Serve's `PredictorDeployment` utility to deploy our trained model. \n",
     "\n",
@@ -636,6 +594,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "-NQDj0rFVUX3"
@@ -648,15 +607,14 @@
     "For each dataset in our stream, we do the following:\n",
     "1. Train on the dataset in Data Parallel fashion. We create a `TorchTrainer`, specify the config for the training loop we defined above, the dataset to train on, and how much we want to scale. `TorchTrainer` also accepts a `checkpoint` arg to continue training from a previously saved checkpoint.\n",
     "2. Get the saved checkpoint from the training run.\n",
-    "3. Test our trained model on a test set containing test data from all the tasks trained on so far.\n",
     "3. After training on each task, we deploy our model so we can query it for predictions.\n",
     "\n",
-    "In this example, the training and test data for each task is well-defined beforehand by the benchmark. For real-world scenarios, this probably will not be the case. It is very likely that the prediction requests after training on one task will become the training data for the next task. \n"
+    "In this example, the training data for each task is well-defined beforehand by the benchmark. For real-world scenarios, this probably will not be the case. It is very likely that the prediction requests after training on one task will become the training data for the next task. \n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": null,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/",
@@ -1390,22 +1348,6 @@
     "  result = trainer.fit()\n",
     "  latest_checkpoint = result.checkpoint\n",
     "\n",
-    "  # **************Batch Prediction**************************\n",
-    "\n",
-    "  # We can do batch prediction on the test data for the tasks seen so far.\n",
-    "  # TODO: Fix type signature in Ray Data\n",
-    "  # TODO: Fix dataset.union when used with empty list.\n",
-    "  if len(all_test_datasets_seen_so_far) > 0:\n",
-    "    full_test_dataset = test_dataset.union(*all_test_datasets_seen_so_far)\n",
-    "  else:\n",
-    "    full_test_dataset = test_dataset\n",
-    "\n",
-    "  all_test_datasets_seen_so_far.append(test_dataset)\n",
-    "\n",
-    "  accuracy_for_this_task = batch_predict(latest_checkpoint, full_test_dataset)\n",
-    "  print(f\"Accuracy for task {task_idx}: {accuracy_for_this_task}\")\n",
-    "  accuracy_for_all_tasks.append(accuracy_for_this_task)\n",
-    "\n",
     "  # *************Model Deployment & Online Inference***************************\n",
     "  \n",
     "  # We can also deploy our model to do online inference with Ray Serve.\n",
@@ -1423,6 +1365,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "ORWpRkPjcPbD"
@@ -1462,328 +1405,7 @@
    ]
   },
   {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "xLLAvsTk8LoV"
-   },
-   "source": [
-    "# [Optional] Step 6: Compare against full training.\n",
-    "\n",
-    "We have now incrementally trained our simple multi-layer perceptron. Let's compare the incrementally trained model via fine tuning against a model that is trained on all the tasks up front.\n",
-    "\n",
-    "Since we are using a naive fine-tuning strategy, we should expect that our incrementally trained model will perform worse than the one that is fully trained! However, there's various other strategies that have been developed and are actively being researched to improve accuracy for incremental training. And overall, incremental/continual learning allows you to train in many real world settings where the entire dataset is not available up front, but new data is arriving at a relatively high rate."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "RNHsEVBHc0p2"
-   },
-   "source": [
-    "Let's first combine all of our datasets for each task into a single, unified dataset"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "pU2fVH068lfF",
-    "outputId": "fd6a3b56-dda1-4fa6-cebd-d0ee8784e698"
-   },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Map Progress (1 actors 1 pending): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  1.37it/s]\n",
-      "Map Progress (1 actors 1 pending): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  1.37it/s]\n",
-      "Map Progress (1 actors 1 pending): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  1.40it/s]\n",
-      "Shuffle Map: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 40.34it/s]\n",
-      "Shuffle Reduce: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 28.99it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "train_stream = permuted_mnist.generate_train_stream()\n",
-    "\n",
-    "# Collect all datasets in the stream into a single dataset.\n",
-    "all_training_datasets = []\n",
-    "for train_dataset in train_stream:\n",
-    "  all_training_datasets.append(train_dataset)\n",
-    "combined_training_dataset = all_training_datasets[0].union(*all_training_datasets[1:])\n",
-    "\n",
-    "\n",
-    "combined_training_dataset = combined_training_dataset.random_shuffle()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "tJ6Oqdgvc5dn"
-   },
-   "source": [
-    "Then, we train a new model on the unified dataset using the same configurations as before."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/",
-     "height": 1000
-    },
-    "id": "PmH9c0-z9KME",
-    "outputId": "653b4dfc-ed47-4307-fa84-e4c4ea3ec354"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/html": [
-       "<div class=\"tuneStatus\">\n",
-       "  <div style=\"display: flex;flex-direction: row\">\n",
-       "    <div style=\"display: flex;flex-direction: column;\">\n",
-       "      <h3>Tune Status</h3>\n",
-       "      <table>\n",
-       "<tbody>\n",
-       "<tr><td>Current time:</td><td>2022-09-23 16:37:13</td></tr>\n",
-       "<tr><td>Running for: </td><td>00:00:25.97        </td></tr>\n",
-       "<tr><td>Memory:      </td><td>19.4/62.7 GiB      </td></tr>\n",
-       "</tbody>\n",
-       "</table>\n",
-       "    </div>\n",
-       "    <div class=\"vDivider\"></div>\n",
-       "    <div class=\"systemInfo\">\n",
-       "      <h3>System Info</h3>\n",
-       "      Using FIFO scheduling algorithm.<br>Resources requested: 0/24 CPUs, 0/0 GPUs, 0.0/32.53 GiB heap, 0.0/16.26 GiB objects\n",
-       "    </div>\n",
-       "    \n",
-       "  </div>\n",
-       "  <div class=\"hDivider\"></div>\n",
-       "  <div class=\"trialStatus\">\n",
-       "    <h3>Trial Status</h3>\n",
-       "    <table>\n",
-       "<thead>\n",
-       "<tr><th>Trial name              </th><th>status    </th><th>loc                  </th><th style=\"text-align: right;\">  iter</th><th style=\"text-align: right;\">  total time (s)</th><th style=\"text-align: right;\">  loss</th><th style=\"text-align: right;\">  _timestamp</th><th style=\"text-align: right;\">  _time_this_iter_s</th></tr>\n",
-       "</thead>\n",
-       "<tbody>\n",
-       "<tr><td>TorchTrainer_971af_00000</td><td>TERMINATED</td><td>10.109.175.190:860035</td><td style=\"text-align: right;\">     4</td><td style=\"text-align: right;\">         22.1282</td><td style=\"text-align: right;\">     0</td><td style=\"text-align: right;\">  1663976231</td><td style=\"text-align: right;\">          0.0924587</td></tr>\n",
-       "</tbody>\n",
-       "</table>\n",
-       "  </div>\n",
-       "</div>\n",
-       "<style>\n",
-       ".tuneStatus {\n",
-       "  color: var(--jp-ui-font-color1);\n",
-       "}\n",
-       ".tuneStatus .systemInfo {\n",
-       "  display: flex;\n",
-       "  flex-direction: column;\n",
-       "}\n",
-       ".tuneStatus td {\n",
-       "  white-space: nowrap;\n",
-       "}\n",
-       ".tuneStatus .trialStatus {\n",
-       "  display: flex;\n",
-       "  flex-direction: column;\n",
-       "}\n",
-       ".tuneStatus h3 {\n",
-       "  font-weight: bold;\n",
-       "}\n",
-       ".tuneStatus .hDivider {\n",
-       "  border-bottom-width: var(--jp-border-width);\n",
-       "  border-bottom-color: var(--jp-border-color0);\n",
-       "  border-bottom-style: solid;\n",
-       "}\n",
-       ".tuneStatus .vDivider {\n",
-       "  border-left-width: var(--jp-border-width);\n",
-       "  border-left-color: var(--jp-border-color0);\n",
-       "  border-left-style: solid;\n",
-       "  margin: 0.5em 1em 0.5em 1em;\n",
-       "}\n",
-       "</style>\n"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m 2022-09-23 16:36:55,188\tINFO config.py:71 -- Setting up process group for: env:// [rank=0, world_size=1]\n",
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m 2022-09-23 16:36:55,399\tINFO train_loop_utils.py:354 -- Moving model to device: cuda:0\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m loss: 2.301066, epoch: 0, iteration: 0\n",
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m loss: 1.869080, epoch: 0, iteration: 500\n",
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m loss: 1.489264, epoch: 0, iteration: 1000\n",
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m loss: 1.646756, epoch: 0, iteration: 1500\n",
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m loss: 1.582330, epoch: 0, iteration: 2000\n",
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m loss: 1.246018, epoch: 0, iteration: 2500\n",
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m loss: 1.035204, epoch: 0, iteration: 3000\n",
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m loss: 0.872962, epoch: 0, iteration: 3500\n",
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m loss: 1.138829, epoch: 0, iteration: 4000\n",
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m loss: 0.753354, epoch: 0, iteration: 4500\n",
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m loss: 0.991935, epoch: 0, iteration: 5000\n",
-      "\u001b[2m\u001b[36m(RayTrainWorker pid=860154)\u001b[0m loss: 0.928292, epoch: 0, iteration: 5500\n"
-     ]
-    },
-    {
-     "data": {
-      "text/html": [
-       "<div class=\"trialProgress\">\n",
-       "  <h3>Trial Progress</h3>\n",
-       "  <table>\n",
-       "<thead>\n",
-       "<tr><th>Trial name              </th><th style=\"text-align: right;\">  _time_this_iter_s</th><th style=\"text-align: right;\">  _timestamp</th><th style=\"text-align: right;\">  _training_iteration</th><th>date               </th><th>done  </th><th>episodes_total  </th><th>experiment_id                   </th><th style=\"text-align: right;\">  experiment_tag</th><th>hostname  </th><th style=\"text-align: right;\">  iterations_since_restore</th><th style=\"text-align: right;\">  loss</th><th>node_ip       </th><th style=\"text-align: right;\">   pid</th><th>should_checkpoint  </th><th style=\"text-align: right;\">  time_since_restore</th><th style=\"text-align: right;\">  time_this_iter_s</th><th style=\"text-align: right;\">  time_total_s</th><th style=\"text-align: right;\">  timestamp</th><th style=\"text-align: right;\">  timesteps_since_restore</th><th>timesteps_total  </th><th style=\"text-align: right;\">  training_iteration</th><th>trial_id   </th><th style=\"text-align: right;\">  warmup_time</th></tr>\n",
-       "</thead>\n",
-       "<tbody>\n",
-       "<tr><td>TorchTrainer_971af_00000</td><td style=\"text-align: right;\">          0.0924587</td><td style=\"text-align: right;\">  1663976231</td><td style=\"text-align: right;\">                    4</td><td>2022-09-23_16-37-11</td><td>True  </td><td>                </td><td>26d685b2612a4752b7d062d1ebfb89f0</td><td style=\"text-align: right;\">               0</td><td>corvus    </td><td style=\"text-align: right;\">                         4</td><td style=\"text-align: right;\">     0</td><td>10.109.175.190</td><td style=\"text-align: right;\">860035</td><td>True               </td><td style=\"text-align: right;\">             22.1282</td><td style=\"text-align: right;\">         0.0941384</td><td style=\"text-align: right;\">       22.1282</td><td style=\"text-align: right;\"> 1663976231</td><td style=\"text-align: right;\">                        0</td><td>                 </td><td style=\"text-align: right;\">                   4</td><td>971af_00000</td><td style=\"text-align: right;\">    0.0034101</td></tr>\n",
-       "</tbody>\n",
-       "</table>\n",
-       "</div>\n",
-       "<style>\n",
-       ".trialProgress {\n",
-       "  display: flex;\n",
-       "  flex-direction: column;\n",
-       "  color: var(--jp-ui-font-color1);\n",
-       "}\n",
-       ".trialProgress h3 {\n",
-       "  font-weight: bold;\n",
-       "}\n",
-       ".trialProgress td {\n",
-       "  white-space: nowrap;\n",
-       "}\n",
-       "</style>\n"
-      ],
-      "text/plain": [
-       "<IPython.core.display.HTML object>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2022-09-23 16:37:13,525\tINFO tune.py:762 -- Total run time: 26.08 seconds (25.96 seconds for the tuning loop).\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Now we do training with the same configurations as before\n",
-    "trainer = TorchTrainer(\n",
-    "            train_loop_per_worker=train_loop_per_worker,\n",
-    "            train_loop_config={\n",
-    "                \"num_epochs\": num_epochs,\n",
-    "                \"learning_rate\": learning_rate,\n",
-    "                \"momentum\": momentum,\n",
-    "                \"batch_size\": batch_size,\n",
-    "            },\n",
-    "            # Have to specify trainer_resources as 0 so that the example works on Colab. \n",
-    "            scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu, trainer_resources={\"CPU\": 0}),\n",
-    "            datasets={\"train\": combined_training_dataset},\n",
-    "            preprocessor=mnist_normalize_preprocessor,\n",
-    "        )\n",
-    "result = trainer.fit()\n",
-    "full_training_checkpoint = result.checkpoint"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "jLaOcmBddRqB"
-   },
-   "source": [
-    "Then, let's test model that was trained on all the tasks up front."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "WC7zV_Cw9TAi",
-    "outputId": "12a86f2b-be90-47b6-e252-25e3199689f9"
-   },
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Map Progress (1 actors 1 pending):   0%|                                                                                                                                                                                                                                                                                                                                                                        | 0/3 [00:01<?, ?it/s]\u001b[2m\u001b[36m(BlockWorker pid=860261)\u001b[0m 2022-09-23 16:37:15,152\tWARNING torch_predictor.py:53 -- You have `use_gpu` as False but there are 1 GPUs detected on host where prediction will only use CPU. Please consider explicitly setting `TorchPredictor(use_gpu=True)` or `batch_predictor.predict(ds, num_gpus_per_worker=1)` to enable GPU prediction.\n",
-      "Map Progress (2 actors 1 pending):  33%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▎                                                                                                                                                                                                                                          | 1/3 [00:03<00:04,  2.45s/it]\u001b[2m\u001b[36m(BlockWorker pid=860289)\u001b[0m 2022-09-23 16:37:16,696\tWARNING torch_predictor.py:53 -- You have `use_gpu` as False but there are 1 GPUs detected on host where prediction will only use CPU. Please consider explicitly setting `TorchPredictor(use_gpu=True)` or `batch_predictor.predict(ds, num_gpus_per_worker=1)` to enable GPU prediction.\n",
-      "Map Progress (2 actors 1 pending): 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:04<00:00,  1.37s/it]\n",
-      "Map_Batches: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 74.29it/s]\n",
-      "Map_Batches: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 134.64it/s]\n",
-      "Shuffle Map: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 3/3 [00:00<00:00, 304.26it/s]\n",
-      "Shuffle Reduce: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00, 108.41it/s]\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Then, we used the fully trained model and do batch prediction on the entire test set.\n",
-    "\n",
-    "# `full_test_dataset` should already contain the combined test datasets.\n",
-    "fully_trained_accuracy = batch_predict(full_training_checkpoint, full_test_dataset)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "Pn5LJ4CUdZgI"
-   },
-   "source": [
-    "Finally, let's compare the accuracies between the incrementally trained model and the fully trained model. We should see that the fully trained model performs better."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "UFhRf_8e-vgA",
-    "outputId": "056ff06f-ff87-4f3a-d740-4cc556bde3dd"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Fully trained model accuracy:  0.8888666666666667\n",
-      "Incrementally trained model accuracy:  0.8439\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[2m\u001b[36m(BlockWorker pid=860324)\u001b[0m 2022-09-23 16:37:18,256\tWARNING torch_predictor.py:53 -- You have `use_gpu` as False but there are 1 GPUs detected on host where prediction will only use CPU. Please consider explicitly setting `TorchPredictor(use_gpu=True)` or `batch_predictor.predict(ds, num_gpus_per_worker=1)` to enable GPU prediction.\n"
-     ]
-    }
-   ],
-   "source": [
-    "print(\"Fully trained model accuracy: \", fully_trained_accuracy)\n",
-    "print(\"Incrementally trained model accuracy: \", accuracy_for_all_tasks[-1])"
-   ]
-  },
-  {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {
     "id": "FuqKePrYe-Fz"
diff --git a/doc/source/ray-air/examples/xgboost_batch_prediction.py b/doc/source/ray-air/examples/xgboost_batch_prediction.py
deleted file mode 100644
index 8c65e9db84a4..000000000000
--- a/doc/source/ray-air/examples/xgboost_batch_prediction.py
+++ /dev/null
@@ -1,42 +0,0 @@
-import ray
-from ray.data.preprocessors import StandardScaler
-from ray.train.batch_predictor import BatchPredictor
-from ray.train.xgboost import XGBoostTrainer, XGBoostPredictor
-from ray.air.config import ScalingConfig
-
-# Split data into train and validation.
-dataset = ray.data.read_csv("s3://anonymous@air-example-data/breast_cancer.csv")
-train_dataset, valid_dataset = dataset.train_test_split(test_size=0.3)
-test_dataset = valid_dataset.drop_columns(["target"])
-
-columns_to_scale = ["mean radius", "mean texture"]
-preprocessor = StandardScaler(columns=columns_to_scale)
-
-trainer = XGBoostTrainer(
-    label_column="target",
-    num_boost_round=20,
-    scaling_config=ScalingConfig(num_workers=2),
-    params={
-        "objective": "binary:logistic",
-        "eval_metric": ["logloss", "error"],
-    },
-    datasets={"train": train_dataset},
-    preprocessor=preprocessor,
-)
-result = trainer.fit()
-
-# You can also create a checkpoint from a trained model using
-# `XGBoostCheckpoint.from_model`.
-
-# import xgboost as xgb
-# from ray.train.xgboost import XGBoostCheckpoint
-# model = xgb.Booster()
-# model.load_model(...)
-# checkpoint = XGBoostCheckpoint.from_model(model, path=".")
-checkpoint = result.checkpoint
-
-batch_predictor = BatchPredictor.from_checkpoint(checkpoint, XGBoostPredictor)
-
-predicted_probabilities = batch_predictor.predict(test_dataset)
-# Call show on the output probabilities to trigger execution.
-predicted_probabilities.show()
diff --git a/doc/source/ray-air/examples/xgboost_example.ipynb b/doc/source/ray-air/examples/xgboost_example.ipynb
index 3473eb5d71e0..6cb37094341a 100644
--- a/doc/source/ray-air/examples/xgboost_example.ipynb
+++ b/doc/source/ray-air/examples/xgboost_example.ipynb
@@ -17,6 +17,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "53d57c1f",
    "metadata": {
@@ -30,19 +31,30 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 21,
    "id": "41f20cc1",
    "metadata": {
     "pycharm": {
      "name": "#%%\n"
     }
    },
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip available: \u001b[0m\u001b[31;49m22.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m23.1.2\u001b[0m\n",
+      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
+     ]
+    }
+   ],
    "source": [
     "!pip install -qU \"ray[tune]\" xgboost_ray"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "d2fe8d4a",
    "metadata": {
@@ -56,27 +68,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 22,
    "id": "7232303d",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "/home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "  from pandas import MultiIndex, Int64Index\n",
-      "FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from typing import Tuple\n",
     "\n",
     "import ray\n",
-    "from ray.train.batch_predictor import BatchPredictor\n",
     "from ray.train.xgboost import XGBoostPredictor\n",
     "from ray.train.xgboost import XGBoostTrainer\n",
     "from ray.air.config import ScalingConfig\n",
@@ -86,6 +85,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "1c75b5ca",
    "metadata": {
@@ -99,7 +99,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 23,
    "id": "37c4f38f",
    "metadata": {},
    "outputs": [],
@@ -112,6 +112,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "9b2850dd",
    "metadata": {
@@ -125,7 +126,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 24,
    "id": "dae8998d",
    "metadata": {
     "pycharm": {
@@ -163,6 +164,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "ce05af87",
    "metadata": {},
@@ -172,7 +174,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 25,
    "id": "5b8076d3",
    "metadata": {
     "pycharm": {
@@ -181,26 +183,37 @@
    },
    "outputs": [],
    "source": [
+    "import pandas as pd\n",
+    "from ray.air import Checkpoint\n",
+    "from ray.data import ActorPoolStrategy\n",
+    "\n",
+    "\n",
+    "class Predict:\n",
+    "\n",
+    "    def __init__(self, checkpoint: Checkpoint):\n",
+    "        self.predictor = XGBoostPredictor.from_checkpoint(checkpoint)\n",
+    "\n",
+    "    def __call__(self, batch: pd.DataFrame) -> pd.DataFrame:\n",
+    "        return self.predictor.predict(batch)\n",
+    "\n",
+    "\n",
     "def predict_xgboost(result: Result):\n",
     "    _, _, test_dataset = prepare_data()\n",
     "\n",
-    "    batch_predictor = BatchPredictor.from_checkpoint(\n",
-    "        result.checkpoint, XGBoostPredictor\n",
-    "    )\n",
-    "\n",
-    "    predicted_labels = (\n",
-    "        batch_predictor.predict(test_dataset)\n",
-    "        .map_batches(lambda df: (df > 0.5).astype(int), batch_format=\"pandas\")\n",
+    "    scores = test_dataset.map_batches(\n",
+    "        Predict, \n",
+    "        fn_constructor_args=[result.checkpoint], \n",
+    "        compute=ActorPoolStrategy(), \n",
+    "        batch_format=\"pandas\"\n",
     "    )\n",
+    "    \n",
+    "    predicted_labels = scores.map_batches(lambda df: (df > 0.5).astype(int), batch_format=\"pandas\")\n",
     "    print(f\"PREDICTED LABELS\")\n",
-    "    predicted_labels.show()\n",
-    "\n",
-    "    shap_values = batch_predictor.predict(test_dataset, pred_contribs=True)\n",
-    "    print(f\"SHAP VALUES\")\n",
-    "    shap_values.show()\n"
+    "    predicted_labels.show()"
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "7e172f66",
    "metadata": {},
@@ -210,7 +223,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 26,
    "id": "0f96d62b",
    "metadata": {
     "pycharm": {
@@ -218,26 +231,71 @@
     }
    },
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2022-06-22 17:28:55,841\tINFO services.py:1477 -- View the Ray dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8270\u001b[39m\u001b[22m\n",
-      "2022-06-22 17:28:58,044\tWARNING read_api.py:260 -- The number of blocks in this dataset (1) limits its parallelism to 1 concurrent tasks. This is much less than the number of available CPU slots in the cluster. Use `.repartition(n)` to increase the number of dataset blocks.\n",
-      "Map_Batches: 100%|██████████| 1/1 [00:00<00:00, 40.28it/s]\n"
-     ]
-    },
     {
      "data": {
       "text/html": [
-       "== Status ==<br>Current time: 2022-06-22 17:29:15 (running for 00:00:16.11)<br>Memory usage on this node: 11.5/31.0 GiB<br>Using FIFO scheduling algorithm.<br>Resources requested: 0/8 CPUs, 0/0 GPUs, 0.0/12.35 GiB heap, 0.0/6.18 GiB objects<br>Result logdir: /home/ubuntu/ray_results/XGBoostTrainer_2022-06-22_17-28-58<br>Number of trials: 1/1 (1 TERMINATED)<br><table>\n",
+       "<div class=\"tuneStatus\">\n",
+       "  <div style=\"display: flex;flex-direction: row\">\n",
+       "    <div style=\"display: flex;flex-direction: column;\">\n",
+       "      <h3>Tune Status</h3>\n",
+       "      <table>\n",
+       "<tbody>\n",
+       "<tr><td>Current time:</td><td>2023-07-06 18:33:25</td></tr>\n",
+       "<tr><td>Running for: </td><td>00:00:06.19        </td></tr>\n",
+       "<tr><td>Memory:      </td><td>14.9/64.0 GiB      </td></tr>\n",
+       "</tbody>\n",
+       "</table>\n",
+       "    </div>\n",
+       "    <div class=\"vDivider\"></div>\n",
+       "    <div class=\"systemInfo\">\n",
+       "      <h3>System Info</h3>\n",
+       "      Using FIFO scheduling algorithm.<br>Logical resource usage: 2.0/10 CPUs, 0/0 GPUs\n",
+       "    </div>\n",
+       "    \n",
+       "  </div>\n",
+       "  <div class=\"hDivider\"></div>\n",
+       "  <div class=\"trialStatus\">\n",
+       "    <h3>Trial Status</h3>\n",
+       "    <table>\n",
        "<thead>\n",
-       "<tr><th>Trial name                </th><th>status    </th><th>loc                  </th><th style=\"text-align: right;\">  iter</th><th style=\"text-align: right;\">  total time (s)</th><th style=\"text-align: right;\">  train-logloss</th><th style=\"text-align: right;\">  train-error</th><th style=\"text-align: right;\">  valid-logloss</th></tr>\n",
+       "<tr><th>Trial name                </th><th>status    </th><th>loc            </th><th style=\"text-align: right;\">  iter</th><th style=\"text-align: right;\">  total time (s)</th><th style=\"text-align: right;\">  train-logloss</th><th style=\"text-align: right;\">  train-error</th><th style=\"text-align: right;\">  valid-logloss</th></tr>\n",
        "</thead>\n",
        "<tbody>\n",
-       "<tr><td>XGBoostTrainer_cc863_00000</td><td>TERMINATED</td><td>172.31.43.110:1493910</td><td style=\"text-align: right;\">   100</td><td style=\"text-align: right;\">         12.5164</td><td style=\"text-align: right;\">       0.005874</td><td style=\"text-align: right;\">            0</td><td style=\"text-align: right;\">       0.078188</td></tr>\n",
+       "<tr><td>XGBoostTrainer_40fed_00000</td><td>TERMINATED</td><td>127.0.0.1:40725</td><td style=\"text-align: right;\">   101</td><td style=\"text-align: right;\">         4.90132</td><td style=\"text-align: right;\">     0.00587595</td><td style=\"text-align: right;\">            0</td><td style=\"text-align: right;\">        0.06215</td></tr>\n",
        "</tbody>\n",
-       "</table><br><br>"
+       "</table>\n",
+       "  </div>\n",
+       "</div>\n",
+       "<style>\n",
+       ".tuneStatus {\n",
+       "  color: var(--jp-ui-font-color1);\n",
+       "}\n",
+       ".tuneStatus .systemInfo {\n",
+       "  display: flex;\n",
+       "  flex-direction: column;\n",
+       "}\n",
+       ".tuneStatus td {\n",
+       "  white-space: nowrap;\n",
+       "}\n",
+       ".tuneStatus .trialStatus {\n",
+       "  display: flex;\n",
+       "  flex-direction: column;\n",
+       "}\n",
+       ".tuneStatus h3 {\n",
+       "  font-weight: bold;\n",
+       "}\n",
+       ".tuneStatus .hDivider {\n",
+       "  border-bottom-width: var(--jp-border-width);\n",
+       "  border-bottom-color: var(--jp-border-color0);\n",
+       "  border-bottom-style: solid;\n",
+       "}\n",
+       ".tuneStatus .vDivider {\n",
+       "  border-left-width: var(--jp-border-width);\n",
+       "  border-left-color: var(--jp-border-color0);\n",
+       "  border-left-style: solid;\n",
+       "  margin: 0.5em 1em 0.5em 1em;\n",
+       "}\n",
+       "</style>\n"
       ],
       "text/plain": [
        "<IPython.core.display.HTML object>"
@@ -250,119 +308,95 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "\u001b[2m\u001b[36m(pid=1493910)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1493910)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(pid=1493910)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1493910)\u001b[0m FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1493910)\u001b[0m FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(XGBoostTrainer pid=1493910)\u001b[0m UserWarning: Dataset 'train' has 1 blocks, which is less than the `num_workers` 2. This dataset will be automatically repartitioned to 2 blocks.\n",
-      "\u001b[2m\u001b[36m(XGBoostTrainer pid=1493910)\u001b[0m UserWarning: Dataset 'valid' has 1 blocks, which is less than the `num_workers` 2. This dataset will be automatically repartitioned to 2 blocks.\n",
-      "\u001b[2m\u001b[36m(XGBoostTrainer pid=1493910)\u001b[0m 2022-06-22 17:29:04,073\tINFO main.py:980 -- [RayXGBoost] Created 2 new actors (2 total actors). Waiting until actors are ready for training.\n",
-      "\u001b[2m\u001b[36m(pid=1494007)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494007)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(pid=1494008)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494008)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(pid=1494009)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494009)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(pid=1494007)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494007)\u001b[0m FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494007)\u001b[0m FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494008)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494008)\u001b[0m FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494008)\u001b[0m FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(_RemoteRayXGBoostActor pid=1494008)\u001b[0m 2022-06-22 17:29:07,324\tWARNING __init__.py:190 -- DeprecationWarning: `ray.worker.get_resource_ids` is a private attribute and access will be removed in a future Ray version.\n",
-      "\u001b[2m\u001b[36m(pid=1494009)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494009)\u001b[0m FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494009)\u001b[0m FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(_RemoteRayXGBoostActor pid=1494009)\u001b[0m 2022-06-22 17:29:07,421\tWARNING __init__.py:190 -- DeprecationWarning: `ray.worker.get_resource_ids` is a private attribute and access will be removed in a future Ray version.\n",
-      "\u001b[2m\u001b[36m(XGBoostTrainer pid=1493910)\u001b[0m 2022-06-22 17:29:07,874\tINFO main.py:1025 -- [RayXGBoost] Starting XGBoost training.\n",
-      "\u001b[2m\u001b[36m(_RemoteRayXGBoostActor pid=1494008)\u001b[0m [17:29:07] task [xgboost.ray]:139731353900128 got new rank 0\n",
-      "\u001b[2m\u001b[36m(_RemoteRayXGBoostActor pid=1494008)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(_RemoteRayXGBoostActor pid=1494009)\u001b[0m [17:29:07] task [xgboost.ray]:140076138558608 got new rank 1\n",
-      "\u001b[2m\u001b[36m(_RemoteRayXGBoostActor pid=1494009)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(_QueueActor pid=1494006)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(_QueueActor pid=1494006)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(_QueueActor pid=1494006)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(_QueueActor pid=1494006)\u001b[0m FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(_QueueActor pid=1494006)\u001b[0m FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Result for XGBoostTrainer_cc863_00000:\n",
-      "  date: 2022-06-22_17-29-09\n",
-      "  done: false\n",
-      "  experiment_id: dc3dac01a34043cfb5751907e2bc648e\n",
-      "  hostname: ip-172-31-43-110\n",
-      "  iterations_since_restore: 1\n",
-      "  node_ip: 172.31.43.110\n",
-      "  pid: 1493910\n",
-      "  should_checkpoint: true\n",
-      "  time_since_restore: 7.967940330505371\n",
-      "  time_this_iter_s: 7.967940330505371\n",
-      "  time_total_s: 7.967940330505371\n",
-      "  timestamp: 1655918949\n",
-      "  timesteps_since_restore: 0\n",
-      "  train-error: 0.017588\n",
-      "  train-logloss: 0.464648\n",
-      "  training_iteration: 1\n",
-      "  trial_id: cc863_00000\n",
-      "  valid-error: 0.081871\n",
-      "  valid-logloss: 0.496374\n",
-      "  warmup_time: 0.004768848419189453\n",
-      "  \n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[2m\u001b[36m(XGBoostTrainer pid=1493910)\u001b[0m 2022-06-22 17:29:14,546\tINFO main.py:1516 -- [RayXGBoost] Finished XGBoost training on training data with total N=398 in 10.52 seconds (6.66 pure XGBoost training time).\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Result for XGBoostTrainer_cc863_00000:\n",
-      "  date: 2022-06-22_17-29-14\n",
-      "  done: true\n",
-      "  experiment_id: dc3dac01a34043cfb5751907e2bc648e\n",
-      "  experiment_tag: '0'\n",
-      "  hostname: ip-172-31-43-110\n",
-      "  iterations_since_restore: 100\n",
-      "  node_ip: 172.31.43.110\n",
-      "  pid: 1493910\n",
-      "  should_checkpoint: true\n",
-      "  time_since_restore: 12.516392230987549\n",
-      "  time_this_iter_s: 0.03008890151977539\n",
-      "  time_total_s: 12.516392230987549\n",
-      "  timestamp: 1655918954\n",
-      "  timesteps_since_restore: 0\n",
-      "  train-error: 0.0\n",
-      "  train-logloss: 0.005874\n",
-      "  training_iteration: 100\n",
-      "  trial_id: cc863_00000\n",
-      "  valid-error: 0.040936\n",
-      "  valid-logloss: 0.078188\n",
-      "  warmup_time: 0.004768848419189453\n",
-      "  \n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "2022-06-22 17:29:15,362\tINFO tune.py:734 -- Total run time: 16.94 seconds (16.08 seconds for the tuning loop).\n"
+      "\u001b[2m\u001b[36m(XGBoostTrainer pid=40725)\u001b[0m The `preprocessor` arg to Trainer is deprecated. Apply preprocessor transformations ahead of time by calling `preprocessor.transform(ds)`. Support for the preprocessor arg will be dropped in a future release.\n",
+      "\u001b[2m\u001b[36m(XGBoostTrainer pid=40725)\u001b[0m Tip: Use `take_batch()` instead of `take() / show()` to return records in pandas or numpy batch format.\n",
+      "\u001b[2m\u001b[36m(XGBoostTrainer pid=40725)\u001b[0m Executing DAG InputDataBuffer[Input] -> AllToAllOperator[Aggregate]\n",
+      "\u001b[2m\u001b[36m(XGBoostTrainer pid=40725)\u001b[0m Execution config: ExecutionOptions(resource_limits=ExecutionResources(cpu=None, gpu=None, object_store_memory=None), locality_with_output=False, preserve_order=False, actor_locality_enabled=True, verbose_progress=False)\n",
+      "\u001b[2m\u001b[36m(XGBoostTrainer pid=40725)\u001b[0m Tip: For detailed progress reporting, run `ray.data.DataContext.get_current().execution_options.verbose_progress = True`\n",
+      "\n",
+      "\u001b[A\n",
+      "\u001b[A\n",
+      "\n",
+      "\u001b[A\u001b[A\n",
+      "\n",
+      "(pid=40725) Running: 0.0/10.0 CPU, 0.0/0.0 GPU, 0.0 MiB/512.0 MiB object_store_memory:   0%|          | 0/14 [00:00<?, ?it/s]\n",
+      "\u001b[A\n",
+      "\u001b[A                                                              \n",
+      "\n",
+      "\u001b[A\u001b[A                                                                      \n",
+      "\n",
+      "\n",
+      "\u001b[2m\u001b[36m(XGBoostTrainer pid=40725)\u001b[0m Executing DAG InputDataBuffer[Input] -> TaskPoolMapOperator[MapBatches(StandardScaler._transform_pandas)]\n",
+      "\n",
+      "\u001b[A\n",
+      "\n",
+      "(pid=40725) Running: 0.0/10.0 CPU, 0.0/0.0 GPU, 0.0 MiB/512.0 MiB object_store_memory:   0%|          | 0/14 [00:01<?, ?it/s]\n",
+      "\u001b[A                                                              \n",
+      "\n",
+      "\u001b[A\u001b[A                                                                       \n",
+      "\n",
+      "\n",
+      "\u001b[2m\u001b[36m(XGBoostTrainer pid=40725)\u001b[0m Execution config: ExecutionOptions(resource_limits=ExecutionResources(cpu=None, gpu=None, object_store_memory=None), locality_with_output=False, preserve_order=False, actor_locality_enabled=True, verbose_progress=False)\n",
+      "\n",
+      "\u001b[A\n",
+      "\n",
+      "(pid=40725) Running: 0.0/10.0 CPU, 0.0/0.0 GPU, 0.0 MiB/512.0 MiB object_store_memory:   0%|          | 0/14 [00:01<?, ?it/s]\n",
+      "\u001b[A                                                              \n",
+      "\n",
+      "\u001b[A\u001b[A                                                                       \n",
+      "\n",
+      "\n",
+      "\u001b[2m\u001b[36m(XGBoostTrainer pid=40725)\u001b[0m Tip: For detailed progress reporting, run `ray.data.DataContext.get_current().execution_options.verbose_progress = True`\n",
+      "\n",
+      "\u001b[A\n",
+      "\n",
+      "(pid=40725) Running: 0.0/10.0 CPU, 0.0/0.0 GPU, 0.0 MiB/512.0 MiB object_store_memory:   0%|          | 0/14 [00:01<?, ?it/s]\n",
+      "\u001b[A                                                              \n",
+      "\n",
+      "\u001b[A\u001b[A                                                                       \n",
+      "\n",
+      "\n",
+      "\u001b[2m\u001b[36m(XGBoostTrainer pid=40725)\u001b[0m Executing DAG InputDataBuffer[Input] -> TaskPoolMapOperator[MapBatches(StandardScaler._transform_pandas)]\n",
+      "\n",
+      "\u001b[A\n",
+      "\n",
+      "(pid=40725) Running: 0.0/10.0 CPU, 0.0/0.0 GPU, 0.0 MiB/512.0 MiB object_store_memory:   0%|          | 0/14 [00:01<?, ?it/s]\n",
+      "\u001b[A                                                              \n",
+      "\n",
+      "\u001b[A\u001b[A                                                                       \n",
+      "\n",
+      "\n",
+      "\u001b[2m\u001b[36m(XGBoostTrainer pid=40725)\u001b[0m Execution config: ExecutionOptions(resource_limits=ExecutionResources(cpu=None, gpu=None, object_store_memory=None), locality_with_output=False, preserve_order=False, actor_locality_enabled=True, verbose_progress=False)\n",
+      "\n",
+      "\u001b[A\n",
+      "\n",
+      "(pid=40725) Running: 0.0/10.0 CPU, 0.0/0.0 GPU, 0.0 MiB/512.0 MiB object_store_memory:   0%|          | 0/14 [00:01<?, ?it/s]\n",
+      "\u001b[A                                                              \n",
+      "\n",
+      "\u001b[A\u001b[A                                                                       \n",
+      "\n",
+      "\n",
+      "\u001b[2m\u001b[36m(XGBoostTrainer pid=40725)\u001b[0m Tip: For detailed progress reporting, run `ray.data.DataContext.get_current().execution_options.verbose_progress = True`\n",
+      "\n",
+      "\u001b[A\n",
+      "\n",
+      "(pid=40725) Running: 0.0/10.0 CPU, 0.0/0.0 GPU, 0.0 MiB/512.0 MiB object_store_memory:   0%|          | 0/14 [00:01<?, ?it/s]\n",
+      "\u001b[A\n",
+      "\n",
+      "                                                                                                                                     \n",
+      "\u001b[A\n",
+      "\n",
+      "\u001b[A\u001b[A\n",
+      "\n",
+      "\u001b[2m\u001b[36m(_RemoteRayXGBoostActor pid=40741)\u001b[0m [18:33:23] task [xgboost.ray]:5022217360 got new rank 1                                   \n",
+      "2023-07-06 18:33:25,975\tINFO tune.py:1148 -- Total run time: 6.20 seconds (6.19 seconds for the tuning loop).\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "{'train-logloss': 0.005874, 'train-error': 0.0, 'valid-logloss': 0.078188, 'valid-error': 0.040936, 'time_this_iter_s': 0.03008890151977539, 'should_checkpoint': True, 'done': True, 'timesteps_total': None, 'episodes_total': None, 'training_iteration': 100, 'trial_id': 'cc863_00000', 'experiment_id': 'dc3dac01a34043cfb5751907e2bc648e', 'date': '2022-06-22_17-29-14', 'timestamp': 1655918954, 'time_total_s': 12.516392230987549, 'pid': 1493910, 'hostname': 'ip-172-31-43-110', 'node_ip': '172.31.43.110', 'config': {}, 'time_since_restore': 12.516392230987549, 'timesteps_since_restore': 0, 'iterations_since_restore': 100, 'warmup_time': 0.004768848419189453, 'experiment_tag': '0'}\n"
+      "{'train-logloss': 0.00587594546605992, 'train-error': 0.0, 'valid-logloss': 0.06215000962556052, 'valid-error': 0.02941176470588235, 'time_this_iter_s': 0.0101318359375, 'should_checkpoint': True, 'done': True, 'training_iteration': 101, 'trial_id': '40fed_00000', 'date': '2023-07-06_18-33-25', 'timestamp': 1688693605, 'time_total_s': 4.901317834854126, 'pid': 40725, 'hostname': 'Balajis-MacBook-Pro-16', 'node_ip': '127.0.0.1', 'config': {}, 'time_since_restore': 4.901317834854126, 'iterations_since_restore': 101, 'experiment_tag': '0'}\n"
      ]
     }
    ],
@@ -371,6 +405,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "7055ad1b",
    "metadata": {},
@@ -380,7 +415,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 27,
    "id": "283b1dba",
    "metadata": {
     "pycharm": {
@@ -392,23 +427,31 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2022-06-22 17:29:16,463\tWARNING read_api.py:260 -- The number of blocks in this dataset (1) limits its parallelism to 1 concurrent tasks. This is much less than the number of available CPU slots in the cluster. Use `.repartition(n)` to increase the number of dataset blocks.\n",
-      "Map_Batches: 100%|██████████| 1/1 [00:00<00:00, 46.14it/s]\n",
-      "Map_Batches:   0%|          | 0/1 [00:00<?, ?it/s]\u001b[2m\u001b[36m(pid=1494373)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494373)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(pid=1494373)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494373)\u001b[0m FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494373)\u001b[0m FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "Map Progress (1 actors 1 pending): 100%|██████████| 1/1 [00:01<00:00,  1.90s/it]\n",
-      "\u001b[2m\u001b[36m(BlockWorker pid=1494373)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "Map_Batches: 100%|██████████| 1/1 [00:00<00:00, 75.10it/s]\n"
+      "2023-07-06 18:33:27,259\tINFO read_api.py:374 -- To satisfy the requested parallelism of 20, each read task output will be split into 20 smaller blocks.\n",
+      "2023-07-06 18:33:28,112\tINFO streaming_executor.py:92 -- Executing DAG InputDataBuffer[Input] -> ActorPoolMapOperator[MapBatches(<lambda>)->MapBatches(Predict)] -> TaskPoolMapOperator[MapBatches(<lambda>)]\n",
+      "2023-07-06 18:33:28,112\tINFO streaming_executor.py:93 -- Execution config: ExecutionOptions(resource_limits=ExecutionResources(cpu=None, gpu=None, object_store_memory=None), locality_with_output=False, preserve_order=False, actor_locality_enabled=True, verbose_progress=False)\n",
+      "2023-07-06 18:33:28,114\tINFO streaming_executor.py:95 -- Tip: For detailed progress reporting, run `ray.data.DataContext.get_current().execution_options.verbose_progress = True`\n",
+      "2023-07-06 18:33:28,150\tINFO actor_pool_map_operator.py:117 -- MapBatches(<lambda>)->MapBatches(Predict): Waiting for 1 pool actors to start...\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "PREDICTED LABELS\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "                                                                                                                        "
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "PREDICTED LABELS\n",
       "{'predictions': 1}\n",
       "{'predictions': 1}\n",
       "{'predictions': 0}\n",
@@ -435,63 +478,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Map_Batches:   0%|          | 0/1 [00:00<?, ?it/s]\u001b[2m\u001b[36m(pid=1494403)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494403)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(pid=1494413)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494413)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(pid=1494403)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494403)\u001b[0m FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494403)\u001b[0m FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494413)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494413)\u001b[0m FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494413)\u001b[0m FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "Map Progress (1 actors 1 pending): 100%|██████████| 1/1 [00:01<00:00,  1.88s/it]"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "SHAP VALUES\n",
-      "{'predictions_0': 0.009930070489645004, 'predictions_1': 0.546318531036377, 'predictions_2': -0.006533853709697723, 'predictions_3': 0.022934239357709885, 'predictions_4': 0.32817941904067993, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.013881205581128597, 'predictions_7': 0.568859875202179, 'predictions_8': -0.27460771799087524, 'predictions_9': 0.013218197971582413, 'predictions_10': 0.009325551800429821, 'predictions_11': 0.04015672579407692, 'predictions_12': 0.11667086184024811, 'predictions_13': 0.9853533506393433, 'predictions_14': 0.05529181659221649, 'predictions_15': -0.005734208971261978, 'predictions_16': -0.0008497871458530426, 'predictions_17': 0.16138489544391632, 'predictions_18': -0.36162295937538147, 'predictions_19': 0.003658014815300703, 'predictions_20': 0.393682062625885, 'predictions_21': 0.6647266149520874, 'predictions_22': 1.7201099395751953, 'predictions_23': 0.35084351897239685, 'predictions_24': 0.4841834604740143, 'predictions_25': 0.013311417773365974, 'predictions_26': 0.8087116479873657, 'predictions_27': 0.5730299353599548, 'predictions_28': 0.1818174123764038, 'predictions_29': 0.1068669781088829, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.009930070489645004, 'predictions_1': 0.39265793561935425, 'predictions_2': -0.007271876558661461, 'predictions_3': 0.022934239357709885, 'predictions_4': -0.06737710535526276, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.04250117391347885, 'predictions_7': 0.7946916222572327, 'predictions_8': 0.30412089824676514, 'predictions_9': 0.013218197971582413, 'predictions_10': 0.009325551800429821, 'predictions_11': 0.04129280149936676, 'predictions_12': -0.012254374101758003, 'predictions_13': 0.2327558547258377, 'predictions_14': 0.05529181659221649, 'predictions_15': 0.6405980587005615, 'predictions_16': -0.0014769809786230326, 'predictions_17': 0.16138489544391632, 'predictions_18': -0.3877210319042206, 'predictions_19': -0.0022315792739391327, 'predictions_20': 0.3209536671638489, 'predictions_21': -0.018735788762569427, 'predictions_22': 1.482913851737976, 'predictions_23': 0.3135913908481598, 'predictions_24': 0.5150958299636841, 'predictions_25': 0.013311417773365974, 'predictions_26': 1.0875523090362549, 'predictions_27': 0.7323897480964661, 'predictions_28': 0.12313760071992874, 'predictions_29': 0.1068669781088829, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': -0.01911582238972187, 'predictions_1': -0.41554388403892517, 'predictions_2': -0.0034923271741718054, 'predictions_3': -0.06306137144565582, 'predictions_4': -0.3192429542541504, 'predictions_5': -0.009094981476664543, 'predictions_6': -0.08258295059204102, 'predictions_7': -0.8088644742965698, 'predictions_8': 0.2545676529407501, 'predictions_9': 0.015295587480068207, 'predictions_10': 0.0021318818908184767, 'predictions_11': -0.003732672892510891, 'predictions_12': -0.054907385259866714, 'predictions_13': -0.9062053561210632, 'predictions_14': -0.029618918895721436, 'predictions_15': 0.4303477704524994, 'predictions_16': -0.007885736413300037, 'predictions_17': 0.17541413009166718, 'predictions_18': -0.3190936744213104, 'predictions_19': 0.0024438181426376104, 'predictions_20': -0.6092430353164673, 'predictions_21': -0.3519248068332672, 'predictions_22': -2.2588469982147217, 'predictions_23': -0.569831371307373, 'predictions_24': -1.0422284603118896, 'predictions_25': -0.031086977571249008, 'predictions_26': -0.5648106932640076, 'predictions_27': -1.0225528478622437, 'predictions_28': -0.1411924958229065, 'predictions_29': -0.1724514663219452, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.009930070489645004, 'predictions_1': 0.5999099016189575, 'predictions_2': -0.007271876558661461, 'predictions_3': 0.022934239357709885, 'predictions_4': 0.31011122465133667, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.013881205581128597, 'predictions_7': 0.5183905363082886, 'predictions_8': -0.27460771799087524, 'predictions_9': -0.015900276601314545, 'predictions_10': 0.009325551800429821, 'predictions_11': 0.04129280149936676, 'predictions_12': 0.11667086184024811, 'predictions_13': 1.0816324949264526, 'predictions_14': -0.01614229381084442, 'predictions_15': -0.23922261595726013, 'predictions_16': -0.0008497871458530426, 'predictions_17': 0.16138489544391632, 'predictions_18': -0.36162295937538147, 'predictions_19': -0.0022315792739391327, 'predictions_20': 0.3189953863620758, 'predictions_21': 0.6754519939422607, 'predictions_22': 1.513157606124878, 'predictions_23': 0.35084351897239685, 'predictions_24': 0.09065212309360504, 'predictions_25': 0.013311417773365974, 'predictions_26': 1.039388656616211, 'predictions_27': 0.5762963891029358, 'predictions_28': 0.1816803514957428, 'predictions_29': 0.1068669781088829, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.009930070489645004, 'predictions_1': 0.41290122270584106, 'predictions_2': -0.007271876558661461, 'predictions_3': 0.04113232344388962, 'predictions_4': 0.4761468768119812, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.04250117391347885, 'predictions_7': 0.6889638304710388, 'predictions_8': 0.30412089824676514, 'predictions_9': -0.00703495554625988, 'predictions_10': 0.009325551800429821, 'predictions_11': 0.04129280149936676, 'predictions_12': -0.012254374101758003, 'predictions_13': 0.8955986499786377, 'predictions_14': 0.05529181659221649, 'predictions_15': 0.6405980587005615, 'predictions_16': -0.0010774275287985802, 'predictions_17': 0.16581103205680847, 'predictions_18': 0.4200459420681, 'predictions_19': 0.003658014815300703, 'predictions_20': 0.31559276580810547, 'predictions_21': 0.08991634845733643, 'predictions_22': 1.374340295791626, 'predictions_23': 0.4179628789424896, 'predictions_24': 0.6432731747627258, 'predictions_25': 0.013311417773365974, 'predictions_26': 1.0250954627990723, 'predictions_27': 0.6740144491195679, 'predictions_28': -0.12222569435834885, 'predictions_29': 0.1068669781088829, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.009930070489645004, 'predictions_1': 0.6223654747009277, 'predictions_2': -0.007271876558661461, 'predictions_3': 0.04113232344388962, 'predictions_4': 0.030091704800724983, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.04250117391347885, 'predictions_7': 0.520163893699646, 'predictions_8': -0.13476626574993134, 'predictions_9': 0.013218197971582413, 'predictions_10': 0.004969821777194738, 'predictions_11': 0.04129280149936676, 'predictions_12': 0.11667086184024811, 'predictions_13': 1.0668600797653198, 'predictions_14': -0.46553897857666016, 'predictions_15': -0.06545835733413696, 'predictions_16': -0.0010774275287985802, 'predictions_17': 0.16138489544391632, 'predictions_18': 0.406240850687027, 'predictions_19': -0.004930071532726288, 'predictions_20': 0.303047239780426, 'predictions_21': 0.7582043409347534, 'predictions_22': 1.4096102714538574, 'predictions_23': 0.43191614747047424, 'predictions_24': 0.47695621848106384, 'predictions_25': 0.013311417773365974, 'predictions_26': 1.0482016801834106, 'predictions_27': 0.5110929608345032, 'predictions_28': -0.12222569435834885, 'predictions_29': 0.1068669781088829, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.009930070489645004, 'predictions_1': 0.5661754012107849, 'predictions_2': -0.0019909553229808807, 'predictions_3': 0.04113232344388962, 'predictions_4': 0.292957067489624, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.013881205581128597, 'predictions_7': 0.6406195759773254, 'predictions_8': -0.27460771799087524, 'predictions_9': -0.007835762575268745, 'predictions_10': 0.009325551800429821, 'predictions_11': 0.04129280149936676, 'predictions_12': -0.012254374101758003, 'predictions_13': 0.21169273555278778, 'predictions_14': -0.02233714796602726, 'predictions_15': -0.9655348658561707, 'predictions_16': -0.0010828523663803935, 'predictions_17': 0.16581103205680847, 'predictions_18': 0.4200459420681, 'predictions_19': -0.0022315792739391327, 'predictions_20': 0.3136344850063324, 'predictions_21': 0.7264875173568726, 'predictions_22': 1.5534507036209106, 'predictions_23': 0.41967567801475525, 'predictions_24': 0.41908225417137146, 'predictions_25': 0.013311417773365974, 'predictions_26': 1.0032241344451904, 'predictions_27': 0.6091798543930054, 'predictions_28': 0.18174558877944946, 'predictions_29': 0.08268986642360687, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.009930070489645004, 'predictions_1': 0.40430790185928345, 'predictions_2': -0.007271876558661461, 'predictions_3': 0.04113232344388962, 'predictions_4': -0.33117765188217163, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.013881205581128597, 'predictions_7': 0.737379252910614, 'predictions_8': -0.27460771799087524, 'predictions_9': 0.007771771401166916, 'predictions_10': 0.009325551800429821, 'predictions_11': 0.04129280149936676, 'predictions_12': -0.012254374101758003, 'predictions_13': 0.18173672258853912, 'predictions_14': -0.02233714796602726, 'predictions_15': 0.6811268329620361, 'predictions_16': -0.0014769809786230326, 'predictions_17': 0.16581103205680847, 'predictions_18': 0.406240850687027, 'predictions_19': -0.0022315792739391327, 'predictions_20': 0.31559276580810547, 'predictions_21': -0.031238339841365814, 'predictions_22': 1.7573171854019165, 'predictions_23': 0.3838556110858917, 'predictions_24': 0.14492939412593842, 'predictions_25': 0.013311417773365974, 'predictions_26': 1.1148661375045776, 'predictions_27': 0.7173758149147034, 'predictions_28': 0.18174558877944946, 'predictions_29': 0.1068669781088829, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': -0.11485510319471359, 'predictions_1': 0.8875605463981628, 'predictions_2': -0.05685592442750931, 'predictions_3': -0.06306137144565582, 'predictions_4': 0.013515152037143707, 'predictions_5': 0.0024520084261894226, 'predictions_6': 0.04250117391347885, 'predictions_7': 0.4688224792480469, 'predictions_8': -0.14307911694049835, 'predictions_9': -0.06901921331882477, 'predictions_10': 0.003290211781859398, 'predictions_11': 0.003781725186854601, 'predictions_12': 0.08078579604625702, 'predictions_13': 0.8216619491577148, 'predictions_14': 0.0463210791349411, 'predictions_15': -0.13255546987056732, 'predictions_16': -0.0010776874842122197, 'predictions_17': 0.21504947543144226, 'predictions_18': -0.27619078755378723, 'predictions_19': -0.0025743553414940834, 'predictions_20': -0.9275118112564087, 'predictions_21': 1.0307989120483398, 'predictions_22': -2.7721312046051025, 'predictions_23': -0.4765593707561493, 'predictions_24': 0.5464398264884949, 'predictions_25': 0.004862718749791384, 'predictions_26': -0.619586706161499, 'predictions_27': -0.2188473790884018, 'predictions_28': 0.10512445122003555, 'predictions_29': 0.1068669781088829, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.0011090459302067757, 'predictions_1': -0.5595121383666992, 'predictions_2': -0.002613282995298505, 'predictions_3': 0.04113232344388962, 'predictions_4': 0.5254822969436646, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.013881205581128597, 'predictions_7': 0.5551156997680664, 'predictions_8': -0.27460771799087524, 'predictions_9': 0.013218197971582413, 'predictions_10': 0.009325551800429821, 'predictions_11': 0.03278094530105591, 'predictions_12': -0.012305313721299171, 'predictions_13': -1.6061651706695557, 'predictions_14': -0.01887715980410576, 'predictions_15': 0.6405980587005615, 'predictions_16': -0.0014769809786230326, 'predictions_17': -0.1776299774646759, 'predictions_18': 0.4914793074131012, 'predictions_19': 0.014222224242985249, 'predictions_20': 0.311679869890213, 'predictions_21': -0.08892179280519485, 'predictions_22': 1.5504939556121826, 'predictions_23': 0.5073927044868469, 'predictions_24': 0.6315706968307495, 'predictions_25': 0.005095706321299076, 'predictions_26': 0.7821602821350098, 'predictions_27': 0.7409825325012207, 'predictions_28': 0.1818174123764038, 'predictions_29': 0.07541733235120773, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': -0.01911582238972187, 'predictions_1': -0.41529545187950134, 'predictions_2': -0.0034923271741718054, 'predictions_3': -0.06306137144565582, 'predictions_4': -0.3192429542541504, 'predictions_5': 0.0166567862033844, 'predictions_6': -0.08258295059204102, 'predictions_7': -0.8088644742965698, 'predictions_8': 0.2545676529407501, 'predictions_9': 0.015658268705010414, 'predictions_10': 0.0021318818908184767, 'predictions_11': -0.003732672892510891, 'predictions_12': -0.054124802350997925, 'predictions_13': -1.0786054134368896, 'predictions_14': -0.03029802441596985, 'predictions_15': 0.46246394515037537, 'predictions_16': -0.007885736413300037, 'predictions_17': -0.168026864528656, 'predictions_18': -0.3190936744213104, 'predictions_19': 0.0024438181426376104, 'predictions_20': -0.6661943197250366, 'predictions_21': -0.2179996222257614, 'predictions_22': -2.689664125442505, 'predictions_23': -0.569831371307373, 'predictions_24': -0.6587631106376648, 'predictions_25': -0.031086977571249008, 'predictions_26': -0.554345428943634, 'predictions_27': -0.9785504937171936, 'predictions_28': -0.14026358723640442, 'predictions_29': 0.08221252262592316, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.009930070489645004, 'predictions_1': 0.46925264596939087, 'predictions_2': -0.007271876558661461, 'predictions_3': 0.04113232344388962, 'predictions_4': 0.49770334362983704, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.04250117391347885, 'predictions_7': 0.7922263741493225, 'predictions_8': 0.30412089824676514, 'predictions_9': -0.020953189581632614, 'predictions_10': 0.009325551800429821, 'predictions_11': 0.04129280149936676, 'predictions_12': 0.10843253135681152, 'predictions_13': 0.21460701525211334, 'predictions_14': -0.01614229381084442, 'predictions_15': 0.6405980587005615, 'predictions_16': -0.0014769809786230326, 'predictions_17': 0.16581103205680847, 'predictions_18': 0.4200459420681, 'predictions_19': -0.004930071532726288, 'predictions_20': 0.31877532601356506, 'predictions_21': -0.01986948773264885, 'predictions_22': 1.556177020072937, 'predictions_23': 0.445791631937027, 'predictions_24': 0.5183241367340088, 'predictions_25': 0.013311417773365974, 'predictions_26': 1.139808177947998, 'predictions_27': 0.5462087988853455, 'predictions_28': -0.12236276268959045, 'predictions_29': 0.1068669781088829, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.009930070489645004, 'predictions_1': 0.4225478768348694, 'predictions_2': -0.007271876558661461, 'predictions_3': 0.022934239357709885, 'predictions_4': 0.38781455159187317, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.013881205581128597, 'predictions_7': 0.6254116892814636, 'predictions_8': -0.28174325823783875, 'predictions_9': 0.008572578430175781, 'predictions_10': 0.009325551800429821, 'predictions_11': 0.022453438490629196, 'predictions_12': 0.09116627275943756, 'predictions_13': 1.01081383228302, 'predictions_14': 0.051560889929533005, 'predictions_15': -0.7526140809059143, 'predictions_16': -0.0010774275287985802, 'predictions_17': 0.13414263725280762, 'predictions_18': -0.36162295937538147, 'predictions_19': -0.004930071532726288, 'predictions_20': 0.21481025218963623, 'predictions_21': -1.7817566394805908, 'predictions_22': 1.5544458627700806, 'predictions_23': 0.37933894991874695, 'predictions_24': -0.5005015134811401, 'predictions_25': 0.013311417773365974, 'predictions_26': 1.2116931676864624, 'predictions_27': 0.44261839985847473, 'predictions_28': -0.1333807408809662, 'predictions_29': 0.12356309592723846, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.009930070489645004, 'predictions_1': 0.383541464805603, 'predictions_2': -0.007271876558661461, 'predictions_3': 0.022934239357709885, 'predictions_4': -0.2877747118473053, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.04250117391347885, 'predictions_7': 0.7181150317192078, 'predictions_8': -0.13476626574993134, 'predictions_9': 0.008572578430175781, 'predictions_10': 0.009325551800429821, 'predictions_11': 0.04129280149936676, 'predictions_12': 0.11661991477012634, 'predictions_13': 1.1911046504974365, 'predictions_14': 0.05529181659221649, 'predictions_15': -0.9263618588447571, 'predictions_16': -0.0010774275287985802, 'predictions_17': 0.16138489544391632, 'predictions_18': -0.3877210319042206, 'predictions_19': -0.0022315792739391327, 'predictions_20': 0.33373206853866577, 'predictions_21': -0.10588698834180832, 'predictions_22': 1.7216718196868896, 'predictions_23': 0.33134832978248596, 'predictions_24': 0.07947567105293274, 'predictions_25': 0.013311417773365974, 'predictions_26': 1.1455278396606445, 'predictions_27': 0.6355651617050171, 'predictions_28': -0.12222569435834885, 'predictions_29': 0.08268986642360687, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.0006543132476508617, 'predictions_1': -0.9144008159637451, 'predictions_2': -0.002613282995298505, 'predictions_3': 0.04113232344388962, 'predictions_4': 0.3871055543422699, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.013881205581128597, 'predictions_7': 0.6403470635414124, 'predictions_8': -0.27460771799087524, 'predictions_9': 0.05645249783992767, 'predictions_10': 0.009325551800429821, 'predictions_11': 0.02168312668800354, 'predictions_12': 0.11661991477012634, 'predictions_13': 1.00325608253479, 'predictions_14': 0.0834161564707756, 'predictions_15': 0.6231592893600464, 'predictions_16': -0.0010774275287985802, 'predictions_17': 0.16138489544391632, 'predictions_18': 0.37865450978279114, 'predictions_19': 0.014222224242985249, 'predictions_20': 0.38761574029922485, 'predictions_21': -0.02124016545712948, 'predictions_22': 1.7526684999465942, 'predictions_23': 0.5112042427062988, 'predictions_24': 0.5034342408180237, 'predictions_25': 0.005095706321299076, 'predictions_26': 0.9831169843673706, 'predictions_27': 0.6794330477714539, 'predictions_28': 0.1802661120891571, 'predictions_29': 0.1068669781088829, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': -0.0415695421397686, 'predictions_1': -1.239405870437622, 'predictions_2': 0.017459018155932426, 'predictions_3': 0.04113232344388962, 'predictions_4': 0.4762759208679199, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.04250117391347885, 'predictions_7': 0.9339500665664673, 'predictions_8': 0.2958080470561981, 'predictions_9': 0.07594440132379532, 'predictions_10': 0.003290211781859398, 'predictions_11': 0.01872287504374981, 'predictions_12': -0.047330088913440704, 'predictions_13': 0.056555308401584625, 'predictions_14': 0.05581967905163765, 'predictions_15': 0.6933314800262451, 'predictions_16': -0.0023022573441267014, 'predictions_17': 0.13856875896453857, 'predictions_18': 0.3069765865802765, 'predictions_19': 0.06924822926521301, 'predictions_20': -0.34636712074279785, 'predictions_21': -0.8682994246482849, 'predictions_22': -0.9693347215652466, 'predictions_23': -0.4783455729484558, 'predictions_24': 0.8846019506454468, 'predictions_25': -0.053959351032972336, 'predictions_26': -0.5963033437728882, 'predictions_27': -0.5447815656661987, 'predictions_28': -0.14065201580524445, 'predictions_29': 0.1068669781088829, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': -0.024941062554717064, 'predictions_1': -1.5298322439193726, 'predictions_2': 0.01525309681892395, 'predictions_3': -0.11309991031885147, 'predictions_4': 0.5172238945960999, 'predictions_5': 0.0024520084261894226, 'predictions_6': 0.04250117391347885, 'predictions_7': 0.5354383587837219, 'predictions_8': 0.2944106459617615, 'predictions_9': -0.0466727539896965, 'predictions_10': 0.003290211781859398, 'predictions_11': 0.0188205074518919, 'predictions_12': -0.051007892936468124, 'predictions_13': -1.3502459526062012, 'predictions_14': -0.033535201102495193, 'predictions_15': 0.10678261518478394, 'predictions_16': -0.0014769809786230326, 'predictions_17': 0.13856875896453857, 'predictions_18': 0.3855004608631134, 'predictions_19': -0.008676297031342983, 'predictions_20': -0.3423697352409363, 'predictions_21': -0.8320394158363342, 'predictions_22': -0.677459716796875, 'predictions_23': -0.5385064482688904, 'predictions_24': 0.49956056475639343, 'predictions_25': 0.005095706321299076, 'predictions_26': 1.6938838958740234, 'predictions_27': 0.842454195022583, 'predictions_28': -0.14346282184123993, 'predictions_29': 0.07541733235120773, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.009930070489645004, 'predictions_1': -0.27544116973876953, 'predictions_2': -0.002613282995298505, 'predictions_3': 0.04113232344388962, 'predictions_4': -0.055784229189157486, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.04250117391347885, 'predictions_7': 0.6722134947776794, 'predictions_8': 0.30412089824676514, 'predictions_9': 0.007771771401166916, 'predictions_10': 0.009325551800429821, 'predictions_11': 0.022453438490629196, 'predictions_12': 0.10843253135681152, 'predictions_13': 0.8989681005477905, 'predictions_14': -0.02233714796602726, 'predictions_15': 0.5417506694793701, 'predictions_16': -0.0010774275287985802, 'predictions_17': 0.16581103205680847, 'predictions_18': 0.406240850687027, 'predictions_19': -0.004930071532726288, 'predictions_20': 0.31123194098472595, 'predictions_21': -0.0816090777516365, 'predictions_22': 2.1133084297180176, 'predictions_23': 0.4506610333919525, 'predictions_24': -0.7193129658699036, 'predictions_25': 0.013311417773365974, 'predictions_26': 1.1290695667266846, 'predictions_27': 0.6277202367782593, 'predictions_28': -0.12229753285646439, 'predictions_29': 0.1477402150630951, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': 0.0006543132476508617, 'predictions_1': -0.8427131175994873, 'predictions_2': -0.002613282995298505, 'predictions_3': 0.04113232344388962, 'predictions_4': -0.2789610028266907, 'predictions_5': 0.004407345782965422, 'predictions_6': 0.04250117391347885, 'predictions_7': 0.7143340706825256, 'predictions_8': 0.30412089824676514, 'predictions_9': 0.0335814394056797, 'predictions_10': 0.009325551800429821, 'predictions_11': -0.11305594444274902, 'predictions_12': -0.018516669049859047, 'predictions_13': 0.10892590880393982, 'predictions_14': -0.08540848642587662, 'predictions_15': -0.09607099741697311, 'predictions_16': -0.0014769809786230326, 'predictions_17': 0.13856875896453857, 'predictions_18': 0.44763222336769104, 'predictions_19': 0.014222224242985249, 'predictions_20': 0.33993175625801086, 'predictions_21': -0.6289445757865906, 'predictions_22': 1.987050175666809, 'predictions_23': 0.5226719379425049, 'predictions_24': -0.6285008788108826, 'predictions_25': 0.003129873890429735, 'predictions_26': 1.1616733074188232, 'predictions_27': 0.8699957132339478, 'predictions_28': -0.12550924718379974, 'predictions_29': 0.1477402150630951, 'predictions_30': 0.5059375166893005}\n",
-      "{'predictions_0': -0.01911582238972187, 'predictions_1': -0.41529545187950134, 'predictions_2': -0.0034923271741718054, 'predictions_3': -0.11309991031885147, 'predictions_4': -0.30164235830307007, 'predictions_5': -0.009094981476664543, 'predictions_6': -0.08258295059204102, 'predictions_7': -0.8088644742965698, 'predictions_8': 0.2545676529407501, 'predictions_9': 0.015658268705010414, 'predictions_10': -0.08612797409296036, 'predictions_11': -0.003732672892510891, 'predictions_12': -0.24946996569633484, 'predictions_13': -0.9515364766120911, 'predictions_14': -0.030668942257761955, 'predictions_15': 0.4303477704524994, 'predictions_16': -0.007885736413300037, 'predictions_17': -0.168026864528656, 'predictions_18': 0.25735917687416077, 'predictions_19': 0.0024438181426376104, 'predictions_20': -0.5780957937240601, 'predictions_21': -0.3737794756889343, 'predictions_22': -2.412389039993286, 'predictions_23': -0.6488395929336548, 'predictions_24': -0.6861838698387146, 'predictions_25': -0.031086977571249008, 'predictions_26': -0.5648106932640076, 'predictions_27': -0.9229058623313904, 'predictions_28': -0.14026358723640442, 'predictions_29': -0.1724514663219452, 'predictions_30': 0.5059375166893005}\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "\u001b[2m\u001b[36m(BlockWorker pid=1494413)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494469)\u001b[0m /home/ubuntu/ray/venv/lib/python3.8/site-packages/xgboost/compat.py:31: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494469)\u001b[0m   from pandas import MultiIndex, Int64Index\n",
-      "\u001b[2m\u001b[36m(pid=1494469)\u001b[0m FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494469)\u001b[0m FutureWarning: pandas.Float64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
-      "\u001b[2m\u001b[36m(pid=1494469)\u001b[0m FutureWarning: pandas.UInt64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n"
+      "\r"
      ]
     }
    ],
@@ -521,7 +508,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.8.10"
+   "version": "3.11.2"
   },
   "vscode": {
    "interpreter": {
diff --git a/doc/source/ray-air/examples/xgboost_starter.py b/doc/source/ray-air/examples/xgboost_starter.py
index 3925e1672b09..d988ca68b41e 100644
--- a/doc/source/ray-air/examples/xgboost_starter.py
+++ b/doc/source/ray-air/examples/xgboost_starter.py
@@ -68,21 +68,3 @@
 best_result = result_grid.get_best_result()
 print("Best result:", best_result)
 # __air_tune_generic_end__
-
-# __air_xgb_batchpred_start__
-from ray.train.batch_predictor import BatchPredictor
-from ray.train.xgboost import XGBoostPredictor
-
-# You can also create a checkpoint from a trained model using
-# `XGBoostCheckpoint.from_model`.
-checkpoint = best_result.checkpoint
-
-batch_predictor = BatchPredictor.from_checkpoint(checkpoint, XGBoostPredictor)
-
-predicted_probabilities = batch_predictor.predict(test_dataset)
-predicted_probabilities.show()
-# {'predictions': 0.9970690608024597}
-# {'predictions': 0.9943051934242249}
-# {'predictions': 0.00334902573376894}
-# ...
-# __air_xgb_batchpred_end__
diff --git a/doc/source/ray-air/getting-started.rst b/doc/source/ray-air/getting-started.rst
index dfdfad46f5e0..5b467cb3230a 100644
--- a/doc/source/ray-air/getting-started.rst
+++ b/doc/source/ray-air/getting-started.rst
@@ -177,30 +177,12 @@ Then use the ``Tuner`` to run the search:
 Batch Inference
 ~~~~~~~~~~~~~~~
 
-After running the steps in :ref:`Training <air-getting-started-training>` or :ref:`Tuning <air-getting-started-tuning>`, use the trained model for scalable batch prediction with a ``BatchPredictor``.
+After running the steps in :ref:`Training <air-getting-started-training>` or
+:ref:`Tuning <air-getting-started-tuning>`, use the trained model for scalable batch
+prediction with :meth:`Dataset.map_batches() <ray.data.Dataset.map_batches>`.
 
-.. tabs::
-
-    .. group-tab:: XGBoost
-
-        .. literalinclude:: examples/xgboost_starter.py
-            :language: python
-            :start-after: __air_xgb_batchpred_start__
-            :end-before: __air_xgb_batchpred_end__
-
-    .. group-tab:: Pytorch
+To learn more, see :ref:`End-to-end: Offline Batch Inference <batch_inference_home>`.
 
-        .. literalinclude:: examples/pytorch_tabular_starter.py
-            :language: python
-            :start-after: __air_pytorch_batchpred_start__
-            :end-before: __air_pytorch_batchpred_end__
-
-    .. group-tab:: Tensorflow
-
-        .. literalinclude:: examples/tf_tabular_starter.py
-            :language: python
-            :start-after: __air_tf_batchpred_start__
-            :end-before: __air_tf_batchpred_end__
 
 Project Status
 --------------
@@ -216,4 +198,4 @@ Next Steps
 - :ref:`air-examples-ref`
 - :ref:`API reference <air-api-ref>`
 - :ref:`Technical whitepaper <whitepaper>`
-- To check how your application is doing, you can use the :ref:`Ray dashboard<observability-getting-started>`. 
+- To check how your application is doing, you can use the :ref:`Ray dashboard<observability-getting-started>`.
diff --git a/doc/source/ray-air/predictors.rst b/doc/source/ray-air/predictors.rst
index 0656bcd4428a..6b87de6f3917 100644
--- a/doc/source/ray-air/predictors.rst
+++ b/doc/source/ray-air/predictors.rst
@@ -138,51 +138,6 @@ Here are some examples:
     :end-before: __configure_batch_predictor_scaling_end__
 
 
-
-Batch Inference Examples
-------------------------
-Below, we provide examples of using common frameworks to do batch inference for different data types:
-
-Tabular
-~~~~~~~
-
-.. tab-set::
-
-    .. tab-item:: XGBoost
-
-        .. literalinclude:: examples/xgboost_batch_prediction.py
-            :language: python
-
-    .. tab-item:: Pytorch
-
-        .. literalinclude:: examples/pytorch_tabular_batch_prediction.py
-            :language: python
-
-    .. tab-item:: Tensorflow
-
-        .. literalinclude:: examples/tf_tabular_batch_prediction.py
-            :language: python
-
-Image
-~~~~~
-
-.. tab-set::
-
-    .. tab-item:: Pytorch
-
-        .. literalinclude:: examples/torch_image_batch_pretrained.py
-            :language: python
-
-
-    .. tab-item:: Tensorflow
-
-        Coming soon!
-
-Text
-~~~~
-
-Coming soon!
-
 Developer Guide: Implementing your own Predictor
 ------------------------------------------------
 
diff --git a/doc/source/ray-overview/examples.rst b/doc/source/ray-overview/examples.rst
index 1ab70176df2b..176e7a6f4876 100644
--- a/doc/source/ray-overview/examples.rst
+++ b/doc/source/ray-overview/examples.rst
@@ -161,13 +161,6 @@ Ray Examples
 
             How to fine-tune a DreamBooth text-to-image model with your own images.
 
-    .. grid-item-card:: :bdg-secondary:`Code example`
-        :class-item: gallery-item data-processing inference
-
-        .. button-ref:: /ray-air/examples/opt_deepspeed_batch_inference
-
-            How to run batch inference on a dataset of texts with a 30B OPT model
-
     .. grid-item-card:: :bdg-secondary:`Code example`
         :class-item: gallery-item training
 
diff --git a/doc/source/ray-overview/getting-started.md b/doc/source/ray-overview/getting-started.md
index 8a0fce79f92a..6ff4267c9bbb 100644
--- a/doc/source/ray-overview/getting-started.md
+++ b/doc/source/ray-overview/getting-started.md
@@ -75,13 +75,10 @@ Run hyperparameter tuning with Ray Tune to find the best model:
 
 `````{dropdown} Use the trained model for Batch prediction
 
-Use the trained model for batch prediction with a ``BatchPredictor``.
+Use the trained model for batch prediction with
+``Dataset.map_batches()``.
 
-```{literalinclude} ../ray-air/examples/xgboost_starter.py
-    :language: python
-    :start-after: __air_xgb_batchpred_start__
-    :end-before: __air_xgb_batchpred_end__
-```
+To learn more, see :ref:`End-to-end: Offline Batch Inference <batch_inference_home>`.
 
 ```{button-ref} air
 :color: primary
diff --git a/doc/source/train/examples/lightning/lightning_cola_advanced.ipynb b/doc/source/train/examples/lightning/lightning_cola_advanced.ipynb
index 488532146c38..f6807aabd53b 100644
--- a/doc/source/train/examples/lightning/lightning_cola_advanced.ipynb
+++ b/doc/source/train/examples/lightning/lightning_cola_advanced.ipynb
@@ -11,7 +11,7 @@
     "\n",
     ":::{note}\n",
     "\n",
-    "This is an advanced example for {class}`LightningTrainer <ray.train.lightning.LightningTrainer>`, which demonstrates how to use LightningTrainer with {ref}`Dataset <data>` and {ref}`Batch Predictor <air-predictors>`. \n",
+    "This is an advanced example for {class}`LightningTrainer <ray.train.lightning.LightningTrainer>`, which demonstrates how to use LightningTrainer with {ref}`Dataset <data>`. \n",
     "\n",
     "If you just want to quickly convert your existing PyTorch Lightning scripts into Ray AIR, you can refer to this starter example:\n",
     "{ref}`Train a Pytorch Lightning Image Classifier <lightning_mnist_example>`.\n",
@@ -22,9 +22,7 @@
     "In particular, we will:\n",
     "- Create Ray Data from the original CoLA dataset.\n",
     "- Define a preprocessor to tokenize the sentences.\n",
-    "- Finetune a BERT model using LightningTrainer.\n",
-    "- Construct a BatchPredictor with the checkpoint and preprocessor.\n",
-    "- Do batch prediction on multiple GPUs, and evaluate the results."
+    "- Finetune a BERT model using LightningTrainer."
    ]
   },
   {
@@ -41,6 +39,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -57,6 +56,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -1369,121 +1369,6 @@
     "result"
    ]
   },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Do Batch Inference with a Saved Checkpoint"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Now that we have fine-tuned the module, we can load the checkpoint into a BatchPredictor and perform fast inference with multiple GPUs. It will distribute the inference workload across multiple workers when calling `predict()` and run prediction on multiple shards of data in parallel. \n",
-    "\n",
-    "You can find more details in [Using Predictors for Inference](air-predictors)."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 17,
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "from ray.train.batch_predictor import BatchPredictor\n",
-    "from ray.train.lightning import LightningCheckpoint, LightningPredictor\n",
-    "\n",
-    "# Use in-memory checkpoint object\n",
-    "checkpoint = result.checkpoint\n",
-    "\n",
-    "# You can also load a checkpoint from disk:\n",
-    "# YOUR_CHECKPOINT_DIR = result.checkpoint.path\n",
-    "# checkpoint = LightningCheckpoint.from_directory(YOUR_CHECKPOINT_DIR)\n",
-    "\n",
-    "batch_predictor = BatchPredictor(\n",
-    "    checkpoint=checkpoint,\n",
-    "    predictor_cls=LightningPredictor,\n",
-    "    use_gpu=True,\n",
-    "    model_class=SentimentModel,\n",
-    "    preprocessor=preprocessor,\n",
-    ")\n",
-    "\n",
-    "# Use 2 GPUs for batch inference\n",
-    "predictions = batch_predictor.predict(\n",
-    "    ray_datasets[\"validation\"],\n",
-    "    feature_columns=[\"input_ids\", \"attention_mask\", \"label\"],\n",
-    "    keep_columns=[\"label\"],\n",
-    "    batch_size=16,\n",
-    "    min_scoring_workers=2,\n",
-    "    max_scoring_workers=2,\n",
-    "    num_gpus_per_worker=1,\n",
-    ")"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "We obtained a Ray dataset containing predictions from `batch_predictor.predict()`. Now we can easily evaluate the results with just a few lines of code:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "# Internally, BatchPredictor calls forward() method of the LightningModule.\n",
-    "# Convert the logits tensor into labels with argmax.\n",
-    "def argmax(batch):\n",
-    "    batch[\"predictions\"] = batch[\"predictions\"].apply(lambda x: np.argmax(x))\n",
-    "    return batch\n",
-    "\n",
-    "\n",
-    "results = predictions.map_batches(argmax, batch_format=\"pandas\").to_pandas()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 23,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "   predictions  label\n",
-      "0            1      1\n",
-      "1            1      1\n",
-      "2            0      1\n",
-      "3            1      1\n",
-      "4            0      0\n",
-      "5            1      0\n",
-      "6            1      0\n",
-      "7            1      1\n",
-      "8            1      1\n",
-      "9            1      1\n",
-      "\n",
-      "{'matthews_correlation': 0.5899314497879129}\n"
-     ]
-    }
-   ],
-   "source": [
-    "matthews_corr = metric.compute(\n",
-    "    predictions=results[\"predictions\"], references=results[\"label\"]\n",
-    ")\n",
-    "print(results.head(10))\n",
-    "print(matthews_corr)"
-   ]
-  },
   {
    "attachments": {},
    "cell_type": "markdown",
@@ -1493,7 +1378,8 @@
     "\n",
     "- {ref}`Fine-tune a Large Language Model with LightningTrainer and FSDP <dolly_lightning_fsdp_finetuning>`\n",
     "- {ref}`Hyperparameter searching with LightningTrainer + Ray Tune. <tune-pytorch-lightning-ref>`\n",
-    "- {ref}`Experiment Tracking with Wandb, CometML, MLFlow, and Tensorboard in LightningTrainer <lightning_experiment_tracking>`"
+    "- {ref}`Experiment Tracking with Wandb, CometML, MLFlow, and Tensorboard in LightningTrainer <lightning_experiment_tracking>`\n",
+    "- {ref}`End-to-end: Offline Batch Inference <batch_inference_home>`"
    ]
   }
  ],
diff --git a/python/ray/train/examples/huggingface/huggingface_basic_language_modeling_example.py b/python/ray/train/examples/huggingface/huggingface_basic_language_modeling_example.py
index 78fd376676f7..548275819a76 100644
--- a/python/ray/train/examples/huggingface/huggingface_basic_language_modeling_example.py
+++ b/python/ray/train/examples/huggingface/huggingface_basic_language_modeling_example.py
@@ -6,7 +6,6 @@
 import argparse
 import tempfile
 
-import pandas as pd
 import torch
 from datasets import load_dataset
 from transformers import (
@@ -19,11 +18,7 @@
 
 import ray
 import ray.data
-from ray.train.batch_predictor import BatchPredictor
-from ray.train.huggingface import (
-    TransformersPredictor,
-    TransformersTrainer,
-)
+from ray.train.huggingface import TransformersTrainer
 from ray.air.config import ScalingConfig
 
 
@@ -128,19 +123,6 @@ def train_function(train_dataset, eval_dataset=None, **config):
     results = trainer.fit()
     print(results.metrics)
 
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_checkpoint)
-    prompt = ["My text: Complete me..."]
-    predictor = BatchPredictor.from_checkpoint(
-        results.checkpoint,
-        TransformersPredictor,
-        task="text-generation",
-        tokenizer=tokenizer,
-    )
-    data = ray.data.from_pandas(pd.DataFrame(prompt, columns=["prompt"]))
-    prediction = predictor.predict(data, num_gpus_per_worker=int(use_gpu))
-
-    print(f"Generated text for prompt '{prompt}': '{prediction.take(1)}'")
-
 
 if __name__ == "__main__":
     # Training settings
diff --git a/python/ray/train/examples/pytorch/torch_regression_example.py b/python/ray/train/examples/pytorch/torch_regression_example.py
index 663f4af5a621..b49fe0558aff 100644
--- a/python/ray/train/examples/pytorch/torch_regression_example.py
+++ b/python/ray/train/examples/pytorch/torch_regression_example.py
@@ -1,7 +1,6 @@
 import argparse
 from typing import Tuple
 
-import numpy as np
 import pandas as pd
 from ray.air.checkpoint import Checkpoint
 
@@ -11,10 +10,8 @@
 import ray
 import ray.train as train
 from ray.air import session
-from ray.air.result import Result
 from ray.data import Dataset
-from ray.train.batch_predictor import BatchPredictor
-from ray.train.torch import TorchPredictor, TorchTrainer
+from ray.train.torch import TorchTrainer
 from ray.air.config import ScalingConfig
 
 
@@ -126,19 +123,6 @@ def train_regression(num_workers=2, use_gpu=False):
     return result
 
 
-def predict_regression(result: Result):
-    batch_predictor = BatchPredictor.from_checkpoint(result.checkpoint, TorchPredictor)
-
-    df = pd.DataFrame(
-        [[np.random.uniform(0, 1, size=100)] for i in range(100)], columns=["x"]
-    )
-    prediction_dataset = ray.data.from_pandas(df)
-
-    predictions = batch_predictor.predict(prediction_dataset, dtype=torch.float)
-
-    return predictions
-
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -170,5 +154,4 @@ def predict_regression(result: Result):
     else:
         ray.init(address=args.address)
         result = train_regression(num_workers=args.num_workers, use_gpu=args.use_gpu)
-    predictions = predict_regression(result)
-    print(predictions.to_pandas())
+    print(result)
diff --git a/python/ray/train/examples/tf/tensorflow_autoencoder_example.py b/python/ray/train/examples/tf/tensorflow_autoencoder_example.py
index c0a91307af6f..678ccaad2416 100644
--- a/python/ray/train/examples/tf/tensorflow_autoencoder_example.py
+++ b/python/ray/train/examples/tf/tensorflow_autoencoder_example.py
@@ -9,8 +9,6 @@
 import tensorflow as tf
 import tensorflow_datasets as tfds
 from ray.data.datasource import SimpleTensorFlowDatasource
-from ray.air.batch_predictor import BatchPredictor
-from ray.air.predictors.integrations.tensorflow import TensorflowPredictor
 from ray.air.result import Result
 from ray.train.tensorflow import TensorflowTrainer
 from ray.train.tensorflow import prepare_dataset_shard
@@ -137,61 +135,6 @@ def train_tensorflow_mnist(
     return results
 
 
-def predict_tensorflow_mnist(result: Result) -> ray.data.Dataset:
-    test_dataset = get_dataset(split_type="test")
-    batch_predictor = BatchPredictor.from_checkpoint(
-        result.checkpoint, TensorflowPredictor, model_definition=build_autoencoder_model
-    )
-
-    predictions = batch_predictor.predict(
-        test_dataset, feature_columns=["image"], dtype=tf.float32
-    )
-
-    pandas_predictions = predictions.to_pandas(float("inf"))
-    print(f"PREDICTIONS\n{pandas_predictions}")
-
-    return pandas_predictions
-
-
-def visualize_tensorflow_mnist_autoencoder(result: Result) -> None:
-    test_dataset = get_dataset(split_type="test")
-    batch_predictor = BatchPredictor.from_checkpoint(
-        result.checkpoint, TensorflowPredictor, model_definition=build_autoencoder_model
-    )
-
-    # test_dataset.
-    predictions = batch_predictor.predict(
-        test_dataset, feature_columns=["image"], dtype=tf.float32
-    )
-
-    pandas_predictions = predictions.to_pandas(float("inf"))
-
-    decoded_imgs = pandas_predictions["predictions"].values
-    x_test = test_dataset.to_pandas(float("inf"))["image"].values
-
-    import matplotlib.pyplot as plt
-
-    n = 10  # How many digits we will display
-    plt.figure(figsize=(20, 4))
-    for i in range(n):
-        # Display original
-        ax = plt.subplot(2, n, i + 1)
-        plt.imshow(np.asarray(x_test[i]).reshape(28, 28))
-        plt.gray()
-        ax.get_xaxis().set_visible(False)
-        ax.get_yaxis().set_visible(False)
-
-        # Display reconstruction
-        ax = plt.subplot(2, n, i + 1 + n)
-        plt.imshow(np.asarray(decoded_imgs[i]).reshape(28, 28))
-        plt.gray()
-        ax.get_xaxis().set_visible(False)
-        ax.get_yaxis().set_visible(False)
-
-    # how to retrieve the folderpath of the checkpoint
-    plt.savefig("test.png")
-
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -219,8 +162,6 @@ def visualize_tensorflow_mnist_autoencoder(result: Result) -> None:
 
     args, _ = parser.parse_known_args()
 
-    import ray
-
     if args.smoke_test:
         # 2 workers, 1 for trainer, 1 for datasets
         num_gpus = args.num_workers if args.use_gpu else 0
@@ -231,6 +172,4 @@ def visualize_tensorflow_mnist_autoencoder(result: Result) -> None:
         result = train_tensorflow_mnist(
             num_workers=args.num_workers, use_gpu=args.use_gpu, epochs=args.epochs
         )
-
-    predict_tensorflow_mnist(result)
-    visualize_tensorflow_mnist_autoencoder(result)
+    print(result)
diff --git a/python/ray/train/examples/tf/tensorflow_regression_example.py b/python/ray/train/examples/tf/tensorflow_regression_example.py
index 5b130a9947b4..80aeafc208d1 100644
--- a/python/ray/train/examples/tf/tensorflow_regression_example.py
+++ b/python/ray/train/examples/tf/tensorflow_regression_example.py
@@ -1,20 +1,13 @@
 import argparse
 
-import numpy as np
-import pandas as pd
 import tensorflow as tf
 
 import ray
 from ray.air import session
 from ray.air.integrations.keras import ReportCheckpointCallback
 from ray.air.result import Result
-from ray.data import Dataset
 from ray.data.preprocessors import Concatenator
-from ray.train.batch_predictor import BatchPredictor
-from ray.train.tensorflow import (
-    TensorflowPredictor,
-    TensorflowTrainer,
-)
+from ray.train.tensorflow import TensorflowTrainer
 from ray.air.config import ScalingConfig
 
 
@@ -75,24 +68,6 @@ def train_tensorflow_regression(num_workers: int = 2, use_gpu: bool = False) ->
     return results
 
 
-def predict_regression(result: Result) -> Dataset:
-    batch_predictor = BatchPredictor.from_checkpoint(
-        result.checkpoint, TensorflowPredictor, model_definition=build_model
-    )
-
-    df = pd.DataFrame(
-        [[np.random.uniform(0, 1, size=100)] for i in range(100)], columns=["x"]
-    )
-    prediction_dataset = ray.data.from_pandas(df)
-
-    predictions = batch_predictor.predict(prediction_dataset, dtype=tf.float32)
-
-    print("PREDICTIONS")
-    predictions.show()
-
-    return predictions
-
-
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
@@ -127,4 +102,4 @@ def predict_regression(result: Result) -> Dataset:
         result = train_tensorflow_regression(
             num_workers=args.num_workers, use_gpu=args.use_gpu
         )
-    predict_regression(result)
+    print(result)
diff --git a/release/release_tests.yaml b/release/release_tests.yaml
index 36320e7ea8a3..5e066064ed75 100644
--- a/release/release_tests.yaml
+++ b/release/release_tests.yaml
@@ -802,7 +802,7 @@
   team: data
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       type: gpu
       runtime_env:
         - RAY_task_oom_retries=50
@@ -904,29 +904,6 @@
     timeout: 4700
     script: python test_myst_doc.py --path lightning-llm-finetuning-7b.ipynb
 
-
-- name: air_example_opt_deepspeed_batch_inference
-  group: AIR examples
-  working_dir: air_examples/opt_deepspeed_batch_inference
-
-  python: "3.9"
-
-  frequency: weekly
-  team: ml
-  cluster:
-    byod:
-      type: gpu
-      pip:
-        - myst-parser==0.15.2
-        - myst-nb==0.13.1
-        - jupytext==1.13.6
-    cluster_env: 30b_deepspeed_env.yaml
-    cluster_compute: 30b_deepspeed_compute.yaml
-
-  run:
-    timeout: 3600
-    script: python test_myst_doc.py --path opt_deepspeed_batch_inference.ipynb
-
   # variations: TODO(jungong): add GCP variation.
 
 
@@ -941,7 +918,7 @@
   frequency: nightly-3x
   team: data
   cluster:
-    byod: 
+    byod:
       type: gpu
     cluster_env: ../testing/cluster_envs/default_cluster_env_nightly_ml_py39.yaml
     cluster_compute: ../testing/compute_configs/gpu/aws.yaml
@@ -2330,7 +2307,7 @@
   team: ml
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       type: gpu
     cluster_env: torch_tune_serve_app_config.yaml
     cluster_compute: gpu_tpl_aws.yaml
@@ -2784,7 +2761,7 @@
   team: serve
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       runtime_env:
         - RLLIB_TEST_NO_JAX_IMPORT=1
     cluster_env: app_config.yaml
@@ -2826,7 +2803,7 @@
   team: serve
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       runtime_env:
         - RLLIB_TEST_NO_JAX_IMPORT=1
     cluster_env: app_config.yaml
@@ -2868,7 +2845,7 @@
   team: serve
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       runtime_env:
         - RLLIB_TEST_NO_JAX_IMPORT=1
     cluster_env: app_config.yaml
@@ -3583,7 +3560,7 @@
   team: ml
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       type: gpu
       post_build_script: byod_alpa_test.sh
     cluster_env: app_config.yaml
@@ -3621,7 +3598,7 @@
   team: ml
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       type: gpu
       post_build_script: byod_alpa_test.sh
     cluster_env: app_config.yaml
@@ -4745,7 +4722,7 @@
   team: rllib
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       type: gpu
       post_build_script: byod_rllib_test.sh
       runtime_env:
@@ -4780,7 +4757,7 @@
   team: rllib
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       type: gpu
       post_build_script: byod_rllib_test.sh
       runtime_env:
@@ -5087,7 +5064,7 @@
   group: core-daily-test
   working_dir: nightly_tests
 
-  stable: false 
+  stable: false
 
   python: "3.8"
   frequency: nightly
@@ -5140,7 +5117,7 @@
   cluster:
     byod:
       runtime_env:
-        - RAY_MAX_LIMIT_FROM_API_SERVER=1000000000 
+        - RAY_MAX_LIMIT_FROM_API_SERVER=1000000000
         - RAY_MAX_LIMIT_FROM_DATA_SOURCE=1000000000
     cluster_env: shuffle/shuffle_with_state_api_app_config.yaml
     cluster_compute: shuffle/shuffle_compute_single.yaml
@@ -5790,7 +5767,7 @@
   team: data
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       type: gpu
     cluster_env: app_config.yaml
     cluster_compute: inference.yaml
@@ -5845,7 +5822,7 @@
   team: data
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       type: gpu
     cluster_env: app_config.yaml
     cluster_compute: single_node_benchmark_compute.yaml
@@ -5904,7 +5881,7 @@
   team: data
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       type: gpu
     cluster_env: app_config.yaml
     cluster_compute: data_ingest_benchmark_compute.yaml
@@ -6208,7 +6185,7 @@
   team: data
   python: "3.8"
   cluster:
-    byod: 
+    byod:
       type: gpu
     cluster_env: app_config.yaml
     cluster_compute: single_node_benchmark_compute.yaml
@@ -6870,4 +6847,4 @@
 
   run:
     timeout: 2400
-    script: python launch_and_verify_cluster.py gcp/example-full.yaml
\ No newline at end of file
+    script: python launch_and_verify_cluster.py gcp/example-full.yaml