open-ai-embed-completions.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "485c4024",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/leemthompo/notebook-tests/blob/main/open-ai-embed-completions.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ec1de10",
   "metadata": {},
   "source": [
    "Learn how to:\n",
    "   - index the OpenAI Wikipedia embeddings dataset into Elasticsearch\n",
    "   - encode a question with the `openai ada-02` model\n",
    "   - perform a kNN semantic search\n",
    "   - send top result to OpenAI chat API to answer the question using that context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46eb7dec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import openai\n",
    "from elasticsearch import Elasticsearch, helpers\n",
    "import getpass\n",
    "from zipfile import ZipFile\n",
    "import os\n",
    "import json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dabb7b43",
   "metadata": {},
   "source": [
    "Download the dataset and extract it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c721163f",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.system('curl -O https://cdn.openai.com/API/examples/data/vector_database_wikipedia_articles_embedded.zip')\n",
    "with ZipFile('vector_database_wikipedia_articles_embedded.zip', 'r') as zip_ref:\n",
    "    zip_ref.extractall()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab274f79",
   "metadata": {},
   "source": [
    "Replace with your local path to the downloaded file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b31d6cdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "file = '/Users/<user>/Downloads/vector_database_wikipedia_articles_embedded.csv'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f07d489",
   "metadata": {},
   "source": [
    "Read the data into a DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6eae7987",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv(file)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e39adc85",
   "metadata": {},
   "source": [
    "Found in the 'Manage Deployment' page"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "994408db",
   "metadata": {},
   "outputs": [],
   "source": [
    "CLOUD_ID = getpass.getpass('Enter Elastic Cloud ID:  ')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5083f2af",
   "metadata": {},
   "source": [
    " Password for the 'elastic' user generated by Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0c5ab90",
   "metadata": {},
   "outputs": [],
   "source": [
    "ELASTIC_API_KEY = getpass.getpass('Enter Elastic password:  ')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "348a4dff",
   "metadata": {},
   "source": [
    "Create the client instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c6416c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = Elasticsearch(\n",
    "    cloud_id= CLOUD_ID,\n",
    "    api_key= ELASTIC_API_KEY\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e697b9c",
   "metadata": {},
   "source": [
    "Define the mapping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e70b1e2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "mapping = {\n",
    "    \"mappings\": {\n",
    "        \"properties\": {\n",
    "            \"url\": {\n",
    "                \"type\": \"text\"\n",
    "            },\n",
    "            \"title\": {\n",
    "                \"type\": \"text\"\n",
    "            },\n",
    "             \"title_vector\": {\n",
    "                \"type\": \"dense_vector\",\n",
    "                \"dims\": 1536,\n",
    "                \"index\": \"true\",\n",
    "                \"similarity\": \"cosine\"\n",
    "            },\n",
    "            \"content_vector\": {\n",
    "                \"type\": \"dense_vector\",\n",
    "                \"dims\": 1536,\n",
    "                \"index\": \"true\",\n",
    "                \"similarity\": \"cosine\"\n",
    "            },\n",
    "            \"text\": {\n",
    "                \"type\": \"text\"\n",
    "            },\n",
    "            \"vector_id\": {\n",
    "                \"type\": \"keyword\"\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5ec7139",
   "metadata": {},
   "source": [
    "  Create an index in Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7970b75c",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    client.indices.create(\n",
    "        index=\"wiki-vector\",\n",
    "        body=mapping,\n",
    "        ignore=400  ignore 400 already exists code\n",
    "    )\n",
    "except:\n",
    "    print('Index already exists')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36c66dfa",
   "metadata": {},
   "source": [
    " Index the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0613efe",
   "metadata": {},
   "outputs": [],
   "source": [
    "actions = [\n",
    "    {\n",
    "        \"_index\": 'wiki-vector',\n",
    "        \"_id\": row['id'],\n",
    "        \"_source\": {\n",
    "            \"url\": row['url'],\n",
    "            \"title\": row['title'],\n",
    "            \"text\": row['text'],\n",
    "            \"title_vector\": json.loads(row['title_vector']),\n",
    "            \"content_vector\": json.loads(row['content_vector']),\n",
    "            \"vector_id\": row['vector_id']\n",
    "        }\n",
    "    }\n",
    "    for _, row in df.iterrows()\n",
    "]\n",
    "\n",
    "helpers.bulk(client, actions)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92724501",
   "metadata": {},
   "source": [
    "Embed user query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e277d8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"who invented the compass?\"\n",
    "\n",
    "openai.api_key = getpass.getpass('Enter OpenAI API key:')\n",
    "\n",
    "response = openai.Embedding.create(\n",
    "    input=query,\n",
    "    model=\"text-embedding-ada-002\"\n",
    ")\n",
    "query_vector = response['data'][0]['embedding']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40a298f9",
   "metadata": {},
   "source": [
    "Perform knn search using the question vector."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23cfc7b1",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cbb02eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = client.search(index=\"wiki-vector\", body={\n",
    "  \"knn\": {\n",
    "    \"field\": \"content_vector\",\n",
    "    \"query_vector\": query_vector,\n",
    "    \"k\": 3,\n",
    "    \"num_candidates\": 100\n",
    "  },\n",
    "  \"_source\": [\"title\", \"text\"]\n",
    "})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d4ae769",
   "metadata": {},
   "source": [
    "Extract the first hit from the response."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b053efe",
   "metadata": {},
   "outputs": [],
   "source": [
    "first_hit = response['hits']['hits'][0]\n",
    "source = first_hit['_source']\n",
    "text = source['text']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91d9bd18",
   "metadata": {},
   "source": [
    "Now we can send the question and the text to OpenAI's chat completion API.\n",
    "The model will generate a response to the question, using the top kNN hit as context.\n",
    "Use the `messages` list to shape your prompt to the gen AI model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16170468",
   "metadata": {},
   "outputs": [],
   "source": [
    "summary = openai.ChatCompletion.create(\n",
    "  model=\"gpt-3.5-turbo\",\n",
    "  messages=[\n",
    "        {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
    "        {\"role\": \"user\", \"content\": \"Answer the following question:\" + query + \"by using the following text:\" + text + \"Please print each sentence on a new line.\"},\n",
    "    ]\n",
    ")\n",
    "\n",
    "choices = summary.choices\n",
    "\n",
    "for choice in choices:\n",
    "    print(\"------------------------------------------------------------\")\n",
    "    print(choice.message.content)\n",
    "    print(\"------------------------------------------------------------\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}