sdk/python/foundation-models/system/inference/text-generation/llama-safe-batch-deployment.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to create an Azure AI Content Safety enabled Llama 2 batch endpoint (Preview)\n",
    "### This notebook will walk you through the steps to create an __Azure AI Content Safety__ enabled __Llama 2__ batch endpoint.\n",
    "### This notebook is under preview\n",
    "### The steps are:\n",
    "1. Create an __Azure AI Content Safety__ resource for moderating the request from user and response from the __Llama 2__ batch endpoint.\n",
    "2. Create a new __Azure AI Content Safety__ enabled __Llama 2__ batch endpoint with a custom score.py which will integrate with the __Azure AI Content Safety__ resource to moderate the response from the __Llama 2__ model and the request from the user, but to make the custom score.py to successfully authenticated to the __Azure AI Content Safety__ resource, for batch inferencing is using __Environment variable__ to pass the access key of the __Azure AI Content Safety__ resource to the custom score.py via environment variable, then the custom score.py can use the key directly to access the Azure AI Content Safety resource, this option is less secure than the first option, 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 Azure AI Content Safety resource.\n",
    "  "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. Prerequisites\n",
    "#### 1.1 Check List:\n",
    "- [x] You have created a new Python virtual environment for this notebook.\n",
    "- [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 Azure AI Content Safety resource using that identity.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.2 Assign variables for the workspace and deployment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The public registry name contains Llama 2 models\n",
    "registry_name = \"azureml-meta\"\n",
    "\n",
    "# Name of the Llama 2 model to be deployed\n",
    "# available_llama_models_text_generation = [\"Llama-2-7b\", \"Llama-2-13b\", \"Llama-2-70b\"]\n",
    "# use the appropriate model name that is suitable for your workload below, this example shows Llama-2-7b\n",
    "model_name = \"Llama-2-7b\"\n",
    "# This notebook has been tested with \"Llama-2-7b version\" \"4\", \"Llama-2-13b\" version 4, and \"Llama-2-70b\" version 4\n",
    "import random\n",
    "\n",
    "endpoint_name = f\"batch-{random.randint(0,10000)}\"  # Replace with your endpoint name\n",
    "deployment_name = \"batch-dep\"  # Replace with your deployment name, lower case only!!!\n",
    "sku_name = \"Standard_ND40rs_v2\"  # Name of the sku(instance type) Check the model-list(can be found in the parent folder(inference)) to get the most optimal sku for your model (Default: Standard_DS2_v2)\n",
    "\n",
    "environment_name = \"llama-model-env\"  # Replace with your environment name\n",
    "compute_name = \"nd40-src\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.3 Install Dependencies(as needed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# uncomment the following lines to 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.23.1\n",
    "# %pip install azure-mgmt-msi==7.0.0\n",
    "# %pip install azure-mgmt-authorization==3.0.0"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.4 All required Imports\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.identity import DefaultAzureCredential, InteractiveBrowserCredential\n",
    "from azure.ai.ml import MLClient\n",
    "from azure.mgmt.cognitiveservices import CognitiveServicesManagementClient\n",
    "from azure.mgmt.cognitiveservices.models import Account, Sku, AccountProperties\n",
    "from IPython.core.display import display, HTML\n",
    "from azure.ai.ml import Input\n",
    "from azure.ai.ml.entities import (\n",
    "    BatchEndpoint,\n",
    "    ModelBatchDeployment,\n",
    "    ModelBatchDeploymentSettings,\n",
    "    Model,\n",
    "    AmlCompute,\n",
    "    Data,\n",
    "    BuildContext,\n",
    "    BatchRetrySettings,\n",
    "    CodeConfiguration,\n",
    "    Environment,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.5 Get credential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    credential = DefaultAzureCredential()\n",
    "    # Check if given credential can get token successfully.\n",
    "    credential.get_token(\"https://management.azure.com/.default\")\n",
    "except Exception as ex:\n",
    "    # Fall back to InteractiveBrowserCredential in case DefaultAzureCredential not work\n",
    "    credential = InteractiveBrowserCredential()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.6 Configure workspace "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    ml_client = MLClient.from_config(credential=credential)\n",
    "except Exception as ex:\n",
    "    # enter details of your AML workspace\n",
    "    subscription_id = \"subscription_id\"\n",
    "    resource_group = \"resource_group\"\n",
    "    workspace = \"workspace\"\n",
    "\n",
    "    # get a handle to the workspace\n",
    "    ml_client = MLClient(\n",
    "        credential,\n",
    "        subscription_id,\n",
    "        resource_group,\n",
    "        workspace,\n",
    "        logging_enable=True,\n",
    "    )\n",
    "\n",
    "\n",
    "subscription_id = ml_client.subscription_id\n",
    "resource_group = ml_client.resource_group_name\n",
    "workspace = ml_client.workspace_name\n",
    "\n",
    "print(f\"Connected to workspace {workspace}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 1.7 Assign variables for Azure Content Safety\n",
    "Currently, Azure AI Content Safety is in a limited set of regions:\n",
    "\n",
    "\n",
    "__NOTE__: before you choose the region to deploy the Azure AI Content Safety, 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": [
    "acs_client = CognitiveServicesManagementClient(credential, subscription_id)\n",
    "\n",
    "\n",
    "# settings for the Azure AI Content Safety resource\n",
    "# we will choose existing AACS resource if it exists, otherwise create a new one\n",
    "# name of azure ai content safety resource, has to be unique\n",
    "import time\n",
    "\n",
    "aacs_name = f\"{endpoint_name}-aacs\"\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 = acs_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\" + sku.tier + \"\\t\" + locations)\n",
    "\n",
    "print(\n",
    "    f\"Choose a new Azure AI Content Safety resource in {aacs_location} with SKU {aacs_sku_name}\"\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Create Azure AI Content Safety"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "# How many seconds to wait between checking the status of an async operation.\n",
    "wait_time = 10\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 = acs_client.accounts.get(resource_group, aacs_name)\n",
    "    print(f\"Found existing Azure AI content safety Account {aacs.name}.\")\n",
    "except:\n",
    "    try:\n",
    "        # check if there is an existing AACS resource within same resource group\n",
    "        aacs = find_acs(acs_client.accounts.list_by_resource_group(resource_group))\n",
    "        print(\n",
    "            f\"Found existing Azure AI content safety Account {aacs.name} in resource group {resource_group}.\"\n",
    "        )\n",
    "    except:\n",
    "        print(f\"Creating Azure AI content safety Account {aacs_name}.\")\n",
    "        acs_client.accounts.begin_create(resource_group, aacs_name, parameters).wait()\n",
    "        print(\"Resource created.\")\n",
    "        aacs = acs_client.accounts.get(resource_group, aacs_name)\n",
    "\n",
    "\n",
    "aacs_endpoint = aacs.properties.endpoint\n",
    "aacs_resource_id = aacs.id\n",
    "print(f\"AACS endpoint is {aacs_endpoint}\")\n",
    "print(f\"AACS ResourceId is {aacs_resource_id}\")\n",
    "\n",
    "aacs_access_key = acs_client.accounts.list_keys(\n",
    "    resource_group_name=resource_group, account_name=aacs.name\n",
    ").key1"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Create Azure AI Content Safety enabled Llama 2 batch endpoint"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.1 Check if Llama 2 model is available in the AML registry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reg_client = MLClient(\n",
    "    credential,\n",
    "    subscription_id=subscription_id,\n",
    "    resource_group_name=resource_group,\n",
    "    registry_name=registry_name,\n",
    ")\n",
    "version_list = list(\n",
    "    reg_client.models.list(model_name)\n",
    ")  # list available versions of the model\n",
    "llama_model = None\n",
    "if len(version_list) == 0:\n",
    "    raise Exception(f\"No model named {model_name} found in registry\")\n",
    "else:\n",
    "    model_version = \"4\"\n",
    "    llama_model = reg_client.models.get(model_name, model_version)\n",
    "    print(\n",
    "        f\"Using model name: {llama_model.name}, version: {llama_model.version}, id: {llama_model.id} for inferencing\"\n",
    "    )"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 3.2 Create environment for Llama 2 endpoint\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    env = 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, build=BuildContext(path=\"./llama-files/docker_env\")\n",
    "    )\n",
    "    ml_client.environments.create_or_update(env)\n",
    "    env = 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}\">\n",
    "                Click here to check env build status in AML studio\n",
    "             </a>\n",
    "             \"\"\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 3.3 Create compute cluster to run batch job on\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.constants import AssetTypes, BatchDeploymentOutputAction\n",
    "\n",
    "if not any(filter(lambda m: m.name == compute_name, ml_client.compute.list())):\n",
    "    compute_cluster = AmlCompute(\n",
    "        name=compute_name,\n",
    "        size=sku_name,\n",
    "        min_instances=0,\n",
    "        max_instances=2,\n",
    "    )\n",
    "    ml_client.compute.begin_create_or_update(compute_cluster).result()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 3.4 Create Llama 2 batch endpoint\n",
    "This step may take a few minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import BatchEndpoint\n",
    "\n",
    "# Check if the endpoint already exists in the workspace\n",
    "try:\n",
    "    endpoint = 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",
    "        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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 3.5 Deploy Llama 2 model\n",
    "This step may take a few minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deployment = ModelBatchDeployment(\n",
    "    name=deployment_name,\n",
    "    endpoint_name=endpoint.name,\n",
    "    model=llama_model,\n",
    "    environment=env,\n",
    "    code_configuration=CodeConfiguration(\n",
    "        code=\"llama-files/score/default\",\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=1,\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",
    "    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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 3.6 Update Batch endpoint to set the default deployment\n",
    "This step may take a few minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint = ml_client.batch_endpoints.get(endpoint_name)\n",
    "endpoint.defaults.deployment_name = deployment.name\n",
    "ml_client.batch_endpoints.begin_create_or_update(endpoint).result()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Prepare to test."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.1 Input data preparation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = \"llama-files/data\"\n",
    "dataset_name = \"input-data-small\"\n",
    "\n",
    "input_data = Data(\n",
    "    path=data_path,\n",
    "    type=AssetTypes.URI_FOLDER,\n",
    "    description=\"A sample of the  dataset for text generation, in CSV file format\",\n",
    "    name=dataset_name,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data = ml_client.data.create_or_update(input_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import sleep\n",
    "\n",
    "print(f\"Waiting for data asset {dataset_name}\", end=\"\")\n",
    "while not any(filter(lambda m: m.name == dataset_name, ml_client.data.list())):\n",
    "    sleep(10)\n",
    "    print(\".\", end=\"\")\n",
    "\n",
    "print(\" [DONE]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data = ml_client.data.get(name=dataset_name, label=\"latest\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "input = Input(type=AssetTypes.URI_FOLDER, path=input_data.id)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.2 Invoke the endpoint\n",
    "\n",
    "Let's now invoke the endpoint for batch scoring job:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "job = ml_client.batch_endpoints.invoke(endpoint_name=endpoint.name, input=input)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.3 Get the details of the invoked job\n",
    "Let us get details and logs of the invoked job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_client.jobs.get(job.name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can wait for the job to finish using the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_client.jobs.stream(job.name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.4 Download the results\n",
    "\n",
    "The deployment creates a child job that executes the scoring. We can get the details of it using the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "scoring_job = list(ml_client.jobs.list(parent_job_name=job.name))[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Job name:\", scoring_job.name)\n",
    "print(\"Job status:\", scoring_job.status)\n",
    "print(\n",
    "    \"Job duration:\",\n",
    "    scoring_job.creation_context.last_modified_at\n",
    "    - scoring_job.creation_context.created_at,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The outputs generated by the deployment job will be placed in an output named `score`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_client.jobs.download(name=scoring_job.name, download_path=\".\", output_name=\"score\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output file above will contain one line for each file, and each line will have multiple arrays corresponding to each line of the file. If you see [], that means ACS has stripped the response from the model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 4.5 Clean up Resources\n",
    "Delete endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_client.batch_endpoints.begin_delete(endpoint.name).result()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}