sdk/python/foundation-models/system/finetune/image-classification/multiclass-classification/hftransformers-fridgeobjects-multiclass-classification.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multi-class Image Classification using transformers specific pipeline component\n",
    "\n",
    "This sample shows how to use `transformers_image_classification_pipeline` component from the `azureml` system registry to fine tune a model for multi-class image classification task using fridgeObjects Dataset. We then deploy the fine tuned model to an online endpoint for real time inference.\n",
    "\n",
    "### Training data\n",
    "We will use the [fridgeObjects](https://cvbp-secondary.z19.web.core.windows.net/datasets/image_classification/fridgeObjects.zip) dataset.\n",
    "\n",
    "### Model\n",
    "We will use the `microsoft-beit-base-patch16-224-pt22k-ft22k` model in this notebook. If you need to fine tune a model that is available on HuggingFace, but not available in `azureml` system registry, you can either register the model and use the registered model or use the `model_name` parameter to instruct the components to pull the model directly from HuggingFace.\n",
    "\n",
    "### Outline\n",
    "1. Install dependencies\n",
    "2. Setup pre-requisites such as compute\n",
    "3. Pick a model to fine tune\n",
    "4. Prepare dataset for finetuning the model\n",
    "5. Submit the fine tuning job using transformers specific image-classification component\n",
    "6. Review training and evaluation metrics\n",
    "7. Register the fine tuned model\n",
    "8. Deploy the fine tuned model for real time inference\n",
    "9. Test deployed end point\n",
    "9. Clean up resources"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Install dependencies\n",
    "Before starting off, if you are running the notebook on Azure Machine Learning Studio or running first time locally, you will need the following packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install azure-ai-ml==1.8.0\n",
    "! pip install azure-identity==1.13.0"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Setup pre-requisites"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.1 Connect to Azure Machine Learning workspace\n",
    "\n",
    "Before we dive in the code, you'll need to connect to your workspace. The workspace is the top-level resource for Azure Machine Learning, providing a centralized place to work with all the artifacts you create when you use Azure Machine Learning.\n",
    "\n",
    "We are using `DefaultAzureCredential` to get access to workspace. `DefaultAzureCredential` should be capable of handling most scenarios. If you want to learn more about other available credentials, go to [set up authentication doc](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-setup-authentication?tabs=sdk), [azure-identity reference doc](https://learn.microsoft.com/en-us/python/api/azure-identity/azure.identity?view=azure-python).\n",
    "\n",
    "Replace `AML_WORKSPACE_NAME`, `RESOURCE_GROUP` and `SUBSCRIPTION_ID` with their respective values in the below cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient\n",
    "from azure.identity import DefaultAzureCredential\n",
    "\n",
    "\n",
    "experiment_name = (\n",
    "    \"AzureML-Train-Finetune-Vision-MultiClass-Samples\"  # can rename to any valid name\n",
    ")\n",
    "\n",
    "credential = DefaultAzureCredential()\n",
    "workspace_ml_client = None\n",
    "try:\n",
    "    workspace_ml_client = MLClient.from_config(credential)\n",
    "    subscription_id = workspace_ml_client.subscription_id\n",
    "    resource_group = workspace_ml_client.resource_group_name\n",
    "    workspace_name = workspace_ml_client.workspace_name\n",
    "except Exception as ex:\n",
    "    print(ex)\n",
    "    # Enter details of your AML workspace\n",
    "    subscription_id = \"SUBSCRIPTION_ID\"\n",
    "    resource_group = \"RESOURCE_GROUP\"\n",
    "    workspace_name = \"AML_WORKSPACE_NAME\"\n",
    "\n",
    "workspace_ml_client = MLClient(\n",
    "    credential, subscription_id, resource_group, workspace_name\n",
    ")\n",
    "registry_ml_client = MLClient(\n",
    "    credential,\n",
    "    subscription_id,\n",
    "    resource_group,\n",
    "    registry_name=\"azureml\",\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.2 Create compute\n",
    "\n",
    "In order to finetune a model on Azure Machine Learning studio, you will need to create a compute resource first. **Creating a compute will take 3-4 minutes.** \n",
    "\n",
    "For additional references, see [Azure Machine Learning in a Day](https://github.com/Azure/azureml-examples/blob/main/tutorials/azureml-in-a-day/azureml-in-a-day.ipynb). "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Create CPU compute for model selection component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import AmlCompute\n",
    "from azure.core.exceptions import ResourceNotFoundError\n",
    "\n",
    "model_import_cluster_name = \"sample-model-import-cluster\"\n",
    "try:\n",
    "    _ = workspace_ml_client.compute.get(model_import_cluster_name)\n",
    "    print(\"Found existing compute target.\")\n",
    "except ResourceNotFoundError:\n",
    "    print(\"Creating a new compute target...\")\n",
    "    compute_config = AmlCompute(\n",
    "        name=model_import_cluster_name,\n",
    "        type=\"amlcompute\",\n",
    "        size=\"Standard_D12_v2\",\n",
    "        idle_time_before_scale_down=120,\n",
    "        min_instances=0,\n",
    "        max_instances=4,\n",
    "    )\n",
    "    workspace_ml_client.begin_create_or_update(compute_config).result()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Create GPU compute for finetune component\n",
    "\n",
    "The list of GPU machines can be found [here](https://learn.microsoft.com/en-us/azure/virtual-machines/sizes-gpu)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "finetune_cluster_name = \"sample-finetune-cluster-gpu\"\n",
    "\n",
    "try:\n",
    "    _ = workspace_ml_client.compute.get(finetune_cluster_name)\n",
    "    print(\"Found existing compute target.\")\n",
    "except ResourceNotFoundError:\n",
    "    print(\"Creating a new compute target...\")\n",
    "    compute_config = AmlCompute(\n",
    "        name=finetune_cluster_name,\n",
    "        type=\"amlcompute\",\n",
    "        size=\"Standard_NC6s_v3\",\n",
    "        idle_time_before_scale_down=120,\n",
    "        min_instances=0,\n",
    "        max_instances=4,\n",
    "    )\n",
    "    workspace_ml_client.begin_create_or_update(compute_config).result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Create GPU compute for model evaluation component\n",
    "\n",
    "The list of GPU machines can be found [here](https://learn.microsoft.com/en-us/azure/virtual-machines/sizes-gpu)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using the same compute cluster for model evaluation as finetuning. If you want to use a different cluster, specify it below\n",
    "model_eval_cluster_name = \"sample-finetune-cluster-gpu\"\n",
    "\n",
    "try:\n",
    "    _ = workspace_ml_client.compute.get(model_eval_cluster_name)\n",
    "    print(\"Found existing compute target.\")\n",
    "except ResourceNotFoundError:\n",
    "    print(\"Creating a new compute target...\")\n",
    "    compute_config = AmlCompute(\n",
    "        name=model_eval_cluster_name,\n",
    "        type=\"amlcompute\",\n",
    "        size=\"Standard_NC6s_v3\",\n",
    "        idle_time_before_scale_down=120,\n",
    "        min_instances=0,\n",
    "        max_instances=4,\n",
    "    )\n",
    "    workspace_ml_client.begin_create_or_update(compute_config).result()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Pick a foundation model to fine tune\n",
    "\n",
    "We will use the `microsoft-beit-base-patch16-224-pt22k-ft22k` model in this notebook. If you need to fine tune a model that is available on HuggingFace, but not available in `azureml` system registry, you can either register the model and use the registered model or use the `model_name` parameter to instruct the components to pull the model directly from HuggingFace.\n",
    "\n",
    "Currently following models are supported:\n",
    "\n",
    "| Model Name | Source |\n",
    "| ------ | ---------- |\n",
    "| [microsoft-beit-base-patch16-224-pt22k-ft22k](https://ml.azure.com/registries/azureml/models/microsoft-beit-base-patch16-224-pt22k-ft22k/version/11) | azureml registry |\n",
    "| [microsoft-swinv2-base-patch4-window12-192-22k](https://ml.azure.com/registries/azureml/models/microsoft-swinv2-base-patch4-window12-192-22k/version/11) | azureml registry |\n",
    "| [facebook-deit-base-patch16-224](https://ml.azure.com/registries/azureml/models/facebook-deit-base-patch16-224/version/10) | azureml registry |\n",
    "| [google-vit-base-patch16-224](https://ml.azure.com/registries/azureml/models/google-vit-base-patch16-224/version/10) | azureml registry |\n",
    "| [Image classification models from Huggingface's Transformer library](https://huggingface.co/models?pipeline_tag=image-classification&library=transformers)| HuggingFace |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "huggingface_model_name = \"microsoft/beit-base-patch16-224-pt22k-ft22k\"\n",
    "\n",
    "aml_registry_model_name = \"microsoft-beit-base-patch16-224-pt22k-ft22k\"\n",
    "foundation_models = registry_ml_client.models.list(aml_registry_model_name)\n",
    "foundation_model = max(foundation_models, key=lambda x: x.version)\n",
    "print(\n",
    "    f\"\\n\\nUsing model name: {foundation_model.name}, version: {foundation_model.version}, id: {foundation_model.id} for fine tuning\"\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Prepare the dataset for fine-tuning the model\n",
    "\n",
    "We will use the [fridgeObjects](https://cvbp-secondary.z19.web.core.windows.net/datasets/image_classification/fridgeObjects.zip) dataset. The fridge object dataset is stored in a directory. There are four different folders inside:\n",
    "- /water_bottle\n",
    "- /milk_bottle\n",
    "- /carton\n",
    "- /can\n",
    "\n",
    "This is the most common data format for multiclass image classification. Each folder's title corresponds to the image label for the images contained inside. \n",
    "\n",
    "#### 4.1 Download the Data\n",
    "We first download and unzip the data locally. By default, the data would be downloaded in `./data` folder in current directory. \n",
    "If you prefer to download the data at a different location, update it in `dataset_parent_dir = ...` in the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import urllib\n",
    "from zipfile import ZipFile\n",
    "\n",
    "# Change to a different location if you prefer\n",
    "dataset_parent_dir = \"./data\"\n",
    "\n",
    "# Create data folder if it doesnt exist.\n",
    "os.makedirs(dataset_parent_dir, exist_ok=True)\n",
    "\n",
    "# Download data\n",
    "download_url = \"https://cvbp-secondary.z19.web.core.windows.net/datasets/image_classification/fridgeObjects.zip\"\n",
    "\n",
    "# Extract current dataset name from dataset url\n",
    "dataset_name = os.path.split(download_url)[-1].split(\".\")[0]\n",
    "# Get dataset path for later use\n",
    "dataset_dir = os.path.join(dataset_parent_dir, dataset_name)\n",
    "\n",
    "# Get the data zip file path\n",
    "data_file = os.path.join(dataset_parent_dir, f\"{dataset_name}.zip\")\n",
    "\n",
    "# Download the dataset\n",
    "urllib.request.urlretrieve(download_url, filename=data_file)\n",
    "\n",
    "# Extract files\n",
    "with ZipFile(data_file, \"r\") as zip:\n",
    "    print(\"extracting files...\")\n",
    "    zip.extractall(path=dataset_parent_dir)\n",
    "    print(\"done\")\n",
    "# Delete zip file\n",
    "os.remove(data_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image\n",
    "\n",
    "sample_image = os.path.join(dataset_dir, \"milk_bottle\", \"99.jpg\")\n",
    "Image(filename=sample_image)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.2 Upload the images to Datastore through an AML Data asset (URI Folder)\n",
    "\n",
    "In order to use the data for training in Azure ML, we upload it to our default Azure Blob Storage of our  Azure ML Workspace."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uploading image files by creating a 'data asset URI FOLDER':\n",
    "\n",
    "from azure.ai.ml.entities import Data\n",
    "from azure.ai.ml.constants import AssetTypes\n",
    "\n",
    "my_data = Data(\n",
    "    path=dataset_dir,\n",
    "    type=AssetTypes.URI_FOLDER,\n",
    "    description=\"Fridge-items images\",\n",
    "    name=\"fridge-items-images-multiclass-classif\",\n",
    ")\n",
    "\n",
    "uri_folder_data_asset = workspace_ml_client.data.create_or_update(my_data)\n",
    "\n",
    "print(uri_folder_data_asset)\n",
    "print(\"\")\n",
    "print(\"Path to folder in Blob Storage:\")\n",
    "print(uri_folder_data_asset.path)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.3 Convert the downloaded data to JSONL\n",
    "\n",
    "For documentation on preparing the datasets beyond this notebook, refer to the [documentation on how to prepare datasets](https://docs.microsoft.com/en-us/azure/machine-learning/how-to-prepare-datasets-for-automl-images).\n",
    "\n",
    "In order to use this data to create an AzureML MLTable, we first need to convert it to the required JSONL format. The following script is creating two `.jsonl` files (one for training and one for validation) in the corresponding MLTable folder. In this example, 20% of the data is kept for validation. For further details on jsonl file used for image classification task in automated ml, please refer to the [data schema documentation for multi-class image classification task](https://learn.microsoft.com/en-us/azure/machine-learning/reference-automl-images-schema#image-classification-binarymulti-class)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "\n",
    "# We will copy each JSONL file within its related MLTable folder\n",
    "training_mltable_path = os.path.join(dataset_parent_dir, \"training-mltable-folder\")\n",
    "validation_mltable_path = os.path.join(dataset_parent_dir, \"validation-mltable-folder\")\n",
    "\n",
    "# Create the folders if they don't exist\n",
    "os.makedirs(training_mltable_path, exist_ok=True)\n",
    "os.makedirs(validation_mltable_path, exist_ok=True)\n",
    "\n",
    "train_validation_ratio = 5\n",
    "\n",
    "# Path to the training and validation files\n",
    "train_annotations_file = os.path.join(training_mltable_path, \"train_annotations.jsonl\")\n",
    "validation_annotations_file = os.path.join(\n",
    "    validation_mltable_path, \"validation_annotations.jsonl\"\n",
    ")\n",
    "\n",
    "# Baseline of json line dictionary\n",
    "json_line_sample = {\n",
    "    \"image_url\": uri_folder_data_asset.path,\n",
    "    \"label\": \"\",\n",
    "}\n",
    "\n",
    "index = 0\n",
    "# Scan each sub directary and generate a jsonl line per image, distributed on train and valid JSONL files\n",
    "with open(train_annotations_file, \"w\") as train_f:\n",
    "    with open(validation_annotations_file, \"w\") as validation_f:\n",
    "        for class_name in os.listdir(dataset_dir):\n",
    "            sub_dir = os.path.join(dataset_dir, class_name)\n",
    "            if not os.path.isdir(sub_dir):\n",
    "                continue\n",
    "            # Scan each sub directary\n",
    "            print(f\"Parsing {sub_dir}\")\n",
    "            for image in os.listdir(sub_dir):\n",
    "                json_line = dict(json_line_sample)\n",
    "                json_line[\"image_url\"] += f\"{class_name}/{image}\"\n",
    "                json_line[\"label\"] = class_name\n",
    "\n",
    "                if index % train_validation_ratio == 0:\n",
    "                    # Validation annotation\n",
    "                    validation_f.write(json.dumps(json_line) + \"\\n\")\n",
    "                else:\n",
    "                    # Train annotation\n",
    "                    train_f.write(json.dumps(json_line) + \"\\n\")\n",
    "                index += 1"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.4 Create MLTable data input\n",
    "\n",
    "Create MLTable data input using the jsonl files created above.\n",
    "\n",
    "For documentation on creating your own MLTable assets for jobs beyond this notebook, please refer to below resources\n",
    "- [MLTable YAML Schema](https://learn.microsoft.com/en-us/azure/machine-learning/reference-yaml-mltable) - covers how to write MLTable YAML, which is required for each MLTable asset.\n",
    "- [Create MLTable data asset](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-create-data-assets?tabs=Python-SDK#create-a-mltable-data-asset) - covers how to create MLTable data asset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_ml_table_file(filename):\n",
    "    \"\"\"Create ML Table definition\"\"\"\n",
    "\n",
    "    return (\n",
    "        \"paths:\\n\"\n",
    "        \"  - file: ./{0}\\n\"\n",
    "        \"transformations:\\n\"\n",
    "        \"  - read_json_lines:\\n\"\n",
    "        \"        encoding: utf8\\n\"\n",
    "        \"        invalid_lines: error\\n\"\n",
    "        \"        include_path_column: false\\n\"\n",
    "        \"  - convert_column_types:\\n\"\n",
    "        \"      - columns: image_url\\n\"\n",
    "        \"        column_type: stream_info\"\n",
    "    ).format(filename)\n",
    "\n",
    "\n",
    "def save_ml_table_file(output_path, mltable_file_contents):\n",
    "    with open(os.path.join(output_path, \"MLTable\"), \"w\") as f:\n",
    "        f.write(mltable_file_contents)\n",
    "\n",
    "\n",
    "# Create and save train mltable\n",
    "train_mltable_file_contents = create_ml_table_file(\n",
    "    os.path.basename(train_annotations_file)\n",
    ")\n",
    "save_ml_table_file(training_mltable_path, train_mltable_file_contents)\n",
    "\n",
    "# Create and save validation mltable\n",
    "validation_mltable_file_contents = create_ml_table_file(\n",
    "    os.path.basename(validation_annotations_file)\n",
    ")\n",
    "save_ml_table_file(validation_mltable_path, validation_mltable_file_contents)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Submit the fine tuning job using `transformers_image_classification_pipeline` component\n",
    " \n",
    "Create the job that uses the `transformers_image_classification_pipeline` component for multi-class image-classification task. Learn more in 5.2 about all the parameters supported for fine tuning."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 5.1 Create component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FINETUNE_PIPELINE_COMPONENT_NAME = \"transformers_image_classification_pipeline\"\n",
    "pipeline_component_transformers_func = registry_ml_client.components.get(\n",
    "    name=FINETUNE_PIPELINE_COMPONENT_NAME, label=\"latest\"\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 5.2 Create arguments to be passed to `transformers_image_classification_pipeline` component\n",
    "\n",
    "The `transformers_image_classification_pipeline` component consists of model selection and finetuning components. The detailed arguments for each component can be found at following README files:\n",
    "- [Model Import Component](../../../docs/component_docs/image_finetune/transformers_model_import_component.md)\n",
    "- [Finetune Component](../../../docs/component_docs/image_finetune/transformers_finetune_component.md)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deepspeed_config_path = \"./deepspeed_configs/zero1.json\"\n",
    "if not os.path.exists(deepspeed_config_path):\n",
    "    print(\"DeepSpeed config file not found\")\n",
    "    deepspeed_config_path = None\n",
    "\n",
    "pipeline_component_args = {\n",
    "    # # Model import args\n",
    "    \"model_family\": \"HuggingFaceImage\",\n",
    "    \"download_from_source\": False,  # True for downloading a model directly from HuggingFace\n",
    "    \"mlflow_model\": foundation_model.id,  # foundation_model.id is provided, only foundation_model gives UserErrorException: only path input is supported now but get: ...\n",
    "    # \"model_name\": huggingface_model_name, # specify the model_name instead of mlflow_model if you want to use a model from the huggingface hub\n",
    "    # # Finetune_args\n",
    "    \"image_width\": -1,  # Default value is -1 which means it would be overwritten by default image width in Hugging Face feature extractor\n",
    "    \"image_height\": -1,  # Default value is -1 which means it would be overwritten by default image height in Hugging Face feature extractor\n",
    "    \"task_name\": \"image-classification\",\n",
    "    \"apply_augmentations\": True,\n",
    "    \"number_of_workers\": 8,\n",
    "    \"deepspeed_config\": deepspeed_config_path,\n",
    "    \"auto_find_batch_size\": False,\n",
    "    \"precision\": \"32\",\n",
    "    \"random_seed\": 42,\n",
    "    \"evaluation_strategy\": \"epoch\",\n",
    "    \"evaluation_steps\": 500,\n",
    "    \"logging_strategy\": \"epoch\",\n",
    "    \"logging_steps\": 500,\n",
    "    \"save_strategy\": \"epoch\",\n",
    "    \"save_steps\": 500,\n",
    "    \"save_total_limit\": -1,\n",
    "    \"early_stopping\": False,\n",
    "    \"early_stopping_patience\": 1,\n",
    "    \"resume_from_checkpoint\": False,\n",
    "    \"save_as_mlflow_model\": True,\n",
    "    # # Uncomment one or more lines below to provide specific values, if you wish you override the autoselected default values.\n",
    "    # \"metric_for_best_model\": \"accuracy\",\n",
    "    # \"number_of_epochs\": 15,\n",
    "    # \"max_steps\": -1,\n",
    "    # \"training_batch_size\": 4,\n",
    "    # \"validation_batch_size\": 4,\n",
    "    # \"learning_rate\": 5e-5,\n",
    "    # \"learning_rate_scheduler\": \"warmup_linear\",\n",
    "    # \"warmup_steps\": 0,\n",
    "    # \"optimizer\": \"adamw_hf\",\n",
    "    # \"weight_decay\": 0.0,\n",
    "    # \"gradient_accumulation_step\": 1,\n",
    "    # \"label_smoothing_factor\": 0.0,\n",
    "    # \"max_grad_norm\": 1.0,\n",
    "    # apply_ort: True,\n",
    "    # apply_deepspeed: True\n",
    "    # # Model evaluation args\n",
    "    # The following parameters map to the dataset fields\n",
    "    # Uncomment one or more lines below to provide specific values, if you wish you override the autoselected default values.\n",
    "    # \"label_column_name\": \"label\",\n",
    "    # \"input_column_names\": \"image_url\",\n",
    "}\n",
    "process_count_per_instance = 1\n",
    "instance_count = 1\n",
    "\n",
    "# Ensure that the user provides only one of mlflow_model or model_name\n",
    "if (\n",
    "    pipeline_component_args.get(\"mlflow_model\") is None\n",
    "    and pipeline_component_args.get(\"model_name\") is None\n",
    "):\n",
    "    raise ValueError(\n",
    "        \"You must specify either mlflow_model or model_name for the model to finetune\"\n",
    "    )\n",
    "if (\n",
    "    pipeline_component_args.get(\"mlflow_model\") is not None\n",
    "    and pipeline_component_args.get(\"model_name\") is not None\n",
    "):\n",
    "    raise ValueError(\n",
    "        \"You must specify ONLY one of mlflow_model and model_name for the model to finetune\"\n",
    "    )\n",
    "elif (\n",
    "    pipeline_component_args.get(\"mlflow_model\") is None\n",
    "    and pipeline_component_args.get(\"model_name\") is not None\n",
    "):\n",
    "    use_model_name = huggingface_model_name\n",
    "elif (\n",
    "    pipeline_component_args.get(\"mlflow_model\") is not None\n",
    "    and pipeline_component_args.get(\"model_name\") is None\n",
    "):\n",
    "    use_model_name = aml_registry_model_name\n",
    "print(f\"Finetuning model {use_model_name}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 5.3 Utility function to create pipeline using `transformers_image_classification_pipeline` component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.dsl import pipeline\n",
    "from azure.ai.ml.entities import PipelineComponent\n",
    "from azure.ai.ml import Input\n",
    "from azure.ai.ml.constants import AssetTypes\n",
    "\n",
    "\n",
    "@pipeline()\n",
    "def create_pipeline_transformers():\n",
    "    \"\"\"Create pipeline.\"\"\"\n",
    "\n",
    "    transformers_pipeline_component: PipelineComponent = pipeline_component_transformers_func(\n",
    "        compute_model_import=model_import_cluster_name,\n",
    "        compute_finetune=finetune_cluster_name,\n",
    "        compute_model_evaluation=model_eval_cluster_name,\n",
    "        training_data=Input(type=AssetTypes.MLTABLE, path=training_mltable_path),\n",
    "        validation_data=Input(type=AssetTypes.MLTABLE, path=validation_mltable_path),\n",
    "        # test data\n",
    "        # Using the same data for validation and test. If you want to use a different dataset for test, specify it below\n",
    "        test_data=Input(type=AssetTypes.MLTABLE, path=validation_mltable_path),\n",
    "        instance_count=instance_count,\n",
    "        process_count_per_instance=process_count_per_instance,\n",
    "        **pipeline_component_args,\n",
    "    )\n",
    "    return {\n",
    "        # Map the output of the fine tuning job to the output of pipeline job so that we can easily register the fine tuned model. Registering the model is required to deploy the model to an online or batch endpoint.\n",
    "        \"trained_model\": transformers_pipeline_component.outputs.mlflow_model_folder,\n",
    "    }"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 5.4 Run the fine tuning job using `transformers_image_classification_pipeline` component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformers_pipeline_object = create_pipeline_transformers()\n",
    "\n",
    "# don't use cached results from previous jobs\n",
    "transformers_pipeline_object.settings.force_rerun = True\n",
    "\n",
    "# set continue on step failure to False\n",
    "transformers_pipeline_object.settings.continue_on_step_failure = False\n",
    "\n",
    "transformers_pipeline_object.display_name = (\n",
    "    use_model_name + \"_transformers_pipeline_component_run_\" + \"multiclass\"\n",
    ")\n",
    "# Don't use cached results from previous jobs\n",
    "transformers_pipeline_object.settings.force_rerun = True\n",
    "\n",
    "print(\"Submitting pipeline\")\n",
    "\n",
    "transformers_pipeline_run = workspace_ml_client.jobs.create_or_update(\n",
    "    transformers_pipeline_object, experiment_name=experiment_name\n",
    ")\n",
    "\n",
    "print(f\"Pipeline created. URL: {transformers_pipeline_run.studio_url}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workspace_ml_client.jobs.stream(transformers_pipeline_run.name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Get metrics from finetune component\n",
    "\n",
    "The model training happens as part of the finetune component. Please follow below steps to extract validation metrics from the run."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 6.1 Initialize MLFlow Client\n",
    "\n",
    "The models and artifacts that are produced by AutoML can be accessed via the MLFlow interface.\n",
    "Initialize the MLFlow client here, and set the backend as Azure ML, via. the MLFlow Client.\n",
    "\n",
    "IMPORTANT - You need to have installed the latest MLFlow packages with:\n",
    "\n",
    "    pip install azureml-mlflow\n",
    "    pip install mlflow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mlflow\n",
    "\n",
    "# Obtain the tracking URL from MLClient\n",
    "MLFLOW_TRACKING_URI = workspace_ml_client.workspaces.get(\n",
    "    name=workspace_ml_client.workspace_name\n",
    ").mlflow_tracking_uri\n",
    "\n",
    "print(MLFLOW_TRACKING_URI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the MLFLOW TRACKING URI\n",
    "mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)\n",
    "print(f\"\\nCurrent tracking uri: {mlflow.get_tracking_uri()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlflow.tracking.client import MlflowClient\n",
    "\n",
    "# Initialize MLFlow client\n",
    "mlflow_client = MlflowClient()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6.2 Get the training and evaluation run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Concat 'tags.mlflow.rootRunId=' and pipeline_job.name in single quotes as filter variable\n",
    "filter = \"tags.mlflow.rootRunId='\" + transformers_pipeline_run.name + \"'\"\n",
    "runs = mlflow.search_runs(\n",
    "    experiment_names=[experiment_name], filter_string=filter, output_format=\"list\"\n",
    ")\n",
    "\n",
    "# Get the training and evaluation runs.\n",
    "# Using a hacky way till 'Bug 2320997: not able to show eval metrics in FT notebooks - mlflow client now showing display names' is fixed\n",
    "for run in runs:\n",
    "    # Check if run.data.metrics.epoch exists\n",
    "    if \"epoch\" in run.data.metrics:\n",
    "        training_run = run\n",
    "    # Else, check if run.data.metrics.accuracy exists\n",
    "    elif \"accuracy\" in run.data.metrics:\n",
    "        evaluation_run = run"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 6.3 Get training metrics\n",
    "\n",
    "Access the results (such as Models, Artifacts, Metrics) of a previously completed run."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "pd.DataFrame(training_run.data.metrics, index=[0]).T"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7. Register the fine tuned model with the workspace\n",
    "\n",
    "We will register the model from the output of the fine tuning job. This will track lineage between the fine tuned model and the fine tuning job. The fine tuning job, further, tracks lineage to the foundation model, data and training code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "# Generating a unique timestamp that can be used for names and versions that need to be unique\n",
    "timestamp = str(int(time.time()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import Model\n",
    "from azure.ai.ml.constants import AssetTypes\n",
    "\n",
    "# Check if the `trained_model` output is available\n",
    "print(\n",
    "    f\"Pipeline job outputs: {workspace_ml_client.jobs.get(transformers_pipeline_run.name).outputs}\"\n",
    ")\n",
    "\n",
    "# Fetch the model from pipeline job output - not working, hence fetching from fine tune child job\n",
    "model_path_from_job = (\n",
    "    f\"azureml://jobs/{transformers_pipeline_run.name}/outputs/trained_model\"\n",
    ")\n",
    "print(f\"Path to register model: {model_path_from_job}\")\n",
    "\n",
    "finetuned_model_name = (\n",
    "    f\"{use_model_name.replace('/', '-')}-fridge-objects-multiclass-classification\"\n",
    ")\n",
    "finetuned_model_description = f\"{use_model_name.replace('/', '-')} fine tuned model for fridge objects multiclass classification\"\n",
    "prepare_to_register_model = Model(\n",
    "    path=model_path_from_job,\n",
    "    type=AssetTypes.MLFLOW_MODEL,\n",
    "    name=finetuned_model_name,\n",
    "    version=timestamp,  # use timestamp as version to avoid version conflict\n",
    "    description=finetuned_model_description,\n",
    ")\n",
    "print(f\"Prepare to register model: \\n{prepare_to_register_model}\")\n",
    "\n",
    "# Register the model from pipeline job output\n",
    "registered_model = workspace_ml_client.models.create_or_update(\n",
    "    prepare_to_register_model\n",
    ")\n",
    "print(f\"Registered model: {registered_model}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 8. Deploy the fine tuned model to an online endpoint\n",
    "Online endpoints give a durable REST API that can be used to integrate with applications that need to use the model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "from azure.ai.ml.entities import ManagedOnlineEndpoint, ManagedOnlineDeployment\n",
    "\n",
    "# Endpoint names need to be unique in a region, hence using timestamp to create unique endpoint name\n",
    "online_endpoint_name = \"hf-mc-fridge-items-\" + datetime.datetime.now().strftime(\n",
    "    \"%m%d%H%M\"\n",
    ")\n",
    "online_endpoint_description = f\"Online endpoint for {registered_model.name}, finetuned for fridge objects multiclass classification\"\n",
    "# Create an online endpoint\n",
    "endpoint = ManagedOnlineEndpoint(\n",
    "    name=online_endpoint_name,\n",
    "    description=online_endpoint_description,\n",
    "    auth_mode=\"key\",\n",
    "    tags={\"foo\": \"bar\"},\n",
    ")\n",
    "workspace_ml_client.begin_create_or_update(endpoint).result()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import OnlineRequestSettings, ProbeSettings\n",
    "\n",
    "deployment_name = \"hf-mc-fridge-mlflow-deploy\"\n",
    "print(registered_model.id)\n",
    "print(online_endpoint_name)\n",
    "print(deployment_name)\n",
    "\n",
    "# Create a deployment\n",
    "demo_deployment = ManagedOnlineDeployment(\n",
    "    name=deployment_name,\n",
    "    endpoint_name=online_endpoint_name,\n",
    "    model=registered_model.id,\n",
    "    instance_type=\"Standard_DS3_V2\",  # Use GPU instance type like STANDARD_NC6s_v3 for faster explanations\n",
    "    instance_count=1,\n",
    "    request_settings=OnlineRequestSettings(\n",
    "        max_concurrent_requests_per_instance=1,\n",
    "        request_timeout_ms=90000,\n",
    "        max_queue_wait_ms=500,\n",
    "    ),\n",
    "    liveness_probe=ProbeSettings(\n",
    "        failure_threshold=49,\n",
    "        success_threshold=1,\n",
    "        timeout=299,\n",
    "        period=180,\n",
    "        initial_delay=180,\n",
    "    ),\n",
    "    readiness_probe=ProbeSettings(\n",
    "        failure_threshold=10,\n",
    "        success_threshold=1,\n",
    "        timeout=10,\n",
    "        period=10,\n",
    "        initial_delay=2000,\n",
    "    ),\n",
    ")\n",
    "workspace_ml_client.online_deployments.begin_create_or_update(demo_deployment).wait()\n",
    "endpoint.traffic = {deployment_name: 100}\n",
    "workspace_ml_client.begin_create_or_update(endpoint).result()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 9. Test the endpoint with sample data\n",
    "\n",
    "We will fetch some sample data from the test dataset and submit to online endpoint for inference. We will then show the display the scored labels alongside the ground truth labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "demo_deployment = workspace_ml_client.online_deployments.get(\n",
    "    name=deployment_name,\n",
    "    endpoint_name=online_endpoint_name,\n",
    ")\n",
    "\n",
    "# Get the details for online endpoint\n",
    "endpoint = workspace_ml_client.online_endpoints.get(name=online_endpoint_name)\n",
    "\n",
    "# Existing traffic details\n",
    "print(endpoint.traffic)\n",
    "\n",
    "# Get the scoring URI\n",
    "print(endpoint.scoring_uri)\n",
    "print(demo_deployment)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create request json\n",
    "import base64\n",
    "import json\n",
    "\n",
    "sample_image = os.path.join(dataset_dir, \"milk_bottle\", \"99.jpg\")\n",
    "\n",
    "\n",
    "def read_image(image_path):\n",
    "    with open(image_path, \"rb\") as f:\n",
    "        return f.read()\n",
    "\n",
    "\n",
    "request_json = {\n",
    "    \"input_data\": {\n",
    "        \"columns\": [\"image\"],\n",
    "        \"data\": [base64.encodebytes(read_image(sample_image)).decode(\"utf-8\")],\n",
    "    }\n",
    "}\n",
    "\n",
    "request_file_name = \"sample_request_data.json\"\n",
    "with open(request_file_name, \"w\") as request_file:\n",
    "    json.dump(request_json, request_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = workspace_ml_client.online_endpoints.invoke(\n",
    "    endpoint_name=online_endpoint_name,\n",
    "    deployment_name=demo_deployment.name,\n",
    "    request_file=request_file_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Visualize detections\n",
    "Now that we have scored a test image, we can visualize the prediction for this image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install matplotlib\n",
    "! pip install seaborn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import json\n",
    "import numpy as np\n",
    "import seaborn as sns\n",
    "from PIL import Image\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.image as mpimg\n",
    "\n",
    "img_np = mpimg.imread(sample_image)\n",
    "img = Image.fromarray(img_np.astype(\"uint8\"), \"RGB\")\n",
    "x, y = img.size\n",
    "\n",
    "# Set a compact figure size\n",
    "fig_width = 12\n",
    "fig_height = 12\n",
    "\n",
    "# Initialize figure and axes\n",
    "fig = plt.figure(figsize=(fig_width, fig_height))\n",
    "gs = fig.add_gridspec(2, 1, height_ratios=[16, 1], hspace=0.2)\n",
    "ax1 = fig.add_subplot(gs[0])\n",
    "ax2 = fig.add_subplot(gs[1])\n",
    "ax1.imshow(img_np)\n",
    "ax1.axis(\"off\")\n",
    "\n",
    "prediction = json.loads(resp)[0]\n",
    "label_index = np.argmax(prediction[\"probs\"])\n",
    "label = prediction[\"labels\"][label_index]\n",
    "conf_score = prediction[\"probs\"][label_index]\n",
    "display_text = f\"{label} ({round(conf_score, 3)})\"\n",
    "print(display_text)\n",
    "\n",
    "color = np.random.rand(3)\n",
    "\n",
    "# Set a stylish color palette\n",
    "sns.set_palette(\"pastel\")\n",
    "\n",
    "# Create the bar plot without x-axis and y-axis markings\n",
    "barplot = sns.barplot(x=[conf_score], y=[label], palette=[color], ax=ax2)\n",
    "ax2.set_xlabel(\"\")  # Remove x-axis label\n",
    "ax2.set_ylabel(\"\")  # Remove y-axis label\n",
    "ax2.set_title(f\"Top Object Scores\", fontsize=12)\n",
    "\n",
    "# Add scores in front of the bars\n",
    "barplot.text(\n",
    "    conf_score + 0.001, 0, f\"{conf_score:.2f}\", va=\"center\", color=\"black\", fontsize=10\n",
    ")\n",
    "\n",
    "# Remove spines and ticks from the bar plot\n",
    "barplot.spines[\"left\"].set_visible(False)\n",
    "barplot.spines[\"top\"].set_visible(False)\n",
    "barplot.spines[\"right\"].set_visible(False)\n",
    "barplot.spines[\"bottom\"].set_visible(False)\n",
    "barplot.tick_params(left=False, top=False, right=False, bottom=False)\n",
    "barplot.xaxis.set_visible(False)  # Remove x-axis\n",
    "barplot.yaxis.grid(False)  # Remove y-axis grid\n",
    "\n",
    "# Customize bar plot lines for a fading effect\n",
    "for container in barplot.containers:\n",
    "    for bar in container:\n",
    "        bar.set_linewidth(0.5)\n",
    "        bar.set_edgecolor(\"gray\")\n",
    "        bar.set_alpha(0.7)\n",
    "\n",
    "# Set plot background color\n",
    "fig.patch.set_facecolor(\"#F7F7F7\")  # Light gray\n",
    "\n",
    "plt.tight_layout()\n",
    "# fig.savefig(\"plot.png\", bbox_inches=\"tight\")\n",
    "plt.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 10. Clean up resources - delete the online endpoint\n",
    "Don't forget to delete the online endpoint, else you will leave the billing meter running for the compute used by the endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workspace_ml_client.online_endpoints.begin_delete(name=online_endpoint_name).wait()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}