vector_search_example/vector_search_example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "- **Docker**  \n",
    "  Required to run Memgraph, as Memgraph is a native Linux application and cannot be installed directly on Windows or macOS.\n",
    "\n",
    "- **pandas**  \n",
    "  A fast, powerful, flexible, and easy-to-use open-source data analysis and manipulation tool, built on top of the Python programming language.\n",
    "\n",
    "- **kagglehub**  \n",
    "  A powerful and flexible library similar to pandas for data manipulation and analysis.\n",
    "\n",
    "- **sentence_transformers**  \n",
    "  A library for state-of-the-art sentence embeddings, similar in purpose to pandas but focused on text data.\n",
    "\n",
    "- **neo4j**  \n",
    "  Used to query Memgraph."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build a Movie Similarity Search Engine with Vector Search in Memgraph\n",
    "\n",
    "In this example, we will demonstrate how vector search can be used to find movies based on their plots or short descriptions. For this, we will use the Wikipedia Movie Plots dataset, available on Kaggle.\n",
    "To get started, launch Memgraph with the `--experimental-enabled=vector-search` flag and the appropriate `--experimental-config flag`.\n",
    "\n",
    "To start Memgraph, run:\n",
    "`docker run -p 7687:7687 -p 7444:7444 memgraph/memgraph:latest --experimental-enabled=vector-search --experimental-config='{\"vector-search\": {\"movies_index\": {\"label\": \"Movie\",\"property\": \"embedding\",\"dimension\": 384,\"capacity\": 100, \"metric\": \"cos\"}}}'`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We created a vector index `movies_index`  which is defined on label `Movie`  and property `embedding` .\n",
    "\n",
    "Firstly let’s load the dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "   Release Year                             Title Origin/Ethnicity  \\\n",
      "0          1901            Kansas Saloon Smashers         American   \n",
      "1          1901     Love by the Light of the Moon         American   \n",
      "2          1901           The Martyred Presidents         American   \n",
      "3          1901  Terrible Teddy, the Grizzly King         American   \n",
      "4          1902            Jack and the Beanstalk         American   \n",
      "\n",
      "                             Director Cast    Genre  \\\n",
      "0                             Unknown  NaN  unknown   \n",
      "1                             Unknown  NaN  unknown   \n",
      "2                             Unknown  NaN  unknown   \n",
      "3                             Unknown  NaN  unknown   \n",
      "4  George S. Fleming, Edwin S. Porter  NaN  unknown   \n",
      "\n",
      "                                           Wiki Page  \\\n",
      "0  https://en.wikipedia.org/wiki/Kansas_Saloon_Sm...   \n",
      "1  https://en.wikipedia.org/wiki/Love_by_the_Ligh...   \n",
      "2  https://en.wikipedia.org/wiki/The_Martyred_Pre...   \n",
      "3  https://en.wikipedia.org/wiki/Terrible_Teddy,_...   \n",
      "4  https://en.wikipedia.org/wiki/Jack_and_the_Bea...   \n",
      "\n",
      "                                                Plot  \n",
      "0  A bartender is working at a saloon, serving dr...  \n",
      "1  The moon, painted with a smiling face hangs ov...  \n",
      "2  The film, just over a minute long, is composed...  \n",
      "3  Lasting just 61 seconds and consisting of two ...  \n",
      "4  The earliest known adaptation of the classic f...  \n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import kagglehub\n",
    "\n",
    "dataset_path = kagglehub.dataset_download(\"jrobischon/wikipedia-movie-plots\")\n",
    "df = pd.read_csv(dataset_path + \"/wiki_movie_plots_deduped.csv\")\n",
    "print(df.head())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This dataset consists of 32,432 movies. To keep the example simple and understandable, we will reduce the dataset size. We will filter movies based on the director."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(9, 8)\n"
     ]
    }
   ],
   "source": [
    "nolan_movies = df[df['Director'] == 'Christopher Nolan']\n",
    "nolan_movies.reset_index(drop=True, inplace=True)\n",
    "print(nolan_movies.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also need a function to compute embeddings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(9, 384)\n"
     ]
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "def compute_embeddings(texts):\n",
    "    model = SentenceTransformer('paraphrase-MiniLM-L6-v2')\n",
    "    return model.encode(texts)\n",
    "    \n",
    "embeddings = compute_embeddings(nolan_movies['Plot'].values)\n",
    "print(embeddings.shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let’s import these movies into Memgraph:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import neo4j\n",
    "\n",
    "driver = neo4j.GraphDatabase.driver(\"bolt://localhost:7687\", auth=(\"\", \"\"))\n",
    "with driver.session() as session:\n",
    "    for index, row in nolan_movies.iterrows():\n",
    "        # remove quotes from the title to avoid parsing issues\n",
    "        title = row[\"Title\"].replace('\"', '')\n",
    "        \n",
    "        embedding = embeddings[index].tolist()\n",
    "        embeddings_str = \",\".join([str(x) for x in embedding])\n",
    "        query = f'CREATE (m:Movie {{title: \"{title}\", year: {row[\"Release Year\"]}, embedding: [{embeddings_str}]}})'\n",
    "        session.run(query)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, we computed the embeddings with `compute_embeddings` function and then we used these embedding vectors to store it in the `embedding` property of each node.\n",
    "\n",
    "After we have imported the data into Memgraph we can start with our experiments!\n",
    "\n",
    "We will define a function which we can use to find most similar movies described with the plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_movie(plot):\n",
    "    embeddings = compute_embeddings([plot])\n",
    "    embeddings_str = \",\".join([str(x) for x in embeddings[0]])\n",
    "    with driver.session() as session:\n",
    "        query = f\"CALL vector_search.search('movies_index', 3, [{embeddings_str}]) yield node, similarity return node.title, similarity\"\n",
    "        result = session.run(query)\n",
    "        for record in result:\n",
    "            print(record)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let’s try to find the Inception, by using the following plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Record node.title='Inception' similarity=0.5250678062438965>\n",
      "<Record node.title='Interstellar' similarity=0.2907602787017822>\n",
      "<Record node.title='The Dark Knight' similarity=0.2784501910209656>\n"
     ]
    }
   ],
   "source": [
    "plot = \"A thief who steals corporate secrets through the use of dream-sharing technology is given the inverse task of planting an idea into the mind of a C.E.O.\"\n",
    "find_movie(plot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, let’s attempt to find Memento:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Record node.title='Memento' similarity=0.37598633766174316>\n",
      "<Record node.title='Insomnia' similarity=0.26347029209136963>\n",
      "<Record node.title='Inception' similarity=0.23714923858642578>\n"
     ]
    }
   ],
   "source": [
    "plot = \"An insurance investigator suffers from anterograde amnesia, leaving him unable to form new memories, and uses notes and tattoos to track down his wife's killer.\"\n",
    "find_movie(plot)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "usr",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}