embedding-document-search.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Search documents with Azure Open AI Embeddings\n",
    "\n",
    "This tutorial will walk you through using the Azure OpenAI embeddings API to perform document search where you'll query a knowledge base to find the most relevant document.\n",
    "\n",
    "## Settings\n",
    "\n",
    "| Setting | Value | \n",
    "| --- | --- |\n",
    "| Model | text-embedding-3-small |\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install additional .NET SDKs\n",
    "\n",
    "You'll need to install 3 additional SDKs from NuGet.\n",
    "\n",
    "### Azure.AI.OpenAI\n",
    "\n",
    "[Azure.AI.OpenAI NuGet package](https://www.nuget.org/packages/Azure.AI.OpenAI/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><div></div><div></div><div><strong>Installed Packages</strong><ul><li><span>Azure.AI.OpenAI, 1.0.0-beta.15</span></li></ul></div></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#r \"nuget: Azure.AI.OpenAI, 1.0.0-beta.15\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utility SDKs\n",
    "\n",
    "* [CSVHelper](https://www.nuget.org/packages/CsvHelper/) => help with reading CSV files\n",
    "* [Microsoft.ML.Tokenizers](https://www.nuget.org/packages/Microsoft.ML.Tokenizers/0.20.1) => help generate tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><div></div><div></div><div><strong>Installed Packages</strong><ul><li><span>CsvHelper, 30.0.1</span></li><li><span>Microsoft.ML.Tokenizers, 0.20.1</span></li></ul></div></div>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#r \"nuget: CsvHelper, 30.0.1\"\n",
    "#r \"nuget: Microsoft.ML.Tokenizers, 0.20.1\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set the `using` statements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "using System;\n",
    "using Azure.AI.OpenAI;\n",
    "using Azure;\n",
    "using Microsoft.ML.Tokenizers;\n",
    "using CsvHelper;\n",
    "using System.Globalization;\n",
    "using System.Text.RegularExpressions;"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configure Azure OpenAI client\n",
    "\n",
    "We'll use the configuration values you set in the **.devcontainer.json** to configure and initialize the `OpenAIClient`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "var AOAI_ENDPOINT = Environment.GetEnvironmentVariable(\"AOAI_ENDPOINT\");\n",
    "var AOAI_KEY = Environment.GetEnvironmentVariable(\"AOAI_KEY\");\n",
    "var AOAI_DEPLOYMENTID = Environment.GetEnvironmentVariable(\"AOAI_DEPLOYMENTID\");"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "var endpoint = new Uri(AOAI_ENDPOINT);\n",
    "var credentials = new AzureKeyCredential(AOAI_KEY);\n",
    "var openAIClient = new OpenAIClient(endpoint, credentials);"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Read the document to be searched into memory\n",
    "\n",
    "The document that is to be searched is a CSV file. It contains legislative bills. \n",
    "\n",
    "First create a class that will model each row of the file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "public class BillSummary\n",
    "{\n",
    "    public string Id { get; set; }\n",
    "    public string Text { get; set; }\n",
    "    public string Title { get; set; }\n",
    "    public string Summary { get; set; }\n",
    "    public int TokenCount { get; set; }\n",
    "    public float[] Embeddings { get; set; }\n",
    "}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, open the CSV file and read certain columns into memory.\n",
    "\n",
    "_(Note: The CSV file was downloaded when the Codespace was started. If you are not using Codespaces, download the [CSV from GitHub](https://raw.githubusercontent.com/Azure-Samples/Azure-OpenAI-Docs-Samples/main/Samples/Tutorials/Embeddings/data/bill_sum_data.csv) and place it in the same directory as this notebook.)_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "List<BillSummary> billSummaries = new ();\n",
    "\n",
    "using (var reader = new System.IO.StreamReader(\"bill_sum_data.csv\"))\n",
    "using (var csv = new CsvReader(reader, CultureInfo.InvariantCulture))\n",
    "{\n",
    "    csv.Read();\n",
    "    csv.ReadHeader();\n",
    "\n",
    "    while (csv.Read())\n",
    "    {\n",
    "        var billSummary = new BillSummary\n",
    "        {\n",
    "            Id = csv.GetField<string>(\"bill_id\"),\n",
    "            Text = csv.GetField<string>(\"text\"),\n",
    "            Title = csv.GetField<string>(\"title\"),\n",
    "            Summary = csv.GetField<string>(\"summary\")\n",
    "        };\n",
    "        billSummaries.Add(billSummary);\n",
    "    }\n",
    "}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before checking the token count or creating the embeddings, we want to cleanup the `Text` property. That involves removing any extra whitespace, replacing \"..\" with \".\" and so on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "// The function to do the normalization\n",
    "string NormalizeText(string s)\n",
    "{\n",
    "    // Replace multiple whitespace characters with a single space\n",
    "    s = Regex.Replace(s, @\"\\s+\", \" \").Trim();\n",
    "\n",
    "    // Remove \". ,\"\n",
    "    s = Regex.Replace(s, @\"\\. ,\", \"\");\n",
    "\n",
    "    // Replace \"..\" with \".\"\n",
    "    s = s.Replace(\"..\", \".\");\n",
    "\n",
    "    // Replace \". .\" with \".\"\n",
    "    s = s.Replace(\". .\", \".\");\n",
    "\n",
    "    // Remove newline characters\n",
    "    s = s.Replace(\"\\n\", \"\");\n",
    "\n",
    "    // Remove leading and trailing spaces\n",
    "    s = s.Trim();\n",
    "\n",
    "    return s;\n",
    "}\n",
    "\n",
    "// Perform the normalization on the Text property\n",
    "billSummaries.ForEach(b => b.Text = NormalizeText(b.Text));"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenize the document's text\n",
    "\n",
    "Because Azure OpenAI's embedding function can only handle 8191 tokens, we need to make sure none of the `Text` fields evaluate to more than that. We'll use **Microsoft.ML.Tokenizers**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "// Function to perform the token count\n",
    "int TokenCount(Tokenizer tokenizer, string s)\n",
    "{\n",
    "    var tokenizerResult = tokenizer.Encode(s);\n",
    "\n",
    "    return tokenizerResult.Tokens.Count();\n",
    "}\n",
    "\n",
    "// In order to tokenize, we need the following\n",
    "var vocabFilePath = @\"vocab.json\";\n",
    "var mergesFilePath = @\"merges.txt\";\n",
    "\n",
    "var tokenizer = new Tokenizer(new Bpe(vocabFilePath, mergesFilePath));\n",
    "\n",
    "// Get a token count for the Text and put it into the TokenCount property\n",
    "billSummaries.ForEach(b => b.TokenCount = TokenCount(tokenizer, b.Text));"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Remove anything that has too many tokens\n",
    "\n",
    "If any of the rows have too many tokens we can't use it and should remove it from the `billSummaries`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "var tooLong = billSummaries.Where(b => b.TokenCount > 8192).ToList();\n",
    "\n",
    "foreach (var bill in tooLong)\n",
    "{\n",
    "    billSummaries.Remove(bill);\n",
    "}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### View the tokens\n",
    "\n",
    "As a side note, let's view the tokens returned for the first row of data. _(This is not necessary for the tutorial, but it's interesting.)_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"dni-plaintext\"><pre>[ S, ECTION, 1, ., SH, ORT, TIT, LE, ., This, Act, may, be, c, ited, as, the, ``, National, Science ... (1712 more) ]</pre></div><style>\r\n",
       ".dni-code-hint {\r\n",
       "    font-style: italic;\r\n",
       "    overflow: hidden;\r\n",
       "    white-space: nowrap;\r\n",
       "}\r\n",
       ".dni-treeview {\r\n",
       "    white-space: nowrap;\r\n",
       "}\r\n",
       ".dni-treeview td {\r\n",
       "    vertical-align: top;\r\n",
       "    text-align: start;\r\n",
       "}\r\n",
       "details.dni-treeview {\r\n",
       "    padding-left: 1em;\r\n",
       "}\r\n",
       "table td {\r\n",
       "    text-align: start;\r\n",
       "}\r\n",
       "table tr { \r\n",
       "    vertical-align: top; \r\n",
       "    margin: 0em 0px;\r\n",
       "}\r\n",
       "table tr td pre \r\n",
       "{ \r\n",
       "    vertical-align: top !important; \r\n",
       "    margin: 0em 0px !important;\r\n",
       "} \r\n",
       "table th {\r\n",
       "    text-align: start;\r\n",
       "}\r\n",
       "</style>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "var firstTokenResult = tokenizer.Encode(billSummaries[0].Text);\n",
    "\n",
    "firstTokenResult.Tokens.Display();"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get the embeddings for each row in the document\n",
    "\n",
    "For every row in the document, this will get the embedding and store it as an array of `float`s in the `AdaV2` property."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "foreach (var bill in billSummaries)\n",
    "{\n",
    "    var embeddingResponse = await openAIClient.GetEmbeddingsAsync(\n",
    "        new ()\n",
    "        {\n",
    "            DeploymentName=AOAI_DEPLOYMENTID,\n",
    "            Input={bill.Text}\n",
    "        });\n",
    "\n",
    "    bill.Embeddings = embeddingResponse.Value.Data[0].Embedding.ToArray();\n",
    "}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get embedding for the incoming query\n",
    "\n",
    "Now we'll get the incoming query, or what we'll want to search for in the overall document, and then get an embedding for that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "var incoming = \"Can I get information about STEM contributions?\";\n",
    "EmbeddingsOptions embeddingOptions = new()\n",
    "{\n",
    "    DeploymentName=AOAI_DEPLOYMENTID,\n",
    "    Input={incoming}\n",
    "};\n",
    "var incomingEmbeddingResponse = await openAIClient.GetEmbeddingsAsync(embeddingOptions);\n",
    "var incomingEmbedding = incomingEmbeddingResponse.Value.Data[0].Embedding.ToArray();"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Find closest match with cosine similarity comparison\n",
    "\n",
    "To find the closest match between the rows in the document and the incoming query, we'll use the cosine similarity to compare the vectors (or arrays) of each embedding."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "float CosineSimilarity (float[] vectorA, float[] vectorB) {\n",
    "    float dotProduct = 0.0f;\n",
    "    float magnitudeA = 0.0f;\n",
    "    float magnitudeB = 0.0f;\n",
    "\n",
    "    for (int i = 0; i < vectorA.Length; i++)\n",
    "    {\n",
    "        dotProduct += vectorA[i] * vectorB[i];\n",
    "        magnitudeA += vectorA[i] * vectorA[i];\n",
    "        magnitudeB += vectorB[i] * vectorB[i];\n",
    "    }\n",
    "\n",
    "    return dotProduct / (float)(Math.Sqrt(magnitudeA) * Math.Sqrt(magnitudeB));\n",
    "}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now use the `CosineSimilarity` function to peform the comparison. Here we are taking the top 4 results for illustration purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [],
   "source": [
    "var topMatches = billSummaries.Select(b => new\n",
    "{\n",
    "    Bill = b,\n",
    "    Similarity = CosineSimilarity(incomingEmbedding, b.Embeddings)\n",
    "}).OrderByDescending(b => b.Similarity).Take(4).ToList();"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, display the top match"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "dotnet_interactive": {
     "language": "csharp"
    },
    "polyglot_notebook": {
     "kernelName": "csharp"
    },
    "vscode": {
     "languageId": "polyglot-notebook"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "To amend the Internal Revenue Code of 1986 to encourage businesses to improve math and science education at elementary and secondary schools."
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "National Science Education Tax Incentive for Businesses Act of 2007 - Amends the Internal Revenue Code to allow a general business tax credit for contributions of property or services to elementary and secondary schools and for teacher training to promote instruction in science, technology, engineering, or mathematics ."
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "topMatches[0].Bill.Title.Display();\n",
    "topMatches[0].Bill.Summary.Display();"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "csharp"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}