{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-tune a Hugging Face Transformers Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VaFMt6AIhYbK"
   },
   "source": [
    "This notebook is based on an official Hugging Face example, [How to fine-tune a model on text classification](https://github.com/huggingface/notebooks/blob/6ca682955173cc9d36ffa431ddda505a048cbe80/examples/text_classification.ipynb). This notebook shows the process of conversion from vanilla HF to Ray Train without changing the training logic unless necessary.\n",
    "\n",
    "This notebook consists of the following steps:\n",
    "1. [Set up Ray](#hf-setup)\n",
    "2. [Load the dataset](#hf-load)\n",
    "3. [Preprocess the dataset with Ray Data](#hf-preprocess)\n",
    "4. [Run the training with Ray Train](#hf-train)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "sQbdfyWQhYbO"
   },
   "source": [
    "Uncomment and run the following line to install all the necessary dependencies. (This notebook is being tested with `transformers==4.19.1`.):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "YajFzmkthYbO"
   },
   "outputs": [],
   "source": [
    "#! pip install \"datasets\" \"transformers>=4.19.0\" \"torch>=1.10.0\" \"mlflow\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "pvSRaEHChYbP"
   },
   "source": [
    "(hf-setup)=\n",
    "## Set up Ray"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "LRdL3kWBhYbQ"
   },
   "source": [
    "Use `ray.init()` to initialize a local cluster. By default, this cluster contains only the machine you are running this notebook on. You can also run this notebook on an [Anyscale](https://www.anyscale.com/) cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "MOsHUjgdIrIW",
    "outputId": "e527bdbb-2f28-4142-cca0-762e0566cbcd",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "import ray\n",
    "\n",
    "ray.init()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "oJiSdWy2hYbR"
   },
   "source": [
    "Check the resources our cluster is composed of. If you are running this notebook on your local machine or Google Colab, you should see the number of CPU cores and GPUs available on your machine."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KlMz0dt9hYbS",
    "outputId": "2d485449-ee69-4334-fcba-47e0ceb63078",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'CPU': 48.0,\n",
      " 'GPU': 4.0,\n",
      " 'accelerator_type:T4': 1.0,\n",
      " 'anyscale/accelerator_shape:4xT4': 1.0,\n",
      " 'anyscale/node-group:head': 1.0,\n",
      " 'anyscale/provider:aws': 1.0,\n",
      " 'anyscale/region:us-west-2': 1.0,\n",
      " 'memory': 206158430208.0,\n",
      " 'node:10.0.114.132': 1.0,\n",
      " 'node:__internal_head__': 1.0,\n",
      " 'object_store_memory': 58913938636.0}\n"
     ]
    }
   ],
   "source": [
    "pprint(ray.cluster_resources())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uS6oeJELhYbS"
   },
   "source": [
    "This notebook fine-tunes a [HF Transformers](https://github.com/huggingface/transformers) model for one of the text classification task of the [GLUE Benchmark](https://gluebenchmark.com/). It runs the training using Ray Train.\n",
    "\n",
    "You can change these two variables to control whether the training, which happens later, uses CPUs or GPUs, and how many workers to spawn. Each worker claims one CPU or GPU. Make sure to not request more resources than the resources present. By default, the training runs with one GPU worker."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "id": "gAbhv9OqhYbT",
    "tags": []
   },
   "outputs": [],
   "source": [
    "use_gpu = True  # set this to False to run on CPUs\n",
    "num_workers = 1  # set this to number of GPUs or CPUs you want to use"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "rEJBSTyZIrIb"
   },
   "source": [
    "## Fine-tune a model on a text classification task"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kTCFado4IrIc"
   },
   "source": [
    "The GLUE Benchmark is a group of nine classification tasks on sentences or pairs of sentences. To learn more, see the [original notebook](https://github.com/huggingface/notebooks/blob/6ca682955173cc9d36ffa431ddda505a048cbe80/examples/text_classification.ipynb).\n",
    "\n",
    "Each task has a name that is its acronym, with `mnli-mm` to indicate that it is a mismatched version of MNLI. Each one has the same training set as `mnli` but different validation and test sets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "id": "YZbiBDuGIrId",
    "tags": []
   },
   "outputs": [],
   "source": [
    "GLUE_TASKS = [\n",
    "    \"cola\",\n",
    "    \"mnli\",\n",
    "    \"mnli-mm\",\n",
    "    \"mrpc\",\n",
    "    \"qnli\",\n",
    "    \"qqp\",\n",
    "    \"rte\",\n",
    "    \"sst2\",\n",
    "    \"stsb\",\n",
    "    \"wnli\",\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4RRkXuteIrIh"
   },
   "source": [
    "This notebook runs on any of the tasks in the list above, with any model checkpoint from the [Model Hub](https://huggingface.co/models) as long as that model has a version with a classification head. Depending on the model and the GPU you are using, you might need to adjust the batch size to avoid out-of-memory errors. Set these three parameters, and the rest of the notebook should run smoothly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "zVvslsfMIrIh",
    "tags": []
   },
   "outputs": [],
   "source": [
    "task = \"cola\"\n",
    "model_checkpoint = \"distilbert-base-uncased\"\n",
    "batch_size = 16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "whPRbBNbIrIl"
   },
   "source": [
    "(hf-load)=\n",
    "### Loading the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "W7QYTpxXIrIl"
   },
   "source": [
    "Use the [HF Datasets](https://github.com/huggingface/datasets) library to download the data and get the metric to use for evaluation and to compare your model to the benchmark. You can do this comparison easily with the `load_dataset` and `load_metric` functions.\n",
    "\n",
    "Apart from `mnli-mm` being special code, you can directly pass the task name to those functions.\n",
    "\n",
    "Run the normal HF Datasets code to load the dataset from the Hub."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 200
    },
    "id": "MwhAeEOuhYbV",
    "outputId": "3aff8c73-d6eb-4784-890a-a419403b5bda",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "actual_task = \"mnli\" if task == \"mnli-mm\" else task\n",
    "datasets = load_dataset(\"glue\", actual_task)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RzfPtOMoIrIu"
   },
   "source": [
    "The `dataset` object itself is a [`DatasetDict`](https://huggingface.co/docs/datasets/package_reference/main_classes.html#datasetdict), which contains one key for the training, validation, and test set, with more keys for the mismatched validation and test set in the special case of `mnli`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "n9qywopnIrJH"
   },
   "source": [
    "(hf-preprocess)=\n",
    "### Preprocessing the data with Ray Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "YVx71GdAIrJH"
   },
   "source": [
    "Before you can feed these texts to the model, you need to preprocess them. Preprocess them with a HF Transformers' `Tokenizer`, which tokenizes the inputs, including converting the tokens to their corresponding IDs in the pretrained vocabulary, and puts them in a format the model expects. It also generates the other inputs that the model requires.\n",
    "\n",
    "To do all of this preprocessing, instantiate your tokenizer with the `AutoTokenizer.from_pretrained` method, which ensures that you:\n",
    "\n",
    "- Get a tokenizer that corresponds to the model architecture you want to use.\n",
    "- Download the vocabulary used when pretraining this specific checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 145
    },
    "id": "eXNLu_-nIrJI",
    "outputId": "f545a7a5-f341-4315-cd89-9942a657aa31",
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ray/anaconda3/lib/python3.9/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "  _torch_pytree._register_pytree_node(\n",
      "/home/ray/anaconda3/lib/python3.9/site-packages/huggingface_hub/file_download.py:795: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cbfed0e37b1f4546a98878cc6090ff6f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer_config.json:   0%|          | 0.00/48.0 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f2992807084545c4b4e6e1b9e476c5b0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "config.json:   0%|          | 0.00/483 [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "cb6c8740578747ab8005379b899bd54f",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "vocab.txt:   0%|          | 0.00/232k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "07f22216bcee452fa0fb94b84c52e8bc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "tokenizer.json:   0%|          | 0.00/466k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vl6IidfdIrJK"
   },
   "source": [
    "Pass `use_fast=True` to the preceding call to use one of the fast tokenizers, backed by Rust, from the HF Tokenizers library. These fast tokenizers are available for almost all models, but if you get an error with the previous call, remove the argument."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "qo_0B1M2IrJM"
   },
   "source": [
    "To preprocess the dataset, you need the names of the columns containing the sentence(s). The following dictionary keeps track of the correspondence task to column names:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "fyGdtK9oIrJM",
    "tags": []
   },
   "outputs": [],
   "source": [
    "task_to_keys = {\n",
    "    \"cola\": (\"sentence\", None),\n",
    "    \"mnli\": (\"premise\", \"hypothesis\"),\n",
    "    \"mnli-mm\": (\"premise\", \"hypothesis\"),\n",
    "    \"mrpc\": (\"sentence1\", \"sentence2\"),\n",
    "    \"qnli\": (\"question\", \"sentence\"),\n",
    "    \"qqp\": (\"question1\", \"question2\"),\n",
    "    \"rte\": (\"sentence1\", \"sentence2\"),\n",
    "    \"sst2\": (\"sentence\", None),\n",
    "    \"stsb\": (\"sentence1\", \"sentence2\"),\n",
    "    \"wnli\": (\"sentence1\", \"sentence2\"),\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "256fOuzjhYbY"
   },
   "source": [
    "Instead of using HF Dataset objects directly, convert them to [Ray Data](https://docs.ray.io/en/latest/data/data.html). Arrow tables back both of them, so the conversion is straightforward. Use the built-in {meth}`~ray.data.from_huggingface` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "2abe2c46da3a440fa7fd6e800562bf61",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Parquet Files Sample 0:   0%|          | 0.00/1.00 [00:00<?, ? file/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0243828d5a244ef1ae881e57d795c8c6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Parquet Files Sample 0:   0%|          | 0.00/1.00 [00:00<?, ? file/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "111662e024f8443fa4c9fd1ab81008fb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Parquet Files Sample 0:   0%|          | 0.00/1.00 [00:00<?, ? file/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'train': Dataset(num_rows=8551, schema={sentence: string, label: int64, idx: int32}),\n",
       " 'validation': Dataset(num_rows=1043, schema={sentence: string, label: int64, idx: int32}),\n",
       " 'test': Dataset(num_rows=1063, schema={sentence: string, label: int64, idx: int32})}"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import ray.data\n",
    "\n",
    "ray_datasets = {\n",
    "    \"train\": ray.data.from_huggingface(datasets[\"train\"]),\n",
    "    \"validation\": ray.data.from_huggingface(datasets[\"validation\"]),\n",
    "    \"test\": ray.data.from_huggingface(datasets[\"test\"]),\n",
    "}\n",
    "ray_datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2C0hcmp9IrJQ"
   },
   "source": [
    "You can then write the function that preprocesses the samples. Feed them to the `tokenizer` with the argument `truncation=True`. This configuration ensures that the `tokenizer` truncates and pads to the longest sequence in the batch, any input longer than what the model selected can handle."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "vc0BSBLIIrJQ",
    "tags": []
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from typing import Dict\n",
    "\n",
    "\n",
    "# Tokenize input sentences\n",
    "def collate_fn(examples: Dict[str, np.array]):\n",
    "    sentence1_key, sentence2_key = task_to_keys[task]\n",
    "    if sentence2_key is None:\n",
    "        outputs = tokenizer(\n",
    "            list(examples[sentence1_key]),\n",
    "            truncation=True,\n",
    "            padding=\"longest\",\n",
    "            return_tensors=\"pt\",\n",
    "        )\n",
    "    else:\n",
    "        outputs = tokenizer(\n",
    "            list(examples[sentence1_key]),\n",
    "            list(examples[sentence2_key]),\n",
    "            truncation=True,\n",
    "            padding=\"longest\",\n",
    "            return_tensors=\"pt\",\n",
    "        )\n",
    "\n",
    "    outputs[\"labels\"] = torch.LongTensor(examples[\"label\"])\n",
    "\n",
    "    # Move all input tensors to GPU\n",
    "    for key, value in outputs.items():\n",
    "        outputs[key] = value.cuda()\n",
    "\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "545PP3o8IrJV"
   },
   "source": [
    "(hf-train)=\n",
    "### Fine-tuning the model with Ray Train"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FBiW8UpKIrJW"
   },
   "source": [
    "Now that the data is ready, download the pretrained model and fine-tune it.\n",
    "\n",
    "Because all of the tasks involve sentence classification, use the `AutoModelForSequenceClassification` class. For more specifics about each individual training component, see the [original notebook](https://github.com/huggingface/notebooks/blob/6ca682955173cc9d36ffa431ddda505a048cbe80/examples/text_classification.ipynb). The original notebook uses the same tokenizer used to encode the dataset in this notebook's preceding example.\n",
    "\n",
    "The main difference when using Ray Train is that you need to define the training logic as a function (`train_func`). You pass this [training function](train-overview-training-function) to the {class}`~ray.train.torch.TorchTrainer` to on every Ray worker. The training then proceeds using PyTorch DDP.\n",
    "\n",
    "\n",
    "```{note}\n",
    "\n",
    "Be sure to initialize the model, metric, and tokenizer within the function. Otherwise, you may encounter serialization errors.\n",
    "\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "TlqNaB8jIrJW",
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-09 15:56:28.075767: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "2025-07-09 15:56:28.124864: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2025-07-09 15:56:28.124884: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2025-07-09 15:56:28.126125: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2025-07-09 15:56:28.133567: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2025-07-09 15:56:29.219640: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
      "/home/ray/anaconda3/lib/python3.9/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "  _torch_pytree._register_pytree_node(\n",
      "/home/ray/anaconda3/lib/python3.9/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "  _torch_pytree._register_pytree_node(\n",
      "comet_ml is installed but `COMET_API_KEY` is not set.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "\n",
    "from evaluate import load\n",
    "from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer\n",
    "\n",
    "import ray.train\n",
    "from ray.train.huggingface.transformers import prepare_trainer, RayTrainReportCallback\n",
    "\n",
    "num_labels = 3 if task.startswith(\"mnli\") else 1 if task == \"stsb\" else 2\n",
    "metric_name = (\n",
    "    \"pearson\"\n",
    "    if task == \"stsb\"\n",
    "    else \"matthews_correlation\"\n",
    "    if task == \"cola\"\n",
    "    else \"accuracy\"\n",
    ")\n",
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "validation_key = (\n",
    "    \"validation_mismatched\"\n",
    "    if task == \"mnli-mm\"\n",
    "    else \"validation_matched\"\n",
    "    if task == \"mnli\"\n",
    "    else \"validation\"\n",
    ")\n",
    "name = f\"{model_name}-finetuned-{task}\"\n",
    "\n",
    "# Calculate the maximum steps per epoch based on the number of rows in the training dataset.\n",
    "# Make sure to scale by the total number of training workers and the per device batch size.\n",
    "max_steps_per_epoch = ray_datasets[\"train\"].count() // (batch_size * num_workers)\n",
    "\n",
    "\n",
    "def train_func(config):\n",
    "    print(f\"Is CUDA available: {torch.cuda.is_available()}\")\n",
    "\n",
    "    metric = load(\"glue\", actual_task)\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)\n",
    "    model = AutoModelForSequenceClassification.from_pretrained(\n",
    "        model_checkpoint, num_labels=num_labels\n",
    "    )\n",
    "\n",
    "    train_ds = ray.train.get_dataset_shard(\"train\")\n",
    "    eval_ds = ray.train.get_dataset_shard(\"eval\")\n",
    "\n",
    "    train_ds_iterable = train_ds.iter_torch_batches(\n",
    "        batch_size=batch_size, collate_fn=collate_fn\n",
    "    )\n",
    "    eval_ds_iterable = eval_ds.iter_torch_batches(\n",
    "        batch_size=batch_size, collate_fn=collate_fn\n",
    "    )\n",
    "\n",
    "    print(\"max_steps_per_epoch: \", max_steps_per_epoch)\n",
    "\n",
    "    args = TrainingArguments(\n",
    "        name,\n",
    "        evaluation_strategy=\"epoch\",\n",
    "        save_strategy=\"epoch\",\n",
    "        logging_strategy=\"epoch\",\n",
    "        per_device_train_batch_size=batch_size,\n",
    "        per_device_eval_batch_size=batch_size,\n",
    "        learning_rate=config.get(\"learning_rate\", 2e-5),\n",
    "        num_train_epochs=config.get(\"epochs\", 2),\n",
    "        weight_decay=config.get(\"weight_decay\", 0.01),\n",
    "        push_to_hub=False,\n",
    "        max_steps=max_steps_per_epoch * config.get(\"epochs\", 2),\n",
    "        disable_tqdm=True,  # declutter the output a little\n",
    "        no_cuda=not use_gpu,  # you need to explicitly set no_cuda if you want CPUs\n",
    "        report_to=\"none\",\n",
    "    )\n",
    "\n",
    "    def compute_metrics(eval_pred):\n",
    "        predictions, labels = eval_pred\n",
    "        if task != \"stsb\":\n",
    "            predictions = np.argmax(predictions, axis=1)\n",
    "        else:\n",
    "            predictions = predictions[:, 0]\n",
    "        return metric.compute(predictions=predictions, references=labels)\n",
    "\n",
    "    trainer = Trainer(\n",
    "        model,\n",
    "        args,\n",
    "        train_dataset=train_ds_iterable,\n",
    "        eval_dataset=eval_ds_iterable,\n",
    "        tokenizer=tokenizer,\n",
    "        compute_metrics=compute_metrics,\n",
    "    )\n",
    "\n",
    "    trainer.add_callback(RayTrainReportCallback())\n",
    "\n",
    "    trainer = prepare_trainer(trainer)\n",
    "\n",
    "    print(\"Starting training\")\n",
    "    trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "CdzABDVcIrJg"
   },
   "source": [
    "With your `train_func` complete, you can now instantiate the {class}`~ray.train.torch.TorchTrainer`. Aside from calling the function, set the `scaling_config`, which controls the amount of workers and resources used, and the `datasets` to use for training and evaluation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "RElw7OgLhYba",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ray.train.torch import TorchTrainer\n",
    "from ray.train import RunConfig, ScalingConfig, CheckpointConfig\n",
    "\n",
    "trainer = TorchTrainer(\n",
    "    train_func,\n",
    "    scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),\n",
    "    datasets={\n",
    "        \"train\": ray_datasets[\"train\"],\n",
    "        \"eval\": ray_datasets[\"validation\"],\n",
    "    },\n",
    "    run_config=RunConfig(\n",
    "        checkpoint_config=CheckpointConfig(\n",
    "            num_to_keep=1,\n",
    "            checkpoint_score_attribute=\"eval_loss\",\n",
    "            checkpoint_score_order=\"min\",\n",
    "        ),\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XvS136zKhYba"
   },
   "source": [
    "Finally, call the `fit` method to start training with Ray Train. Save the `Result` object to a variable so you can access metrics and checkpoints."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "uNx5pyRlIrJh",
    "outputId": "8496fe4f-f1c3-48ad-a6d3-b16a65716135",
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-09 15:56:32,564\tINFO tune.py:616 -- [output] This uses the legacy output and progress reporter, as Jupyter notebooks are not supported by the new engine, yet. For more information, please see https://github.com/ray-project/ray/issues/36949\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Status ==\n",
      "Current time: 2025-07-09 15:56:32 (running for 00:00:00.11)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/node-group:head, 0.0/1.0 accelerator_type:T4, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/region:us-west-2)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 PENDING)\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m /home/ray/anaconda3/lib/python3.9/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m   _torch_pytree._register_pytree_node(\n",
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m 2025-07-09 15:56:36.371154: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m 2025-07-09 15:56:36.418819: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m 2025-07-09 15:56:36.418845: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m 2025-07-09 15:56:36.420083: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m 2025-07-09 15:56:36.427078: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m 2025-07-09 15:56:37.464124: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Status ==\n",
      "Current time: 2025-07-09 15:56:37 (running for 00:00:05.13)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/node-group:head, 0.0/1.0 accelerator_type:T4, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/region:us-west-2)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 PENDING)\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m /home/ray/anaconda3/lib/python3.9/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m   _torch_pytree._register_pytree_node(\n",
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m /home/ray/anaconda3/lib/python3.9/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m   _torch_pytree._register_pytree_node(\n",
      "\u001b[36m(TrainTrainable pid=41390)\u001b[0m comet_ml is installed but `COMET_API_KEY` is not set.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Status ==\n",
      "Current time: 2025-07-09 15:56:42 (running for 00:00:10.18)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/node-group:head, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 accelerator_type:T4, 0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 anyscale/accelerator_shape:4xT4)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m Setting up process group for: env:// [rank=0, world_size=1]\n",
      "\u001b[36m(TorchTrainer pid=41390)\u001b[0m Started distributed worker processes: \n",
      "\u001b[36m(TorchTrainer pid=41390)\u001b[0m - (node_id=f67b5f412a227b4c6b3ddd85d6f5b1eecd0bd0917efa8f9cd4b5e4da, ip=10.0.114.132, pid=41521) world_rank=0, local_rank=0, node_rank=0\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m /home/ray/anaconda3/lib/python3.9/site-packages/transformers/utils/generic.py:441: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m   _torch_pytree._register_pytree_node(\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m 2025-07-09 15:56:44.730942: I tensorflow/core/util/port.cc:113] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m 2025-07-09 15:56:44.779207: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m 2025-07-09 15:56:44.779230: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m 2025-07-09 15:56:44.780437: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m 2025-07-09 15:56:44.787541: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m 2025-07-09 15:56:45.863740: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m /home/ray/anaconda3/lib/python3.9/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m   _torch_pytree._register_pytree_node(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Status ==\n",
      "Current time: 2025-07-09 15:56:47 (running for 00:00:15.21)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/node-group:head, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 accelerator_type:T4, 0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 anyscale/accelerator_shape:4xT4)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m /home/ray/anaconda3/lib/python3.9/site-packages/transformers/utils/generic.py:309: UserWarning: torch.utils._pytree._register_pytree_node is deprecated. Please use torch.utils._pytree.register_pytree_node instead.\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m   _torch_pytree._register_pytree_node(\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m comet_ml is installed but `COMET_API_KEY` is not set.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m Is CUDA available: True\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m /home/ray/anaconda3/lib/python3.9/site-packages/huggingface_hub/file_download.py:795: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m   warnings.warn(\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m Some weights of DistilBertForSequenceClassification were not initialized from the model checkpoint at distilbert-base-uncased and are newly initialized: ['classifier.bias', 'classifier.weight', 'pre_classifier.bias', 'pre_classifier.weight']\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m /home/ray/anaconda3/lib/python3.9/site-packages/ray/data/iterator.py:436: RayDeprecationWarning: Passing a function to `iter_torch_batches(collate_fn)` is deprecated in Ray 2.47. Please switch to using a callable class that inherits from `ArrowBatchCollateFn`, `NumpyBatchCollateFn`, or `PandasBatchCollateFn`.\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m   warnings.warn(\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m /home/ray/anaconda3/lib/python3.9/site-packages/accelerate/accelerator.py:432: FutureWarning: Passing the following arguments to `Accelerator` is deprecated and will be removed in version 1.0 of Accelerate: dict_keys(['dispatch_batches', 'split_batches']). Please pass an `accelerate.DataLoaderConfiguration` instead: \n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m dataloader_config = DataLoaderConfiguration(dispatch_batches=None, split_batches=False)\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m   warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m max_steps_per_epoch:  534\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m Starting training\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "77e95e89b2094af1be48278c9e7d65d2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "(pid=41621) Running 0: 0.00 row [00:00, ? row/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e4e962af7830421da3dc3a58d85333e1",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "(pid=41621) - ReadParquet->SplitBlocks(96) 1: 0.00 row [00:00, ? row/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b3a0c80878d46bfa496425a6c750405",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "(pid=41621) - split(1, equal=True) 2: 0.00 row [00:00, ? row/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(SplitCoordinator pid=41621)\u001b[0m Registered dataset logger for dataset train_23_0\n",
      "\u001b[36m(SplitCoordinator pid=41621)\u001b[0m Starting execution of Dataset train_23_0. Full logs are in /tmp/ray/session_2025-07-09_15-09-59_163606_3385/logs/ray-data\n",
      "\u001b[36m(SplitCoordinator pid=41621)\u001b[0m Execution plan of Dataset train_23_0: InputDataBuffer[Input] -> TaskPoolMapOperator[ReadParquet] -> OutputSplitter[split(1, equal=True)]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Status ==\n",
      "Current time: 2025-07-09 15:56:52 (running for 00:00:20.23)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 accelerator_type:T4, 0.0/1.0 anyscale/node-group:head)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m /tmp/ipykernel_40967/133795194.py:24: 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:206.)\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m [rank0]:[W reducer.cpp:1389] Warning: find_unused_parameters=True was specified in DDP constructor, but did not find any unused parameters in the forward pass. This flag results in an extra traversal of the autograd graph every iteration,  which can adversely affect performance. If your model indeed never has any unused parameters in the forward pass, consider turning this flag off. Note that this warning may be a false positive if your model has flow control causing later iterations to have unused parameters. (function operator())\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Status ==\n",
      "Current time: 2025-07-09 15:56:57 (running for 00:00:25.25)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 accelerator_type:T4, 0.0/1.0 anyscale/node-group:head)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n",
      "== Status ==\n",
      "Current time: 2025-07-09 15:57:02 (running for 00:00:30.27)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 anyscale/node-group:head, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 accelerator_type:T4)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n",
      "== Status ==\n",
      "Current time: 2025-07-09 15:57:07 (running for 00:00:35.29)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 anyscale/node-group:head, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 accelerator_type:T4)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n",
      "== Status ==\n",
      "Current time: 2025-07-09 15:57:12 (running for 00:00:40.32)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/node-group:head, 0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 accelerator_type:T4)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n",
      "== Status ==\n",
      "Current time: 2025-07-09 15:57:17 (running for 00:00:45.34)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/node-group:head, 0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 accelerator_type:T4)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(SplitCoordinator pid=41621)\u001b[0m ✔️  Dataset train_23_0 execution finished in 28.21 seconds\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m {'loss': 0.5441, 'learning_rate': 9.9812734082397e-06, 'epoch': 0.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(SplitCoordinator pid=41622)\u001b[0m Registered dataset logger for dataset eval_24_0\n",
      "\u001b[36m(SplitCoordinator pid=41622)\u001b[0m Starting execution of Dataset eval_24_0. Full logs are in /tmp/ray/session_2025-07-09_15-09-59_163606_3385/logs/ray-data\n",
      "\u001b[36m(SplitCoordinator pid=41622)\u001b[0m Execution plan of Dataset eval_24_0: InputDataBuffer[Input] -> TaskPoolMapOperator[ReadParquet] -> OutputSplitter[split(1, equal=True)]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7afaf8757b8e49ff8ba1432fc50fb053",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "(pid=41622) Running 0: 0.00 row [00:00, ? row/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1c0abd88c1694102bd779fca0ea61804",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "(pid=41622) - ReadParquet->SplitBlocks(96) 1: 0.00 row [00:00, ? row/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "a5cd7456953145c195d74c23ed8badd7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "(pid=41622) - split(1, equal=True) 2: 0.00 row [00:00, ? row/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Status ==\n",
      "Current time: 2025-07-09 15:57:22 (running for 00:00:50.36)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 accelerator_type:T4, 0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/node-group:head)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m {'eval_loss': 0.51453697681427, 'eval_matthews_correlation': 0.37793570732654813, 'eval_runtime': 1.8456, 'eval_samples_per_second': 565.126, 'eval_steps_per_second': 35.761, 'epoch': 0.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-09 15:57:26,970\tWARNING experiment_state.py:206 -- Experiment state snapshotting has been triggered multiple times in the last 5.0 seconds and may become a bottleneck. A snapshot is forced if `CheckpointConfig(num_to_keep)` is set, and a trial has checkpointed >= `num_to_keep` times since the last snapshot.\n",
      "You may want to consider increasing the `CheckpointConfig(num_to_keep)` or decreasing the frequency of saving checkpoints.\n",
      "You can suppress this warning by setting the environment variable TUNE_WARN_EXCESSIVE_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S to a smaller value than the current threshold (5.0). Set it to 0 to completely suppress this warning.\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/ray/ray_results/TorchTrainer_2025-07-09_15-56-32/TorchTrainer_f5114_00000_0_2025-07-09_15-56-32/checkpoint_000000)\n",
      "\u001b[36m(SplitCoordinator pid=41622)\u001b[0m ✔️  Dataset eval_24_0 execution finished in 1.73 seconds\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "99ce23489c4a46ec9f621589bd83ff9a",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "(pid=41621) Running 0: 0.00 row [00:00, ? row/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c77d0ec4db2e40b982a32d3171238ac8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "(pid=41621) - ReadParquet->SplitBlocks(96) 1: 0.00 row [00:00, ? row/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e20adeec80504e68875294b053e0c1c0",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "(pid=41621) - split(1, equal=True) 2: 0.00 row [00:00, ? row/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Status ==\n",
      "Current time: 2025-07-09 15:57:27 (running for 00:00:55.36)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 accelerator_type:T4, 0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/node-group:head)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n",
      "== Status ==\n",
      "Current time: 2025-07-09 15:57:32 (running for 00:01:00.38)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 accelerator_type:T4, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/node-group:head)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n",
      "== Status ==\n",
      "Current time: 2025-07-09 15:57:38 (running for 00:01:05.41)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 accelerator_type:T4, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/node-group:head)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n",
      "== Status ==\n",
      "Current time: 2025-07-09 15:57:43 (running for 00:01:10.43)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/node-group:head, 0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 accelerator_type:T4, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/provider:aws)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n",
      "== Status ==\n",
      "Current time: 2025-07-09 15:57:48 (running for 00:01:15.45)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/node-group:head, 0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 accelerator_type:T4, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/provider:aws)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n",
      "== Status ==\n",
      "Current time: 2025-07-09 15:57:53 (running for 00:01:20.47)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/node-group:head, 0.0/1.0 accelerator_type:T4)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(SplitCoordinator pid=41621)\u001b[0m ✔️  Dataset train_23_1 execution finished in 26.58 seconds\n",
      "\u001b[36m(SplitCoordinator pid=41621)\u001b[0m Registered dataset logger for dataset train_23_1\n",
      "\u001b[36m(SplitCoordinator pid=41621)\u001b[0m Starting execution of Dataset train_23_1. Full logs are in /tmp/ray/session_2025-07-09_15-09-59_163606_3385/logs/ray-data\n",
      "\u001b[36m(SplitCoordinator pid=41621)\u001b[0m Execution plan of Dataset train_23_1: InputDataBuffer[Input] -> TaskPoolMapOperator[ReadParquet] -> OutputSplitter[split(1, equal=True)]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m {'loss': 0.3864, 'learning_rate': 0.0, 'epoch': 1.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[36m(SplitCoordinator pid=41622)\u001b[0m Registered dataset logger for dataset eval_24_1\n",
      "\u001b[36m(SplitCoordinator pid=41622)\u001b[0m Starting execution of Dataset eval_24_1. Full logs are in /tmp/ray/session_2025-07-09_15-09-59_163606_3385/logs/ray-data\n",
      "\u001b[36m(SplitCoordinator pid=41622)\u001b[0m Execution plan of Dataset eval_24_1: InputDataBuffer[Input] -> TaskPoolMapOperator[ReadParquet] -> OutputSplitter[split(1, equal=True)]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3853079fdd524064890b0dfccb41aa9b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "(pid=41622) Running 0: 0.00 row [00:00, ? row/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7398954babdc47798b8da28a2adfc080",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "(pid=41622) - ReadParquet->SplitBlocks(96) 1: 0.00 row [00:00, ? row/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "37889407c41e4a5782e8f74b02a401a3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "(pid=41622) - split(1, equal=True) 2: 0.00 row [00:00, ? row/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m {'eval_loss': 0.5683005452156067, 'eval_matthews_correlation': 0.45115517656589194, 'eval_runtime': 1.6027, 'eval_samples_per_second': 650.77, 'eval_steps_per_second': 41.18, 'epoch': 1.5}\n",
      "== Status ==\n",
      "Current time: 2025-07-09 15:57:58 (running for 00:01:25.49)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/node-group:head, 0.0/1.0 accelerator_type:T4)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 RUNNING)\n",
      "\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-09 15:57:59,354\tWARNING experiment_state.py:206 -- Experiment state snapshotting has been triggered multiple times in the last 5.0 seconds and may become a bottleneck. A snapshot is forced if `CheckpointConfig(num_to_keep)` is set, and a trial has checkpointed >= `num_to_keep` times since the last snapshot.\n",
      "You may want to consider increasing the `CheckpointConfig(num_to_keep)` or decreasing the frequency of saving checkpoints.\n",
      "You can suppress this warning by setting the environment variable TUNE_WARN_EXCESSIVE_EXPERIMENT_CHECKPOINT_SYNC_THRESHOLD_S to a smaller value than the current threshold (5.0). Set it to 0 to completely suppress this warning.\n",
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/ray/ray_results/TorchTrainer_2025-07-09_15-56-32/TorchTrainer_f5114_00000_0_2025-07-09_15-56-32/checkpoint_000001)\n",
      "\u001b[36m(SplitCoordinator pid=41622)\u001b[0m ✔️  Dataset eval_24_1 execution finished in 1.49 seconds\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36m(RayTrainWorker pid=41521)\u001b[0m {'train_runtime': 66.7725, 'train_samples_per_second': 255.914, 'train_steps_per_second': 15.995, 'train_loss': 0.4653928092356478, 'epoch': 1.5}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-07-09 15:58:00,649\tINFO tune.py:1009 -- Wrote the latest version of all result files and experiment state to '/home/ray/ray_results/TorchTrainer_2025-07-09_15-56-32' in 0.0022s.\n",
      "2025-07-09 15:58:00,651\tINFO tune.py:1041 -- Total run time: 88.09 seconds (88.03 seconds for the tuning loop).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Status ==\n",
      "Current time: 2025-07-09 15:58:00 (running for 00:01:28.04)\n",
      "Using FIFO scheduling algorithm.\n",
      "Logical resource usage: 1.0/48 CPUs, 1.0/4 GPUs (0.0/1.0 anyscale/region:us-west-2, 0.0/1.0 anyscale/provider:aws, 0.0/1.0 anyscale/accelerator_shape:4xT4, 0.0/1.0 anyscale/node-group:head, 0.0/1.0 accelerator_type:T4)\n",
      "Result logdir: /tmp/ray/session_2025-07-09_15-09-59_163606_3385/artifacts/2025-07-09_15-56-32/TorchTrainer_2025-07-09_15-56-32/driver_artifacts\n",
      "Number of trials: 1/1 (1 TERMINATED)\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "result = trainer.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "4cnWqUWmhYba"
   },
   "source": [
    "You can use the returned `Result` object to access metrics and the Ray Train `Checkpoint` associated with the last iteration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "AMN5qjUwhYba",
    "outputId": "7b754c36-c58b-4ff4-d7a8-63ec9764bd0c"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Result(\n",
       "  metrics={'loss': 0.3864, 'learning_rate': 0.0, 'epoch': 1.5, 'step': 1068, 'eval_loss': 0.5683005452156067, 'eval_matthews_correlation': 0.45115517656589194, 'eval_runtime': 1.6027, 'eval_samples_per_second': 650.77, 'eval_steps_per_second': 41.18},\n",
       "  path='/home/ray/ray_results/TorchTrainer_2025-07-09_15-56-32/TorchTrainer_f5114_00000_0_2025-07-09_15-56-32',\n",
       "  filesystem='local',\n",
       "  checkpoint=Checkpoint(filesystem=local, path=/home/ray/ray_results/TorchTrainer_2025-07-09_15-56-32/TorchTrainer_f5114_00000_0_2025-07-09_15-56-32/checkpoint_000001)\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## See also\n",
    "\n",
    "* {doc}`Ray Train Examples <../../examples>` for more use cases\n",
    "* {ref}`Ray Train User Guides <train-user-guides>` for how-to guides\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "huggingface_text_classification.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.23"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}