demo-python/code/vector-quantization-and-storage/vector-quantization-and-storage.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reducing storage usage for vectors on Azure AI Search\n",
    "\n",
    "This code demonstrates how to use the following features to reduce vector storage on Azure AI Search.\n",
    "\n",
    "+ Use smaller \"narrow\" data types instead of `Edm.Single`. Types such as `Edm.Half` reduce storage overhead.\n",
    "+ Disable storing vectors used in the query response. Vectors returned in a query response are stored separately from the vectors used during queries.\n",
    "+ Quantizing vectors. Use built-in scalar quantization to quantize embeddings to `Edm.Int8` without any reduction in query performance. Information loss from quantization can be compensated for using the original unquantized embeddings and oversampling.\n",
    "\n",
    "### Prerequisites\n",
    "\n",
    "+ An Azure subscription.\n",
    " \n",
    "+ Azure AI Search, any tier, but we recommend Basic or higher for this workload. [Enable semantic ranker](https://learn.microsoft.com/azure/search/semantic-how-to-enable-disable) if you want to run a hybrid query with semantic ranking."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Set up a Python virtual environment in Visual Studio Code\n",
    "\n",
    "1. Open the Command Palette (Ctrl+Shift+P).\n",
    "1. Search for **Python: Create Environment**.\n",
    "1. Select **Venv**.\n",
    "1. Select a Python interpreter. Choose 3.10 or later.\n",
    "\n",
    "It can take a minute to set up. If you run into problems, see [Python environments in VS Code](https://code.visualstudio.com/docs/python/environments)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Install packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install -r vector-quantization-and-storage-requirements.txt --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load .env file (Copy .env-sample to .env and update accordingly)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "from azure.identity import DefaultAzureCredential\n",
    "from azure.core.credentials import AzureKeyCredential\n",
    "import os\n",
    "\n",
    "load_dotenv(override=True) # take environment variables from .env.\n",
    "\n",
    "# Variables not used here do not need to be updated in your .env file\n",
    "endpoint = os.environ[\"AZURE_SEARCH_SERVICE_ENDPOINT\"]\n",
    "credential = AzureKeyCredential(os.environ[\"AZURE_SEARCH_ADMIN_KEY\"]) if len(os.environ[\"AZURE_SEARCH_ADMIN_KEY\"]) > 0 else DefaultAzureCredential()\n",
    "base_index_name = os.environ[\"AZURE_SEARCH_INDEX\"]\n",
    "embedding_dimensions = int(os.getenv(\"AZURE_OPENAI_EMBEDDING_DIMENSIONS\", 3072))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load embeddings\n",
    "\n",
    "Load the embeddings from a precomputed file. These embeddings use [text-embedding-3-large](https://learn.microsoft.com/azure/ai-services/openai/concepts/models#embeddings) with 3072 dimensions. The chunks are from the sample data in the document folder, chunked using the [Split Skill](https://learn.microsoft.com/azure/search/cognitive-search-skill-textsplit)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "from lib.embeddings import content_path\n",
    "\n",
    "with open(content_path, \"r\") as f:\n",
    "    chunks = json.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create indexes\n",
    "\n",
    "To demonstrate the storage impact of the different options, the following code creates indexes that use each option, and another index that combines all the options together"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to define the indexes on the search service\n",
    "\n",
    "from azure.search.documents.indexes import SearchIndexClient\n",
    "from azure.search.documents.indexes.models import (\n",
    "    SimpleField,\n",
    "    SearchFieldDataType,\n",
    "    SearchableField,\n",
    "    VectorSearch,\n",
    "    HnswAlgorithmConfiguration,\n",
    "    VectorSearchProfile,\n",
    "    SemanticConfiguration,\n",
    "    SemanticPrioritizedFields,\n",
    "    SemanticField,\n",
    "    SemanticSearch,\n",
    "    SearchIndex,\n",
    "    SearchField,\n",
    "    ScalarQuantizationCompressionConfiguration\n",
    ")\n",
    "\n",
    "def create_index(index_name, dimensions, use_compression=False, use_float16=False, use_stored=True):\n",
    "    if use_float16:\n",
    "        vector_type = \"Collection(Edm.Half)\"\n",
    "    else:\n",
    "        vector_type = \"Collection(Edm.Single)\"\n",
    "\n",
    "    # Vector fields that aren't stored can never be returned in the response\n",
    "    fields = [\n",
    "        SimpleField(name=\"id\", type=SearchFieldDataType.String, key=True, sortable=True, filterable=True),\n",
    "        SearchableField(name=\"title\", type=SearchFieldDataType.String),\n",
    "        SearchableField(name=\"chunk\", type=SearchFieldDataType.String),\n",
    "        SearchField(name=\"embedding\", type=vector_type, searchable=True, stored=use_stored, vector_search_dimensions=dimensions, vector_search_profile_name=\"myHnswProfile\")\n",
    "    ]\n",
    "\n",
    "    if use_compression:\n",
    "        compression_name = \"myCompression\"\n",
    "        compression_configurations = [\n",
    "            ScalarQuantizationCompressionConfiguration(name=compression_name)\n",
    "        ]\n",
    "    else:\n",
    "        compression_name = None\n",
    "        compression_configurations = None\n",
    "    \n",
    "    vector_search = VectorSearch(\n",
    "        algorithms=[\n",
    "            HnswAlgorithmConfiguration(name=\"myHnsw\")\n",
    "        ],\n",
    "        profiles=[\n",
    "            VectorSearchProfile(name=\"myHnswProfile\", algorithm_configuration_name=\"myHnsw\", compression_configuration_name=compression_name)\n",
    "        ],\n",
    "        compressions=compression_configurations\n",
    "    )\n",
    "\n",
    "    semantic_config = SemanticConfiguration(\n",
    "        name=\"my-semantic-config\",\n",
    "        prioritized_fields=SemanticPrioritizedFields(\n",
    "            title_field=SemanticField(field_name=\"title\"),\n",
    "            content_fields=[SemanticField(field_name=\"chunk\")]\n",
    "        )\n",
    "    )\n",
    "    semantic_search = SemanticSearch(configurations=[semantic_config])\n",
    "\n",
    "    return SearchIndex(name=index_name, fields=fields, vector_search=vector_search, semantic_search=semantic_search)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created indexes\n"
     ]
    }
   ],
   "source": [
    "# Create indexes to compare storage usage\n",
    "# The baseline index does not use any options\n",
    "\n",
    "baseline_index = f\"{base_index_name}-baseline\"\n",
    "compression_index = f\"{base_index_name}-compression\"\n",
    "narrow_index = f\"{base_index_name}-narrow\"\n",
    "no_stored_index = f\"{base_index_name}-no-stored\"\n",
    "all_index = f\"{base_index_name}-all-options\"\n",
    "\n",
    "search_index_client = SearchIndexClient(endpoint, credential)\n",
    "search_index_client.create_or_update_index(\n",
    "    create_index(compression_index, embedding_dimensions, use_compression=True))\n",
    "search_index_client.create_or_update_index(\n",
    "    create_index(narrow_index, embedding_dimensions, use_float16=True))\n",
    "search_index_client.create_or_update_index(\n",
    "    create_index(no_stored_index, embedding_dimensions, use_stored=False))\n",
    "search_index_client.create_or_update_index(\n",
    "    create_index(all_index, embedding_dimensions, use_compression=True, use_float16=True, use_stored=False))\n",
    "search_index_client.create_or_update_index(\n",
    "    create_index(baseline_index, embedding_dimensions))\n",
    "print(\"Created indexes\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to upload the embeddings to each index\n",
    "\n",
    "import json\n",
    "from lib.embeddings import content_path\n",
    "from azure.search.documents import SearchIndexingBufferedSender\n",
    "\n",
    "def upload_embeddings(index_name):\n",
    "    with open(content_path, \"r\") as f:\n",
    "        content = json.load(f)\n",
    "    \n",
    "    with SearchIndexingBufferedSender(endpoint, index_name, credential) as client:\n",
    "        client.upload_documents(content)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Uploaded embeddings\n"
     ]
    }
   ],
   "source": [
    "upload_embeddings(compression_index)\n",
    "upload_embeddings(narrow_index)\n",
    "upload_embeddings(no_stored_index)\n",
    "upload_embeddings(all_index)\n",
    "upload_embeddings(baseline_index)\n",
    "\n",
    "print(\"Uploaded embeddings\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check storage sizes\n",
    "\n",
    "Find the new storage size in MB to demonstrate how the various options affect storage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "****************************************\n",
      "Index Name: compressiontest-baseline\n",
      "Storage Size: 21.3613MB\n",
      "Vector Size: 4.8277MB\n",
      "****************************************\n",
      "Index Name: compressiontest-compression\n",
      "Storage Size: 17.7604MB\n",
      "Vector Size: 1.2242MB\n",
      "****************************************\n",
      "Index Name: compressiontest-narrow\n",
      "Storage Size: 16.5567MB\n",
      "Vector Size: 2.4254MB\n",
      "****************************************\n",
      "Index Name: compressiontest-no-stored\n",
      "Storage Size: 10.9224MB\n",
      "Vector Size: 4.8277MB\n",
      "****************************************\n",
      "Index Name: compressiontest-all-options\n",
      "Storage Size: 4.9192MB\n",
      "Vector Size: 1.2242MB\n"
     ]
    }
   ],
   "source": [
    "# Please note - there may be delays in finding index statistics after document upload\n",
    "# Index statistics is not a real time API\n",
    "# See https://learn.microsoft.com/en-us/rest/api/searchservice/preview-api/get-index-statistics for more information\n",
    "\n",
    "def bytes_to_mb(bytes):\n",
    "    return round(bytes / (1024 * 1024), 4)\n",
    "\n",
    "def find_storage_size_mb(index_name):\n",
    "    response = search_index_client.get_index_statistics(index_name)\n",
    "    return bytes_to_mb(response[\"storage_size\"]), bytes_to_mb(response[\"vector_index_size\"])\n",
    "\n",
    "index_sizes = [(find_storage_size_mb(index_name), index_name) for index_name in [compression_index, baseline_index, no_stored_index, narrow_index, all_index]]\n",
    "index_sizes.sort(key=lambda item: item[0][0], reverse=True)\n",
    "\n",
    "for ((storage_size, vector_size), index_name) in index_sizes:\n",
    "    print(\"*\" * 40)\n",
    "    print(f\"Index Name: {index_name}\\nStorage Size: {storage_size}MB\\nVector Size: {vector_size}MB\")\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}