notebooks/GenAI/GCP_GenAI_Huggingface.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "92021b22-3fbf-4489-9e88-aea4d73f3529",
   "metadata": {},
   "source": [
    "# Finetuning and Deploying Hugging Face Models on Vertex AI"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40d0c590",
   "metadata": {},
   "source": [
    "## Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfbc2ea5-1aca-4322-a15c-7bcbb3925ef6",
   "metadata": {},
   "source": [
    "This tutorial will focus on utilizing Hugging Face which is a repository for user to share and download machine learning models, datasets, and demos. For this tutorial we will load in a model and dataset from Hugging Face and train and test our model before deploying it on Vertex AI. The model we will be deploying is Flan T5 and the datasets is [ccdv/pubmed-summarization](https://HuggingFace.co/datasets/ccdv/pubmed-summarization). Steps will show how to hypertune a model locally and how to launch our custom training job on Vertex AI Training, these steps are based on Keras NLP Tutorials for [abstractive summarization](https://keras.io/examples/nlp/t5_hf_summarization/).\n",
    "\n",
    "You may be wondering why are we training a pretrained model? The reason for this is because we are fine tuning our pretrained model for optimal performance on a particular application, in our case summarizing scientific documents. This is not a necessary step anymore as new methods have been made to enhance model performance like zero-shot learning which we will go over in our next tutorial."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1705116a",
   "metadata": {},
   "source": [
    "## Prerequisites \n",
    "For this tutorial it is recommended to use 1 GPU to speed up processes, this notebooks was run using the machinetype n1-highcpu-8 (8 vCPUs, 7.199 GB RAM) on Tensorflow. Visit the following tutorial to set up notebooks that utilize: GPUs [Spinning up a Vertex AI Notebook](/docs/vertexai.md)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5870770c",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "+ Learning the basics of Hugging Face models\n",
    "+ Learn how to deploy a Hugging Face model and fine tune it"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "565a354d",
   "metadata": {},
   "source": [
    "## Get Started"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ab3668c-58c0-489a-aa4f-6f7e045b450f",
   "metadata": {},
   "source": [
    "### Install packages"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "161dea96-601e-4194-a285-51b0c4403e9d",
   "metadata": {},
   "source": [
    "Hugging Face **transformers** are an open-source framework that allows you to utilize APIs and tools to download pretrained models, set hyperparameters, tokenize datasets, and further tune them to suite your needs. Here we are updating Vertex AI as well as installing the transformers package and **datasets** so that we can have access to Hugging Face datasets and as a bonus we are adding the S3 feature to help download datasets that may already be in a S3 bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6e5884b-ac90-42d4-aafd-d34d5495d24d",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install \"transformers\" \"datasets\" \"rouge_score\" \"evaluate\" \"keras_nlp\" \"tf_keras\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "560867d9",
   "metadata": {},
   "source": [
    "## Get Started"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9b287a7-3035-48c1-8520-a7d3de4f925b",
   "metadata": {},
   "source": [
    "### Download your dataset from Hugging Face"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd58ed15-7e8f-4f24-b68f-4a5b5fcd5cf2",
   "metadata": {},
   "source": [
    "We will be downloading Hugging Face dataset 'ccdv/pubmed-summarization' which contains the full article and their abstracts which will help train our model to summarize scientific articles. Once the dataset is loaded we'll split the data into train, test, and validation datasets. Since these are large datasets we will only be using 5% of dataset to help our process run faster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f59e17e-c006-45ee-be0b-766774f9d420",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# load dataset\n",
    "train, test, validation = load_dataset(\"ccdv/pubmed-summarization\", split=[\"train[:5%]\", \"test[:5%]\", \"validation[:5%]\" ])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51d2756c-d280-4aec-a6fb-416dd91c3ba5",
   "metadata": {},
   "source": [
    "Lets list the feaures of one of our datasets to determine what we will need to tokenize in a later step. this dataset features are 'article' and 'abstract'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "760c9128-793a-4bed-a127-b92ef496e33b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3446d1b6-af0c-4127-93c0-c47fde142546",
   "metadata": {},
   "source": [
    "### Finetuning our Model Locally"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed6ddff1-2636-4e3b-88ee-e3c86c584245",
   "metadata": {},
   "source": [
    "Now that we have our datasets we can upload our model which will be the small version of Flan T5."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b112574b-0e33-4c31-b12a-f1839024ea44",
   "metadata": {},
   "source": [
    "\n",
    "**Flan T5** is a text-to-text generation model and an advancement to the original T5 model and can be run on both CPUs and GPUs. **Text-to-text** is a method of creating text by using a neural network to generate new text from a given input. These T5 models can be fine-tuned for various zero shot NLP tasks that we have seen and heard of before: text classification, summarization, translation, and question-answering. Text-to-text is not to be confused by text2text generation which is a earlier version of T5 that is designed specifically for sequence-to-sequence tasks, such as machine translation and text generation and is limited to these task where as T5 models are more flexible due to the wider range of NPL tasks they can execute.\n",
    "\n",
    "Because it is a seq2seq class model we will be using the transformer **TFAutoModelForSeq2Seq** (specifically for tensorflow models) to help find a load our pretrained model architecture. Then we will assign an **AutoTokenizer** to preprocess the text of our inputs (the test, train, validation datasets) into an array of numbers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfd433b3-9790-4a10-ac08-6c90c194d8b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#model name\n",
    "CHECKPOINT = \"google/flan-t5-small\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1988cbcb-4bec-4aa2-a356-a211584ceacb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TFAutoModelForSeq2SeqLM, AutoTokenizer\n",
    "\n",
    "model = TFAutoModelForSeq2SeqLM.from_pretrained(CHECKPOINT)\n",
    "tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f6ca0419-0075-4f62-becf-b859312cea22",
   "metadata": {},
   "source": [
    "Now that we have loaded the architecture of our model and configured it to tokenize our inputs we can now implement a tokenization functions to start processing our datasets.\n",
    "Since we are using a T5 model we will have to prefix the inputs with \"summarize:\" to know which task to perform. We create a preprocess function to append the prefix to each row within the \"article\" column of our dataset labeling them as inputs. The inputs are then tokenized, limited by a set max length, and truncated.\n",
    "\n",
    "A similar process is done for the \"abstract\" column within our dataset except we do not add the prefix and we labels them as **labels**.\n",
    "\n",
    "**What is Truncating?**\n",
    "\n",
    "Our group of inputs or batch will usually be different lengths which makes it hard to be converted to fixed-size tensors. To fix this problem **truncation** removes tokens ensure longer sequences will have the same length as the longest sequence in the batch which we have set to be **1024** for our inputs and **128** for our labels.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f101c309-f214-4b3f-b77b-d55491e48a59",
   "metadata": {},
   "outputs": [],
   "source": [
    "prefix = \"summarize: \"\n",
    "\n",
    "def preprocess_function(examples):\n",
    "    inputs = [prefix + doc for doc in examples[\"article\"]]\n",
    "    model_inputs = tokenizer(inputs, max_length=1024, truncation=True)\n",
    "\n",
    "    labels = tokenizer(text_target=\n",
    "            examples[\"abstract\"], max_length=128, truncation=True\n",
    "        )\n",
    "\n",
    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "\n",
    "    return model_inputs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "491efc2a-1679-4044-9f53-4aff27329856",
   "metadata": {},
   "source": [
    "Now that we have our tokenized function the next step is to implement the **map** function to iterate the function **preprocess_function** over our loaded datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e58eb58-a655-4e2b-8665-b4b770bc87a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenized_train = train.map(preprocess_function, batched=True)\n",
    "\n",
    "tokenized_test = test.map(preprocess_function, batched=True)\n",
    "\n",
    "tokenized_validation = validation.map(preprocess_function, batched=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41bc1576-44dd-4604-b5db-6c57b711096a",
   "metadata": {},
   "source": [
    "Lets look at the structure of one of our new tokenized datasets you should see 3 new features (**'input_ids', 'attention_mask', 'labels'**) making 5 features total:\n",
    "\n",
    "- **input_ids:** As our inputs are being tokenized an ID is assigned for each token, meaning as each text is broken up into sequences (which can be words or subwords) and converted to tokens within our dataset they are assign an ID.\n",
    "- **attention_masks:** Tokens that should be ignored by the model usually represented by a 0. Masking can be done when some sequences are not the same length so they can not belong in the same tensor and need to be padded.\n",
    "- **labels:** The new name of the abstract column that has been tokenized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80a25bc8-00db-4b8d-9b68-d52c5d6ca7fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tokenized_train)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "285d972c-6aa3-4072-a405-86e7dd82904e",
   "metadata": {},
   "source": [
    "DataCollators are objects that dynamically pads the inputs and the labels in our batches, reverse to truncating **padding** adds a special padding token to ensure shorter sequences will have the same length as the longest sequence in the batch which a gain we set in out preprocess_function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "875ef33d-5ef3-4b07-b1de-6d471743a8ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForSeq2Seq\n",
    "data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=CHECKPOINT, return_tensors=\"tf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b48b037-d666-417b-b28e-88b715a1083c",
   "metadata": {},
   "source": [
    "Then the last step will be to set our data format to be suitable for Tensorflow using the function **'prepare_tf_dataset()'** by automatically inspecting your model and keep only the features that are necessary. As you can see there are only 2 of our features left represented in the dataset: **input_ids and attention_mask**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fcac5a8-912f-461f-bfab-990e472c01ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf_train_set = model.prepare_tf_dataset(\n",
    "    tokenized_train,\n",
    "    shuffle=True,\n",
    "    batch_size=10,\n",
    "    collate_fn=data_collator,\n",
    "    \n",
    ")\n",
    "\n",
    "tf_test_set = model.prepare_tf_dataset(\n",
    "    tokenized_test,\n",
    "    shuffle=False,\n",
    "    batch_size=10,\n",
    "    collate_fn=data_collator,\n",
    "    \n",
    ")\n",
    "\n",
    "tf_validation_set = model.prepare_tf_dataset(\n",
    "    tokenized_validation,\n",
    "    shuffle=False,\n",
    "    batch_size=10,\n",
    "    collate_fn=data_collator,\n",
    "    \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11aaf028-c713-4064-84cc-f699df3151ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "print (tf_train_set)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7ad7839-4967-40fb-aa26-39cea71fa085",
   "metadata": {},
   "source": [
    "**Learning rate** controls how much the model will change in response to the estimated error each time the model weights are updated. Too small of a learning rate could result very slow training process that could eventually get stuck, whereas a value too large may result in an unstable training process. Setting the **weight decay** helps to avoid overfitting, weights small, and avoid exploding gradient. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50ed6068-763a-46a6-8aed-4862f84413a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AdamWeightDecay\n",
    "optimizer = AdamWeightDecay(learning_rate=2e-5, weight_decay_rate=0.01)\n",
    "model.compile(optimizer=optimizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd6a770a-e149-4970-a915-1362f544ef40",
   "metadata": {},
   "source": [
    "Using the function metric_fn will help us calculate the **ROUGE** score between the ground-truth and predictions while training. ROUGE stands for **Recall-Oriented Understudy for Gisting Evaluation** this metric compares a reference sentence with what our model produces see if there is overlap if there is it calculates the precision and recall using the overlap.\n",
    "\n",
    "As an example say our model produced a sentence like so:\n",
    "\n",
    "**'the cat was found under the bed'**\n",
    "\n",
    "but the reference sentence normally written by a human is:\n",
    "\n",
    "**'the cat was under the bed'**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9c6b45b-7349-4965-938b-3a334ced3882",
   "metadata": {},
   "outputs": [],
   "source": [
    "import keras_nlp\n",
    "\n",
    "rouge_l = keras_nlp.metrics.RougeL()\n",
    "\n",
    "\n",
    "def metric_fn(eval_predictions):\n",
    "    predictions, labels = eval_predictions\n",
    "    decoded_predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)\n",
    "    for label in labels:\n",
    "        label[label < 0] = tokenizer.pad_token_id  # Replace masked label tokens\n",
    "    decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "    result = rouge_l(decoded_labels, decoded_predictions)\n",
    "    # We will print only the F1 score, you can use other aggregation metrics as well\n",
    "    result = {\"RougeL\": result[\"f1_score\"]}\n",
    "\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b57e9aa9-9d7c-41d0-8ab1-5d79277c969d",
   "metadata": {},
   "source": [
    "We will use the validation dataset for calculating our ROUGE score. While our ROUGE score is being calculated and our training is running its best to set up a **callback system**. A callback is an object that can perform actions at various stages of training and helps to write logs after every batch of training to monitor your metrics, periodically save your model to disk, and if need be do early stopping. Here we are using Keras call back system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06e014b1-e9d2-4d9f-a149-c6c0381f7407",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers.keras_callbacks import KerasMetricCallback\n",
    "metric_callback = KerasMetricCallback(\n",
    "    metric_fn, eval_dataset=tf_validation_set, predict_with_generate=True, use_xla_generation=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cd1b16ea-0e8a-4c78-875e-e3edea6cf043",
   "metadata": {},
   "source": [
    "Before we start to train our model the last step will be to set how many batches of training we should do, the number of iterations is called **epochs**, we will set ours to 3. Now we can start to train our model using the function **'fit'** and save our artifacts to a directory. The artifact that holds our model will be a file named **tf_model.h5**. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8fd0c64-4d85-4b5e-86fe-538c7dc65da7",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(x=tf_train_set, validation_data=tf_test_set, epochs=3, callbacks=metric_callback)\n",
    "\n",
    "model.save_pretrained('saved_model')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1c35de0-4cdf-4c2b-9c0d-c2272b71b362",
   "metadata": {},
   "source": [
    "### Testing the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31f7f62a-ca17-4dcc-b25d-825572ee1630",
   "metadata": {},
   "source": [
    "Here we will use a sample text that we want our model to summarize."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "980d0053-0c3b-4d0d-91b9-9dd6e6dd3e64",
   "metadata": {},
   "outputs": [],
   "source": [
    "text = \"Severe acute respiratory syndrome coronavirus 2 (SARS-CoV-2) is a \\\n",
    "highly transmissible and pathogenic coronavirus that emerged in late 2019 and has \\\n",
    "caused a pandemic of acute respiratory disease, named ‘coronavirus disease 2019’ (COVID-19), \\\n",
    "which threatens human health and public safety. In this Review, we describe the basic virology of \\\n",
    "SARS-CoV-2, including genomic characteristics and receptor use, highlighting its key difference \\\n",
    "from previously known coronaviruses. We summarize current knowledge of clinical, epidemiological and \\\n",
    "pathological features of COVID-19, as well as recent progress in animal models and antiviral treatment \\\n",
    "approaches for SARS-CoV-2 infection. We also discuss the potential wildlife hosts and zoonotic origin \\\n",
    "of this emerging virus in detail.\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3156ef96-20ab-46ff-ae76-57d6ea54f0ff",
   "metadata": {},
   "source": [
    "To predict the following tokenizes the text to gather the inputs, then uses **generate()** generate sequences of token ids for our model. We then decode our output to translate our tokenized output into text."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb34d7c2-a815-4c9d-bd8b-495aca8b2d02",
   "metadata": {},
   "source": [
    "Below you will see that we have provided a paragraph about SARS-CoV-2 as our output, we also have some parameters that we specify to further tune our model to get a concise summary of what our text is about.\n",
    "\n",
    "- **Max_Length:** Max number of words to generate.\n",
    "- **Num_Return_Sequences:** Number of different outputs to generate. For our example we want one sentence or sequence.\n",
    "- **Temperature:** Controls randomness, higher values increase diversity meaning a more unique response make the model to think harder. Must be a number from 0 to 1.\n",
    "- **Top_p (nucleus):** The cumulative probability cutoff for token selection. Lower values mean sampling from a smaller, more top-weighted nucleus. Must be a number from 0 to 1.\n",
    "- **Top_k**: Sample from the k most likely next tokens at each step. Lower k focuses on higher probability tokens. This means the model choses the most probable words. Lower values eliminate fewer coherent words."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc2206c7-1bbf-41eb-8c63-abb17752d00d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)\n",
    "inputs = tokenizer.encode(text, return_tensors=\"tf\")\n",
    "\n",
    "from transformers import TFAutoModelForSeq2SeqLM\n",
    "\n",
    "model = TFAutoModelForSeq2SeqLM.from_pretrained(\"saved_model\")\n",
    "\n",
    "outputs = model.generate(inputs, \n",
    "                         max_length=1000,\n",
    "                         num_return_sequences = 1,\n",
    "                         do_sample=True, \n",
    "                         temperature = 0.6,\n",
    "                         top_k = 50, \n",
    "                         top_p = 0.95,)\n",
    "\n",
    "tokenizer.decode(outputs[0], skip_special_tokens=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af0f0064-d9b4-430c-90ae-d2e390d1b78c",
   "metadata": {},
   "source": [
    "### Optional: Summarizing PDF Files"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1abe4b0b-8ee2-48f1-afed-0cd2796717a2",
   "metadata": {},
   "source": [
    "The process of summarizing scientific PDF files is relatively the same except that we first need to extract the text from the PDF. To do so lets download a PDF file from PubMed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d69aa008-80c0-4a19-aa0f-8f5798673c47",
   "metadata": {},
   "outputs": [],
   "source": [
    "! wget --user-agent=\"Chrome\" https://www.ncbi.nlm.nih.gov/pmc/articles/PMC7784226/pdf/12248_2020_Article_532.pdf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6555b048-a76c-4f10-b13c-4ac93f436fe8",
   "metadata": {},
   "source": [
    "We'll be downloading some tools that help us extract only the text from our pdf file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1347b3cd-5ce0-44c9-864d-a688bcacb1d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"fitz\" \"PyMuPDF\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3331b05f-d00e-4ff2-be24-edea732e4af2",
   "metadata": {},
   "source": [
    "Now we can make a function **extract_text_from_pdf** to extract the text from the pdf and save it as a variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d77b6ffe-90e1-4a01-aa52-9cf93a9c5c85",
   "metadata": {},
   "outputs": [],
   "source": [
    "import fitz\n",
    "def extract_text_from_pdf(pdf_path):\n",
    "    doc = fitz.open(pdf_path)\n",
    "    text = ''\n",
    "    for page in doc:\n",
    "        text += page.get_text()\n",
    "    return text\n",
    "\n",
    "text_pdf=extract_text_from_pdf('12248_2020_Article_532.pdf')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10e761c6-da93-4ca0-827e-a3715c20eb10",
   "metadata": {},
   "source": [
    "Finally we'll follow the same steps we did before to encode our inputs, pass it to our model, and then decode our output. Notice how we increased the max_length of what is expected of our input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a1f5dbd-6a9f-4533-a12e-8a6c4073df74",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)\n",
    "inputs = tokenizer.encode(text_pdf, max_length=1000, truncation=True, return_tensors=\"tf\")\n",
    "\n",
    "from transformers import TFAutoModelForSeq2SeqLM\n",
    "\n",
    "model = TFAutoModelForSeq2SeqLM.from_pretrained(\"saved_model\")\n",
    "\n",
    "outputs = model.generate(inputs, max_new_tokens=100, do_sample=False)\n",
    "\n",
    "tokenizer.decode(outputs[0], skip_special_tokens=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "257fb45f-7752-481d-a1f6-f3eeb7655fac",
   "metadata": {},
   "source": [
    "### Finetuning our Model via Vertex AI Training API"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ac841f6-c65e-4ebf-8c42-3030e2f92cb0",
   "metadata": {},
   "source": [
    "### Setting up our Datasets for Training "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fff825cf-86fb-4777-885b-9e981be831b7",
   "metadata": {},
   "source": [
    "Although we have our datasets saved locally inorder to utilize the Vertex AI Training API we will need to store our datasets in a bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3f49896-b2c1-47e6-a7cc-aca7753bb6c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# load dataset\n",
    "train, test, validation = load_dataset(\"ccdv/pubmed-summarization\", split=[\"train[:5%]\", \"test[:5%]\", \"validation[:5%]\" ])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b91fe7c-5970-45c1-9401-1db3206a8ce9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#load in the storage package and name our bucket\n",
    "from google.cloud import storage\n",
    "BUCKET='flan-t5-model-resources'\n",
    "client = storage.Client()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0066ad72-e451-41c0-b30a-c3a7dfa5f17c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create bucket\n",
    "bucket = client.bucket(BUCKET)\n",
    "bucket.create()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4fc7335e-cfa7-4b0e-8880-85edfa573772",
   "metadata": {},
   "source": [
    "Convert our datasets to csv and upload to our bucket in one step!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bfbbd92-4b2c-4e5c-95f8-d4e645a6ab24",
   "metadata": {},
   "outputs": [],
   "source": [
    "from io import BytesIO\n",
    "\n",
    "#convert train dataset to csv and push to GCS bucket\n",
    "csv_buffer = BytesIO()\n",
    "train.to_csv(csv_buffer)\n",
    "client = storage.Client()\n",
    "bucket = client.get_bucket(BUCKET)\n",
    "bucket.blob('train.csv').upload_from_file(csv_buffer, 'text/csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbf3f68f-acc8-4086-9b89-be0d3eacf898",
   "metadata": {},
   "outputs": [],
   "source": [
    "#convert test dataset to csv and push to GCS bucket\n",
    "csv_buffer = BytesIO()\n",
    "test.to_csv(csv_buffer)\n",
    "client = storage.Client()\n",
    "bucket = client.get_bucket(BUCKET)\n",
    "bucket.blob('test.csv').upload_from_file(csv_buffer, 'text/csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea1773e1-dfe2-46b5-a63d-782101d79096",
   "metadata": {},
   "outputs": [],
   "source": [
    "#convert validation dataset to csv and push to GCS bucket\n",
    "csv_buffer = BytesIO()\n",
    "validation.to_csv(csv_buffer)\n",
    "client = storage.Client()\n",
    "bucket = client.get_bucket(BUCKET)\n",
    "bucket.blob('validation.csv').upload_from_file(csv_buffer, 'text/csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f18630a7-109c-4f53-9233-1842f5c27029",
   "metadata": {},
   "source": [
    "Here we will be saving the location of our datasets be used when we execute the training of our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebc1bc39-a554-473b-949a-d9588f6e7fb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# save train_dataset to s3\n",
    "training_input_path = f'gs://{BUCKET}/train.csv'\n",
    "\n",
    "# save test_dataset to s3\n",
    "test_input_path = f'gs://{BUCKET}/test.csv'\n",
    "\n",
    "validation_input_path = f'gs://{BUCKET}/validation.csv'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9204b6dc-8f6e-407e-8c68-a036a6a5b7c9",
   "metadata": {},
   "source": [
    "### Training our Model via Vertex AI Training API"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f873f2f-90b8-4566-96f5-37a23a2294e1",
   "metadata": {},
   "source": [
    "To train our model on Vertex AI Training API you must first create a custom AI job, this is done by creating a autopkg that holds your requirements.txt and task.py files is a specific structure like so: "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8eeafae2-a698-4a52-a4f1-cae550245d0b",
   "metadata": {},
   "source": [
    "```\n",
    "autopkg-summarizer /\n",
    "    + requirements.txt\n",
    "    + trainer/\n",
    "        + task.py\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48d49de7-6d86-411e-9e6e-104763ae36e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Creates the following directories and files\n",
    "!mkdir autopkg-summarizer\n",
    "!touch autopkg-summarizer/requirements.txt\n",
    "!mkdir autopkg-summarizer/trainer\n",
    "!touch autopkg-summarizer/trainer/task.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7bd953b5-8d21-4c07-adcc-23604d5d0279",
   "metadata": {},
   "source": [
    "Add your requirements.txt file by adding the packages below:\n",
    "```\n",
    "nltk\n",
    "transformers\n",
    "keras_nlp\n",
    "datasets\n",
    "rouge_score\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "443f0e45-dbe6-4ce9-b0ad-96a5fe52a455",
   "metadata": {},
   "source": [
    "To create our training script we will be adding all the steps that we ran from the 'Finetuning our Model Locally' section of this tutorial to a file named task.py:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14ab5a7d-36bf-449d-b2cd-4f7e107c75de",
   "metadata": {},
   "source": [
    "```\n",
    "import nltk\n",
    "import argparse\n",
    "from datasets import load_dataset\n",
    "#import evaluate\n",
    "import numpy as np\n",
    "from transformers import create_optimizer, AdamWeightDecay, TFAutoModelForSeq2SeqLM, AutoTokenizer, DataCollatorForSeq2Seq, set_seed\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from transformers.keras_callbacks import KerasMetricCallback\n",
    "import keras_nlp\n",
    "\n",
    "def get_args():\n",
    "    '''Parses args.'''\n",
    "    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)\n",
    "    parser.add_argument(\n",
    "        '--model_name_or_path',\n",
    "        required=True,\n",
    "        type=str,\n",
    "        help='name of model or path to load into tokenizer and class')\n",
    "    parser.add_argument(\n",
    "        '--train_file',\n",
    "        required=True,\n",
    "        type=str,\n",
    "        help='train dataset in csv or json format')\n",
    "    parser.add_argument(\n",
    "        '--test_file',\n",
    "        required=True,\n",
    "        type=str,\n",
    "        help='test dataset in csv or json format')\n",
    "    parser.add_argument(\n",
    "        '--validation_file',\n",
    "        required=True,\n",
    "        type=str,\n",
    "        help='validation dataset in csv or json format used to calculate ROUGE score')\n",
    "    parser.add_argument(\n",
    "        '--text_column',\n",
    "        required=True,\n",
    "        type=str,\n",
    "        help='The name of the column in the datasets containing the full texts (for summarization)')\n",
    "    parser.add_argument(\n",
    "        '--summary_column',\n",
    "        required=True,\n",
    "        type=str,\n",
    "        help='The name of the column in the datasets containing the abstracts or summary of the full text')\n",
    "    parser.add_argument(\n",
    "        '--num_train_epochs',\n",
    "        required=False,\n",
    "        type=int,\n",
    "        default=3,\n",
    "        help='number of complete passes through the training dataset')\n",
    "    parser.add_argument(\n",
    "        '--source_prefix',\n",
    "        required=False,\n",
    "        type=str,\n",
    "        help='A prefix to add before every source text (needed for T5 models)')\n",
    "    parser.add_argument(\n",
    "        '--inputs_max_length',\n",
    "        required=False,\n",
    "        type=int,\n",
    "        default=1024,\n",
    "        help='max token length for model inputs')\n",
    "    parser.add_argument(\n",
    "        '--labels_max_length',\n",
    "        required=False,\n",
    "        type=int,\n",
    "        default=128,\n",
    "        help='max token length for model labels or targets')\n",
    "    parser.add_argument(\n",
    "        '--batch_size',\n",
    "        required=False,\n",
    "        type=int,\n",
    "        default=10,\n",
    "        help='max token length for model labels or targets')\n",
    "    parser.add_argument(\n",
    "        '--output_dir',\n",
    "        required=True,\n",
    "        type=str,\n",
    "        help='bucket to store saved model, include gs://')\n",
    "    \n",
    "    args = parser.parse_args()\n",
    "    return args\n",
    "\n",
    "def main():\n",
    "    \n",
    "    args = get_args() \n",
    " \n",
    "    checkpoint = args.model_name_or_path\n",
    "    \n",
    "    tokenizer = AutoTokenizer.from_pretrained(checkpoint)\n",
    "    \n",
    "    text = args.text_column\n",
    "    summary = args.summary_column\n",
    "    inputs_max_length = args.inputs_max_length\n",
    "    labels_max_length = args.labels_max_length\n",
    "    prefix = args.source_prefix \n",
    "    \n",
    "    model = TFAutoModelForSeq2SeqLM.from_pretrained(checkpoint) \n",
    "    \n",
    "    data_files = {'train':args.train_file, 'test':args.test_file, 'validation':args.validation_file}\n",
    "    extension = args.train_file.split(\".\")[-1]\n",
    "    \n",
    "    raw_datasets = load_dataset(\n",
    "        extension,\n",
    "        data_files=data_files)\n",
    "    \n",
    "    raw_datasets = raw_datasets.filter(lambda x: x[text] is not None) \n",
    "    \n",
    "    train = raw_datasets[\"train\"]\n",
    "    test = raw_datasets[\"test\"]\n",
    "    validation = raw_datasets[\"validation\"]\n",
    "         \n",
    "    def preprocess_function(examples):\n",
    "        \n",
    "        inputs = [prefix + doc for doc in examples[text]]\n",
    "        model_inputs = tokenizer(inputs, max_length=inputs_max_length, truncation=True)\n",
    "\n",
    "    #    labels = tokenizer(text_target=examples[\"abstract\"], max_length=128, truncation=True)\n",
    "\n",
    "        labels = tokenizer(text_target=\n",
    "                examples[summary], max_length=labels_max_length, truncation=True\n",
    "            )\n",
    "\n",
    "        model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
    "        return model_inputs\n",
    "    \n",
    "    tokenized_train = train.map(preprocess_function, batched=True)\n",
    "    tokenized_test = test.map(preprocess_function, batched=True)\n",
    "    tokenized_validation = validation.map(preprocess_function, batched=True)\n",
    "    \n",
    "    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=checkpoint, return_tensors=\"tf\")\n",
    "\n",
    "    optimizer = AdamWeightDecay(learning_rate=2e-5, weight_decay_rate=0.01)\n",
    "    model.compile(optimizer=optimizer)\n",
    "\n",
    "    tf_train_set = model.prepare_tf_dataset(\n",
    "        tokenized_train,\n",
    "        shuffle=True,\n",
    "        batch_size=args.batch_size,\n",
    "        collate_fn=data_collator\n",
    "    )\n",
    "\n",
    "    tf_test_set = model.prepare_tf_dataset(\n",
    "        tokenized_test,\n",
    "        shuffle=False,\n",
    "        batch_size=args.batch_size,\n",
    "        collate_fn=data_collator\n",
    "    )\n",
    "    \n",
    "    tf_validation_set = model.prepare_tf_dataset(\n",
    "        tokenized_validation,\n",
    "        shuffle=False,\n",
    "        batch_size=args.batch_size,\n",
    "        collate_fn=data_collator\n",
    "    )   \n",
    "    \n",
    "    def metric_fn(eval_predictions):\n",
    "        predictions, labels = eval_predictions\n",
    "        decoded_predictions = tokenizer.batch_decode(predictions, skip_special_tokens=True)\n",
    "        for label in labels:\n",
    "            label[label < 0] = tokenizer.pad_token_id  # Replace masked label tokens\n",
    "        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)\n",
    "        result = rouge_l(decoded_labels, decoded_predictions)\n",
    "        # We will print only the F1 score, you can use other aggregation metrics as well\n",
    "        result = {\"RougeL\": result[\"f1_score\"]}\n",
    "\n",
    "        return result\n",
    "    \n",
    "    rouge_l = keras_nlp.metrics.RougeL()\n",
    "\n",
    "    metric_callback = KerasMetricCallback(\n",
    "        metric_fn, eval_dataset=tf_validation_set, predict_with_generate=True, use_xla_generation=True)\n",
    "\n",
    "\n",
    "    model.fit(x=tf_train_set, validation_data=tf_test_set, epochs=args.num_train_epochs, callbacks=metric_callback)\n",
    "    model.save(f'{args.output_dir}/saved_model_artifacts_tf')\n",
    "    model.save_pretrained(f'{args.output_dir}/saved_model_hf_tf')\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    main()\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8775dba1-c47b-4375-9587-6fec561bc5f9",
   "metadata": {},
   "source": [
    "### Hyperparameters (for the training script and custom AI job)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a66d0c47-f6df-4b79-87a6-a637b04ebc87",
   "metadata": {},
   "source": [
    "The first step to training our model other than setting up our datasets is to set our **hyperparameters**. Hyperparameters depend on your training script and for this one we need to identify our model, the location of our train and test files, etc. \n",
    "\n",
    "The batch_size, inputs_max_length, num_train_epochs, and labels_max_length already have defualts setting same as the ones we used in the first section of this tutorial!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "764980e6-9bc1-4715-b540-9e254b12f1f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#to view options and defaults you can run the command below\n",
    "!python autopkg-summarizer/trainer/task.py --help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b21c8c79-1709-4052-8522-ae332cfec934",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Parameters for task.py script\n",
    "CHECKPOINT = \"google/flan-t5-small\"\n",
    "train_file=training_input_path\n",
    "test_file=test_input_path\n",
    "validation_file=validation_input_path\n",
    "text_column=\"article\"\n",
    "summary_column=\"abstract\"\n",
    "source_prefix=\"summarize: \" \n",
    "output_dir= f'gs://{BUCKET}'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e830bb4a-854e-412d-93ec-3059faf603d6",
   "metadata": {},
   "source": [
    "For custom AI we need to set the machine type, the accelerator for GPUs, and prebuilt docker image that will run our training. See here for more available containers: https://cloud.google.com/vertex-ai/docs/training/pre-built-containers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09392ddd-aa9d-4358-95a6-3e64fa1692ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Parameters for custom AI job\n",
    "display_name='flan-t5-training-tf'\n",
    "BASE_GPU_IMAGE_tf='us-docker.pkg.dev/vertex-ai/training/tf-gpu.2-12.py310:latest'\n",
    "machine_type='n1-standard-4'\n",
    "accelerator_type='NVIDIA_TESLA_V100'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a6a7856-624d-4229-8a5d-cfd263a84033",
   "metadata": {},
   "source": [
    "### Submit Custom AI Training Job"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16c88abb-2a4f-4475-ad58-1d69ca31c449",
   "metadata": {},
   "source": [
    "Finally we can submit our training via a custom job! It will first deploy the container that we specified and then submit our model for training. This custom job can take 15 - 20 min using our sample datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "252d8e16-5b3d-409b-bc86-9da0ce996f72",
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud ai custom-jobs create \\\n",
    "--region=us-central1 \\\n",
    "--display-name=$display_name \\\n",
    "--args=--model_name_or_path=$CHECKPOINT \\\n",
    "--args=--train_file=$train_file \\\n",
    "--args=--test_file=$test_file \\\n",
    "--args=--validation_file=$validation_file \\\n",
    "--args=--text_column=$text_column \\\n",
    "--args=--summary_column=$summary_column \\\n",
    "--args=--output_dir=gs://$BUCKET \\\n",
    "--args=--source_prefix=$source_prefix \\\n",
    "--worker-pool-spec=machine-type=$machine_type,replica-count=1,accelerator-type=$accelerator_type,executor-image-uri=$BASE_GPU_IMAGE_tf,local-package-path=autopkg-summarizer,python-module=trainer.task"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "403fe77f-990c-4518-ad3a-0aac3d2c8b92",
   "metadata": {},
   "source": [
    "Once you start training the output from the command line should show you the command to use to view the progress of your training via the command `gcloud ai custom-jobs stream-logs <`. You can also monitor and view logs on the console by going to `Vertex AI > Training > Custom Jobs`\n",
    "select your custom job and click on \"View Logs\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc7b81af-d424-427d-a37a-ac5da197567e",
   "metadata": {},
   "source": [
    "### Deploy the Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d772dc95-95e9-40f2-a5c5-dc782d6f7e14",
   "metadata": {},
   "source": [
    "### Upload the Model to Vertex AI's Model Registry"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dff127a-4d14-4fa1-a22a-75662eccce02",
   "metadata": {},
   "source": [
    "Once our model is done training you should see a model_save.pd file in your bucket. We will need this inorder to upload our model to the Model Registry. Here we are specifiying a prebuilt docker image that will run our predictions, the name of our model and the directory in our bucket that holds our **model_save.pd** file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26042230-dc95-4b6c-bd32-bf3596e5de52",
   "metadata": {},
   "outputs": [],
   "source": [
    "TF_PREDICTION_IMAGE_URI_RUNTIME = 'us-docker.pkg.dev/vertex-ai-restricted/prediction/tf_opt-gpu.2-12:latest'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5a3ef3b-8080-4e2d-bb8f-7e2f22c59e05",
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.cloud import aiplatform as vertexai\n",
    "from google.cloud import aiplatform\n",
    "\n",
    "#give your model a name\n",
    "MODEL_DISPLAY_NAME = \"summarizer-tf-runtime\"\n",
    "MODEL_DESCRIPTION = \"summarizes scientific texts and pdfs\" #optional\n",
    "\n",
    "#add your project ID and location\n",
    "project='<PROJECT_ID>'\n",
    "location='<LOCATION>'\n",
    "\n",
    "vertexai.init(project=project, location=location, staging_bucket=BUCKET)\n",
    "\n",
    "\n",
    "model = aiplatform.Model.upload(\n",
    "    display_name=MODEL_DISPLAY_NAME,\n",
    "    description=MODEL_DESCRIPTION,\n",
    "    serving_container_image_uri=TF_PREDICTION_IMAGE_URI_RUNTIME,\n",
    "    serving_container_args=[\"--allow_precompilation\", \"--allow_compression\", \"--use_tfrt\"],\n",
    "    artifact_uri=f'gs://{BUCKET}/saved_model_artifacts_tf', #directory where our artifacts are in our bucket\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d9772a7-f00d-4633-aa87-3861fa5dec79",
   "metadata": {},
   "source": [
    "### Create a Endpoint and Deploy it to our Model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26122187-da1c-4d26-b1a0-ec1bd403cb19",
   "metadata": {},
   "source": [
    "A **endpoint** is how the user of the model can communicate with the model. A single model endpoint responds by returning a single inference from at least one model. It can take 20 min or more to establish a endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74a2c3dd-0e34-4049-804b-940c9a440570",
   "metadata": {},
   "outputs": [],
   "source": [
    "ENDPOINT_DISPLAY_NAME = \"summarizer-endpoint\" \n",
    "endpoint = aiplatform.Endpoint.create(display_name=ENDPOINT_DISPLAY_NAME)\n",
    "\n",
    "model_endpoint = model.deploy(\n",
    "    endpoint=endpoint,\n",
    "    deployed_model_display_name=MODEL_DISPLAY_NAME,\n",
    "    machine_type=\"n1-standard-8\",\n",
    "    accelerator_type=\"NVIDIA_TESLA_V100\",\n",
    "    accelerator_count=1,\n",
    "    traffic_percentage=100,\n",
    "    deploy_request_timeout=1200,\n",
    "    sync=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dee7811-559f-4bdc-b56e-b932b0831c0f",
   "metadata": {},
   "source": [
    "Here we are creating a endpoint and deploying our model to said endpoint. We are deploying our endpoint using 1 GPU which can take 20min to run, feel free to try out other machine types that utilize more GPUs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf29655a",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "Here you learned how to deploy and fine tune a Hugging Face model on Google Cloud"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1544dbe2-8f06-43e8-9b2c-e9d57332e00e",
   "metadata": {},
   "source": [
    "## Clean Up"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b7a6d9a-3d0c-425e-a8f3-cb3160f1ee3b",
   "metadata": {},
   "source": [
    "**Warning:** Once you are done don't forget to delete your endpoint, model, buckets, and shutdown or delete your Vertex AI notebook to avoid additional charges!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24d0fea3-fd4a-4735-b6d5-b910239b5ffa",
   "metadata": {},
   "source": [
    "First we will delete our custom job. The command below will list custom jobs allowing you to gather the job id from the field called **'name:projects/<PROJECT_ID>/locations/us-central1/customJobs/<JOB_ID>'**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9721f52d-040f-4dc7-808e-8d1ffb5efb4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud ai custom-jobs list --project=$project --region=$location"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6e10933-3d84-41fd-8785-fa801b97bfb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.cloud import aiplatform\n",
    "custom_job_id='<Custom_Job_ID_from_List>'\n",
    "\n",
    "def delete_custom_job_sample(custom_job_id: str,\n",
    "    project: str = project,\n",
    "    location: str = location,\n",
    "    api_endpoint: str = f'{location}-aiplatform.googleapis.com',\n",
    "    timeout: int = 300,\n",
    "):\n",
    "    # The AI Platform services require regional API endpoints.\n",
    "    client_options = {\"api_endpoint\": api_endpoint}\n",
    "    # Initialize client that will be used to create and send requests.\n",
    "    # This client only needs to be created once, and can be reused for multiple requests.\n",
    "    client = aiplatform.gapic.JobServiceClient(client_options=client_options)\n",
    "    name = client.custom_job_path(\n",
    "        project=project, location=location, custom_job=custom_job_id\n",
    "    )\n",
    "    response = client.delete_custom_job(name=name)\n",
    "    print(\"Long running operation:\", response.operation.name)\n",
    "    delete_custom_job_response = response.result(timeout=timeout)\n",
    "    print(\"delete_custom_job_response:\", delete_custom_job_response)\n",
    "    \n",
    "delete_custom_job_sample(custom_job_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a02ba4b-1c8f-4fcc-93b9-3e2a6121b59b",
   "metadata": {},
   "source": [
    "Now we will undeploy our model, delete endpoints, and delete finally our model!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfc2276e-8ab2-4c80-9721-26153ea80d63",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_endpoint.undeploy_all()\n",
    "model_endpoint.delete()\n",
    "model.delete()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8fc98dd1-b4e5-4ab2-a6ce-e46b5a23ab5d",
   "metadata": {},
   "source": [
    "Delete custom container stored in Custom Registry or Artifacr Registry. List the images to gather the tag id."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "588af684-4c7e-43d5-a1f5-5510157aa40f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#list the containers\n",
    "!gcloud container images list-tags gcr.io/$project/cloudai-autogenerated/$display_name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "635cc519-230d-48a0-b9b7-c350c2d62ac4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Save the tag ID\n",
    "tag_id='<TAG_ID>'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "946fa12b-ad77-4d19-a556-c926309a14c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#delete \n",
    "!gcloud container images delete gcr.io/$project/cloudai-autogenerated/$display_name:$tag_id --force-delete-tags --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd4a6cfd-9567-4da3-8d1b-6a4207442680",
   "metadata": {},
   "source": [
    "And finally delete our bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85824daa-66a5-4303-8b17-2565863a2844",
   "metadata": {},
   "outputs": [],
   "source": [
    "!gcloud storage rm --recursive gs://$BUCKET/"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}