sdk/python/generative-ai/rag/notebooks/qa_data_generation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "345e6aa4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install azure-ai-ml\n",
    "%pip install -U 'azureml-rag[faiss]>=0.1.14'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "038912d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If `import win32file` fails with a DLL error then run the following and restart kernel:\n",
    "# %pip uninstall -y pywin32\n",
    "# %conda install -y --force-reinstall pywin32"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "12eb9e1f",
   "metadata": {},
   "source": [
    "# QA Data Generation\n",
    "\n",
    "QA Data Generation is a part of RAG (Retrieval Augemented Generation) creation process where the autogenerated QA dataset is used:\n",
    "\n",
    "1. To get the best prompt for RAG\n",
    "2. To get evaluation metrics for RAG\n",
    "\n",
    "This notebook shows you how to create a QA dataset from your data (Git repo). We run just the components needed for QA Data Generation and not for the full RAG creation flow."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1a99955c",
   "metadata": {},
   "source": [
    "## Get client for AzureML Workspace\n",
    "\n",
    "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. In this section we will connect to the workspace in which the job will be run.\n",
    "\n",
    "Enter your Workspace details below, running this still will write a `workspace.json` file to the current folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63178816",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile workspace.json\n",
    "{\n",
    "    \"subscription_id\": \"<subscription_id>\",\n",
    "    \"resource_group\": \"<resource_group_name>\",\n",
    "    \"workspace_name\": \"<workspace_name>\"\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fb5c8f0",
   "metadata": {},
   "source": [
    "`MLClient` is how you interact with AzureML"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1af37c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.identity import DefaultAzureCredential, InteractiveBrowserCredential\n",
    "from azure.ai.ml import MLClient\n",
    "from azureml.core import Workspace\n",
    "\n",
    "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()\n",
    "\n",
    "try:\n",
    "    ml_client = MLClient.from_config(credential=credential, path=\"workspace.json\")\n",
    "except Exception as ex:\n",
    "    raise Exception(\n",
    "        \"Failed to create MLClient from config file. Please modify and then run the above cell with your AzureML Workspace details.\"\n",
    "    ) from ex\n",
    "    # ml_client = MLClient(\n",
    "    #     credential=credential,\n",
    "    #     subscription_id=\"\",\n",
    "    #     resource_group_name=\"\",\n",
    "    #     workspace_name=\"\"\n",
    "    # )\n",
    "\n",
    "ws = Workspace(\n",
    "    subscription_id=ml_client.subscription_id,\n",
    "    resource_group=ml_client.resource_group_name,\n",
    "    workspace_name=ml_client.workspace_name,\n",
    ")\n",
    "print(ml_client)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "30906d39",
   "metadata": {},
   "source": [
    "## Azure OpenAI\n",
    "\n",
    "We recommend using gpt-35-turbo model to get good quality QAs. [Follow these instructions](https://learn.microsoft.com/en-us/azure/cognitive-services/openai/how-to/create-resource?pivots=web-portal) to setup an Azure OpenAI Instance and deploy the model. Once you have the model deployed in AOAI you can specify your Model name and Deployment name below.\n",
    "\n",
    "We will use the automatically created `Default_AzureOpenAI` connection, change `aoai_connection_name` to use your own."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab3f1c33",
   "metadata": {},
   "outputs": [],
   "source": [
    "aoai_connection_name = \"Default_AzureOpenAI\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "129ac0b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azureml.rag.utils.connections import (\n",
    "    get_connection_by_name_v2,\n",
    "    create_connection_v2,\n",
    ")\n",
    "\n",
    "try:\n",
    "    aoai_connection = get_connection_by_name_v2(ws, aoai_connection_name)\n",
    "except Exception as ex:\n",
    "    # Create New Connection\n",
    "    # Modify the details below to match the `Endpoint` and API key of your AOAI resource, these details can be found in Azure Portal\n",
    "    raise RuntimeError(\n",
    "        \"Have you entered your AOAI resource details below? If so, delete me!\"\n",
    "    )\n",
    "    aoai_connection = create_connection_v2(\n",
    "        workspace=ws,\n",
    "        name=aoai_connection,\n",
    "        category=\"AzureOpenAI\",\n",
    "        # 'Endpoint' from Azure OpenAI resource overview\n",
    "        target=\"https://<endpoint_name>.openai.azure.com/\",\n",
    "        auth_type=\"ApiKey\",\n",
    "        credentials={\n",
    "            # Either `Key` from the `Keys and Endpoint` tab of your Azure OpenAI resource, will be stored in your Workspace associated Azure Key Vault.\n",
    "            \"key\": \"<api-key>\"\n",
    "        },\n",
    "        metadata={\"ApiType\": \"azure\", \"ApiVersion\": \"2023-05-15\"},\n",
    "    )\n",
    "\n",
    "aoai_connection_id = aoai_connection[\"id\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73c4ece4",
   "metadata": {},
   "source": [
    "Now that your Workspace has a connection to Azure OpenAI we will make sure the `gpt-35-turbo` model has been deployed ready for inference. This cell will fail if there is not deployment for the embeddings model, [follow these instructions](https://learn.microsoft.com/azure/cognitive-services/openai/how-to/create-resource?pivots=web-portal#deploy-a-model) to deploy a model with Azure OpenAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1e5cc13",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azureml.rag.utils.deployment import infer_deployment\n",
    "\n",
    "aoai_completion_model_name = \"gpt-35-turbo\"\n",
    "\n",
    "try:\n",
    "    aoai_completion_deployment_name = infer_deployment(\n",
    "        aoai_connection, aoai_completion_model_name\n",
    "    )\n",
    "except Exception as e:\n",
    "    print(\n",
    "        f\"Deployment name in AOAI workspace for model '{aoai_completion_model_name}' is not found.\"\n",
    "    )\n",
    "    print(\n",
    "        f\"Please create a deployment for this model by following the deploy instructions on the resource page for '{aoai_connection['properties']['target']}' in Azure Portal.\"\n",
    "    )\n",
    "\n",
    "print(\n",
    "    f\"Deployment name in AOAI workspace for model '{aoai_completion_model_name}' is '{aoai_completion_deployment_name}'\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff40ff04",
   "metadata": {},
   "source": [
    "Finally we will combine the deployment and model information into a uri form which the AzureML embeddings components expect as input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1dfe810e",
   "metadata": {},
   "outputs": [],
   "source": [
    "llm_completion_config = f'{{\"type\":\"azure_open_ai\",\"model_name\":\"{aoai_completion_model_name}\",\"deployment_name\":\"{aoai_completion_deployment_name}\",\"temperature\":0,\"max_tokens\":\"1500\"}}'"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "56878876",
   "metadata": {},
   "source": [
    "### Setup Pipeline\n",
    "\n",
    "AzureML [Pipelines](https://learn.microsoft.com/azure/machine-learning/concept-ml-pipelines?view=azureml-api-2) connect together multiple [Components](https://learn.microsoft.com/azure/machine-learning/concept-component?view=azureml-api-2). Each Component defines inputs, code that consumes the inputs and outputs produced from the code. Pipelines themselves can have inputs, and outputs produced by connecting together individual sub Components.\n",
    "To process your data for embedding and indexing we will chain together multiple components each performing their own step of the workflow.\n",
    "\n",
    "The Components are published to a [Registry](https://learn.microsoft.com/azure/machine-learning/how-to-manage-registries?view=azureml-api-2&tabs=cli), `azureml`, which should have access to by default, it can be accessed from any Workspace.\n",
    "In the below cell we get the Component Definitions from the `azureml` registry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2a3752a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ml_registry = MLClient(credential=credential, registry_name=\"azureml\")\n",
    "\n",
    "validate_deployments_component = ml_registry.components.get(\n",
    "    \"llm_rag_validate_deployments\", label=\"latest\"\n",
    ")\n",
    "# Clones git repository to output folder of pipeline, by default this will be on the default Workspace Datastore `workspaceblobstore`\n",
    "git_clone_component = ml_registry.components.get(\"llm_rag_git_clone\", label=\"latest\")\n",
    "# Walks input folder according to provided glob pattern (all files by default: '**/*') and attempts to open them, extract text chunks and further chunk if necessary to fir within provided `chunk_size`.\n",
    "crack_and_chunk_component = ml_registry.components.get(\n",
    "    \"llm_rag_crack_and_chunk\", label=\"latest\"\n",
    ")\n",
    "# Generates test data for QA model training and evaluation\n",
    "data_generation_component = ml_registry.components.get(\n",
    "    \"llm_rag_qa_data_generation\", label=\"latest\"\n",
    ")\n",
    "# Takes a uri to a storage location where test data is stored and registers it as a uri_file Data asset in the AzureML Workspace.\n",
    "register_qa_data_component = ml_registry.components.get(\n",
    "    \"llm_rag_register_qa_data_asset\", label=\"latest\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c2836fb",
   "metadata": {},
   "source": [
    "Each Component has documentation which provides an overall description of the Components purpose and each of the inputs/outputs.\n",
    "For example we can see understand what `data_generation_component` does by inspecting the Component definition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fc995a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(data_generation_component)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eaf3a251",
   "metadata": {},
   "source": [
    "Below a Pipeline is built by defining a python function which chains together the above components inputs and outputs. Arguments to the function are inputs to the Pipeline itself and the return value is a dictionary defining the outputs of the Pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53285c16",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import Output\n",
    "from azure.ai.ml.dsl import pipeline\n",
    "from azure.ai.ml.entities._job.pipeline._io import PipelineInput\n",
    "from typing import Optional\n",
    "\n",
    "\n",
    "def use_automatic_compute(\n",
    "    component, instance_count=1, instance_type=\"Standard_DS12_v2\"\n",
    "):\n",
    "    \"\"\"Configure input `component` to use automatic compute with `instance_count` and `instance_type`.\n",
    "\n",
    "    This avoids the need to provision a compute cluster to run the component.\n",
    "    \"\"\"\n",
    "    component.set_resources(\n",
    "        instance_count=instance_count,\n",
    "        instance_type=instance_type,\n",
    "        properties={\"compute_specification\": {\"automatic\": True}},\n",
    "    )\n",
    "    return component\n",
    "\n",
    "\n",
    "def optional_pipeline_input_provided(input: Optional[PipelineInput]):\n",
    "    \"\"\"Checks if optional pipeline inputs are provided.\"\"\"\n",
    "    return input is not None and input._data is not None\n",
    "\n",
    "\n",
    "def use_aoai_connection(component, aoai_connection_id, custom_env: str = None):\n",
    "    \"\"\"Configure input `component` to use AOAI connection.\"\"\"\n",
    "    if optional_pipeline_input_provided(aoai_connection_id):\n",
    "        if custom_env is not None:\n",
    "            component.environment_variables[custom_env] = aoai_connection_id\n",
    "        else:\n",
    "            component.environment_variables[\n",
    "                \"AZUREML_WORKSPACE_CONNECTION_ID_AOAI\"\n",
    "            ] = aoai_connection_id\n",
    "\n",
    "\n",
    "# If you have an existing compute cluster you want to use instead of automatic compute, uncomment the following line, replace `dedicated_cpu_compute` with the name of your cluster.\n",
    "# Also comment out the `component.set_resources` line in `use_automatic_compute` above and the `default_compute='serverless'` line below.\n",
    "# @pipeline(compute=dedicated_cpu_compute)\n",
    "@pipeline(default_compute=\"serverless\")\n",
    "def qa_data_generation(\n",
    "    git_url: str,\n",
    "    asset_name: str,\n",
    "    llm_completion_config: str,\n",
    "    data_source_glob: str = None,\n",
    "    data_source_url: str = None,\n",
    "    document_path_replacement_regex: str = None,\n",
    "    aoai_connection_id: str = None,\n",
    "):\n",
    "    \"\"\"Pipeline for generating QA data from a git repository.\"\"\"\n",
    "\n",
    "    validate_deployments = validate_deployments_component(\n",
    "        llm_config=llm_completion_config,\n",
    "        check_completion=\"True\",\n",
    "        check_embeddings=\"False\",\n",
    "    )\n",
    "    use_automatic_compute(validate_deployments)\n",
    "    use_aoai_connection(\n",
    "        validate_deployments,\n",
    "        aoai_connection_id,\n",
    "        custom_env=\"AZUREML_WORKSPACE_CONNECTION_ID_AOAI_COMPLETION\",\n",
    "    )\n",
    "\n",
    "    git_clone = git_clone_component(\n",
    "        git_repository=git_url,\n",
    "    )\n",
    "    use_automatic_compute(git_clone)\n",
    "\n",
    "    crack_and_chunk = crack_and_chunk_component(\n",
    "        input_data=git_clone.outputs.output_data,\n",
    "        input_glob=data_source_glob,\n",
    "        chunk_size=1024,\n",
    "        data_source_url=data_source_url,\n",
    "        document_path_replacement_regex=document_path_replacement_regex,\n",
    "    )\n",
    "    use_automatic_compute(crack_and_chunk)\n",
    "\n",
    "    # QA Data Generation\n",
    "    data_generation = data_generation_component(\n",
    "        input_data=crack_and_chunk.outputs.output_chunks,\n",
    "        deployment_validation=validate_deployments.outputs.output_data,\n",
    "        llm_config=llm_completion_config,\n",
    "        dataset_size=10,  # Number of QAs to be generated\n",
    "    )\n",
    "    use_automatic_compute(data_generation)\n",
    "    use_aoai_connection(data_generation, aoai_connection_id)\n",
    "\n",
    "    register_qa_data = register_qa_data_component(\n",
    "        storage_uri=data_generation.outputs.output_data,\n",
    "        asset_name=asset_name,\n",
    "        register_output=True,\n",
    "    )\n",
    "    use_automatic_compute(register_qa_data)\n",
    "\n",
    "    return {\n",
    "        \"qa_data\": data_generation.outputs.output_data,\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74957019",
   "metadata": {},
   "source": [
    "The settings below show how the different git and data_source parameters can be set to processes only the AzureML documentation from the larger AzureDocs git repo, and ensure the source url for each document is processed to link to the publicly hosted URL instead of the git url."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cd4b8bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "git_url = \"https://github.com/MicrosoftDocs/azure-docs/\"\n",
    "data_source_url = \"https://learn.microsoft.com/en-us/azure\"\n",
    "data_source_glob = \"articles/machine-learning/**/*\"\n",
    "# This regex is used to remove the 'articles' folder from the source url put in each files metadata in the index.\n",
    "document_path_replacement_regex = r'{\"match_pattern\": \"(.*)/articles/(.*)(\\\\.[^.]+)$\", \"replacement_pattern\": \"\\\\1/\\\\2\"}'\n",
    "asset_name = \"azure_docs_ml_qa_data\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f5d18c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import Input\n",
    "from azure.ai.ml.entities import UserIdentityConfiguration\n",
    "\n",
    "pipeline_job = qa_data_generation(\n",
    "    git_url=git_url,\n",
    "    data_source_glob=data_source_glob,\n",
    "    data_source_url=data_source_url,\n",
    "    document_path_replacement_regex=document_path_replacement_regex,\n",
    "    aoai_connection_id=aoai_connection_id,\n",
    "    llm_completion_config=llm_completion_config,\n",
    "    asset_name=asset_name,\n",
    ")\n",
    "\n",
    "# By default AzureML Pipelines will reuse the output of previous component Runs when inputs have not changed.\n",
    "# If you want to rerun the Pipeline every time each time so that any changes to upstream data sources are processed uncomment the below line.\n",
    "# pipeline_job.settings.force_rerun = True # Rerun each time so that git_clone isn't cached, if intent is to ingest latest data."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "67f89a25",
   "metadata": {},
   "source": [
    "### Submit Pipeline\n",
    "\n",
    "**In case of any errors see [TROUBLESHOOT.md](../../TROUBLESHOOT.md).**\n",
    "\n",
    "The output of each step in the pipeline can be inspected via the Workspace UI, click the link under 'Details Page' after running the below cell. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52d9b2be",
   "metadata": {},
   "outputs": [],
   "source": [
    "running_pipeline_job = ml_client.jobs.create_or_update(\n",
    "    pipeline_job, experiment_name=\"qa_data_generation\"\n",
    ")\n",
    "running_pipeline_job"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ed144c98",
   "metadata": {},
   "source": [
    "### Review generated QA data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "174bd7de",
   "metadata": {},
   "outputs": [],
   "source": [
    "import fsspec\n",
    "import pandas as pd\n",
    "\n",
    "qa_data = ml_client.data.get(f\"{asset_name}-test-data\", label=\"latest\")\n",
    "with fsspec.open(qa_data.path) as f:\n",
    "    df = pd.read_json(f, lines=True)\n",
    "df.head(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d236cf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "for qa_type in [\"TOPIC\", \"FACTUAL\", \"BOOLEAN\"]:\n",
    "    print(f\"{qa_type} Question Answers:\")\n",
    "    for _, row in df[df[\"qaType\"] == qa_type][:2].iterrows():\n",
    "        print(\"Q:\", row[\"question\"])\n",
    "        print(\"A:\", row[\"answer\"])\n",
    "        print()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e598a8a4",
   "metadata": {},
   "source": [
    "### Review token usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f97a1a98",
   "metadata": {},
   "outputs": [],
   "source": [
    "child_runs = ml_client.jobs.list(parent_job_name=running_pipeline_job.name)\n",
    "child_runs = list(child_runs)\n",
    "data_generation_run = child_runs[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "369793ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from azureml.core import Run\n",
    "\n",
    "run = Run.get(ws, data_generation_run.name)\n",
    "metrics = run.get_metrics()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbed9c85",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Tokens used: {metrics['total_tokens']}\")\n",
    "print(f\"Model used: {metrics['llm_model_name']}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "69b9772c",
   "metadata": {},
   "source": [
    "Given the token usage and the model you can compute cost using the pricing here: https://azure.microsoft.com/en-us/pricing/details/cognitive-services/openai-service/."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}