fine_tuning_using_ray.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4d15a75",
   "metadata": {},
   "outputs": [],
   "source": [
    "#! pip install \"datasets\" \"transformers>=4.19.0\" \"torch>=1.10.0\" \"mlflow\" \"ray[air]>=1.13\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08640c03",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pprint import pprint\n",
    "import ray\n",
    "\n",
    "ray.init()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a3a3dcf",
   "metadata": {},
   "outputs": [],
   "source": [
    "pprint(ray.cluster_resources())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb609fae",
   "metadata": {},
   "outputs": [],
   "source": [
    "use_gpu = True  # set this to False to run on CPUs\n",
    "num_workers = 1  # set this to number of GPUs/CPUs you want to use"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b63c1cbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "GLUE_TASKS = [\"cola\", \"mnli\", \"mnli-mm\", \"mrpc\", \"qnli\", \"qqp\", \"rte\", \"sst2\", \"stsb\", \"wnli\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88c0580d",
   "metadata": {},
   "outputs": [],
   "source": [
    "task = \"cola\"\n",
    "model_checkpoint = \"distilbert-base-uncased\"\n",
    "batch_size = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91ce68da",
   "metadata": {},
   "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": "code",
   "execution_count": null,
   "id": "88ea4ead",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_metric\n",
    "\n",
    "def load_metric_fn():\n",
    "    return load_metric('glue', actual_task)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "56f39bae",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfdae727",
   "metadata": {},
   "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": "code",
   "execution_count": null,
   "id": "3e437537",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ray.data\n",
    "\n",
    "ray_datasets = ray.data.from_huggingface(datasets)\n",
    "ray_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "edacbccc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from ray.data.preprocessors import BatchMapper\n",
    "\n",
    "def preprocess_function(examples: pd.DataFrame):\n",
    "    # if we only have one column, we are inferring.\n",
    "    # no need to tokenize in that case. \n",
    "    if len(examples.columns) == 1:\n",
    "        return examples\n",
    "    examples = examples.to_dict(\"list\")\n",
    "    sentence1_key, sentence2_key = task_to_keys[task]\n",
    "    if sentence2_key is None:\n",
    "        ret = tokenizer(examples[sentence1_key], truncation=True)\n",
    "    else:\n",
    "        ret = tokenizer(examples[sentence1_key], examples[sentence2_key], truncation=True)\n",
    "    # Add back the original columns\n",
    "    ret = {**examples, **ret}\n",
    "    return pd.DataFrame.from_dict(ret)\n",
    "\n",
    "batch_encoder = BatchMapper(preprocess_function, batch_format=\"pandas\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f63a2498",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForSequenceClassification, TrainingArguments, Trainer\n",
    "import numpy as np\n",
    "import torch\n",
    "\n",
    "num_labels = 3 if task.startswith(\"mnli\") else 1 if task==\"stsb\" else 2\n",
    "metric_name = \"pearson\" if task == \"stsb\" else \"matthews_correlation\" if task == \"cola\" else \"accuracy\"\n",
    "model_name = model_checkpoint.split(\"/\")[-1]\n",
    "validation_key = \"validation_mismatched\" if task == \"mnli-mm\" else \"validation_matched\" if task == \"mnli\" else \"validation\"\n",
    "name = f\"{model_name}-finetuned-{task}\"\n",
    "\n",
    "def trainer_init_per_worker(train_dataset, eval_dataset = None, **config):\n",
    "    print(f\"Is CUDA available: {torch.cuda.is_available()}\")\n",
    "    metric = load_metric_fn()\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)\n",
    "    model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint, num_labels=num_labels)\n",
    "    args = TrainingArguments(\n",
    "        name,\n",
    "        evaluation_strategy=\"epoch\",\n",
    "        save_strategy=\"epoch\",\n",
    "        logging_strategy=\"epoch\",\n",
    "        learning_rate=config.get(\"learning_rate\", 2e-5),\n",
    "        per_device_train_batch_size=batch_size,\n",
    "        per_device_eval_batch_size=batch_size,\n",
    "        num_train_epochs=config.get(\"epochs\", 2),\n",
    "        weight_decay=config.get(\"weight_decay\", 0.01),\n",
    "        push_to_hub=False,\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",
    "    )\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_dataset,\n",
    "        eval_dataset=eval_dataset,\n",
    "        tokenizer=tokenizer,\n",
    "        compute_metrics=compute_metrics\n",
    "    )\n",
    "\n",
    "    print(\"Starting training\")\n",
    "    return trainer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32e07524",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.train.huggingface import HuggingFaceTrainer\n",
    "from ray.air.config import RunConfig, ScalingConfig, CheckpointConfig\n",
    "from ray.air.integrations.mlflow import MLflowLoggerCallback\n",
    "\n",
    "trainer = HuggingFaceTrainer(\n",
    "    trainer_init_per_worker=trainer_init_per_worker,\n",
    "    scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu),\n",
    "    datasets={\n",
    "        \"train\": ray_datasets[\"train\"],\n",
    "        \"evaluation\": ray_datasets[validation_key],\n",
    "    },\n",
    "    run_config=RunConfig(\n",
    "        callbacks=[MLflowLoggerCallback(experiment_name=name)],\n",
    "        checkpoint_config=CheckpointConfig(\n",
    "            num_to_keep=1,\n",
    "            checkpoint_score_attribute=\"eval_loss\",\n",
    "            checkpoint_score_order=\"min\",\n",
    "        ),\n",
    "    ),\n",
    "    preprocessor=batch_encoder,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e85cb089",
   "metadata": {},
   "outputs": [],
   "source": [
    "result = trainer.fit()\n",
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6f7cfc7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray import tune\n",
    "from ray.tune import Tuner\n",
    "from ray.tune.schedulers.async_hyperband import ASHAScheduler\n",
    "\n",
    "tune_epochs = 4\n",
    "tuner = Tuner(\n",
    "    trainer,\n",
    "    param_space={\n",
    "        \"trainer_init_config\": {\n",
    "            \"learning_rate\": tune.grid_search([2e-5, 2e-4, 2e-3, 2e-2]),\n",
    "            \"epochs\": tune_epochs,\n",
    "        }\n",
    "    },\n",
    "    tune_config=tune.TuneConfig(\n",
    "        metric=\"eval_loss\",\n",
    "        mode=\"min\",\n",
    "        num_samples=1,\n",
    "        scheduler=ASHAScheduler(\n",
    "            max_t=tune_epochs,\n",
    "        )\n",
    "    ),\n",
    "    run_config=RunConfig(\n",
    "        checkpoint_config=CheckpointConfig(num_to_keep=1, checkpoint_score_attribute=\"eval_loss\", checkpoint_score_order=\"min\")\n",
    "    ),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "69b0d2a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "tune_results = tuner.fit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a10e13e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tune_results.get_dataframe().sort_values(\"eval_loss\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbab6e29",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ray.train.huggingface import HuggingFacePredictor\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=HuggingFacePredictor,\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\"]]), num_gpus_per_worker=int(use_gpu))\n",
    "prediction.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d68ed0d3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}