LegalQABloom.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5cf09ad4",
   "metadata": {},
   "source": [
    "Fine-tuning bloomz-560m using a small legal QA dataset (portuguese)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0960dfcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import torch\n",
    "import json\n",
    "from transformers import BloomTokenizerFast, BloomForCausalLM, TrainingArguments, Trainer\n",
    "from datasets import load_dataset\n",
    "from accelerate import Accelerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7373463a",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "37a524c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading bloomz model and tokenizer \n",
    "tokenizer = BloomTokenizerFast.from_pretrained(\"bigscience/bloomz-560m\")\n",
    "model = BloomForCausalLM.from_pretrained(\"bigscience/bloomz-560m\").to(\"cuda\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6dcc88a7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Found cached dataset json (/home/dvianna/.cache/huggingface/datasets/json/default-5329563cfdcc119a/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "72d6861ac782416eb3159478f82c6bb7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Loading dataset prompts.json built using de portuguese legalQA dataset\n",
    "dataset = load_dataset(\"json\", data_files=\"prompts.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2a21c46d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['text'],\n",
       "        num_rows: 395\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ef27414c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# prepare the data for training\n",
    "def prepare_train_data(data):\n",
    "    # prompt + completion\n",
    "    #text_input = data['prompt'] + ' ' + data['completion']\n",
    "    text_input = data['text']\n",
    "    # tokenize the input (prompt + completion) text\n",
    "    tokenized_input = tokenizer(text_input, return_tensors='pt', padding=True)\n",
    "    # generative models: labels are the same as the input\n",
    "    tokenized_input['labels'] = tokenized_input['input_ids']\n",
    "    return tokenized_input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d0978847",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading cached processed dataset at /home/dvianna/.cache/huggingface/datasets/json/default-5329563cfdcc119a/0.0.0/e347ab1c932092252e717ff3f949105a4dd28b27e842dd53157d2f72e276c2e4/cache-3a16b40a0956a8e0.arrow\n"
     ]
    }
   ],
   "source": [
    "train_dataset = dataset['train'].map(prepare_train_data, batched=True, remove_columns=[\"text\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "66248e9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# setting arguments to be used during training\n",
    "training_arguments = TrainingArguments(\n",
    "    'LegalQA-bloom-560m',\n",
    "    learning_rate=2e-5,\n",
    "    per_device_train_batch_size=2,\n",
    "    num_train_epochs=2,\n",
    "    weight_decay=0.01,\n",
    "    fp16=True,\n",
    "    optim=\"adafactor\",\n",
    "    gradient_accumulation_steps=4,\n",
    "    gradient_checkpointing=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c1ba254a",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model = model,\n",
    "    args = training_arguments,\n",
    "    train_dataset = train_dataset\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d92e67b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='98' max='98' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [98/98 04:14, Epoch 1/2]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=98, training_loss=5.802616041533801, metrics={'train_runtime': 257.7121, 'train_samples_per_second': 3.065, 'train_steps_per_second': 0.38, 'total_flos': 2008259929718784.0, 'train_loss': 5.802616041533801, 'epoch': 1.98})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6ccfd00b",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.save_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90a55e0c",
   "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
}