{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2eec5cc39a59"
      },
      "outputs": [],
      "source": [
        "# Copyright 2024 Google LLC\n",
        "#\n",
        "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "#     https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JAPoU8Sm5E6e"
      },
      "source": [
        "# Anomaly Detection of Infrastructure Logs using Gemini and BigQuery Vector Search\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-infra-logs.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fembeddings%2Fuse-cases%2Foutlier-detection%2Fbq-vector-search-outlier-detection-infra-logs.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-infra-logs.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\" alt=\"Vertex AI logo\"><br> Open in Vertex AI Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/bigquery/import?url=https://github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-infra-logs.ipynb\">\n",
        "      <img src=\"https://www.gstatic.com/images/branding/gcpiconscolors/bigquery/v1/32px.svg\" alt=\"BigQuery Studio logo\"><br> Open in BigQuery Studio\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-infra-logs.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n",
        "\n",
        "<div style=\"clear: both;\"></div>\n",
        "\n",
        "<b>Share to:</b>\n",
        "\n",
        "<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-infra-logs.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-infra-logs.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-infra-logs.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-infra-logs.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
        "</a>\n",
        "\n",
        "<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-infra-logs.ipynb\" target=\"_blank\">\n",
        "  <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
        "</a>            "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "84f0f73a0f76"
      },
      "source": [
        "| | |\n",
        "|-|-|\n",
        "| Author: | [Roy Arsan](https://github.com/rarsan) |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tvgnzT1CKxrO"
      },
      "source": [
        "## Overview\n",
        "\n",
        "This notebook guides you through processing large volumes of logs using Gemini, vector embeddings and [BigQuery Vector Search](https://cloud.google.com/bigquery/docs/vector-search-intro) to perform anomaly detection. We evaluate the accuracy using publicly available labelled [HDFS](https://github.com/logpai/loghub) dataset, and contrast it with other common log anomaly detection techniques.\n",
        "\n",
        "\n",
        "As part of this notebook you will learn how to:\n",
        "1. Use SQL for preprocessing raw logs at scale, interact with LLMs and do vector search analysis\n",
        "1. Use Gemini 2.0 to translate a log sequence into simple natural language summary directly from BigQuery\n",
        "1. Use a text embedding model to generate a vector embedding for each log summary directly from BigQuery, and store all embeddings in BigQuery vector index for fast lookup\n",
        "1. Tune and run vector search-based anomaly detection using BigQuery vector search\n",
        "1. Evaluate performance and compare against popular unsupervised and semi-supervised outlier detection ML algorithms\n",
        "\n",
        "\n",
        "This notebook demonstrates how using off-the-shelf Gemini 2.0 + Text Embeddings + BigQuery vector search yields comparable results to custom pre-trained language models including DeepLog, an [LSTM deep neural network](https://dl.acm.org/doi/10.1145/3133956.3134015).\n",
        "\n",
        "We also compare vector search-based outlier detection with common [scikit outlier detection](https://scikit-learn.org/stable/modules/outlier_detection.html#) algorithms (OneClassSVM, LocalOutlierFactor) using the same generated embeddings and we found that a custom SQL logic using BigQuery vector search to be more accurate and flexible for this use case and dataset.\n",
        "\n",
        "For performance results, see `Summary` section at the bottom.\n",
        "\n",
        " You can use and tweak this notebook to evaluate for your own outlier detection use case, with your own dataset and target sensitivity level."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DW8bwUaS2JQj"
      },
      "source": [
        "### Data pipeline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Rt10T8iZrNSc"
      },
      "source": [
        "![anomaly-detection-data-pipeline.png](https://storage.googleapis.com/github-repo/generative-ai/embeddings/use-cases/outlier-detection/bq-vector-search-outlier-detection-infra-logs/anomaly-detection-data-pipeline.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sh4i937UcI2q"
      },
      "source": [
        "### Task, metrics and dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YWviBrbU7cJ_"
      },
      "source": [
        "Outlier detection, specifically novelty detection, is about detecting all anomalies including 'unknown' types of anomalies (e.g. system failures, new cyberattacks, etc.). Therefore supervised approaches where you train the model on both normal and abnormal logs are simply not applicable. Therefore, the approach described in this notebook is compared against leading and popular unsupervised and semi-supervised techniques. **Recall** is our primary objective for outlier detection (minimize false negatives) while also keeping reasonably high **precision** (minimize false positives).\n",
        "\n",
        "We use the open-source labelled HDFS dataset from [Loghub]((https://github.com/logpai/loghub) which is freely accessible and commonly used in log analysis evaluations, including anomaly detection. HDFS dataset provides 11M+ system log lines (\\~575k log sessions) collected from a real Hadoop cluster on 200 nodes, and labelled by Hadoop domain experts to identify the runtime anomalies. For the purpose of this notebook, we use a subset of the HDFS dataset (70k log sessions or ~12%) taking into consideration default BigQuery ML quota limits (e.g. 200 rpm and 72,000 rows per job with `gemini-2.0-flash` at the time of this writing). This way you can run this notebook as is in your own Google Cloud project in a timely fashion without necessarily requesting quota increases."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "61RBz8LLbxCR"
      },
      "source": [
        "## Getting started"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "No17Cw5hgx12"
      },
      "source": [
        "### Install Vertex AI SDK and other required packages\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tFy3H3aPgx12"
      },
      "outputs": [],
      "source": [
        "%pip install --upgrade --user --quiet \\\n",
        "    google-cloud-aiplatform \\\n",
        "    google-cloud-bigquery \\\n",
        "    scikit-learn==1.3.1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which restarts the current kernel.\n",
        "\n",
        "The restart might take a minute or longer. After it's restarted, continue to the next step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Wait until it's finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dmWOrTJ3gx13"
      },
      "source": [
        "### Authenticate your notebook environment (Colab only)\n",
        "\n",
        "If you're running this notebook on Google Colab, run the cell below to authenticate your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NyKGtVQjgx13"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth\n",
        "\n",
        "    auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DF4l8DTdWgPY"
      },
      "source": [
        "### Set Google Cloud project information and initialize Vertex AI SDK\n",
        "\n",
        "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nqwi-5ufWp_B"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"[project-id]\"  # @param {type:\"string\"}\n",
        "REGION = \"us-central1\"  # @param {type:\"string\"}\n",
        "LOCATION = \"US\"  # @param {type:\"string\"}\n",
        "\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=REGION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5303c05f7aa6"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "01qn1OZeh_EP"
      },
      "outputs": [],
      "source": [
        "import bigframes.pandas as bpd\n",
        "from google.cloud import bigquery\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "# Import necessary libraries\n",
        "import numpy as np\n",
        "import seaborn as sns\n",
        "from sklearn.manifold import TSNE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jLc02A9ch_EP"
      },
      "outputs": [],
      "source": [
        "# initialize bigquery client\n",
        "client = bigquery.Client()\n",
        "\n",
        "# initialize bigquery dataframes client\n",
        "bpd.close_session()\n",
        "bpd.options.bigquery.project = PROJECT_ID\n",
        "bpd.options.bigquery.location = LOCATION"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_68282UqlxrR"
      },
      "outputs": [],
      "source": [
        "DATASET = \"vs_logs_demo\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aF5qTVGRh_EP"
      },
      "source": [
        "### Create a new dataset"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L6OQ55Axh_EP"
      },
      "outputs": [],
      "source": [
        "%%bigquery\n",
        "CREATE SCHEMA `vs_logs_demo` OPTIONS (LOCATION = 'US');"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AxdBl98uUM09"
      },
      "source": [
        "### Create Cloud resource connection for Vertex AI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qsJ_dwvG-28r"
      },
      "source": [
        "You need to create a Cloud resource connection to enable BigQuery to interact with Vertex AI services such as LLM remote models."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7EZlKHyyUMNY"
      },
      "outputs": [],
      "source": [
        "CONNECTION_NAME = \"bq-llm-connection\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BVbxqd9C_Tki"
      },
      "outputs": [],
      "source": [
        "!bq mk --connection --location={LOCATION} --project_id={PROJECT_ID} \\\n",
        "    --connection_type=CLOUD_RESOURCE {CONNECTION_NAME}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pQphx8hz_cyj"
      },
      "source": [
        "Retrieve the connection service account ID"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5hgJatLm_eFF"
      },
      "outputs": [],
      "source": [
        "!bq show --connection {PROJECT_ID}.{LOCATION}.{CLOUD_RESOURCE_CONN}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1YVxe9B3_hhe"
      },
      "source": [
        "Copy the service account ID in the following parameter:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XPqGhAt6_jUU"
      },
      "outputs": [],
      "source": [
        "# Copy the value of `serviceAccountId` field in last cell's output\n",
        "CONNECTION_SA_ID = \"[bq-connection-service-account]\"  # @param {type:\"string\"}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wlsjW4-X_vOw"
      },
      "source": [
        "Give that service account access to Vertex AI"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oKYvvcZM_wMu"
      },
      "outputs": [],
      "source": [
        "!gcloud projects add-iam-policy-binding {PROJECT_ID} --member='serviceAccount:{CONNECTION_SA_ID}' --role='roles/aiplatform.user' --condition=None --no-user-output-enabled"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E0qF8NR7_0o3"
      },
      "source": [
        "If your notebook environment does not have `setIamPolicy` permission on the project, you (or your project's administrator) can alternatively grant  permissions to that cloud connection's service account using Cloud Console following these steps: https://cloud.google.com/bigquery/docs/generate-text-tutorial-gemini#grant-permissions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HysosuDkh_Ec"
      },
      "source": [
        "## Import logs dataset into BigQuery"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iVPIbKUzXx5l"
      },
      "source": [
        "The HDFS dataset was originally developed by Wei Xu et al. in [Detecting Large-Scale System Problems by Mining Console Logs](https://people.eecs.berkeley.edu/~jordan/papers/xu-etal-sosp09.pdf) and made available in [Wei Xu's website](https://people.iiis.tsinghua.edu.cn/~weixu/sospdata.html). The raw logs are also available at [loghub](https://github.com/logpai/loghub) thanks to the work of Zhu et al. in [Loghub: A Large Collection of System Log Datasets for AI-driven Log Analytics](https://arxiv.org/abs/2008.06448). To facilitate log ingestion and analysis in this notebook, the logs are expected to be converted to structured format (csv) using a standard log parser such as [this one](https://github.com/logpai/logparser/blob/main/logparser/Drain/README.md) from [logparser](https://github.com/logpai/logparser).\n",
        "\n",
        "For the purpose of this notebook, the pre-parsed logs are available as .csv files in a GCS bucket:\n",
        "\n",
        "*   Full dataset: `HDFS.log_structured.csv` (2.9 GB)\n",
        "*   Limited 2k dataset: `HDFS_2k.log_structured.csv` (< 1 MB)\n",
        "\n",
        "You will import the full dataset in this notebook, but subsequently process ~15% of the data."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aRBDEawbh_Ef"
      },
      "source": [
        "### Load logs from GCS"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2zn4ds3gh_Ef"
      },
      "outputs": [],
      "source": [
        "gcs_bucket = \"gs://csa-datasets-public/hdfs/\"\n",
        "gcs_filename = \"HDFS.log_structured.csv\"  # @param {type:\"string\"}\n",
        "\n",
        "gcs_uri = gcs_bucket + gcs_filename\n",
        "\n",
        "\n",
        "def camelCase(col):\n",
        "    return col[0].lower() + col[1:]\n",
        "\n",
        "\n",
        "df = bpd.read_csv(gcs_uri, engine=\"bigquery\")\n",
        "# lowercase first letter of each column name\n",
        "df.columns = [camelCase(col) for col in df.columns]\n",
        "print(df.info())\n",
        "\n",
        "df.to_gbq(\"vs_logs_demo.hdfs_full_structured\", if_exists=\"replace\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CUseo2PAh_Ef"
      },
      "source": [
        "### Group logs by sessions"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mgfkFhITXVm-"
      },
      "source": [
        "Grouping logs by session ID or other unique identifiers allows us to correlate events that belong to the same user or transaction, providing context and a clear chronological sequence of actions. This is essential for understanding user behavior, troubleshooting issues, and identifying patterns that might otherwise be lost in the noise of a massive, undifferentiated log stream.\n",
        "\n",
        "In our case, since HDFS logs are records of events related to file system operations at the block level, such as block replication and deletion, we need to group HDFS logs by unique `blockId`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D-fsGL0Uh_Eg"
      },
      "outputs": [],
      "source": [
        "%%bigquery\n",
        "CREATE OR REPLACE TABLE `vs_logs_demo.hdfs_full_sessionized` AS (\n",
        "  SELECT\n",
        "    blockId,\n",
        "    -- Capture start and end time\n",
        "    MIN(timestamp) AS ts_start,\n",
        "    MAX(timestamp) AS ts_end,\n",
        "    COUNT(*) AS numEvents,\n",
        "    -- Concatenate sequence of events\n",
        "    STRING_AGG(CONCAT(\n",
        "     component, \" \", level, \" \", content\n",
        "    ), '\\n' ORDER BY timestamp ASC) AS eventSequence,\n",
        "    -- Concatenate sequence of event templates\n",
        "    STRING_AGG(CONCAT(\n",
        "     component, \" \", level, \" \", eventTemplate\n",
        "    ), '\\n' ORDER BY timestamp ASC) AS eventTemplateSequence,\n",
        "    -- For debug only\n",
        "    ARRAY_AGG(STRUCT(\n",
        "        timestamp, level, component, content, eventId, eventTemplate\n",
        "    ) ORDER BY timestamp ASC) AS sequence,\n",
        "  FROM (\n",
        "    SELECT\n",
        "      PARSE_TIMESTAMP('%y%m%d %H%M%S', CONCAT(LPAD(CAST(date AS STRING), 6, '0'), ' ', LPAD(CAST(time AS STRING), 6, '0'))) AS timestamp,\n",
        "      REGEXP_EXTRACT(content, r'(blk_-?[0-9]+)') AS blockId,\n",
        "      * EXCEPT(date, time)\n",
        "    FROM `vs_logs_demo.hdfs_full_structured`\n",
        "  )\n",
        "  GROUP BY\n",
        "    blockId\n",
        "  ORDER BY\n",
        "    ts_start ASC, blockId ASC\n",
        ");"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DhANs7B7h_Eg"
      },
      "source": [
        "### Label log sequences"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xjRjHbkLawlP"
      },
      "source": [
        "This is a labelled dataset where each `blockId` has a binary label: either 'Normal' or 'Anomaly'. Let's retrieve the labels CSV file and store it in separate BigQuery table `hdfs_labels`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sbTbkLv5h_Eg"
      },
      "outputs": [],
      "source": [
        "gcs_uri = \"gs://csa-datasets-public/hdfs/anomaly_label.csv\"\n",
        "df_labels = bpd.read_csv(gcs_uri, engine=\"bigquery\", header=0)\n",
        "df_labels.columns = [\"blockId\", \"label\"]\n",
        "print(df_labels.info())\n",
        "\n",
        "df_labels.to_gbq(\"vs_logs_demo.hdfs_labels\", if_exists=\"replace\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ADUU9kcJh_Eg"
      },
      "source": [
        "Let's correlate the log sequences with the labels table, and create a new labelled table for log sequences by setting a new boolean column `abnormal`. We also capture the sequential order of rows in `sessionId` which will come in handy later on for splitting the dataset and batching calls to Vertex AI API given its rate limits."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iFmql46fh_Eh"
      },
      "outputs": [],
      "source": [
        "%%bigquery\n",
        "CREATE OR REPLACE TABLE `vs_logs_demo.hdfs_full_labelled` AS (\n",
        "SELECT\n",
        "  -- Capture sequential order of sessions\n",
        "  ROW_NUMBER() OVER(ORDER BY logs.ts_start ASC) AS sessionId,\n",
        "  logs.* EXCEPT(sequence),\n",
        "  CASE\n",
        "    WHEN labels.label = 'Anomaly' THEN TRUE\n",
        "    ELSE FALSE\n",
        "  END AS abnormal\n",
        "FROM `vs_logs_demo.hdfs_full_sessionized` logs\n",
        "LEFT JOIN `vs_logs_demo.hdfs_labels` labels USING (blockId)\n",
        ")\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A52AUFlzh_Eh"
      },
      "source": [
        "### Exploratory data analysis"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TgKum7IZh_Eh"
      },
      "outputs": [],
      "source": [
        "df = bpd.read_gbq(\"vs_logs_demo.hdfs_full_labelled\")\n",
        "df.shape"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nUtFdr5lh_Eh"
      },
      "outputs": [],
      "source": [
        "df.head(5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XwpnSSbjh_Eh"
      },
      "source": [
        "Let's see the distribution of log sequence length, by displaying summary stats for `numEvents` numeric column"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VZq7oaj9h_Ei"
      },
      "outputs": [],
      "source": [
        "df.describe()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OYm6XxT6VI1e"
      },
      "source": [
        "Let's the sample the data and get a breakdown of lengths of log sequences"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "RxNtzI0Ph_Ei"
      },
      "outputs": [],
      "source": [
        "bpd.options.sampling.enable_downsampling = True  # enable downsampling\n",
        "bpd.options.sampling.max_download_size = 200  # download only 200 mb of data"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dQ-ODwuHh_Ei"
      },
      "outputs": [],
      "source": [
        "local_df = df.to_pandas(\n",
        "    max_download_size=200, sampling_method=\"uniform\", random_state=101\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9yrQLumPh_Ej"
      },
      "outputs": [],
      "source": [
        "sns.countplot(x=\"abnormal\", data=local_df)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9u6ADdKwh_Ej"
      },
      "outputs": [],
      "source": [
        "plt.figure(figsize=(10, 6))\n",
        "\n",
        "# Plot with Seaborn (hue creates separate bars for each boolean value)\n",
        "sns.countplot(\n",
        "    data=local_df,\n",
        "    x=\"numEvents\",\n",
        "    hue=\"abnormal\",\n",
        "    palette=[\"blue\", \"red\"],\n",
        "    order=local_df[\"numEvents\"].value_counts().iloc[:15].index,\n",
        ")\n",
        "plt.title(\"Distribution of log sequence length by abnormality\")\n",
        "plt.xlabel(\"Events#/Sequence\")\n",
        "plt.ylabel(\"Count\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gzFpxV9Mh_Ej"
      },
      "source": [
        "## Create and store embeddings in vector index"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U1tK1ZSGh_Ej"
      },
      "source": [
        "### Explain log sequences\n",
        "Prompt Gemini 2.0 to translate log sequences into natural language"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l7A220hCh_Ek"
      },
      "source": [
        "#### Create the remote model for Gemini 2.0 in BigQuery"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8c6ac9ace0f4"
      },
      "outputs": [],
      "source": [
        "#"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6PJ71LkVh_Ek"
      },
      "outputs": [],
      "source": [
        "%%bigquery\n",
        "CREATE OR REPLACE MODEL `vs_logs_demo.gemini_1_5_flash`\n",
        "REMOTE WITH CONNECTION `us.bq-llm-connection`\n",
        "OPTIONS (endpoint = 'gemini-2.0-flash')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tejGKDOtcgOy"
      },
      "source": [
        "#### Prepare prompt"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "w7Rqy_tJc_aE"
      },
      "outputs": [],
      "source": [
        "# Prompt for Gemini\n",
        "prompt_prefix = f\"\"\"\\\n",
        "You're an IT expert and administrator of a Hadoop Distributed File System or HDFS cluster.\n",
        "You are provided with a sequence of HDFS logs in chronological order, where each log line contains the log source component, followed by the log severity and the log content itself.\n",
        "Describe succinctly in plain english the events corresponding to this sequence of logs.\n",
        "Do not repeat the raw logs in your explanation. Do not speculate, explain or make stuff up.\n",
        "Be concise in your answer and combine similar events together.\n",
        "Output must be formatted as an ordered numbered list and without any preamble.\n",
        "\"\"\"\n",
        "\n",
        "# One-shot example where the input is the raw log sequence for one session,\n",
        "# and the output is the corresponding log explanation in natural language.\n",
        "example_input = f\"\"\"\\\n",
        "dfs.DataNode$DataXceiver INFO Receiving block blk_4970672161979262840 src: /10.251.111.130:54916 dest: /10.251.111.130:50010\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.allocateBlock: /user/root/rand5/_temporary/_task_200811101024_0011_m_000844_0/part-00844. blk_4970672161979262840\n",
        "dfs.DataNode$DataXceiver INFO Receiving block blk_4970672161979262840 src: /10.251.66.102:49145 dest: /10.251.66.102:50010\n",
        "dfs.DataNode$DataXceiver INFO Receiving block blk_4970672161979262840 src: /10.251.111.130:34614 dest: /10.251.111.130:50010\n",
        "dfs.DataNode$PacketResponder INFO PacketResponder 1 for block blk_4970672161979262840 terminating\n",
        "dfs.DataNode$PacketResponder INFO PacketResponder 0 for block blk_4970672161979262840 terminating\n",
        "dfs.DataNode$PacketResponder INFO PacketResponder 2 for block blk_4970672161979262840 terminating\n",
        "dfs.DataNode$PacketResponder INFO Received block blk_4970672161979262840 of size 67108864 from /10.251.111.130\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.addStoredBlock: blockMap updated: 10.251.111.130:50010 is added to blk_4970672161979262840 size 67108864\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.addStoredBlock: blockMap updated: 10.251.66.102:50010 is added to blk_4970672161979262840 size 67108864\n",
        "dfs.DataNode$PacketResponder INFO Received block blk_4970672161979262840 of size 67108864 from /10.251.66.102\n",
        "dfs.DataNode$PacketResponder INFO Received block blk_4970672161979262840 of size 67108864 from /10.251.111.130\n",
        "dfs.FSNamesystem WARN BLOCK* NameSystem.addStoredBlock: Redundant addStoredBlock request received for blk_4970672161979262840 on 10.251.194.147:50010 size 67108864\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.addStoredBlock: blockMap updated: 10.251.194.147:50010 is added to blk_4970672161979262840 size 67108864\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.delete: blk_4970672161979262840 is added to invalidSet of 10.251.66.102:50010\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.delete: blk_4970672161979262840 is added to invalidSet of 10.251.111.130:50010\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.delete: blk_4970672161979262840 is added to invalidSet of 10.251.194.147:50010\n",
        "dfs.FSDataset INFO Deleting block blk_4970672161979262840 file /mnt/hadoop/dfs/data/current/subdir59/blk_4970672161979262840\n",
        "dfs.FSDataset INFO Deleting block blk_4970672161979262840 file /mnt/hadoop/dfs/data/current/subdir9/blk_4970672161979262840\n",
        "dfs.FSDataset INFO Deleting block blk_4970672161979262840 file /mnt/hadoop/dfs/data/current/subdir31/blk_4970672161979262840\n",
        "\"\"\"\n",
        "\n",
        "example_output = f\"\"\"\\\n",
        "1. The NameNode allocates a new block for a file in the user directory.\n",
        "2. Three DataNodes are receiving the block.\n",
        "3. The DataNodes packet responders for that block are terminating.\n",
        "4. The DataNodes receives the block.\n",
        "5. The NameNode updates block map after receiving confirmation that the block has been received by the three DataNodes.\n",
        "6. The NameNode reports a warning that it received a redundant confirmation from one of the DataNodes.\n",
        "7. The NameNode deletes the block and marks it as invalid on all three DataNodes.\n",
        "8. The DataNodes delete the block from their local storage.\n",
        "\"\"\"\n",
        "\n",
        "# Escape newlines before subsequent use in f-strings\n",
        "prompt_prefix = prompt_prefix.replace(\"\\n\", \"\\\\n\")\n",
        "example_input = example_input.replace(\"\\n\", \"\\\\n\")\n",
        "example_output = example_output.replace(\"\\n\", \"\\\\n\")\n",
        "\n",
        "params = {\n",
        "    \"prompt_prefix\": prompt_prefix,\n",
        "    \"example_input\": example_input,\n",
        "    \"example_output\": example_output,\n",
        "}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "J53myNZ5cMBn"
      },
      "source": [
        "Let's save the prompt and one-shot example in a new BigQuery table `hdfs_prompts`"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i8t9yIgRr9n-"
      },
      "outputs": [],
      "source": [
        "sql = f\"\"\"\\\n",
        "CREATE OR REPLACE TABLE `vs_logs_demo.hdfs_prompts` AS (\n",
        "  SELECT * FROM UNNEST([STRUCT(\n",
        "    'Explain HDFS log sequences' AS task,\n",
        "    \"{prompt_prefix}\" AS prompt_prefix,\n",
        "    \"{example_input}\" AS example_input,\n",
        "    \"{example_output}\" AS example_output\n",
        "  )])\n",
        ");\n",
        "\"\"\"\n",
        "\n",
        "query_job = client.query(sql)\n",
        "print(query_job.result())  # Wait for the job to complete.\n",
        "print(f\"Created Table {PROJECT_ID}.vs_logs_demo.hdfs_prompts\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2VnrZSeKZfHJ"
      },
      "source": [
        "As an example, let's take the following 'normal' raw log sequence for specific blockID `blk_-3750780870143089898`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wwdVnGOnZxeT"
      },
      "outputs": [],
      "source": [
        "dfs.DataNode$DataXceiver INFO Receiving block blk_-3750780870143089898 src: /10.251.67.211:46933 dest: /10.251.67.211:50010\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.allocateBlock: /user/root/randtxt5/_temporary/_task_200811101024_0012_m_000993_0/part-00993. blk_-3750780870143089898\n",
        "dfs.DataNode$DataXceiver INFO Receiving block blk_-3750780870143089898 src: /10.251.42.207:51200 dest: /10.251.42.207:50010\n",
        "dfs.DataNode$DataXceiver INFO Receiving block blk_-3750780870143089898 src: /10.251.67.211:50346 dest: /10.251.67.211:50010\n",
        "dfs.DataNode$PacketResponder INFO PacketResponder 0 for block blk_-3750780870143089898 terminating\n",
        "dfs.DataNode$PacketResponder INFO PacketResponder 2 for block blk_-3750780870143089898 terminating\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.addStoredBlock: blockMap updated: 10.251.42.207:50010 is added to blk_-3750780870143089898 size 67108864\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.addStoredBlock: blockMap updated: 10.251.67.211:50010 is added to blk_-3750780870143089898 size 67108864\n",
        "dfs.DataNode$PacketResponder INFO Received block blk_-3750780870143089898 of size 67108864 from /10.251.67.211\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.addStoredBlock: blockMap updated: 10.251.194.129:50010 is added to blk_-3750780870143089898 size 67108864\n",
        "dfs.DataNode$PacketResponder INFO PacketResponder 1 for block blk_-3750780870143089898 terminating\n",
        "dfs.DataNode$PacketResponder INFO Received block blk_-3750780870143089898 of size 67108864 from /10.251.67.211\n",
        "dfs.DataNode$PacketResponder INFO Received block blk_-3750780870143089898 of size 67108864 from /10.251.42.207\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.delete: blk_-3750780870143089898 is added to invalidSet of 10.251.42.207:50010\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.delete: blk_-3750780870143089898 is added to invalidSet of 10.251.67.211:50010\n",
        "dfs.FSNamesystem INFO BLOCK* NameSystem.delete: blk_-3750780870143089898 is added to invalidSet of 10.251.194.129:50010\n",
        "dfs.FSNamesystem INFO BLOCK* ask 10.251.194.129:50010 to delete  blk_-3750780870143089898 blk_8839488796462659348 blk_-1076638448786103147 blk_-2607303556568476338 blk_-872077376619694315 blk_3658484321727696675 blk_-8783295298979970351 blk_7253228695508151434 blk_8353593257930202451 blk_6896918372378359715 blk_8476431461885450664 blk_-453468701147780383 blk_2747276825479054342 blk_7831811674826419739 blk_-4676554322964835078 blk_838168686993939059 blk_-3667727247990911503 blk_554950885642981306 blk_3690236354878191248 blk_3594952792786303214 blk_710500522210229744 blk_-641656950218902981 blk_-8914315030958124984 blk_8473337661129085162 blk_8747565120239750968 blk_8483881408521191191 blk_970871787748196992 blk_-1238776083272375823 blk_-4355328787228390393 blk_-6232254141637974415 blk_2851445563013723058 blk_4729371515425302371 blk_-8138185632733985405 blk_3530286662170963267 blk_6705596331043244121 blk_5856355648225187587 blk_-5271003563714979606 blk_6099159980310485523 blk_5770736739131705773 blk_-4630334698677080743 blk_-316813592070079363 blk_-8376103172399604209 blk_-5004999163141797571 blk_-1184480050082523451 blk_-8644451221914479826 blk_-4421425625682059690 blk_8309579147878024429 blk_3172294788463323663 blk_-426829215056729544 blk_-785413567218674273 blk_2115841663373157801 blk_2855785622370264073 blk_8386172561326404060 blk_6747728035734009764 blk_7206820189825343828 blk_9222274141331127342 blk_-7791132919862479353 blk_-1813922770106404412 blk_-6794746648243431237 blk_-3114372317175128819 blk_-5375459594742433201 blk_6912573532770191892 blk_124792821150689426 blk_8787609558519171611 blk_-5450607639120667562 blk_769133640763536215 blk_4689006757714287009 blk_-2322954586860795404 blk_-3083400559019231217 blk_-2398840409229247956 blk_4606567481735284177 blk_-1098865217826643562 blk_-3075042532557633130 blk_6737487517130970246 blk_3299232526334609785 blk_-1866148958272946201 blk_4803758272177687030 blk_2068490022276187676 blk_-4724682685867818017 blk_-1311156078103024510 blk_-1781420834572645132 blk_-8680808247265383818 blk_7865294689332645011 blk_-7864388093305159311 blk_-7259235578951211339 blk_4437793449639312068 blk_2972756848238658967 blk_160665971665268866 blk_7862577900146719312 blk_2141823066215925124 blk_4044494381360076828 blk_868975602193282101 blk_-8493544210473294093 blk_-3909434520726038447 blk_-2397800658299122228 blk_-6929912433144478107 blk_-2022835339724037507 blk_-251514397381670680 blk_-4491187364929180193 blk_-4766279731310845414\n",
        "dfs.FSDataset INFO Deleting block blk_-3750780870143089898 file /mnt/hadoop/dfs/data/current/subdir32/blk_-3750780870143089898\n",
        "dfs.FSDataset INFO Deleting block blk_-3750780870143089898 file /mnt/hadoop/dfs/data/current/subdir16/blk_-3750780870143089898\n",
        "dfs.FSDataset INFO Deleting block blk_-3750780870143089898 file /mnt/hadoop/dfs/data/current/subdir55/blk_-3750780870143089898"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fbR3E9SCaOsD"
      },
      "source": [
        "The Gemini prompt above converts those raw log entries into a sequence of events in natural language which captures concisely the execution path including detailed semantics information while excluding specific identifiers:\n",
        "\n",
        "```\n",
        "1. The NameNode allocates a new block for a file in the user directory.\n",
        "2. Three DataNodes are receiving the block.\n",
        "3. The NameNode updates block map after receiving confirmation that the block has been received by the three DataNodes.\n",
        "4. The NameNode asks one of the DataNodes to delete the block.\n",
        "5. The DataNodes delete the block from their local storage.\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GA6jBUfF59yX"
      },
      "source": [
        "#### Run on subset of logs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OPRRTIHjNm86"
      },
      "source": [
        "The following **may take 6 hours or less** depending on your Vertex AI quota of requests per minute for `gemini-2.0-flash` base model. The following query goes over a small subset of the dataset: 70k log sequences will be translated into natural language summary. At the time of this writing, the default rate limit for `gemini-2.0-flash` in `us-central1` region is **200 requests/min**. To speed up this step, you can [request quota increase](https://cloud.google.com/vertex-ai/generative-ai/docs/quotas#view-quotas-in-console) for Vertex AI endpoint for `gemini-2.0-flash` for your specific region and project. If you do increase that Vertex AI quota, send an email to bqml-feedback@google.com to adjust and increase your BigQuery ML quota for `ML.GENERATE_TEXT` calls to that base model, since BigQuery ML rate-limits the calls to Vertex AI endpoint accordingly.\n",
        "\n",
        "The following SQL code firsts create the table `hdfs_full_explained`, then creates a procedure to generate content from Gemini and inserts into the destination table. You call the procedure iteratively to process data in batches and avoid losing processed results due to runtime error like query timeout or quota exhaustion, both of which are unlikely with smaller batches."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LoCQyX7sMhvX"
      },
      "outputs": [],
      "source": [
        "%%bigquery\n",
        "CREATE OR REPLACE TABLE `vs_logs_demo.hdfs_full_explained`\n",
        "(\n",
        "  sessionId INT64,\n",
        "  blockId STRING,\n",
        "  abnormal BOOL,\n",
        "  numEvents INT64,\n",
        "  eventSequence STRING,\n",
        "  eventTemplateSequence STRING,\n",
        "  status STRING,\n",
        "  response STRING,\n",
        "  input_tokens INT64,\n",
        "  output_tokens INT64,\n",
        "  total_token_count INT64\n",
        ")\n",
        "PARTITION BY RANGE_BUCKET(sessionId, GENERATE_ARRAY(0, 600000, 10000))\n",
        "CLUSTER BY abnormal, numEvents;"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ye-yBze9f8Kh"
      },
      "outputs": [],
      "source": [
        "params[\"prompt_prefix\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o8jZ9EnLMs46"
      },
      "outputs": [],
      "source": [
        "%%bigquery\n",
        "CREATE OR REPLACE PROCEDURE `vs_logs_demo.explain_hdfs_logs_test`(\n",
        "  start_row INT64,\n",
        "  batch_size INT64\n",
        ")\n",
        "BEGIN\n",
        "  INSERT INTO `vs_logs_demo.hdfs_full_explained`\n",
        "  SELECT\n",
        "    sessionId, blockId, abnormal, numEvents, eventSequence, eventTemplateSequence,\n",
        "    ml_generate_text_status AS status,\n",
        "    LAX_STRING(ml_generate_text_result.candidates[0].content.parts[0].text) AS response,\n",
        "    LAX_INT64(ml_generate_text_result.usage_metadata.prompt_token_count) AS input_tokens,\n",
        "    LAX_INT64(ml_generate_text_result.usage_metadata.candidates_token_count) AS output_tokens,\n",
        "    LAX_INT64(ml_generate_text_result.usage_metadata.total_token_count) AS total_token_count\n",
        "  FROM\n",
        "    ML.GENERATE_TEXT(\n",
        "      MODEL `vs_logs_demo.gemini_1_5_flash`,\n",
        "      (\n",
        "        SELECT\n",
        "          -- A prompt for HDFS logs explanation with one-shot example.\n",
        "          sessionId, blockId, abnormal, numEvents, eventSequence, eventTemplateSequence,\n",
        "          CONCAT(llm.prompt_prefix,\n",
        "            \"Here is one-shot example:\\n\",\n",
        "            \"Log sequence:\\n\", llm.example_input,\n",
        "            \"Log sequence summary:\\n\", llm.example_output,\n",
        "            \"Here is the input:\\n\",\n",
        "            \"Log sequence: \", eventSequence, \"\\n\",\n",
        "            \"Answer: \"\n",
        "        ) AS prompt\n",
        "        FROM `vs_logs_demo.hdfs_full_labelled`, `vs_logs_demo.hdfs_prompts` AS llm\n",
        "        WHERE\n",
        "          llm.task = \"Explain HDFS log sequences\"\n",
        "          AND sessionId BETWEEN start_row AND start_row + batch_size -1\n",
        "      ),\n",
        "      STRUCT(\n",
        "        0.2 AS temperature,\n",
        "        2048 AS max_output_tokens, -- limit response to ~1400 words\n",
        "        FALSE AS flatten_json_output\n",
        "      )\n",
        "    );\n",
        "END;"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Msl7KJBi90Nc"
      },
      "outputs": [],
      "source": [
        "%%bigquery\n",
        "-- Process 70k sequences with batches of 20k at a time to stay well within #rows/job quota\n",
        "-- for ML.GENERATE_TEXT (72k with Gemini 2.0, 21.6k with Gemini 2.0)\n",
        "-- https://cloud.google.com/bigquery/quotas#cloud_ai_service_functions\n",
        "BEGIN\n",
        "  CALL `vs_logs_demo.explain_hdfs_logs`(1, 20000);\n",
        "  CALL `vs_logs_demo.explain_hdfs_logs`(20001, 20000);\n",
        "  CALL `vs_logs_demo.explain_hdfs_logs`(40001, 20000);\n",
        "  CALL `vs_logs_demo.explain_hdfs_logs`(60001, 10000);\n",
        "END;"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0H5-RmNwZGBp"
      },
      "source": [
        "To process even larger datasets, you can use the procedure iteratively  sequentially while accounting for the quota limit of 72k rows per job.\n",
        "Note the cumulative time limit for a multi-statement query in BigQuery is 24 hours. If the above code times out (unlikely with small 20k batches for a total of 70k), simply comment out the successfully completed batches before running it again so it can resume after that batch."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sdGQq2Rwh_Ek"
      },
      "source": [
        "### Generate embeddings for log sequences"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EjbNVI1Fh_El"
      },
      "source": [
        "#### Create the remote model for generating text embeddings"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "01yw0xUoh_El"
      },
      "outputs": [],
      "source": [
        "%%bigquery\n",
        "CREATE OR REPLACE MODEL `vs_logs_demo.text_embedding`\n",
        "REMOTE WITH CONNECTION `us.bq-llm-connection`\n",
        "OPTIONS (endpoint = 'text-embedding-005')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gZnn6PJXU5sY"
      },
      "source": [
        "#### Generate embeddings"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EbUr8Zvbddwn"
      },
      "source": [
        "Let's generate vector embedding for each log sequence to capture the semantic meaning of the Gemini-produced sequence of operations over each `blockId`.\n",
        "\n",
        "The following **will take several minutes** depending on your quota of requests per minute for `text-embedding` base model. At the time of this writing, the default rate limit for `text-embedding` in `us-central1` region is 1500 requests/min.\n",
        "\n",
        "The following query processes the entire `hdfs_full_explained` table. Unlike the previous section, where `ML.GENERATE_TEXT` was run in iterative batches, we run `ML.GENERATE_TEXT_EMBEDDING` as one job given the higher limit of 2,700,000 rows per job, that is well above our dataset of 70,000 session records.\n",
        "\n",
        "\n",
        "**Important Note:** we use `CLUSTERING` embedding task type which is best suited for (clustering-based) anomaly detection task, as opposed to other tasks like document retrieval (RAG) or semantic similarity search."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_qD7SJeeh_El"
      },
      "outputs": [],
      "source": [
        "%%bigquery\n",
        "CREATE OR REPLACE TABLE `vs_logs_demo.hdfs_full_embeddings` AS (\n",
        "  SELECT\n",
        "    * EXCEPT(ml_generate_embedding_result, ml_generate_embedding_status, ml_generate_embedding_statistics),\n",
        "    ml_generate_embedding_status AS status,\n",
        "    LAX_INT64(ml_generate_embedding_statistics.token_count) AS token_count,\n",
        "    LAX_BOOL(ml_generate_embedding_statistics.truncated) AS truncated,\n",
        "    ml_generate_embedding_result AS embeddings\n",
        "  FROM\n",
        "    ML.GENERATE_EMBEDDING(\n",
        "      MODEL `vs_logs_demo.text_embedding`,\n",
        "      (\n",
        "        SELECT\n",
        "          # sessionId, timeInterval, numEvents, sequence, abnormal,\n",
        "          sessionId, blockId, abnormal, numEvents, eventSequence, eventTemplateSequence,\n",
        "          response as content\n",
        "        FROM `vs_logs_demo.hdfs_full_explained`\n",
        "        -- skip records where previous ml.generate_text step failed due to quota exhaustion\n",
        "        WHERE status=\"\"\n",
        "        ORDER BY sessionId ASC\n",
        "        LIMIT 70000\n",
        "      ),\n",
        "      STRUCT(\n",
        "          TRUE AS flatten_json_output,\n",
        "          'CLUSTERING' as task_type\n",
        "      )\n",
        "    )\n",
        ")\n",
        "PARTITION BY RANGE_BUCKET(sessionId, GENERATE_ARRAY(0, 600000, 10000))\n",
        "CLUSTER BY abnormal, numEvents;"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q3hRmLiY2O0w"
      },
      "source": [
        "### Create vector index"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bH9FNwlfbUmv"
      },
      "outputs": [],
      "source": [
        "%%bigquery\n",
        "CREATE VECTOR INDEX index_treeah_1000 ON `vs_logs_demo.hdfs_full_embeddings`(embeddings)\n",
        "# STORING (abnormal, sessionId, numEvents)\n",
        "OPTIONS (index_type = 'TREE_AH',\n",
        "  distance_type = 'COSINE',\n",
        "  tree_ah_options = '{\"leaf_node_embedding_count\": 1000}');"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EGO4QPhZh_Em"
      },
      "source": [
        "Inspect vector index progress, and total storage (chargeable) from INFORMATION_SCHEMA.VECTOR_INDEXES view:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IOmgcRWfh_Em"
      },
      "outputs": [],
      "source": [
        "%%bigquery\n",
        "SELECT\n",
        "  table_name,\n",
        "  index_name,\n",
        "  coverage_percentage,\n",
        "  unindexed_row_count,\n",
        "  last_refresh_time,\n",
        "  total_logical_bytes,\n",
        "  total_storage_bytes\n",
        "FROM `vs_logs_demo.INFORMATION_SCHEMA.VECTOR_INDEXES`\n",
        "WHERE index_status = 'ACTIVE';"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ul3FaK8Ph_Em"
      },
      "source": [
        "## Visualize clusters of log sequences"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c6qQx9zHlXoJ"
      },
      "source": [
        "As log sequences are projected into the 768-dimension vector space, similar log sequences are clustered as \"neighbors\" in the sense they share a similar sequence of what actually occurred.\n",
        "\n",
        "Let's do a first-order evaluation to see how the clustering is performing qualitatively. Specifically, we're looking at how well-defined the clusters are and how much separation there is between clusters."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XYdXKXd1fzoy"
      },
      "source": [
        "We use dimensionality reduction for visualizing in 2D space the clusters of HDFS log sequences embeddings by going from 768 to 2 dimension vectors."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PIjQeAU78HeB"
      },
      "source": [
        "### Visualize embeddings"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nM03efDp8GPk"
      },
      "outputs": [],
      "source": [
        "client = bigquery.Client()\n",
        "\n",
        "# sql = f\"\"\"\n",
        "# SELECT *\n",
        "# FROM `vs_logs_demo.hdfs_full_embeddings`\n",
        "# TABLESAMPLE SYSTEM (10 PERCENT)\n",
        "# \"\"\"\n",
        "# df = client.query(sql).to_dataframe()\n",
        "\n",
        "table_id = \"vs_logs_demo.hdfs_full_embeddings\"\n",
        "df = client.list_rows(table_id).to_dataframe()\n",
        "df.head(2)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JsVt6qXph_Eo"
      },
      "outputs": [],
      "source": [
        "print(df.iloc[70].content)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4PCddtSs8GPl"
      },
      "outputs": [],
      "source": [
        "# Reduce dimensionality using t-SNE. Will take 5 minutes with 10% of the data\n",
        "tsne = TSNE(n_components=2, random_state=0)\n",
        "embeddings_array = np.array(df[\"embeddings\"].to_list(), dtype=np.float32)\n",
        "embeddings_2d = tsne.fit_transform(embeddings_array)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wEx5o4a07Dxd"
      },
      "outputs": [],
      "source": [
        "# Mapping boolean values to colors\n",
        "color_map = {True: \"red\", False: \"blue\"}\n",
        "colors = df[\"abnormal\"].map(color_map)\n",
        "\n",
        "# Create a scatter plot of the embeddings\n",
        "plt.figure(figsize=(10, 8))\n",
        "plt.scatter(embeddings_2d[:, 0], embeddings_2d[:, 1], c=colors, s=25, alpha=0.7)\n",
        "\n",
        "plt.title(\"Log Embeddings Clusters\")\n",
        "plt.xlabel(\"TSNE-1\")\n",
        "plt.ylabel(\"TSNE-2\")\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ecf8102c85e7"
      },
      "source": [
        "Results should be rendered in a scatter plot as follows:"
      ]
    },
    {
      "attachments": {
        "830d2fb9-2eb8-481f-b3ff-47129a741321.png": {
          "image/png": "[encoded data, 261416 characters]"
        }
      },
      "cell_type": "markdown",
      "metadata": {
        "id": "394a7829ce06"
      },
      "source": [
        "![image.png](attachment:830d2fb9-2eb8-481f-b3ff-47129a741321.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EWDzewJyiCHc"
      },
      "source": [
        "We can visually see many blue clusters (30-50), depicting different normal types of log sequences. The red dots are the abnormal log sequences that ultimately should be identified as outliers. While many of the red dots are clustered together, some are scattered and at the periphery of blue clusters."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BVy4K8xnh_Ew"
      },
      "source": [
        "## Detect outliers with BigQuery vector search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dyCvM3ssh_E0"
      },
      "source": [
        "Let's implement a simple outlier detection logic leveraging BQ built-in optimized vector search using simple SQL."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YTp-jko3ln7B"
      },
      "source": [
        "### Split dataset into training and test"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7CBZnNDOmkHH"
      },
      "source": [
        "First, let's split the dataset into training and test datasets in BigQuery.\n",
        "\n",
        "We use a session cutoff to split the first 50,000 sessions (ordered chronologically) to use for 'training' and then a subsequent 20,000 sessions for testing. Since this is semi-supervised learning on the negative class only (normal logs), we use only the normal logs sequences in the first 50,000 sessions.\n",
        "\n",
        "In the context of BigQuery vector search, the vector index uses a k-means algorithm to cluster these records as they are ingested. That process is managed by BigQuery during and after index creation. We still writing the training and test dataset in separate table to facilitate testing of other ML models in subsequent section."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IWM6FBBDYadc"
      },
      "outputs": [],
      "source": [
        "session_id_cutoff = 50000  # Split train and eval datasets\n",
        "num_sessions = 20000  # Number of sessions to use for testing"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gVH6LbLBpRuw"
      },
      "source": [
        "Retrieve and write test and training datasets in separate tables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1D4PINTvTe26"
      },
      "outputs": [],
      "source": [
        "# Train dataset\n",
        "sql = f\"\"\"\n",
        "SELECT * FROM `vs_logs_demo.hdfs_full_embeddings`\n",
        "WHERE\n",
        "  abnormal=FALSE\n",
        "  AND sessionId <= {session_id_cutoff}\n",
        "\"\"\"\n",
        "y_train = bpd.read_gbq(sql, index_col=\"sessionId\")\n",
        "\n",
        "# Persist the train dataset in a separate table to be used later for other models e.g. scikit-learn algor\n",
        "y_train.to_gbq(\"vs_logs_demo.hdfs_full_embeddings_train\", if_exists=\"replace\")\n",
        "\n",
        "# Test dataset\n",
        "sql = f\"\"\"\n",
        "SELECT\n",
        "  sessionId, blockId, embeddings,\n",
        "  CAST(abnormal AS INT64) AS label\n",
        "FROM `vs_logs_demo.hdfs_full_embeddings`\n",
        "WHERE sessionId > {session_id_cutoff}\n",
        "ORDER BY sessionId ASC\n",
        "LIMIT {num_sessions}\n",
        "\"\"\"\n",
        "y_test = bpd.read_gbq(sql, index_col=\"sessionId\")\n",
        "\n",
        "# Persist the test dataset in a separate table to be used several times with different vector search configurations\n",
        "y_test.to_gbq(\"vs_logs_demo.hdfs_full_embeddings_test\", if_exists=\"replace\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hRS72B88h_E3"
      },
      "source": [
        "### Predict anomalies with vector search"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aY4PKVQy2MDJ"
      },
      "source": [
        "The premise behind outlier detection is to identify log sequences that do not belong to any known 'normal' blue clusters. Given a log sequence test vector, if there is a sufficient number of 'normal' neighboring base vectors within a certain distance, then that log sequence is considered 'normal'. Otherwise, it is considered an outlier or an anomaly."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d2iVXlbdWFSm"
      },
      "source": [
        " To get the best performance in terms of outlier detection, the following procedure uses a brute force search to maximize recall of top-k closest 'normal' neighbors for every query vector."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TEkbws0mh_E3"
      },
      "outputs": [],
      "source": [
        "def predict_anomalies(n_neighbors, distance_threshold):\n",
        "    sql = f\"\"\"\n",
        "    SELECT\n",
        "      query.blockId AS blockId,\n",
        "      ANY_VALUE(query.sessionId) AS sessionId,\n",
        "      COUNTIF(distance < {distance_threshold}) AS similar_normal_instances,\n",
        "      IF(COUNTIF(distance < {distance_threshold}) >= {n_neighbors}, 0, 1) as predicted\n",
        "    FROM VECTOR_SEARCH(\n",
        "      (SELECT * FROM `vs_logs_demo.hdfs_full_embeddings` WHERE abnormal=FALSE AND sessionId <= {session_id_cutoff}), # TreeAH index\n",
        "      'embeddings',\n",
        "      TABLE `vs_logs_demo.hdfs_full_embeddings_test`,\n",
        "      --(SELECT blockId, sessionId, embeddings FROM `vs_logs_demo.hdfs_full_embeddings` WHERE sessionId > {session_id_cutoff} ORDER BY sessionId ASC LIMIT {num_sessions}),\n",
        "      top_k => {n_neighbors}, # Increase top_k value to 3x if using TreeAH with post-filtering\n",
        "      distance_type => 'COSINE',\n",
        "      options => '{{\"use_brute_force\":true}}'\n",
        "      --options => '{{\"fraction_lists_to_search\":0.1}}'\n",
        "    )\n",
        "    GROUP BY\n",
        "      query.blockId\n",
        "    \"\"\"\n",
        "    return bpd.read_gbq(sql, index_col=\"sessionId\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cZsxyyQnYq6s"
      },
      "source": [
        "**Note**: While the above procedure runs a brute force job for best outlier detection performance, you can still experiment with using the vector index by setting `fraction_lists_to_search` to optimize speed and cost while trading off some recall. For example, you could set it to 0.1 to only search only 10% of the clusters, those with the closest centroids, thereby reducing total latency and slot usage. Since we're using a TreeAH index, the prefilters on the base table are treated as post-filters and may reduce the top-k result set. In that case, make sure to multiply top_k value by say 3x to have a sufficient result set post-filtering."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T91-nLT25rpH"
      },
      "source": [
        "Let's run the batch prediction: this will compare 20,000 test records against 47,770 'normal' base records, essentially doing 955 Millions distance calculations of 768-dimentional vectors. The following completes in **only 1 minute** given the small base table and sufficient parallel shards."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jv4n_N51h_E3"
      },
      "outputs": [],
      "source": [
        "n_neighbors = 10  # Number of neighbors to compare with\n",
        "dist_threshold = 0.004\n",
        "# Neighbor within this distance is a match\n",
        "\n",
        "# Batch predict on the test set\n",
        "y_pred = predict_anomalies(n_neighbors, dist_threshold)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "welzwxwC6gW6"
      },
      "source": [
        "Let's see how much data was processed and the total time it took to predict anomalies across 20,000 test records. Replace `insert-job-id` with the BigQuery job ID returned in the last cell output."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7ZaKLkSHh_E4"
      },
      "outputs": [],
      "source": [
        "vs_job_id = \"[insert-job-id]\"\n",
        "\n",
        "client = bigquery.Client()\n",
        "job = client.get_job(vs_job_id)\n",
        "\n",
        "print(f\"Job ID: {vs_job_id}\")\n",
        "print(\n",
        "    f\"Total Bytes Processed: {round(job.total_bytes_processed / (1024 * 1024) , 2) } MB\"\n",
        ")\n",
        "print(f\"Slot-hours Used: {round(job.slot_millis / (1000 * 60 * 60), 2)}\")\n",
        "\n",
        "start_time = job.started\n",
        "end_time = job.ended\n",
        "execution_time = (end_time - start_time).total_seconds()\n",
        "print(f\"Query Execution Time: {execution_time:.2f} seconds\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2fhhcrCF6-gC"
      },
      "source": [
        "Make sure the predicted result set `y_pred` is of the same size as the test dataset `y_test`, that is 20,000 records:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_J7BqRzJh_E4"
      },
      "outputs": [],
      "source": [
        "print(y_test.shape)\n",
        "print(y_pred.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qmZ1waubh_E5"
      },
      "source": [
        "### Evaluate performance"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kUxPZ0x-ASfL"
      },
      "outputs": [],
      "source": [
        "from sklearn.metrics import (\n",
        "    ConfusionMatrixDisplay,\n",
        "    accuracy_score,\n",
        "    classification_report,\n",
        "    confusion_matrix,\n",
        "    f1_score,\n",
        "    recall_score,\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AlrYy0Mch_E5"
      },
      "outputs": [],
      "source": [
        "cnf_matrix = confusion_matrix(\n",
        "    y_test[\"label\"].to_pandas(), y_pred[\"predicted\"].to_pandas()\n",
        ")\n",
        "print(cnf_matrix)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EcXfacuEAkf9"
      },
      "outputs": [],
      "source": [
        "disp = ConfusionMatrixDisplay(\n",
        "    confusion_matrix=cnf_matrix, display_labels=[\"normal\", \"abnormal\"]\n",
        ")\n",
        "disp.plot()\n",
        "plt.show()"
      ]
    },
    {
      "attachments": {
        "3823d730-4998-48f5-86c8-874384d16f68.png": {
          "image/png": "[encoded data, 32084 characters]"
        }
      },
      "cell_type": "markdown",
      "metadata": {
        "id": "1b7d5522cd39"
      },
      "source": [
        "You should see the following confusion matrix rendered:\n",
        "\n",
        "![image.png](attachment:3823d730-4998-48f5-86c8-874384d16f68.png)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "88Nw7619Akf-"
      },
      "outputs": [],
      "source": [
        "target_names = [\"normal\", \"abnormal\"]\n",
        "report = classification_report(\n",
        "    y_test[\"label\"].to_pandas(),\n",
        "    y_pred[\"predicted\"].to_pandas(),\n",
        "    target_names=target_names,\n",
        ")\n",
        "print(\"Classification Report:\\n\", report)"
      ]
    },
    {
      "attachments": {
        "9f900d6f-09c7-479a-b93d-c441e1a97dec.png": {
          "image/png": "[encoded data, 17752 characters]"
        }
      },
      "cell_type": "markdown",
      "metadata": {
        "id": "19c10f1027b6"
      },
      "source": [
        "You should observe the following classification metrics:\n",
        "\n",
        "![image.png](attachment:9f900d6f-09c7-479a-b93d-c441e1a97dec.png)    "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7d67bcafbbf8"
      },
      "source": [
        "In particular, results for 'abnormal' show 92% recall and 90% precision:\n",
        "\n",
        "|  | precision |   recall | f1-score | support |\n",
        "|---|---|---|---|---|\n",
        "| normal    |   1.00   |   1.00  |    1.00   |  19169 |\n",
        "| abnormal    |   0.90    |  0.92  |    0.91   |    831 |"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GZJ6u6Bih_E5"
      },
      "source": [
        "#### Grid search to tune hyper-parameters\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1aw_kVg0h_E6"
      },
      "outputs": [],
      "source": [
        "from sklearn.metrics import f1_score, precision_score, recall_score\n",
        "\n",
        "# Define the parameter grid for BQ vector search (adjust values as needed)\n",
        "param_grid = {\n",
        "    \"n_neighbors\": [6, 10, 12, 15],  # Number of neighbors to consider\n",
        "    \"dist_threshold\": [0.004, 0.006, 0.008, 0.01],  # Distance threshold for a match\n",
        "}\n",
        "\n",
        "# Initialize variables to track the best score and parameters\n",
        "best_f1 = 0\n",
        "best_recall = 0\n",
        "best_params = {}\n",
        "\n",
        "y_test_pd = y_test[\"label\"].to_pandas()\n",
        "\n",
        "# Manually iterate through parameter combinations\n",
        "for n_neighbors in param_grid[\"n_neighbors\"]:\n",
        "    for dist_threshold in param_grid[\"dist_threshold\"]:\n",
        "        # No need to train any model - simply do a vector search\n",
        "\n",
        "        # Make predictions on the test set\n",
        "        y_pred_df = predict_anomalies(n_neighbors, dist_threshold)\n",
        "        y_pred_pd = y_pred_df[\"predicted\"].to_pandas()\n",
        "\n",
        "        # Calculate F1-score\n",
        "        f1 = f1_score(y_test_pd, y_pred_pd)\n",
        "        # Calculate Recall\n",
        "        recall = recall_score(y_test_pd, y_pred_pd)\n",
        "        # Calculate Precision\n",
        "        precision = precision_score(y_test_pd, y_pred_pd)\n",
        "\n",
        "        # Update f1 score and parameters if current combination is better\n",
        "        if f1 > best_f1:\n",
        "            best_f1 = f1\n",
        "            best_params = {\"n_neighbors\": n_neighbors, \"dist_threshold\": dist_threshold}\n",
        "\n",
        "        print(\n",
        "            f\"Parameters: n_neighbors={n_neighbors}, dist_threshold={dist_threshold}, Recall: {recall}, Precision: {precision}, F1-Score: {f1}\"\n",
        "        )\n",
        "\n",
        "# Print the best parameters and score\n",
        "print(\"Best Parameters:\", best_params)\n",
        "print(\"Best F1-Score:\", best_f1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X4VepQhgh_Er"
      },
      "source": [
        "## Detect outliers with scikit-learn algorithms"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KqZSOjAE8hL3"
      },
      "source": [
        "As a baseline comparison, we'll evaluate popular [outlier detection algorithms](https://scikit-learn.org/stable/modules/outlier_detection.html#overview-of-outlier-detection-methods) implemented in scikit-learn, using the same logs embeddings that were generated and stored in BigQuery. [One-Class SVM](https://scikit-learn.org/stable/modules/outlier_detection.html#novelty-detection) and [Local Outlier Factor](https://scikit-learn.org/stable/modules/outlier_detection.html#novelty-detection-with-local-outlier-factor) are both semi-supervised learning methods suitable for novelty detection, and are available in the `svm` and `neighbors` modules respectively:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iA2joe7Uh_Er"
      },
      "outputs": [],
      "source": [
        "from sklearn import neighbors, svm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "25svuP1jovxf"
      },
      "source": [
        "### Use binary classifiers (semi-supervised on negative class only)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oAUh7rFPovxg"
      },
      "source": [
        "We will use the same training and test datasets persisted in BigQuery tables"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g2djOrYJxqsa"
      },
      "outputs": [],
      "source": [
        "# read a BigQuery table to a BigQuery DataFrame\n",
        "df_train = bpd.read_gbq(\n",
        "    \"vs_logs_demo.hdfs_full_embeddings_train\", index_col=\"sessionId\"\n",
        ")\n",
        "\n",
        "# X_train = list(df_train['embeddings'].to_list())\n",
        "X_train = df_train[[\"embeddings\"]].to_list()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gtwOE_j5sdt_"
      },
      "outputs": [],
      "source": [
        "df_test = bpd.read_gbq(\"vs_logs_demo.hdfs_full_embeddings_test\", index_col=\"sessionId\")\n",
        "\n",
        "X_test = df_test[\"embeddings\"].to_list()\n",
        "y_test = df_test[\"label\"].to_list()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YxLPlWO5ovxg"
      },
      "outputs": [],
      "source": [
        "print(np.array(X_train).shape)\n",
        "print(np.array(X_test).shape)\n",
        "print(np.array(y_test).shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O2GoQSEQovxg"
      },
      "source": [
        "#### Use Novelty Detection with One-Class SVM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3FzknanWovxg"
      },
      "outputs": [],
      "source": [
        "# Train a single-class classifier (e.g., One Class SVM)\n",
        "clf = svm.OneClassSVM(\n",
        "    nu=0.02, kernel=\"rbf\", gamma=\"auto\"\n",
        ")  # Adjust parameters as needed (highest recall, then f1-score)\n",
        "# clf = svm.OneClassSVM(nu=0.01, kernel=\"poly\", gamma='scale')  # Adjust parameters as needed (highest f1-score)\n",
        "clf.fit(X_train)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BJJW8JJtovxg"
      },
      "outputs": [],
      "source": [
        "# Make predictions on the test set (consider outliers as the positive class)\n",
        "y_pred = clf.predict(X_test)\n",
        "y_pred = [1 if y == -1 else 0 for y in y_pred]  # Convert -1 (outlier) to 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Qmp0u021ovxh"
      },
      "outputs": [],
      "source": [
        "cnf_matrix = confusion_matrix(y_test, y_pred)\n",
        "cnf_matrix"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gUEOX79Sovxh"
      },
      "outputs": [],
      "source": [
        "disp = ConfusionMatrixDisplay(\n",
        "    confusion_matrix=cnf_matrix, display_labels=[\"normal\", \"abnormal\"]\n",
        ")\n",
        "disp.plot()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V24LufrGovxh"
      },
      "outputs": [],
      "source": [
        "# Evaluate the classifier\n",
        "accuracy = accuracy_score(y_test, y_pred)\n",
        "target_names = [\"normal\", \"abnormal\"]\n",
        "report = classification_report(y_test, y_pred, target_names=target_names)\n",
        "\n",
        "print(f\"Accuracy: {accuracy}\")\n",
        "print(\"Classification Report:\\n\", report)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aNNvclFYovxh"
      },
      "source": [
        "Let's do some naive hyperparameter tuning"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y4UmwsRgovxh"
      },
      "outputs": [],
      "source": [
        "# Define the parameter grid for hyperparameter tuning\n",
        "param_grid = {\n",
        "    \"nu\": [0.01, 0.02],  # Regularization parameter (controls the fraction of outliers)\n",
        "    \"kernel\": [\"linear\", \"rbf\", \"poly\"],  # Kernel type\n",
        "    \"gamma\": [\"scale\", \"auto\"],  # Kernel coefficient for 'rbf', 'poly', and 'sigmoid'\n",
        "    # 'gamma': ['scale', 'auto', 0.1, 1, 10]  # Kernel coefficient for 'rbf', 'poly', and 'sigmoid'\n",
        "}\n",
        "\n",
        "# Initialize variables to track the best score and parameters\n",
        "best_f1_score = 0\n",
        "best_params = {}\n",
        "\n",
        "# Manually iterate through parameter combinations\n",
        "for nu in param_grid[\"nu\"]:\n",
        "    for kernel in param_grid[\"kernel\"]:\n",
        "        for gamma in param_grid[\"gamma\"]:\n",
        "            # Create and train the OneClassSVM classifier\n",
        "            clf = svm.OneClassSVM(nu=nu, kernel=kernel, gamma=gamma)\n",
        "            clf.fit(X_train)\n",
        "\n",
        "            # Make predictions on the test set\n",
        "            y_pred = clf.predict(X_test)\n",
        "            y_pred = [1 if y == -1 else 0 for y in y_pred]\n",
        "\n",
        "            # Calculate F1-score\n",
        "            f1 = f1_score(y_test, y_pred)\n",
        "            # Calculate Recall\n",
        "            recall = recall_score(y_test, y_pred)\n",
        "            # Calculate Precision\n",
        "            precision = precision_score(y_test_pd, y_pred_pd)\n",
        "\n",
        "            # Update best score and parameters if current combination is better\n",
        "            if f1 > best_f1_score:\n",
        "                best_f1_score = f1\n",
        "                best_params = {\"nu\": nu, \"kernel\": kernel, \"gamma\": gamma}\n",
        "\n",
        "            print(\n",
        "                f\"Parameters: nu={nu}, kernel='{kernel}', gamma='{gamma}', Recall: {recall}, Precision: {precision}, F1-Score: {f1}\"\n",
        "            )\n",
        "\n",
        "# Print the best parameters and score\n",
        "print(\"Best Parameters:\", best_params)\n",
        "print(\"Best F1-Score:\", best_f1_score)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ar0_xsdOovxh"
      },
      "source": [
        "#### Use Novelty Detection with Local Outlier Factor"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cXRKK_Z1ovxh"
      },
      "outputs": [],
      "source": [
        "# lof = neighbors.LocalOutlierFactor(novelty=True, contamination=0.3, n_neighbors=12)\n",
        "lof = neighbors.LocalOutlierFactor(novelty=True, contamination=\"auto\", n_neighbors=5)\n",
        "lof.fit(X_train)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4MytWnctovxh"
      },
      "outputs": [],
      "source": [
        "# Make predictions on the test set (consider outliers as the positive class)\n",
        "y_pred = lof.predict(X_test)\n",
        "y_pred = [1 if y == -1 else 0 for y in y_pred]  # Convert -1 (outlier) to 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p0NWdGtcovxi"
      },
      "outputs": [],
      "source": [
        "cnf_matrix = confusion_matrix(y_test, y_pred)\n",
        "cnf_matrix"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Asx1Yemyovxi"
      },
      "outputs": [],
      "source": [
        "disp = ConfusionMatrixDisplay(\n",
        "    confusion_matrix=cnf_matrix, display_labels=[\"normal\", \"abnormal\"]\n",
        ")\n",
        "disp.plot()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Ze54m7dovxi"
      },
      "outputs": [],
      "source": [
        "# Evaluate the novel detection\n",
        "accuracy = accuracy_score(y_test, y_pred)\n",
        "target_names = [\"normal\", \"abnormal\"]\n",
        "report = classification_report(y_test, y_pred, target_names=target_names)\n",
        "\n",
        "print(f\"Accuracy: {accuracy}\")\n",
        "print(\"Classification Report:\\n\", report)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "849b1e62a910"
      },
      "source": [
        "Let's do some naive hyperparameter tuning"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dC0C1Fv7ovxi"
      },
      "outputs": [],
      "source": [
        "# Define the parameter grid for LocalOutlierFactor (adjust values as needed)\n",
        "param_grid = {\n",
        "    \"n_neighbors\": [5, 10, 15, 20],  # Number of neighbors to consider\n",
        "    \"contamination\": [\n",
        "        0.01,\n",
        "        0.02,\n",
        "        0.1,\n",
        "        0.2,\n",
        "        \"auto\",\n",
        "    ],  # Estimated proportion of outliers in the data set\n",
        "    \"novelty\": [True],  # Must be True for anomaly detection\n",
        "}\n",
        "\n",
        "# Initialize variables to track the best score and parameters\n",
        "best_f1_score = 0\n",
        "best_recall = 0\n",
        "best_params = {}\n",
        "\n",
        "# Manually iterate through parameter combinations\n",
        "for n_neighbors in param_grid[\"n_neighbors\"]:\n",
        "    for contamination in param_grid[\"contamination\"]:\n",
        "        # Create and train the OneClassSVM classifier\n",
        "        clf = neighbors.LocalOutlierFactor(\n",
        "            novelty=True, contamination=contamination, n_neighbors=n_neighbors\n",
        "        )\n",
        "        clf.fit(X_train)\n",
        "\n",
        "        # Make predictions on the test set\n",
        "        y_pred = clf.predict(X_test)\n",
        "        y_pred = [1 if y == -1 else 0 for y in y_pred]\n",
        "\n",
        "        # Calculate F1-score\n",
        "        f1 = f1_score(y_test, y_pred)\n",
        "        # Calculate Recall\n",
        "        recall = recall_score(y_test, y_pred)\n",
        "        # Calculate Precision\n",
        "        precision = precision_score(y_test, y_pred)\n",
        "\n",
        "        # Update best score and parameters if current combination is better\n",
        "        if f1 > best_f1_score:\n",
        "            best_f1_score = f1\n",
        "            best_params = {\"n_neighbors\": n_neighbors, \"contamination\": contamination}\n",
        "\n",
        "        print(\n",
        "            f\"Parameters: n_neighbors={n_neighbors}, contamination={contamination}, Recall: {recall}, Precision: {precision}, F1-Score: {f1}\"\n",
        "        )\n",
        "\n",
        "# Print the best parameters and score\n",
        "print(\"Best Parameters:\", best_params)\n",
        "print(\"Best F1-Score:\", best_recall)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ihfuDUFXm-lO"
      },
      "source": [
        "## Summary"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SlEDYC1PglVH"
      },
      "source": [
        "Here's the summary of the HDFS logs outlier detection you did in this notebook:\n",
        "\n",
        "| Outlier Detection Method | Recall | Precision | F1-score |\n",
        "|---|---|---|---|\n",
        "| **Embeddings + BQ Vector Search** | **0.92** | **0.90** | **0.91** |\n",
        "| Embeddings + LocalOutlierFactor | 0.85 | 0.73 | 0.78 |\n",
        "| Embeddings + OneClassSVM | 0.81 | 0.66 | 0.73 |\n",
        "\n",
        "These metrics show how **BigQuery vector search** can perform similar, or even better than common novelty or outlier detection algorithms. This demonstrates the robustness of its underlying native k-means clustering. These results stem from a combination of factors:\n",
        "\n",
        "- First, Gemini 2.0 effectively summarizes logs sequence including any abnormal activity.\n",
        "- Second, the vector embeddings generated by the text embedding model are well suited for clustering similar normal log sequences together, capturing nuances and representing most abherrations into outliers.\n",
        "- Finally, vector search offers speed and flexibility in fine-tuning distance and neighbor count parameters to optimize the balance between precision and recall.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vK-9z_T_gSxa"
      },
      "source": [
        "Now, let's compare with other leading methods such as DeepLog, N-gram, IM and PCA which were evaluated on HDFS logs dataset by [DeepLog: Anomaly Detection and Diagnosis from System Logs through Deep Learning](https://dl.acm.org/doi/10.1145/3133956.3134015) paper:\n",
        "\n",
        "| Outlier Detection Method | Recall | Precision | F1-score |\n",
        "|---|---|---|---|\n",
        "| DeepLog (RNN using LSTM) | 0.96 | 0.95 | 0.96 |\n",
        "| N-gram | 0.96 | 0.92 | 0.94 |\n",
        "| **Embeddings* + BQ Vector Search** | **0.92** | **0.90** | **0.91** |\n",
        "| Invariants Mining (offline) | 0.95 | 0.88 | 0.91 |\n",
        "| PCA (offline) | 0.67 | 0.98 | 0.79 |\n",
        "\n",
        "\\* Given notebook constraints, this method was evaluated using a test subset of 20k log sessions (equivalent to 400,000 log entries) as opposed to the entire HDFS dataset. This method was also evaluated without incremental updates to the index/model.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0ojJDnje0Hv_"
      },
      "source": [
        "Some observations:\n",
        "\n",
        "* Similar to N-gram and DeepLog, vector search-based outlier detection only requires a very small training set, about 1-10% of normal logs.\n",
        "\n",
        "* Unlike N-gram and DeepLog, vector search-based outlier detection **requires neither training nor serving a custom model**. Vector search results can improve with newly ingested data (e.g. normal log entries that were false positives) without necessarily requiring model retraining, or re-indexing in the case of vector index. However, re-indexing is recommended when a large dataset is added after initial index creation, to help handle data skew.\n",
        "\n",
        "* Unlike IM and PCA offline techniques, vector search-based outlier detection can be applied in a continuous fashion with streaming logs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IL6KstPsuuBZ"
      },
      "source": [
        "### Next steps\n",
        "\n",
        "This notebook provides a step-by-step walkthrough of applying LLMs, vector embeddings and BigQuery vector search for anomaly detection in semi-structured logs using familiar SQL.  Use and extend this notebook to experiment with embeddings and evaluate vector search accuracy, speed and cost using the publicly available HDFS logs dataset or your own dataset.\n",
        "\n",
        "The performance achieved with HDFS logs dataset is satisfactory and incurs much lower overhead compared to other alternatives given the use of off-the-shelf embedding and Gemini models. More exploratory work can be done to further evaluate and optimize recall, from fine-tuning the embedding model itself to optimizing the vector search-based anomaly detection logic (e.g. distance normalization, dynamic vs fixed threshold, etc.). That also includes index-time settings (e.g. `num_lists` for vector index) and search-time settings (e.g. `normalization_type`, `fraction_lists_to_search` for vector search). Results will vary depending on your dataset and your specific configurations."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2a4e033321ad"
      },
      "source": [
        "## Cleaning up"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "21f2a7432655"
      },
      "outputs": [],
      "source": [
        "# Delete BigQuery dataset. Uncomment and run the command below if you want to delete the BigQuery set.\n",
        "# from google.cloud import bigquery\n",
        "# Do this only if the dataset is created for this demo.\n",
        "# dataset = f\"{PROJECT_ID}.{DATASET_ID}\"\n",
        "# dataset_object = bigquery.Dataset(dataset)\n",
        "# client.delete_dataset(dataset_object, delete_contents=True, not_found_ok=True)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "rYn2yZ90cRG5"
      ],
      "name": "bq-vector-search-outlier-detection-infra-logs.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
