Neo4j Vector Embedding and Search.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "981caecc-e884-420f-9968-ae1e2c5b0d7b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download and install Python packages needed for this Jupyter Notebook\n",
    "\n",
    "!pip install openai tiktoken neo4j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fcf7654d-382a-42ab-8cfa-cfc6e7e7d2a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import openai\n",
    "from neo4j import GraphDatabase\n",
    "import tiktoken\n",
    "\n",
    "openai.api_key = os.getenv(\"OPENAI_API_KEY\")\n",
    "\n",
    "# Uncomment the line below for verification that your OpenAI API key was successfully imported\n",
    "#openai.Model.list()\n",
    "\n",
    "# Your OpenAI API key can also be manually entered using the command below\n",
    "#openai.api_key = \"Your Open AI key between the quotation marks\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf335cc3-d9bb-476e-9c1b-543befe793b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The block connect the Jupyter Notebook to your Neo4j Database\n",
    "# Note: Your Neo4j Database must be running and accepting connections\n",
    "# Note: This example is for connecting to a local instance of Neo4j\n",
    "# More information on interfacing with can be found at\n",
    "# https://neo4j.com/docs/python-manual/current/connect/\n",
    "\n",
    "uri = 'bolt://localhost:7687'\n",
    "username = 'neo4j'\n",
    "password = 'password'\n",
    "driver = GraphDatabase.driver(uri, auth=(username, password))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc00c79b-b188-422c-b96a-02aded389a25",
   "metadata": {},
   "outputs": [],
   "source": [
    "#This block calulates the dollar cost of obtaining Vector Embeddings for \n",
    "#abstracts within the database that do not currently have an embedding\n",
    "\n",
    "def num_tokens_from_string(string: str, encoding_name: str) -> int:\n",
    "    \"\"\"Returns the number of tokens in a text string.\"\"\"\n",
    "    encoding = tiktoken.get_encoding(encoding_name)\n",
    "    num_tokens = len(encoding.encode(string))\n",
    "    return num_tokens\n",
    "\n",
    "records, summary, keys = driver.execute_query(\n",
    "    \"MATCH (n:Work) WHERE n.embedding IS NULL AND n.abstract IS NOT NULL \\\n",
    "    RETURN n.abstract, n.id\")\n",
    "\n",
    "abstract_count = 0 \n",
    "token_count = 0 \n",
    "\n",
    "# Calculate number of tokens for each record \n",
    "for record in records:  \n",
    "    abstract = record.data('n.abstract').get('n.abstract')\n",
    "    token_count += num_tokens_from_string(abstract, \"cl100k_base\")\n",
    "\n",
    "print(\"It will cost approximately - $\" + str(round((token_count/1000)*0.0001,3))\\\n",
    "      + \" Dollars to process abstracts into Vector Embeddings, assuming a cost \"+\n",
    "      \"of $0.0001 per 1000 tokens\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c9a2f36-e2c1-45c8-947b-eb2ae39071d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This block identifies Work nodes with abstracts that lack vector embeddings\n",
    "# The abstract and id for these Work nodes are returned \n",
    "# The abstract is sent to OpenAI to return the vector embedding which is added\n",
    "# as an 'embedding' property to the respective Work node\n",
    "\n",
    "records_for_node_embedding_start, summary_for_node_embedding_start, \\\n",
    "    keys_for_node_embedding_start = driver.execute_query(\\\n",
    "    \"MATCH (n:Work) WHERE n.embedding IS NULL AND n.abstract <> '' \\\n",
    "    RETURN COUNT(n)\")\n",
    "eligible_node_without_embedding = records_for_node_embedding_start[0][0]\n",
    "\n",
    "print(\"Nodes with abstracts that do not have vector embeddings - \" + \\\n",
    "      str(eligible_node_without_embedding))\n",
    "print(\"These will be processed in batches of 10\")\n",
    "\n",
    "# Returns Work node id and abstract for Work nodes with abstracts\n",
    "# that do not have vector embeddings\n",
    "# This is done in batches of 10\n",
    "      \n",
    "while eligible_node_without_embedding > 0:\n",
    "    print(\"Processing embedding for node number - \" + \\\n",
    "          str(eligible_node_without_embedding))\n",
    "    records, summary, keys = driver.execute_query( \\\n",
    "        \"MATCH (n:Work) WHERE n.embedding IS NULL AND n.abstract <> '' \\\n",
    "        AND n.abstract IS NOT NULL RETURN n.abstract, n.id LIMIT 10\")\n",
    "     \n",
    "    # Send each record to OpenAI to create a Vector Embedding\n",
    "    # NOTE: embedding is not captured as a variable in Python, only Neo4j\n",
    "    for record in records:  \n",
    "\n",
    "        node_id = record.data('n.id').get('n.id')\n",
    "        abstract = record.data('n.abstract').get('n.abstract')\n",
    " \n",
    "        neo4j_record, summary, keys = driver.execute_query(\"MATCH \\\n",
    "            (w:Work {id: $id}) WITH w \\\n",
    "            CALL db.create.setVectorProperty(w, 'embedding', $embedding) \\\n",
    "            YIELD node RETURN node\", \\\n",
    "            id = node_id, \\\n",
    "            embedding = openai.Embedding.create( \\\n",
    "            input= abstract, model=\\\n",
    "            \"text-embedding-ada-002\")[\"data\"][0][\"embedding\"])\n",
    "\n",
    "        # Note: If you have a new OpenAI account, the number of queries per minute \n",
    "        # will be throttled. This can be addressed by including a sleep timer\n",
    "        # The sleep function will need to be imported - from time import sleep \n",
    "        #sleep(41.2)\n",
    "\n",
    "    eligible_node_without_embedding -= 10\n",
    "\n",
    "    if eligible_node_without_embedding > 0:\n",
    "        print(\"Remaining nodes without embeddings - \" + \\\n",
    "              str(eligible_node_without_embedding))\n",
    "    else:\n",
    "        print (\"Vector embeddings have been added for all Work Nodes\") "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0585bf91-f944-4e25-8c16-8946630277b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This block runs the command needed to create the Vector Search Index\n",
    "# More information is available at https://neo4j.com/docs/cypher-manual/current/indexes-for-vector-search/\n",
    "# Note: Only run this block a single time per database as only one index can be created\n",
    "\n",
    "neo4j_record, summary, keys = driver.execute_query(\\\n",
    "    \"CALL db.index.vector.createNodeIndex(\\\n",
    "    'abstract-embeddings', 'Work', 'embedding', 1536, 'cosine')\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0f92da2-0ce8-4a9d-a530-d4f6055b4b16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This block will find the closest 25 matches to the Work selected based on \n",
    "# the cosine similiarities of the vector embeddings for each abstract\n",
    "\n",
    "# This will search across all Works, regardless of original source if\n",
    "# metadata was imported using the data schema model provided by the\n",
    "# Jupyter Notebooks available at \n",
    "# https://github.com/vtmike2015/Graph-Based-Literature-Review-Tool\n",
    "\n",
    "# Users need to input the Id of the Work of interest to find the closest\n",
    "# match. This is the 'id' property for a Work node, regardless of source\n",
    "\n",
    "work_id = input('Please enter the Id of the Work you want to the closest 25 matches for')\n",
    "\n",
    "records, summary, keys  = driver.execute_query(\\\n",
    "    \"MATCH (w:Work) WHERE w.id = '\" + work_id + \\\n",
    "    \"' CALL db.index.vector.queryNodes('abstract-embeddings', 25, w.embedding)\\\n",
    "    YIELD node AS similarAbstract, score MATCH (n:Work)<-[:WROTE]-(a:Author) \\\n",
    "    WHERE n.id = similarAbstract.id AND w.id <> similarAbstract.id \\\n",
    "    RETURN w.id, w.display_name, w.abstract, score, n.id, n.display_name, \\\n",
    "    COLLECT(a.display_name) AS authors, n.publication_year, n.abstract, n.source LIMIT 25\"\n",
    "    )\n",
    "\n",
    "original_work_id = records[0].data('w.id').get('w.id')\n",
    "original_work_title = records[0].data('w.display_name').get('w.display_name')\n",
    "original_work_abstract = records[0].data('w.abstract').get('w.abstract')\n",
    "\n",
    "print(\"You selected '\" + original_work_title + \"' with Id - \" + original_work_id)\n",
    "print(\"This is the original work's abstract - \" + original_work_abstract)\n",
    "print(\"\")\n",
    "\n",
    "counter = 1\n",
    "\n",
    "for record in records:\n",
    "    \n",
    "    node_id = record.data('n.id').get('n.id')\n",
    "    score = record.data('score').get('score')\n",
    "    title = record.data('n.display_name').get('n.display_name')\n",
    "    abstract = record.data('n.abstract').get('n.abstract') \n",
    "    authors = record.data('authors').get('authors')\n",
    "    source = record.data('n.source').get('n.source')\n",
    "    \n",
    "    print(\"Article Match Number \" + str(counter) + \" is - \" + title)\n",
    "    print(\"The Similarity score is \" + str(round(score,5)))\n",
    "    print(\"The Source is \" + source)\n",
    "    print(\"The ID is - \" + node_id)\n",
    "    print(\"The authors are - \" + str(authors))\n",
    "    print(\"This is the abstract - \" + abstract)\n",
    "    print(\"\")\n",
    "    counter += 1 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58ce032d-67db-4fcd-a4e0-c0dd2c127e38",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}