Lab 7 - Semantic Search/02-semantic-search.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "19ae6874-7585-4a0b-848c-965f639def41",
   "metadata": {},
   "source": [
    "# Semantic Search\n",
    "Semantic search is [defined as \"search with meaning.\"](https://en.wikipedia.org/wiki/Semantic_search) It is key for effective knowledge retrieval.\n",
    "\n",
    "As opposed to traditional, lexical, search which finds matches based on keywords, semantic search seeks to improve search quality and accuracy by understanding search intent and pulling results that match the user’s contextual meaning.  Semantic search is often used in reference to text embedding and vector similarity search, but this is just one implementation aspect of it. Knowledge graph and symbolic query logic can also play a critical role in making semantic search a reality. \n",
    "\n",
    "If all you care about is analyzing a set of documents on a file system, then sure, vector indexing and search may be sufficient.  However, once you need to retrieve and make inferences about people, places, and things connected to those documents, Knowledge graph becomes key. \n",
    "\n",
    "If documents are the entities of interest, for example: \"find all documents that talk about pharma related things,\" then text embeddings with vector similarity search suffices.  But, what if we want second or third-order entities related to the documents?  For example, \"find investors who are most focused on pharma related strategies,\" how would we efficiently search for them at scale in an enterprise setting?\n",
    "\n",
    "This is what we demonstrate below.  We will also show how you can use graph relationships and Graph Data Science algorithms to further improve search results, especially in common scenarios where the presence of text data is inconsistent or sparse."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4026b768-0406-4ce1-a1bb-ed79b64701a3",
   "metadata": {},
   "source": [
    "## Connect to Neo4j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50798482-b5d5-4e96-adb2-851a3855f82e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# username is neo4j by default\n",
    "NEO4J_USERNAME = 'neo4j'\n",
    "\n",
    "# You will need to change these to match your credentials\n",
    "NEO4J_URI = 'neo4j+s://6688b25b.databases.neo4j.io'\n",
    "NEO4J_PASSWORD = '_kogrNk53u8oTk5be55kmit1kHGdhZj98yJlG-VYSR'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5aa11282-874f-4126-832d-9a2f772989a5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphdatascience import GraphDataScience\n",
    "\n",
    "gds = GraphDataScience(\n",
    "    NEO4J_URI,\n",
    "    auth = (NEO4J_USERNAME, NEO4J_PASSWORD),\n",
    "    aura_ds = True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ded4103-008d-4f93-82c9-918add3688eb",
   "metadata": {},
   "source": [
    "## Neo4j Vector Index\n",
    "We will need to create a vector index for similarity search on Document nodes. Neo4j offers a vector index that enables Approximate Nearest Neighbor Search (ANN). Let's creat an index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f08fe029-fd4f-476a-b510-7dbcdb3d892d",
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher(\"CALL db.index.vector.createNodeIndex('document-embeddings', 'Document', 'textEmbedding', 768, 'cosine')\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27e3300f-e799-4ac1-b236-a23c300ae5e2",
   "metadata": {},
   "source": [
    "You can see that vector index has been created using `SHOW INDEXES`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b9c4634-bfee-4d73-a197-ed4345e729bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher(''' \n",
    "SHOW INDEXES YIELD name, type, labelsOrTypes, properties, options\n",
    "WHERE type = \"VECTOR\"\n",
    "''')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "313f706a-08c3-4fc0-9c50-293d588414ae",
   "metadata": {},
   "source": [
    "## Deep Semantic Search with Knowledge Graph\n",
    "Now that we have an index.  Let’s use it in action.\n",
    "In this case, we will Answer The Question - \"What are the companies associated with energy, oil and gas?\". \n",
    "Remember we do not have documents on investment managers, just companies, and there can be multiple documents for each company."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a00ab005-bc0f-49b7-8cf9-e6ecf5cbe298",
   "metadata": {},
   "source": [
    "## Dense Retrieval using Neo4j"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c13411b2-37ae-4d80-a147-9e32904d9012",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.docstore.document import Document\n",
    "from typing import (\n",
    "    List,\n",
    "    Tuple,\n",
    ")\n",
    "\n",
    "def to_df(results: List[Tuple[Document, float]]):\n",
    "    return pd.DataFrame({\n",
    "        \"doc_id\": [r[0].metadata.get('documentId') for r in results],\n",
    "        \"score\": [r[1] for r in results],\n",
    "        \"text\": [r[0].page_content for r in results]\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48299f48-f956-4648-8739-9c990811fa55",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.vectorstores import Neo4jVector\n",
    "from langchain.embeddings.vertexai import VertexAIEmbeddings\n",
    "from langchain.vectorstores.neo4j_vector import SearchType\n",
    "import pandas as pd\n",
    "\n",
    "db = Neo4jVector.from_existing_index(\n",
    "    url = NEO4J_URI,\n",
    "    username = NEO4J_USERNAME,\n",
    "    password = NEO4J_PASSWORD,\n",
    "    embedding = VertexAIEmbeddings(),\n",
    "    index_name = 'document-embeddings',\n",
    "    keyword_index_name = 'doc_text',\n",
    "    search_type = SearchType.VECTOR\n",
    ")\n",
    "query = 'What are the companies associated with energy, oil and gas?'\n",
    "results = db.similarity_search_with_score(query, k=10)\n",
    "res_df = to_df(results)\n",
    "res_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6842196c-c83a-4e07-889e-e4775ef1ff26",
   "metadata": {},
   "source": [
    "Lets take a look at the first result and then understand whether the chunk has the information we asked for."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "365985b9-eb2b-47cd-8869-4f31351a23ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_df['text'][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8f4a919-4a40-4b37-a565-73111324fa31",
   "metadata": {},
   "source": [
    "## Re-rank results\n",
    "If you explore the results, you may find some irrelevant results. The retriever has to be efficient for large document collections with millions of entries. However, it might return irrelevant candidates.\n",
    "A re-ranker based on a Cross-Encoder can substantially improve the final results for the user. The query and a possible document is passed simultaneously to transformer network, which then outputs a single score between 0 and 1 indicating how relevant the document is for the given query.\n",
    "\n",
    "We will use a Cross-encoder model named `ms-marco-MiniLM-L-6-v2` from SentenceTransformers library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29a717ff-9df5-47fb-9962-f296187faae2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rerank_results(res_df):\n",
    "    from sentence_transformers import SentenceTransformer, CrossEncoder, util\n",
    "    cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')\n",
    "    passages = []\n",
    "    \n",
    "    # Score all retrieved passages with the cross_encoder\n",
    "    cross_inp = [[query, f\"Text: {hit['text']}\"] for _,hit in res_df[:].iterrows()]\n",
    "    cross_scores = cross_encoder.predict(cross_inp)\n",
    "\n",
    "    # Sort results by the cross-encoder scores\n",
    "    res_df['cross-score'] = cross_scores\n",
    "\n",
    "    # Output of top-10 hits from re-ranker\n",
    "    print(\"Top-10 Cross-Encoder Re-ranker hits\")\n",
    "    hits = res_df.sort_values('cross-score', ascending=False).reset_index(drop=True)\n",
    "    for _,hit in hits.iterrows():\n",
    "        print(\"\\t{:.3f}\\t{}\".format(hit['cross-score'], hit['doc_id']))\n",
    "    return hits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5851d3bb-d93f-41fc-9315-8d8dc34f4a72",
   "metadata": {},
   "outputs": [],
   "source": [
    "ranked_results = rerank_results(res_df)\n",
    "ranked_results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3479218e-c372-4296-8058-af73eb85096b",
   "metadata": {},
   "source": [
    "As seen above, the cross-encoder finds this passage to be more relevant to the query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c38a08dc-b443-48ab-8e8e-11eb9bc53062",
   "metadata": {},
   "outputs": [],
   "source": [
    "ranked_results['text'][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f1d9d37-7f27-4683-9bb4-c21db62257f8",
   "metadata": {},
   "source": [
    "## Hybrid Search with Neo4j\n",
    "Vector Search is not the only solution. Often you will find a hybrid approach - a combination of Vector and Fulltext search perform better. You can do both in Neo4j.\n",
    "\n",
    "To do a Fulltext Search, lets first create a fulltext index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e107afba-2ae8-4da8-88b1-e5fe7236ced4",
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher(\"CREATE FULLTEXT INDEX doc_text IF NOT EXISTS FOR (n:Document) ON EACH [n.text]\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c75544b-ec3f-4eaa-b57d-bde5258d68b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher(''' \n",
    "SHOW INDEXES YIELD name, type, labelsOrTypes, properties, options\n",
    "WHERE type = \"FULLTEXT\"\n",
    "''')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c512be4c-9e55-4b92-9d7c-bb8f2eb2a583",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.embeddings.vertexai import VertexAIEmbeddings\n",
    "from langchain.text_splitter import CharacterTextSplitter\n",
    "from langchain.vectorstores import Neo4jVector\n",
    "from langchain.vectorstores.neo4j_vector import SearchType\n",
    "\n",
    "db = Neo4jVector.from_existing_index(\n",
    "    url = NEO4J_URI,\n",
    "    username = NEO4J_USERNAME,\n",
    "    password = NEO4J_PASSWORD,\n",
    "    embedding = VertexAIEmbeddings(),\n",
    "    index_name = 'document-embeddings',\n",
    "    keyword_index_name = 'doc_text',\n",
    "    search_type = SearchType.HYBRID # Hybrid search type\n",
    ")\n",
    "\n",
    "results = db.similarity_search_with_score(query, k=10)\n",
    "res_df = to_df(results)\n",
    "res_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "daf3c617-4915-4357-999e-31c346031842",
   "metadata": {},
   "source": [
    "Lets re-rank our results for relevancy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae38378f-8288-4481-94ba-2fb7fa5083a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "ranked_results = rerank_results(res_df)\n",
    "ranked_results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19387fbd-0144-49cf-928c-6c2c25a8caf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "ranked_results['text'][4]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "541c6aad-05b4-4391-9f27-b668a82ae910",
   "metadata": {},
   "source": [
    "We don't have much documents from other companies chunked. But once you have it, you can see additional results with Hybrid search. The re-ranker can then help rank these results accordingly. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07198bd2-2a8d-432a-8d7d-94f617cc0969",
   "metadata": {},
   "source": [
    "## Semantic Search with Multi-hop retrieval using Knowledge Graph"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7502f4a3-5ca8-4771-b637-af58ea07f1ac",
   "metadata": {},
   "source": [
    "First, let's create a database connection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8bbab47-d50a-4827-aeb4-42d4c9c612af",
   "metadata": {},
   "outputs": [],
   "source": [
    "db = Neo4jVector(\n",
    "    url = NEO4J_URI,\n",
    "    username = NEO4J_USERNAME,\n",
    "    password = NEO4J_PASSWORD,\n",
    "    embedding = VertexAIEmbeddings()\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91518685-9430-4f52-8e6b-d7e4b61a7a25",
   "metadata": {},
   "source": [
    "Now let's use that query vector to search for companies.  Remember, companies have multiple documents so we will need to use a graph traversal on top of a document lookup to find which companies are most similar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9428f7f-5aea-445d-a2ee-203c05a4d4f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_vector = db.embedding.embed_query(text=query)\n",
    "query_vector[0:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26e84306-7097-4ea7-ab34-166ef2421f2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Search for similar companies\n",
    "res_df = db.query(\"\"\"\n",
    "CALL db.index.vector.queryNodes('document-embeddings', 10, $queryVector)\n",
    "YIELD node AS similarDocuments, score\n",
    "\n",
    "MATCH (similarDocuments)<-[:HAS]-(c:Company)\n",
    "RETURN c.companyName as companyName, avg(score) AS score\n",
    "ORDER BY score DESC LIMIT 10\n",
    "\"\"\", params =  {'queryVector': query_vector})\n",
    "res_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "085bb735-27de-4726-b23f-df0762ff28e6",
   "metadata": {},
   "source": [
    "The above result is based on the limited set of documents we have in the DB. The scores are based on Average similarity scores across chunks"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64e2d9f1-f2f6-4438-98ec-f04be86972c3",
   "metadata": {},
   "source": [
    "Now let's take this one step further and find investment managers who are most heavily focused in Energy.  This will involve a bit more Cypher for a 2-hop traversal. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c36b2196-554e-43ca-a52c-a3719d12c146",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Search for managers with significiant investments in area\n",
    "res_df = gds.run_cypher(\"\"\"\n",
    "CALL db.index.vector.queryNodes('document-embeddings', 1000, $queryVector)\n",
    "YIELD node AS similarDocuments, score\n",
    "MATCH (similarDocuments)<-[:HAS]-(c:Company)\n",
    "WITH c, avg(score) AS score ORDER BY score LIMIT 100\n",
    "MATCH (c)<-[r:OWNS]-(m:Manager)\n",
    "WITH m, r.value as value, score*r.value as weightedScore\n",
    "WITH m.managerName AS managerName, sum(weightedScore) AS aggScore, sum(value) AS aggValue\n",
    "RETURN managerName, aggScore/aggValue AS score ORDER BY score DESC LIMIT 1000\n",
    "\n",
    "\"\"\", params =  {'queryVector': query_vector})\n",
    "res_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "458e5b2c-8bf6-4755-8e3c-9eaf06bc4096",
   "metadata": {},
   "source": [
    "And we can see that our top result is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20f01f64-f15a-4181-9ec3-2adb326e239d",
   "metadata": {},
   "outputs": [],
   "source": [
    "top_manager = res_df['managerName'][0]\n",
    "top_manager"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "930804d4-d3ec-48a7-b8c6-92ab8a64ef04",
   "metadata": {},
   "source": [
    "## Expanding Available Data for Knowledge Retrieval\n",
    "Not every element in your data will have rich text data, and further, much like we only have 10K documents for some companies, your use cases may also have incomplete, unevenly distributed text data. \n",
    "\n",
    "We can check our top result investment manager to this. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10aefd2b-4100-4be3-a985-1ceabcd2c8ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher('''\n",
    "MATCH (m:Manager {managerName: $managerName})-[:OWNS]->(c:Company)-[:HAS]->(d:Document)\n",
    "WITH m, count(DISTINCT c) AS ownedCompaniesWithDocs\n",
    "MATCH (m:Manager {managerName: $managerName})-[:OWNS]->(c:Company)\n",
    "RETURN m.managerName AS managerName, ownedCompaniesWithDocs, count(DISTINCT c) AS totalOwnedCompanies\n",
    "''', params =  {'managerName':top_manager})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "532553dd-682e-45cf-96d8-7cecf58068a2",
   "metadata": {},
   "source": [
    "This manager has significantly more other companies they own without documents.  We can use Graph Data Science Node Similarity to find the managers that have the most overlap to this one which should give us other Energy companies that we missed due to sparse text data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "252ff6a0-eae7-4dcc-aa39-fbb7a37b1114",
   "metadata": {},
   "outputs": [],
   "source": [
    "g, _ = gds.graph.project('proj', ['Company', 'Manager'], {'OWNS':{'properties':['value']}})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23119a94-d0d0-4c60-b319-f2d94805a144",
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.nodeSimilarity.write(g, writeRelationshipType = 'SIMILAR', writeProperty = 'score', relationshipWeightProperty = 'value')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "acd39e18-fc3d-4734-bbe8-e1862516fc29",
   "metadata": {},
   "outputs": [],
   "source": [
    "g.drop()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "addcf2b2-a3b3-41a2-bc5d-d9630b0689cb",
   "metadata": {},
   "source": [
    "And now we can pull back other relevant results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ec2ed6f-6d34-414f-82b3-9cf28602d4e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher('''\n",
    "MATCH (m0:Manager {managerName: $managerName})-[r:SIMILAR]->(m:Manager)\n",
    "RETURN DISTINCT m.managerName AS managerName, r.score AS score\n",
    "ORDER BY score DESC LIMIT 10\n",
    "''', params =  {'managerName':top_manager})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a0a43fc-72c5-4a40-914f-01e7f3001c56",
   "metadata": {},
   "source": [
    "## Clean Up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f548f8e-e48f-4d5f-b2fd-2251280fc8f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher('MATCH (M:Manager)-[s:SIMILAR]->() DELETE s')"
   ]
  }
 ],
 "metadata": {
  "environment": {
   "kernel": "conda-root-py",
   "name": "workbench-notebooks.m111",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/workbench-notebooks:m111"
  },
  "kernelspec": {
   "display_name": "Python (Local)",
   "language": "python",
   "name": "local-base"
  },
  "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": 5
}