{ "cells": [ { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "# Fine-Tune a Generative AI Model for Dialogue Summarization" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "In this notebook we will see how to fine tune an existing LLM from HuggingFace for enhanced dialogue summarization. We will be using the [Flan-T5](https://huggingface.co/docs/transformers/model_doc/flan-t5) model as it provides a high quality instruction tuned model at various sizes. Flan-T5 can summarize text out of the box, but in this notebook we will see how fine-tuning on a high quality dataset can improve its performance for a specific task. Specifically, we will be using the [DialogSum](https://huggingface.co/datasets/knkarthick/dialogsum) dataset from HuggingFace which contains chunks of dialogue and associated summarizations of the dialogue." ] }, { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "## Setup\n", "\n", "First up, lets make sure we install some libraries which are needed for this notebook. After the installation, we will import the necessary packages for the notebook" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\u001b[33mWARNING: The directory '/root/.cache/pip' or its parent directory is not owned or is not writable by the current user. The cache has been disabled. Check the permissions and owner of that directory. If executing pip with sudo, you should use sudo's -H flag.\u001b[0m\u001b[33m\n", "\u001b[0m\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n", "pytest-astropy 0.8.0 requires pytest-cov>=2.0, which is not installed.\n", "pytest-astropy 0.8.0 requires pytest-filter-subpackage>=0.1, which is not installed.\n", "sagemaker 2.144.0 requires importlib-metadata<5.0,>=1.4.0, but you have importlib-metadata 6.1.0 which is incompatible.\n", "sagemaker 2.144.0 requires PyYAML==5.4.1, but you have pyyaml 6.0 which is incompatible.\n", "docker-compose 1.29.2 requires PyYAML<6,>=3.10, but you have pyyaml 6.0 which is incompatible.\u001b[0m\u001b[31m\n", "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n", "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n", "\u001b[33mWARNING: The directory '/root/.cache/pip' or its parent directory is not owned or is not writable by the current user. The cache has been disabled. Check the permissions and owner of that directory. If executing pip with sudo, you should use sudo's -H flag.\u001b[0m\u001b[33m\n", "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n", "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n", "\u001b[33mWARNING: The directory '/root/.cache/pip' or its parent directory is not owned or is not writable by the current user. The cache has been disabled. Check the permissions and owner of that directory. If executing pip with sudo, you should use sudo's -H flag.\u001b[0m\u001b[33m\n", "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n", "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n", "\u001b[33mWARNING: The directory '/root/.cache/pip' or its parent directory is not owned or is not writable by the current user. The cache has been disabled. Check the permissions and owner of that directory. If executing pip with sudo, you should use sudo's -H flag.\u001b[0m\u001b[33m\n", "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n", "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n", "\u001b[33mWARNING: The directory '/root/.cache/pip' or its parent directory is not owned or is not writable by the current user. The cache has been disabled. Check the permissions and owner of that directory. If executing pip with sudo, you should use sudo's -H flag.\u001b[0m\u001b[33m\n", "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n", "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n", "\u001b[33mWARNING: The directory '/root/.cache/pip' or its parent directory is not owned or is not writable by the current user. The cache has been disabled. Check the permissions and owner of that directory. If executing pip with sudo, you should use sudo's -H flag.\u001b[0m\u001b[33m\n", "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n", "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n", "\u001b[33mWARNING: The directory '/root/.cache/pip' or its parent directory is not owned or is not writable by the current user. The cache has been disabled. Check the permissions and owner of that directory. If executing pip with sudo, you should use sudo's -H flag.\u001b[0m\u001b[33m\n", "\u001b[0m\u001b[33mWARNING: Running pip as the 'root' user can result in broken permissions and conflicting behaviour with the system package manager. It is recommended to use a virtual environment instead: https://pip.pypa.io/warnings/venv\u001b[0m\u001b[33m\n", "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n" ] } ], "source": [ "%pip install transformers==4.27.2 --quiet\n", "%pip install torch==1.13.1 --quiet\n", "%pip install py7zr --quiet\n", "%pip install datasets --quiet\n", "%pip install sentencepiece --quiet\n", "%pip install evaluate --quiet\n", "%pip install rouge_score --quiet" ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "tags": [] }, "outputs": [], "source": [ "from transformers import AutoTokenizer, T5ForConditionalGeneration, TrainingArguments, Trainer, GenerationConfig\n", "from datasets import load_dataset\n", "import datasets\n", "import torch\n", "import time\n", "import evaluate\n", "import numpy as np\n", "DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'" ] }, { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "# Load Flan-T5 Model\n", "\n", "We can load the pre-trained Flan-T5 model directly from HuggingFace. Notice that we will be using the [base version](https://huggingface.co/google/flan-t5-base) of flan. This model version has ~247 million model parameters which makes it small compared to other LLMs. For higher quality results, we recommend looking into the larger versions of this model." ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "tags": [] }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "c671e91bb5ff42e78021315f86d5520a", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading (…)okenizer_config.json: 0%| | 0.00/2.54k [00:00, ?B/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "97babcc3a18d48468a8a064f89ef949d", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading spiece.model: 0%| | 0.00/792k [00:00, ?B/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "44405f2d472b492a8a5e533e80f844db", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading (…)/main/tokenizer.json: 0%| | 0.00/2.42M [00:00, ?B/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "b8e7da5661094d54906618dc00ae3311", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading (…)cial_tokens_map.json: 0%| | 0.00/2.20k [00:00, ?B/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "62d4cccb4c8243e597bb45a3ecb76237", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading (…)lve/main/config.json: 0%| | 0.00/1.40k [00:00, ?B/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "cd0819cd2c6f4a4ba23972d969c761d0", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading pytorch_model.bin: 0%| | 0.00/990M [00:00, ?B/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "aea7f5553f4c4912b4d989c586f8b77b", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading (…)neration_config.json: 0%| | 0.00/147 [00:00, ?B/s]" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "tokenizer = AutoTokenizer.from_pretrained(\"google/flan-t5-base\")\n", "model = T5ForConditionalGeneration.from_pretrained(\"google/flan-t5-base\")" ] }, { "cell_type": "code", "execution_count": 8, "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Total Number of Model Parameters: 247577856\n" ] } ], "source": [ "params = sum(p.numel() for p in model.parameters())\n", "print(f'Total Number of Model Parameters: {params}')" ] }, { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "# Load Dataset\n", "\n", "The DialogSum dataset can also be loaded directly from HuggingFace. There are ~15k examples of dialogue in this dataset with associated human summarizations of these datasets" ] }, { "cell_type": "code", "execution_count": 9, "metadata": { "tags": [] }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "d8155b5aa4f64fbeae82fe57508be3f3", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading readme: 0%| | 0.00/4.58k [00:00, ?B/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "Downloading and preparing dataset csv/knkarthick--dialogsum to /root/.cache/huggingface/datasets/knkarthick___csv/knkarthick--dialogsum-6d41e9a7b96e340e/0.0.0/6954658bab30a358235fa864b05cf819af0e179325c740e4bc853bcc7ec513e1...\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "66474d04ea5d4f6abbc9467fa7fe18cf", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading data files: 0%| | 0/3 [00:00, ?it/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "66079d80a74b4c07969ab9abe4fcb473", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading data: 0%| | 0.00/11.3M [00:00, ?B/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "0262c5259f7d4e1c802355700e8669dd", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading data: 0%| | 0.00/1.35M [00:00, ?B/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "79ac8bd008c14f5c8acf9539bda6c2b9", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Downloading data: 0%| | 0.00/442k [00:00, ?B/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "45856718f7d94249801c8f290c161cb6", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Extracting data files: 0%| | 0/3 [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": [ "Generating train split: 0 examples [00:00, ? examples/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Generating test split: 0 examples [00:00, ? examples/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Generating validation split: 0 examples [00:00, ? examples/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "name": "stdout", "output_type": "stream", "text": [ "Dataset csv downloaded and prepared to /root/.cache/huggingface/datasets/knkarthick___csv/knkarthick--dialogsum-6d41e9a7b96e340e/0.0.0/6954658bab30a358235fa864b05cf819af0e179325c740e4bc853bcc7ec513e1. Subsequent calls will reuse this data.\n" ] }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "152b145cafb74038b1893486c3c745ed", "version_major": 2, "version_minor": 0 }, "text/plain": [ " 0%| | 0/3 [00:00, ?it/s]" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "dataset = load_dataset(\"knkarthick/dialogsum\")" ] }, { "cell_type": "code", "execution_count": 11, "metadata": { "tags": [] }, "outputs": [ { "data": { "text/plain": [ "DatasetDict({\n", " train: Dataset({\n", " features: ['id', 'dialogue', 'summary', 'topic'],\n", " num_rows: 12460\n", " })\n", " test: Dataset({\n", " features: ['id', 'dialogue', 'summary', 'topic'],\n", " num_rows: 1500\n", " })\n", " validation: Dataset({\n", " features: ['id', 'dialogue', 'summary', 'topic'],\n", " num_rows: 500\n", " })\n", "})" ] }, "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ "dataset" ] }, { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "# Test the Model with Zero-Shot Prompts Before Tuning\n", "\n", "In the example below, we highlight how the summarization capability of the model is lacking compared to the baseline summary provided in the dataset. You can see that the model struggles to summarize the dialogue compared to the baseline summary, but it does pull out some important information from the text which indicates the model can be fine tuned to the task at hand." ] }, { "cell_type": "code", "execution_count": 15, "metadata": { "tags": [] }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Prompt:\n", "--------------------------\n", "Summarize the following conversation.\n", "\n", "#Person1#: What time is it, Tom?\n", "#Person2#: Just a minute. It's ten to nine by my watch.\n", "#Person1#: Is it? I had no idea it was so late. I must be off now.\n", "#Person2#: What's the hurry?\n", "#Person1#: I must catch the nine-thirty train.\n", "#Person2#: You've plenty of time yet. The railway station is very close. It won't take more than twenty minutes to get there.\n", "\n", "Summary:\n", "--------------------------\n", "\n", "Original Response: The train is about to leave.\n", "Baseline Summary : #Person1# is in a hurry to catch a train. Tom tells #Person1# there is plenty of time.\n" ] } ], "source": [ "ind = 40\n", "diag = dataset['test'][ind]['dialogue']\n", "summary = dataset['test'][ind]['summary']\n", "\n", "prompt = f'Summarize the following conversation.\\n\\n{diag}\\n\\nSummary:'\n", "input_ids = tokenizer(prompt, return_tensors=\"pt\").input_ids\n", "\n", "original_outputs = model.to('cpu').generate(input_ids, GenerationConfig(max_new_tokens=200))\n", "original_text_output = tokenizer.decode(original_outputs[0], skip_special_tokens=True)\n", "\n", "print(f'Prompt:\\n--------------------------\\n{prompt}\\n--------------------------')\n", "print(f'\\nOriginal Response: {original_text_output}')\n", "print(f'Baseline Summary : {summary}')" ] }, { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "# Preprocessing\n", "\n", "To preprocess the dataset, we need to append a useful prompt to the start and end of each dialogue set then tokenize the words with HuggingFace. The output dataset will be ready for fine tuning in the next step." ] }, { "cell_type": "code", "execution_count": 16, "metadata": { "tags": [] }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Map: 0%| | 0/12460 [00:00, ? examples/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Map: 0%| | 0/1500 [00:00, ? examples/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Map: 0%| | 0/500 [00:00, ? examples/s]" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "def tokenize_function(example):\n", " prompt = 'Summarize the following conversation.\\n\\n'\n", " end_prompt = '\\n\\nSummary: '\n", " inp = [prompt + i + end_prompt for i in example[\"dialogue\"]]\n", " example['input_ids'] = tokenizer(inp, padding=\"max_length\", truncation=True, return_tensors=\"pt\").input_ids\n", " example['labels'] = tokenizer(example[\"summary\"], padding=\"max_length\", truncation=True, return_tensors=\"pt\").input_ids\n", " return example\n", "\n", "tokenized_datasets = dataset.map(tokenize_function, batched=True)\n", "tokenized_datasets = tokenized_datasets.remove_columns(['id', 'topic', 'dialogue', 'summary',])" ] }, { "cell_type": "code", "execution_count": 17, "metadata": { "tags": [] }, "outputs": [ { "data": { "text/plain": [ "DatasetDict({\n", " train: Dataset({\n", " features: ['input_ids', 'labels'],\n", " num_rows: 12460\n", " })\n", " test: Dataset({\n", " features: ['input_ids', 'labels'],\n", " num_rows: 1500\n", " })\n", " validation: Dataset({\n", " features: ['input_ids', 'labels'],\n", " num_rows: 500\n", " })\n", "})" ] }, "execution_count": 17, "metadata": {}, "output_type": "execute_result" } ], "source": [ "tokenized_datasets" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Fine Tuning with HuggingFace Trainer\n", "\n", "Now that the dataset is preprocessed, we can utilize the built-in HuggingFace `Trainer` class to fine tune our model to the task at hand. Please note that training this full model takes a few hours on a GPU, so for the sake of time, a checkpoint for a model which has been trained on 10 epochs without downsampling has been provided. If you have time to experiment on fully training the model yourself, please see the inline comments for how to change up the code. If you are looking to train on a GPU machine, we have used a `ml.g5.xlarge` instance for the checkpoint provided as a place to start." ] }, { "cell_type": "code", "execution_count": 18, "metadata": { "tags": [] }, "outputs": [ { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Filter: 0%| | 0/12460 [00:00, ? examples/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Filter: 0%| | 0/1500 [00:00, ? examples/s]" ] }, "metadata": {}, "output_type": "display_data" }, { "data": { "application/vnd.jupyter.widget-view+json": { "model_id": "", "version_major": 2, "version_minor": 0 }, "text/plain": [ "Filter: 0%| | 0/500 [00:00, ? examples/s]" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "# for the sake of time in the lab, we will subsample our dataset\n", "# if you want to take the time to train a model fully, feel free to alter this subsampling to create a larger dataset\n", "# the line below can be completely removed to remove the subsampling\n", "tokenized_datasets = tokenized_datasets.filter(lambda example, indice: indice % 100 == 0, with_indices=True)\n", "\n", "output_dir = f'./diag-summary-training-{str(int(time.time()))}'\n", "training_args = TrainingArguments(\n", " output_dir=output_dir,\n", " evaluation_strategy=\"epoch\",\n", " save_strategy=\"epoch\",\n", " learning_rate=1e-5,\n", " num_train_epochs=1,\n", " # num_train_epochs=10, # Use a higher number of epochs when you are not in the lab and have more time to experiment\n", " per_device_train_batch_size=4,\n", " per_device_eval_batch_size=4,\n", " weight_decay=0.01,\n", ")\n", "\n", "trainer = Trainer(\n", " model=model,\n", " args=training_args,\n", " train_dataset=tokenized_datasets['train'],\n", " eval_dataset=tokenized_datasets['validation']\n", ")" ] }, { "cell_type": "code", "execution_count": 19, "metadata": { "tags": [] }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/opt/conda/lib/python3.7/site-packages/transformers/optimization.py:395: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n", " FutureWarning,\n" ] }, { "data": { "text/html": [ "\n", "
Epoch | \n", "Training Loss | \n", "Validation Loss | \n", "
---|
"
],
"text/plain": [
" Shutting down your kernel for this notebook to release resources. Shutting down your kernel for this notebook to release resources.