sdk/python/foundation-models/system/inference/text-to-image/safe-image-text-to-image-batch-deployment.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create an Azure AI Content Safety (AACS) enabled image to image batch endpoint (Preview)\n",
    "### This notebook is under preview.\n",
    "\n",
    "### Steps to create an __AACS__ enabled __image to image__ batch endpoint\n",
    "1. Create a __AACS__ enabled image-to-image batch endpoint with a custom [score_batch.py](./aacs-scoring-files/score/score_batch.py) script. This will integrate the batch endpoint with the AACS resource to moderate the response from the __image-to-image__ model and the request from the user.\n",
    "2. Create a new __AACS__ enabled __image-to-image__ batch endpoint with a custom [score_batch.py](./aacs-scoring-files/score/score_batch.py) which will integrate with the __AACS__ resource to moderate the response from the __image-to-image__ model and the request from the user. To make the custom [score_batch.py](./aacs-scoring-files/score/score_batch.py) to successfully authenticated to the __AACS__ resource, use __Environment variable__ to pass the access key of the __AACS__ resource to the custom [score_batch.py](./aacs-scoring-files/score/score_batch.py) via environment variable. The custom [score_batch.py](./aacs-scoring-files/score/score_batch.py) can use the key directly to access the AACS resource. This option is less secure, if someone in your org has access to the endpoint, he/she can get the access key from the environment variable and use it to access the AACS resource.\n",
    "\n",
    "### Task\n",
    "`image to image` task takes an original image, a text prompt as input. The model generates an image by modifying the original image.\n",
    "\n",
    " \n",
    "### Model\n",
    "Models that can perform the `image to image` task are tagged with `image-to-image`. We will use the `stabilityai-stable-diffusion-xl-refiner-1-0` model in this notebook. If you opened this notebook from a specific model card, remember to replace the specific model name.\n",
    "\n",
    "### Outline\n",
    "1. Setup pre-requisties\n",
    "2. Create AACS resource\n",
    "3. Create AACS enabled image to image batch endpoint\n",
    "4. Prepare data for inference - using a folder of csv files with prompt, and image columns\n",
    "5. Test the endpoint - using csv files\n",
    "6. Clean up resources - delete the endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Setup pre-requisites\n",
    "* Check List\n",
    "* Install dependencies\n",
    "* Connect to AzureML Workspace. Learn more at [set up SDK authentication](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-setup-authentication?tabs=sdk). Replace  `<WORKSPACE_NAME>`, `<RESOURCE_GROUP>` and `<SUBSCRIPTION_ID>` below.\n",
    "* Connect to `azureml` system registry"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> [x] The identity you are using to execute this notebook(yourself or your VM) need to have the __Contributor__ role on the resource group where the AML Workspace your specified is located, because this notebook will create an AACS resource using that identity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install the required packages\n",
    "%pip install azure-identity==1.13.0\n",
    "%pip install azure-mgmt-cognitiveservices==13.4.0\n",
    "%pip install azure-ai-ml==1.8.0\n",
    "%pip install azure-mgmt-msi==7.0.0\n",
    "%pip install azure-mgmt-authorization==3.0.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient\n",
    "from azure.identity import DefaultAzureCredential, InteractiveBrowserCredential\n",
    "\n",
    "try:\n",
    "    credential = DefaultAzureCredential()\n",
    "    credential.get_token(\"https://management.azure.com/.default\")\n",
    "except Exception as ex:\n",
    "    credential = InteractiveBrowserCredential()\n",
    "\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",
    "workspace_ml_client = MLClient(\n",
    "    credential, subscription_id, resource_group, workspace_name\n",
    ")\n",
    "\n",
    "print(f\"Connected to workspace {workspace_name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The models, fine tuning pipelines and environments are available in the AzureML system registry, \"azureml\"\n",
    "\n",
    "registry_name = \"azureml\"\n",
    "\n",
    "registry_ml_client = MLClient(\n",
    "    credential,\n",
    "    subscription_id,\n",
    "    resource_group,\n",
    "    registry_name=registry_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Create AACS resource"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.1 Assign variables for Azure Content Safety\n",
    "Currently, AACS is available in a limited set of regions:\n",
    "\n",
    "__NOTE__: before you choose the region to deploy the AACS, please be aware that your data will be transferred to the region you choose and by selecting a region outside your current location, you may be allowing the transmission of your data to regions outside your jurisdiction. It is important to note that data protection and privacy laws may vary between jurisdictions. Before proceeding, we strongly advise you to familiarize yourself with the local laws and regulations governing data transfer and ensure that you are legally permitted to transmit your data to an overseas location for processing. By continuing with the selection of a different region, you acknowledge that you have understood and accepted any potential risks associated with such data transmission. Please proceed with caution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The severity level that will trigger response be blocked\n",
    "# Please reference Azure AI content documentation for more details\n",
    "# https://learn.microsoft.com/en-us/azure/cognitive-services/content-safety/concepts/harm-categories\n",
    "content_severity_threshold = \"2\"\n",
    "\n",
    "# If you choose environment variables for authentication of AACS resource, then assign empty (\"\") value to uai_name\n",
    "uai_name = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from uuid import uuid4\n",
    "from azure.mgmt.cognitiveservices import CognitiveServicesManagementClient\n",
    "\n",
    "aacs_client = CognitiveServicesManagementClient(credential, subscription_id)\n",
    "\n",
    "\n",
    "# settings for the Azure AI Content Safety (AACS) resource\n",
    "# we will choose existing AACS resource if it exists, otherwise create a new one\n",
    "# name of AACS resource, has to be unique\n",
    "\n",
    "aacs_name = f\"aacs-image-to-image-{str(uuid4())[:8]}\"\n",
    "available_aacs_locations = [\"east us\", \"west europe\"]\n",
    "\n",
    "# create a new Cognitive Services Account\n",
    "kind = \"ContentSafety\"\n",
    "aacs_sku_name = \"S0\"\n",
    "aacs_location = available_aacs_locations[0]\n",
    "\n",
    "\n",
    "print(\"Available SKUs:\")\n",
    "aacs_skus = aacs_client.resource_skus.list()\n",
    "print(\"SKU Name\\tSKU Tier\\tLocations\")\n",
    "for sku in aacs_skus:\n",
    "    if sku.kind == \"ContentSafety\":\n",
    "        locations = \",\".join(sku.locations)\n",
    "        print(sku.name + \"\\t\\t\" + sku.tier + \"\\t\\t\" + locations)\n",
    "\n",
    "print(f\"Choose a new AACS resource in {aacs_location} with SKU {aacs_sku_name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 2.2 Create AACS Resource"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.mgmt.cognitiveservices.models import Account, Sku, AccountProperties\n",
    "\n",
    "parameters = Account(\n",
    "    sku=Sku(name=aacs_sku_name),\n",
    "    kind=kind,\n",
    "    location=aacs_location,\n",
    "    properties=AccountProperties(\n",
    "        custom_sub_domain_name=aacs_name, public_network_access=\"Enabled\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "\n",
    "def find_acs(accounts):\n",
    "    return next(\n",
    "        x\n",
    "        for x in accounts\n",
    "        if x.kind == \"ContentSafety\"\n",
    "        and x.location == aacs_location\n",
    "        and x.sku.name == aacs_sku_name\n",
    "    )\n",
    "\n",
    "\n",
    "try:\n",
    "    # check if AACS exists\n",
    "    aacs = aacs_client.accounts.get(resource_group, aacs_name)\n",
    "    print(f\"Found existing AACS Account {aacs.name}.\")\n",
    "except:\n",
    "    try:\n",
    "        # check if there is an existing AACS resource within same resource group\n",
    "        aacs = find_acs(aacs_client.accounts.list_by_resource_group(resource_group))\n",
    "        print(\n",
    "            f\"Found existing AACS Account {aacs.name} in resource group {resource_group}.\"\n",
    "        )\n",
    "    except:\n",
    "        print(f\"Creating AACS Account {aacs_name}.\")\n",
    "        aacs_client.accounts.begin_create(resource_group, aacs_name, parameters).wait()\n",
    "        print(\"Resource created.\")\n",
    "        aacs = aacs_client.accounts.get(resource_group, aacs_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aacs_endpoint = aacs.properties.endpoint\n",
    "aacs_resource_id = aacs.id\n",
    "aacs_name = aacs.name\n",
    "print(\n",
    "    f\"AACS name is {aacs.name} .\\nUse this name in UAI preparation notebook to create UAI.\"\n",
    ")\n",
    "print(f\"AACS endpoint is {aacs_endpoint}\")\n",
    "print(f\"AACS ResourceId is {aacs_resource_id}\")\n",
    "\n",
    "aacs_access_key = aacs_client.accounts.list_keys(\n",
    "    resource_group_name=resource_group, account_name=aacs.name\n",
    ").key1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Create AACS enabled image to image batch endpoint\n",
    "\n",
    "#### 3.1 Check if image to image model is available in the AML registry\n",
    "\n",
    "Browse models in the Model Catalog in the AzureML Studio, filtering by the image-to-image task. In this example, we use the `stabilityai-stable-diffusion-xl-refiner-1-0 model`. If you have opened this notebook for a different model, replace the model name accordingly. This is a pre-trained model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Name of the image to image model to be deployed\n",
    "model_name = \"stabilityai-stable-diffusion-xl-refiner-1-0\"\n",
    "\n",
    "try:\n",
    "    model = registry_ml_client.models.get(model_name, label=\"latest\")\n",
    "    print(\n",
    "        f\"Using model name: {model.name}, version: {model.version}, id: {model.id} for inference.\"\n",
    "    )\n",
    "except:\n",
    "    raise Exception(\n",
    "        f\"No model named {model_name} found in registry. \"\n",
    "        \"Please check model name in Azure model catalog.\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 3.2 Create environment for image-to-image endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import Environment, BuildContext\n",
    "from IPython.core.display import display, HTML\n",
    "\n",
    "environment_name = \"imageto-image-model-env\"  # Replace with your environment name\n",
    "\n",
    "try:\n",
    "    env = workspace_ml_client.environments.get(environment_name, label=\"latest\")\n",
    "    print(\"---Environment already exists---\")\n",
    "except:\n",
    "    print(\"---Creating environment---\")\n",
    "    env = Environment(\n",
    "        name=environment_name,\n",
    "        build=BuildContext(path=\"./scoring-files/docker_env\"),\n",
    "    )\n",
    "    workspace_ml_client.environments.create_or_update(env)\n",
    "    env = workspace_ml_client.environments.get(environment_name, label=\"latest\")\n",
    "    print(\"---Please use link below to check build status---\")\n",
    "\n",
    "\n",
    "display(\n",
    "    HTML(\n",
    "        f\"\"\"\n",
    "             <a href=\"https://ml.azure.com/environments/{environment_name}/version/{env.version}?wsid=/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.MachineLearningServices/workspaces/{workspace_name}\">\n",
    "                Click here to check env build status in AML studio\n",
    "             </a>\n",
    "             \"\"\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 3.3 Create compute cluster to run batch job on\n",
    "\n",
    "Use the model card from the AzureML system registry to check the minimum required inferencing SKU, referenced as size below. If you already have a sufficient compute cluster that you wish to use, you can simply define the name in `compute_name` in the following code block. Otherwise, the below snippet will create a new compute cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import AmlCompute\n",
    "\n",
    "sku_name = \"STANDARD_NC6S_V3\"  # Name of the sku(compute instance type)\n",
    "compute_name = \"gpu-compute\"  # Replace with your compute name\n",
    "\n",
    "if not any(\n",
    "    filter(lambda m: m.name == compute_name, workspace_ml_client.compute.list())\n",
    "):\n",
    "    compute_cluster = AmlCompute(\n",
    "        name=compute_name,\n",
    "        size=sku_name,\n",
    "        min_instances=0,\n",
    "        max_instances=2,\n",
    "    )\n",
    "    workspace_ml_client.compute.begin_create_or_update(compute_cluster).result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 3.4 Deploy the model to a batch endpoint\n",
    "\n",
    "Batch endpoints are endpoints that are used to do batch inferencing on large volumes of data over a period of time. The endpoints receive pointers to data and run jobs asynchronously to process the data in parallel on compute clusters. Batch endpoints store outputs to a data store for further analysis. For more information on batch endpoints and deployments, see <a href=\"https://learn.microsoft.com/en-us/azure/machine-learning/concept-endpoints?view=azureml-api-2#what-are-batch-endpoints\" target=\"_blank\"> What are batch endpoints?</a> In this sub-section, we will cover the following items:\n",
    "\n",
    "* Create a batch endpoint.\n",
    "* Create a batch deployment.\n",
    "* Set the deployment as default. Doing so allows invoking the endpoint without specifying the deployment's name."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Create a batch endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import BatchEndpoint\n",
    "\n",
    "# Endpoint names need to be unique in a region,\n",
    "# hence using uuid (first 8 character) to create unique endpoint name\n",
    "\n",
    "endpoint_name = (\n",
    "    f\"safe-image-to-image-{str(uuid4())[:8]}\"  # Replace with your endpoint name\n",
    ")\n",
    "\n",
    "# Check if the endpoint already exists in the workspace\n",
    "try:\n",
    "    endpoint = workspace_ml_client.batch_endpoints.get(endpoint_name)\n",
    "    print(\"---Endpoint already exists---\")\n",
    "except:\n",
    "    # Create an batch endpoint if it doesn't exist\n",
    "\n",
    "    # Define the endpoint\n",
    "    endpoint = BatchEndpoint(name=endpoint_name, description=\"Test endpoint for model\")\n",
    "\n",
    "    # Trigger the endpoint creation\n",
    "    try:\n",
    "        workspace_ml_client.begin_create_or_update(endpoint).wait()\n",
    "        print(\"\\n---Endpoint created successfully---\\n\")\n",
    "    except Exception as err:\n",
    "        raise RuntimeError(\n",
    "            f\"Endpoint creation failed. Detailed Response:\\n{err}\"\n",
    "        ) from err"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Deploy image to image model\n",
    "This step may take a few minutes.\n",
    "\n",
    "\n",
    "__Note__: `mini_batch_size` is the number of CSV files processed by the model in a single mini_batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import (\n",
    "    ModelBatchDeploymentSettings,\n",
    "    CodeConfiguration,\n",
    "    BatchRetrySettings,\n",
    "    ModelBatchDeployment,\n",
    ")\n",
    "\n",
    "from azure.ai.ml.constants import BatchDeploymentOutputAction\n",
    "\n",
    "\n",
    "deployment_name = \"image-to-image-deploy\"\n",
    "\n",
    "deployment = ModelBatchDeployment(\n",
    "    name=deployment_name,\n",
    "    endpoint_name=endpoint.name,\n",
    "    model=model,\n",
    "    environment=env,\n",
    "    code_configuration=CodeConfiguration(\n",
    "        code=\"scoring-files/score\",\n",
    "        scoring_script=\"score_batch.py\",\n",
    "    ),\n",
    "    compute=compute_name,\n",
    "    settings=ModelBatchDeploymentSettings(\n",
    "        instance_count=1,\n",
    "        max_concurrency_per_instance=1,\n",
    "        mini_batch_size=2,\n",
    "        output_action=BatchDeploymentOutputAction.APPEND_ROW,\n",
    "        output_file_name=\"predictions.csv\",\n",
    "        retry_settings=BatchRetrySettings(max_retries=3, timeout=3000),\n",
    "        logging_level=\"info\",\n",
    "        environment_variables={\n",
    "            \"CONTENT_SAFETY_ENDPOINT\": aacs_endpoint,\n",
    "            \"CONTENT_SAFETY_KEY\": aacs_access_key,\n",
    "        },\n",
    "    ),\n",
    ")\n",
    "# Trigger the deployment creation\n",
    "try:\n",
    "    workspace_ml_client.begin_create_or_update(deployment).wait()\n",
    "    print(\"\\n---Deployment created successfully---\\n\")\n",
    "except Exception as err:\n",
    "    raise RuntimeError(\n",
    "        f\"Deployment creation failed. Detailed Response:\\n{err}\"\n",
    "    ) from err"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Update Batch endpoint to set the default deployment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = workspace_ml_client.batch_endpoints.get(endpoint_name)\n",
    "endpoint.defaults.deployment_name = deployment.name\n",
    "workspace_ml_client.batch_endpoints.begin_create_or_update(endpoint).result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Prepare data for inference - using a folder of csv files with prompt and image columns\n",
    "\n",
    "The CSV files should consist of 3 columns namely, \n",
    "* `image`: Original image to be used as base image for image to image task. It should either be in base64 format or publicly accessible URL.\n",
    "* `prompt`: A text prompt.\n",
    "\n",
    "We provide the text prompts in a csv file starting from the first row of a column named \"prompt\". We provide the input image as base64 string starting from the first row of the column \"image\". We provide the input image as base64 string starting from the first row of the column \"image\".\n",
    "\n",
    "The deployment in the \"Deploy image to image model\" section above takes the argument mini_batch_size, which is the number of CSV files processed by the model in a single mini_batch. To limit the number of prompts processed in each mini_batch we split the dataset into multiple csv files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prepare input csv file\n",
    "\n",
    "import base64\n",
    "import os\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "def read_image(image_path: str) -> bytes:\n",
    "    \"\"\"Reads an image from a file path into a byte array.\"\"\"\n",
    "    with open(image_path, \"rb\") as f:\n",
    "        return f.read()\n",
    "\n",
    "\n",
    "base_image1 = \"inpainting_data/images/dog_on_bench.png\"\n",
    "base_image2 = \"inpainting_data/images/teapot.png\"\n",
    "\n",
    "dataset_parent_dir = \"inpainting_data/batch_data\"\n",
    "os.makedirs(dataset_parent_dir, exist_ok=True)\n",
    "\n",
    "input_data = {\n",
    "    \"columns\": [\"image\", \"prompt\"],\n",
    "    \"data\": [\n",
    "        {\n",
    "            \"image\": base64.encodebytes(read_image(base_image1)).decode(\"utf-8\"),\n",
    "            \"prompt\": \"a dog with severed leg and bleeding profusely from deep laceration to the lower extremities, exposing tissues and nerve.\",\n",
    "        },\n",
    "        {\n",
    "            \"image\": base64.encodebytes(read_image(base_image2)).decode(\"utf-8\"),\n",
    "            \"prompt\": \"A small flower featuring a blend of pink and purple colors.\",\n",
    "        },\n",
    "    ],\n",
    "}\n",
    "pd.DataFrame(**input_data).to_csv(\n",
    "    os.path.join(dataset_parent_dir, \"input1.csv\"), index=False\n",
    ")\n",
    "\n",
    "input_data = {\n",
    "    \"columns\": [\"image\", \"prompt\"],\n",
    "    \"data\": [\n",
    "        {\n",
    "            \"image\": base64.encodebytes(read_image(base_image1)).decode(\"utf-8\"),\n",
    "            \"prompt\": \"Pikachu, cinematic, digital art, sitting on bench\",\n",
    "        },\n",
    "        {\n",
    "            \"image\": base64.encodebytes(read_image(base_image2)).decode(\"utf-8\"),\n",
    "            \"prompt\": \"dead body killed with a big dagger\",\n",
    "        },\n",
    "    ],\n",
    "}\n",
    "pd.DataFrame(**input_data).to_csv(\n",
    "    os.path.join(dataset_parent_dir, \"input2.csv\"), index=False\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read all the csvs in the data folder into a pandas dataframe\n",
    "import glob\n",
    "import os\n",
    "import pandas as pd\n",
    "\n",
    "# Specify the folder where your CSV files are located\n",
    "dataset_parent_dir = \"inpainting_data/batch_data\"\n",
    "\n",
    "# Use glob to get a list of CSV files in the folder\n",
    "csv_files = glob.glob(os.path.join(dataset_parent_dir, \"*.csv\"))\n",
    "\n",
    "# Read all CSV files into a single DataFrame using pd.concat\n",
    "batch_df = pd.concat((pd.read_csv(file) for file in csv_files), ignore_index=True)\n",
    "\n",
    "# Now, 'batch_df' contains all the data from the CSV files in the folder\n",
    "print(batch_df.head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "# Specify the folder where your CSV files should be saved\n",
    "processed_dataset_parent_dir = \"inpainting_data/processed_batch_data\"\n",
    "os.makedirs(processed_dataset_parent_dir, exist_ok=True)\n",
    "batch_input_file = \"batch_input.csv\"\n",
    "\n",
    "# Divide this into files of <x> rows each\n",
    "batch_size_per_predict = 2\n",
    "for i in range(0, len(batch_df), batch_size_per_predict):\n",
    "    j = i + batch_size_per_predict\n",
    "    batch_df[i:j].to_csv(\n",
    "        os.path.join(processed_dataset_parent_dir, str(i) + batch_input_file)\n",
    "    )\n",
    "\n",
    "# Check out the first and last file name created\n",
    "input_paths = sorted(Path(processed_dataset_parent_dir).iterdir(), key=os.path.getmtime)\n",
    "input_files = [os.path.basename(path) for path in input_paths]\n",
    "print(f\"{input_files[0]} to {str(i)}{batch_input_file}.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Register folder containing csv files in AML as data asset to use in batch job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import Data\n",
    "from azure.ai.ml.constants import AssetTypes\n",
    "\n",
    "dataset_name = \"inpainting-data\"\n",
    "input = Data(\n",
    "    name=dataset_name,\n",
    "    description=\"A sample of the dataset for image generation for batch deployment, in CSV file format\",\n",
    "    type=AssetTypes.URI_FOLDER,\n",
    "    path=processed_dataset_parent_dir,\n",
    ")\n",
    "workspace_ml_client.data.create_or_update(input)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 5. Test the endpoint - using csv files\n",
    "\n",
    "Invoke the batch endpoint with the input parameter pointing to the directory containing one or more csv files containing the batch inference input. This creates a pipeline job using the default deployment in the endpoint. Wait for the job to complete.\n",
    "\n",
    "__Note__: If job failed with Out of Memory Error then please try splitting your input into smaller csv files or decreasing mini_batch_size for the deployment."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from azure.ai.ml import Input\n",
    "\n",
    "job = None\n",
    "input = Input(path=dataset_parent_dir, type=AssetTypes.URI_FOLDER)\n",
    "num_retries = 3\n",
    "for i in range(num_retries):\n",
    "    try:\n",
    "        job = workspace_ml_client.batch_endpoints.invoke(\n",
    "            endpoint_name=endpoint.name, input=input\n",
    "        )\n",
    "        break\n",
    "    except Exception as e:\n",
    "        if i == num_retries - 1:\n",
    "            raise e\n",
    "        else:\n",
    "            print(\"Endpoint invocation failed. Retrying after 5 seconds...\")\n",
    "            time.sleep(5)\n",
    "if job is not None:\n",
    "    workspace_ml_client.jobs.stream(job.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "__Note__: If the job failed with error Assertion Error (The actual length exceeded max length 100 MB) then please consider dividing input csv file into multiple csv files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "scoring_job = list(workspace_ml_client.jobs.list(parent_job_name=job.name))[0]\n",
    "\n",
    "workspace_ml_client.jobs.download(\n",
    "    name=scoring_job.name,\n",
    "    download_path=\".\",\n",
    "    output_name=\"score\",\n",
    ")\n",
    "\n",
    "predictions_file = os.path.join(\"named-outputs\", \"score\", \"predictions.csv\")\n",
    "\n",
    "# Load the batch predictions file with no headers into a dataframe and set your column names\n",
    "score_df = pd.read_csv(\n",
    "    predictions_file,\n",
    "    header=None,\n",
    "    names=[\n",
    "        \"row_number_per_file\",\n",
    "        \"image_file_name\",\n",
    "        \"nsfw_content_detected\",\n",
    "        \"input_csv_name\",\n",
    "    ],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(score_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. Clean up resources - delete the endpoint\n",
    "Batch endpoints use compute resources only when jobs are submitted. You can keep the batch endpoint for your reference without worrying about compute bills, or choose to delete the endpoint. If you created your compute cluster to have zero minimum instances and scale down soon after being idle, you won't be charged for an unused compute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workspace_ml_client.batch_endpoints.begin_delete(name=endpoint_name).result()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
    "language_info": {
     "name": "ipython"
    }
   },
 "nbformat": 4,
 "nbformat_minor": 1
}