fine_tuning_pipeline.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ccf955f-2e73-4f68-8a78-7c2b336394eb",
   "metadata": {
    "is_executing": true
   },
   "outputs": [],
   "source": [
    "from kfp import dsl, compiler\n",
    "\n",
    "@dsl.component(packages_to_install=[\"gcsfs\", \"transformers\", \"datasets==2.16\", \"evaluate==0.4.3\", \"accelerate\", \"scikit-learn\", \"kubeflow-training\"])\n",
    "def start_distributed_training(bucket: str, dataset_file: str, output_model_name: str, gproject: str) -> str:\n",
    "    import os\n",
    "    import gcsfs\n",
    "    import numpy as np\n",
    "    from datasets import load_dataset\n",
    "    from datasets.distributed import split_dataset_by_node\n",
    "    from transformers import (\n",
    "        AutoModelForSequenceClassification,\n",
    "        AutoTokenizer,\n",
    "        Trainer,\n",
    "        TrainingArguments,\n",
    "    )\n",
    "    from kubeflow.training import TrainingClient\n",
    "    import torch\n",
    "\n",
    "    def train_func(parameters):\n",
    "        import os\n",
    "        import gcsfs\n",
    "        import numpy as np\n",
    "        from datasets import load_dataset\n",
    "        from datasets.distributed import split_dataset_by_node\n",
    "        from transformers import (\n",
    "            AutoModelForSequenceClassification,\n",
    "            AutoTokenizer,\n",
    "            Trainer,\n",
    "            TrainingArguments,\n",
    "        )\n",
    "        from kubeflow.training import TrainingClient\n",
    "        import torch\n",
    "        import evaluate\n",
    "    \n",
    "        model_name = parameters['MODEL_NAME']\n",
    "        storage_options= parameters['STORAGE_OPTIONS'] \n",
    "        dataset = load_dataset(\"json\", data_files=f'gs://{parameters[\"BUCKET\"]}/{parameters[\"DATASET_FILE\"]}', storage_options=storage_options)\n",
    "        ds = dataset[\"train\"].train_test_split(test_size=0.2)\n",
    "        \n",
    "        labels = [label for label in ds['train'].features.keys() if label not in ['body', 'title']]\n",
    "        id2label = {idx:label for idx, label in enumerate(labels)}\n",
    "        label2id = {label:idx for idx, label in enumerate(labels)}\n",
    "        \n",
    "    \n",
    "        print(\"-\" * 40)\n",
    "        print(\"Download BERT Model\")\n",
    "        model = AutoModelForSequenceClassification.from_pretrained(\"bert-base-uncased\", \n",
    "                                                               problem_type=\"multi_label_classification\", \n",
    "                                                               num_labels=len(labels),\n",
    "                                                               id2label=id2label,\n",
    "                                                               label2id=label2id)\n",
    "        tokenizer = AutoTokenizer.from_pretrained(\"bert-base-cased\")\n",
    "        \n",
    "        # [2] Preprocess dataset.        \n",
    "        def preprocess_data(example):\n",
    "          text = f'{example[\"title\"]}\\n{example[\"body\"]}'\n",
    "          # encode them\n",
    "          encoding = tokenizer(text, padding=True, truncation=True)\n",
    "        \n",
    "          lbls = [0. for i in range(len(labels))]\n",
    "          for label in labels:\n",
    "            if label in example and example[label] == True:\n",
    "              label_id = label2id[label]\n",
    "              lbls[label_id] = 1.\n",
    "        \n",
    "          encoding[\"labels\"] = lbls  \n",
    "          return encoding\n",
    "        \n",
    "        # Map Yelp review dataset to BERT tokenizer.\n",
    "        print(\"-\" * 40)\n",
    "        print(\"Map dataset to BERT Tokenizer\")\n",
    "        encoded_dataset = ds.map(preprocess_data, remove_columns=ds['train'].column_names)\n",
    "        encoded_dataset.set_format(\"torch\")\n",
    "        \n",
    "        # Distribute train and test datasets between PyTorch workers.\n",
    "        # Every worker will process chunk of training data.\n",
    "        # RANK and WORLD_SIZE will be set by Kubeflow Training Operator.\n",
    "        RANK = int(os.environ[\"RANK\"])\n",
    "        WORLD_SIZE = int(os.environ[\"WORLD_SIZE\"])\n",
    "        distributed_ds_train = split_dataset_by_node(\n",
    "            encoded_dataset[\"train\"],\n",
    "            rank=RANK,\n",
    "            world_size=WORLD_SIZE,\n",
    "        )\n",
    "        distributed_ds_test = split_dataset_by_node(\n",
    "            encoded_dataset[\"test\"],\n",
    "            rank=RANK,\n",
    "            world_size=WORLD_SIZE,\n",
    "        )\n",
    "        \n",
    "        # Evaluate accuracy.    \n",
    "        clf_metrics = evaluate.combine([\"accuracy\", \"f1\", \"precision\", \"recall\"])\n",
    "    \n",
    "        def sigmoid(x):\n",
    "           return 1/(1 + np.exp(-x))\n",
    "        \n",
    "        def compute_metrics(eval_pred):\n",
    "           predictions, labels = eval_pred\n",
    "           predictions = sigmoid(predictions)\n",
    "           predictions = (predictions > 0.5).astype(int).reshape(-1)\n",
    "           return clf_metrics.compute(predictions=predictions, references=labels.astype(int).reshape(-1))\n",
    "    \n",
    "    \n",
    "        batch_size = 3\n",
    "        metric_name = \"f1\"\n",
    "        args = TrainingArguments(\n",
    "            f\"{model_name}\",\n",
    "            evaluation_strategy = \"epoch\",\n",
    "            save_strategy = \"epoch\",\n",
    "            learning_rate=2e-5,\n",
    "            per_device_train_batch_size=batch_size,\n",
    "            per_device_eval_batch_size=batch_size,\n",
    "            num_train_epochs=5,\n",
    "            weight_decay=0.01,\n",
    "            load_best_model_at_end=True,\n",
    "            metric_for_best_model=metric_name,\n",
    "            #push_to_hub=True,\n",
    "        )\n",
    "    \n",
    "        trainer = Trainer(\n",
    "            model=model,\n",
    "            args=args,\n",
    "            train_dataset=distributed_ds_train,\n",
    "            eval_dataset=distributed_ds_test,\n",
    "            tokenizer=tokenizer,\n",
    "            compute_metrics=compute_metrics,\n",
    "        )\n",
    "        \n",
    "        print(\"-\" * 40)\n",
    "        print(f\"Start Distributed Training. RANK: {RANK} WORLD_SIZE: {WORLD_SIZE}\")\n",
    "        \n",
    "        trainer.train()\n",
    "        \n",
    "        print(\"-\" * 40)\n",
    "        print(\"Training is complete\")\n",
    "\n",
    "        # Export trained model to GCS from the worker with RANK = 0 (master).\n",
    "        if RANK == 0:\n",
    "            trainer.save_model(f\"./{model_name}\")\n",
    "            fs = gcsfs.GCSFileSystem(**storage_options)\n",
    "            files = ['config.json', 'model.safetensors', 'special_tokens_map.json', 'tokenizer_config.json', 'tokenizer.json', 'training_args.bin', 'vocab.txt']\n",
    "            for f in files: \n",
    "                fs.put(f'{model_name}/{f}', f'{parameters[\"BUCKET\"]}/{model_name}/{f}')\n",
    "        \n",
    "        print(\"-\" * 40)\n",
    "        print(\"Model export complete\")\n",
    "        \n",
    "    job_name = \"training-pipeline-job\"\n",
    "    # Create PyTorchJob\n",
    "    TrainingClient().create_job(\n",
    "        name=job_name,\n",
    "        train_func=train_func,\n",
    "        parameters={\n",
    "            \"BUCKET\": bucket,\n",
    "            \"STORAGE_OPTIONS\": {\"project\": gproject, \"token\": \"google_default\"},\n",
    "            \"MODEL_NAME\": output_model_name,\n",
    "            \"DATASET_FILE\": dataset_file\n",
    "        },\n",
    "        num_workers=2,  # Number of PyTorch workers to use.\n",
    "        resources_per_worker={\n",
    "            \"cpu\": \"3\",\n",
    "            \"memory\": \"10G\",\n",
    "            \"gpu\": \"1\",\n",
    "        },\n",
    "        packages_to_install=[\n",
    "            \"gcsfs\",\n",
    "            \"transformers\",\n",
    "            \"datasets==2.16\",\n",
    "            \"evaluate\",\n",
    "            \"accelerate\",\n",
    "            \"scikit-learn\",\n",
    "            \"kubeflow-training\"\n",
    "        ],  # PIP packages will be installed during PyTorchJob runtime.\n",
    "    )\n",
    "    # Wait until PyTorchJob has Running condition.\n",
    "    job = TrainingClient().wait_for_job_conditions(\n",
    "        job_name,\n",
    "        expected_conditions={\"Running\"},\n",
    "    )\n",
    "    return \"job is running\"\n",
    "\n",
    "@dsl.pipeline\n",
    "def training_pipeline(bucket: str, dataset_file: str, output_model_name: str, gproject: str) -> str:\n",
    "    training_task = start_distributed_training(bucket=bucket, dataset_file=dataset_file, output_model_name=output_model_name, gproject=gproject)\n",
    "    return training_task.output\n",
    "\n",
    "compiler.Compiler().compile(training_pipeline, 'training_pipeline.yaml')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [],
   "metadata": {
    "collapsed": false
   },
   "id": "544738a655b827d"
  }
 ],
 "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}