Lab 7 - Semantic Search/01-text-embedding.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "19ae6874-7585-4a0b-848c-965f639def41",
   "metadata": {},
   "source": [
    "# Text Embedding\n",
    "In this notebook, we generate 10-K filings text embeddings with the Vertex AI [textembedding-gecko](https://cloud.google.com/vertex-ai/docs/generative-ai/model-reference/text-embeddings) model. A 10-K is a comprehensive report filed annually by a publicly-traded company about its financial performance and is required by the SEC.  Unstructured text from 10-K filings has been extracted using a parser beforehand.\n",
    "\n",
    "In this notebook, we will:\n",
    "1. Get 10-K filings unstructured text from a Google storage bucket\n",
    "2. specifically select Item 1 from the 10K which describes the business of the company: who and what the company does, what subsidiaries it owns, and what markets it operates in. \n",
    "3. Chunk the text into natural sections using NLTK (to avoid input token limits)\n",
    "4. Save text with embeddings to csv to stage for loading into graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8de69a3c-2817-4ab6-9c02-a2cf500feb73",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --user tabulate sentence-transformers\n",
    "%pip install --user altair"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0386daec-95ed-40b9-8222-efa65197f426",
   "metadata": {},
   "source": [
    "Be sure to restart the kernel after you run the pip command."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e1b0998-b1f5-43f6-a278-9242e1b2b71e",
   "metadata": {},
   "source": [
    "## Get 10-K Filings from Google Cloud Storage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4286ddb9-950d-4bee-ba7b-c717e3ee1ada",
   "metadata": {},
   "outputs": [],
   "source": [
    "from google.cloud import storage\n",
    "\n",
    "storage_client = storage.Client()\n",
    "storage_client.bucket('neo4j-datasets').blob('hands-on-lab/form10k.zip').download_to_filename('/home/jupyter/form10k.zip')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09449433-e3b9-4946-84ba-05a1ab0c47fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir /home/jupyter/form10k\n",
    "!unzip -qq -n '/home/jupyter/form10k.zip' -d /home/jupyter/form10k"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c64dbe14-2aaa-4df0-97cc-6e03d47dd655",
   "metadata": {},
   "source": [
    "## 10-K Filings Exploration and Chunking\n",
    "Let's open one file to understand its contents.  It is actually a json file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3d86432-2679-4f0a-a33c-c16064769d51",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "with open('/home/jupyter/form10k/0001830197-22-000038.txt') as f:\n",
    "    f10_k = json.load(f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "772f8c51-b2fc-4704-85c0-12cd3b5a8c1f",
   "metadata": {},
   "source": [
    "We are interested in Item 1 specifically. \n",
    "\n",
    "Item 1 describes the business of the company: who and what the company does, what subsidiaries it owns, and what markets it operates in. It may also include recent events, competition, regulations, and labor issues. (Some industries are heavily regulated, and have complex labor requirements, which have significant effects on the business.) Other topics in this section may include special operating costs, seasonal factors, or insurance matters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76566ea2-d2da-4dae-b196-d66c9570c8d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(f10_k['item1'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f63be173-75e7-4017-9d61-8fe683ea199e",
   "metadata": {},
   "source": [
    "This text has the ability to exceed token limits for `textembedding-gecko`.  Also the quality of embeddings can go down if the text gets to large. As such we should find some way to chunk the text up into seperate sections for embedding.\n",
    "\n",
    "Below is a way to do this with Langchain's `RecursiveCharacterTextSplitter` ([doc](https://python.langchain.com/docs/modules/data_connection/document_transformers/text_splitters/recursive_text_splitter)) which takes into account of Chunk overlaps. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04061873-2b51-40b2-98ea-8cfd46e6fabd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "\n",
    "text = f10_k['item1']\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size = 2000,\n",
    "    chunk_overlap  = 15,\n",
    "    length_function = len,\n",
    "    is_separator_regex = False,\n",
    ")\n",
    "docs = text_splitter.split_text(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92d48aca-8acd-4ea2-ae2b-e8e0e18f2c94",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(docs[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ede16b2-a881-48b5-b8ea-f3e4de6f0022",
   "metadata": {},
   "source": [
    "## Get 10-K Text Embeddings with Vertex AI\n",
    "Now that we understand our data and how to chunk it.  Let's Generate embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e8ff3f6-fba3-4d86-b153-a637575fdbec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from vertexai.language_models import TextEmbeddingModel\n",
    "EMBEDDING_MODEL = TextEmbeddingModel.from_pretrained(\"textembedding-gecko@001\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e56acd5-e781-4ce7-b25f-49abdd9de8d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We will need a chunking utility to stay within token limits as we loop through files\n",
    "def chunks(xs, n=3):\n",
    "    n = max(1, n)\n",
    "    return [xs[i:i + n] for i in range(0, len(xs), n)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee20c6ac-6793-4de6-aaac-f83bb2178d59",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "def create_text_embedding_entries(input_text:str, company_name: str, cusip: str):\n",
    "    text_splitter = RecursiveCharacterTextSplitter(\n",
    "        chunk_size = 2000,\n",
    "        chunk_overlap  = 15,\n",
    "        length_function = len,\n",
    "        is_separator_regex = False,\n",
    "    )\n",
    "    docs = text_splitter.split_text(input_text)\n",
    "    res = []\n",
    "    seq_id = -1\n",
    "    for d in chunks(docs):\n",
    "        embeddings = EMBEDDING_MODEL.get_embeddings(d)\n",
    "        \n",
    "        # throttle so we don't blow through the quota.\n",
    "        time.sleep(1)\n",
    "        \n",
    "        for i in range(len(d)):\n",
    "            seq_id += 1\n",
    "            res.append({'companyName': company_name, 'cusip': cusip, 'seqId': seq_id, 'contextId': company_name + str(seq_id), 'textEmbedding': embeddings[i].values, 'text': d[i]})\n",
    "    return res"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bfe0a9f0-d7bb-404c-b899-fd497602cebc",
   "metadata": {},
   "source": [
    "Due to Quota Limitations, lets only do 5 form 10k files out of the 95 we have"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f3006c7-fa97-4282-bb31-6352ab0ed6a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "file_names = os.listdir('/home/jupyter/form10k/')[0:5]\n",
    "len(file_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7cfb362-264d-412e-a6c6-0e372fde8943",
   "metadata": {},
   "source": [
    "This cell takes about 15 minutes to run.  That's largely down to us throttling so we don't exceed the quota on our free account.  If you're using an enterprise account, you won't need to throttle like this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8497e11-d4f5-4bb3-80de-8cc538bd16ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "count = 0\n",
    "embedding_entries = []\n",
    "for file_name in file_names:\n",
    "    if '.txt' in file_name:\n",
    "        count += 1\n",
    "        if count % 5 == 0:\n",
    "            print(f'Parsed {count} of {len(file_names)}')\n",
    "        with open('/home/jupyter/form10k/' + file_name) as f:\n",
    "            f10_k = json.load(f)\n",
    "        embedding_entries.extend(create_text_embedding_entries(f10_k['item1'], f10_k['companyName'], f10_k['cusip']))\n",
    "len(embedding_entries)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e446e651-9134-4634-a527-4f2745c7d623",
   "metadata": {},
   "source": [
    "## Save 10-K Documents with Embeddings\n",
    "We will save these locally to use in graph loading, in the next part."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54ca3c83-fdc8-4d30-bacd-4c0b0328655f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "edf = pd.DataFrame(embedding_entries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "971cd795-4097-415f-83d6-f4caa651218a",
   "metadata": {},
   "outputs": [],
   "source": [
    "edf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05e64c29-2139-4694-ba7f-751bbd8c9daf",
   "metadata": {},
   "source": [
    "Provide your Neo4j credentials.  We need the DB conection URL, the username (probably `neo4j`), and your password."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58743d29-c611-4927-b61a-2b05721a9ae9",
   "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": "916b25a2-4dc8-454d-8fb6-06862b115ecc",
   "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",
    ")\n",
    "gds.set_database('neo4j')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f323e2e-7288-4887-b104-35be62dcae99",
   "metadata": {},
   "source": [
    "Remember to create indexes. We will be merging 10K documents by `companyName`. In a production setting, we would want to use a better identifier here (like we did with cusip for Company) However, this should suffice for our intents and purposes as we are just getting acquainted to learning about semantic search."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a57ef3d0-83f8-4d6d-9824-2f5206c36f78",
   "metadata": {},
   "outputs": [],
   "source": [
    "gds.run_cypher('CREATE INDEX company_name IF NOT EXISTS FOR (n:Company) ON (n.companyName)')\n",
    "gds.run_cypher('CREATE CONSTRAINT unique_document_id IF NOT EXISTS FOR (n:Document) REQUIRE (n.documentId) IS NODE KEY')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96be481f-7be8-4e70-b441-98c9d629b56b",
   "metadata": {},
   "source": [
    "Due to the size of the documents we will want to transform the dataframe into a list of dict that we can chunk up and insert via parameterized query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a0e68d63-b162-457b-8a60-aec01828aa82",
   "metadata": {},
   "outputs": [],
   "source": [
    "emb_entries = edf.to_dict(orient='records')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eecde179-7b86-46b8-a3dd-ae75b902472a",
   "metadata": {},
   "outputs": [],
   "source": [
    "total = len(emb_entries)\n",
    "count = 0\n",
    "for d in chunks(emb_entries, 100):\n",
    "    gds.run_cypher('''\n",
    "    UNWIND $records AS record\n",
    "    MATCH(c:Company {cusip:record.cusip})\n",
    "    MERGE(b:Document {documentId:record.contextId})\n",
    "    SET b.documentType = 'FORM_10K_ITEM1', b.seqId = record.seqId, b.textEmbedding = record.textEmbedding, b.text = record.text\n",
    "    MERGE(c)-[:HAS]->(b)\n",
    "    RETURN count(b) as cnt\n",
    "    ''', params = {'records':d})\n",
    "    count += len(d)\n",
    "    print(f'loaded {count} of {total}')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d75200e1-d4d3-4aaf-aae5-b989a99a855c",
   "metadata": {},
   "source": [
    "## Check Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45754ddd-c62c-4d4f-9e80-0648e99933db",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check node count\n",
    "gds.run_cypher('MATCH(doc:Document) RETURN count(doc)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33bf4496-6783-4ec3-be91-2dad92bc477b",
   "metadata": {},
   "source": [
    "Note that we were only getting 10-K docs for a minority of companies. It should be fine for this, but in a more rigorous setting, you may want to try and pull more.  There are likely a few factors attributing to this. \n",
    "\n",
    "1. We used company names to search EDGAR which resulted in many misses and dups which were discarded. In a more rigorous setting, we would investigate other endpoints and use more parsing to extract EDGAR cik keys for exact matching companies when pulling forms.\n",
    "\n",
    "2. Company names are not consistent across form13 filings, so even if we successfully pull on one version of a company name, we may not be able to merge it into the graph via the one company name represented there. \n",
    "\n",
    "3. Not all companies in the dataset are obligated to file 10-Ks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1034a8eb-cd6d-4a09-913b-d1d5e97c1bda",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check count and percentage of companies with 10-K docs.  Note it is the minority\n",
    "gds.run_cypher('''\n",
    "MATCH(b:Company)\n",
    "WITH b, count{(b)-[:HAS]->(d:Document)} AS docCount\n",
    "WITH count(b) AS total, sum(toInteger(docCount > 0)) AS numWithDocs\n",
    "RETURN total, numWithDocs, round(100*toFloat(numWithDocs)/toFloat(total), 2) As PercWithDocs\n",
    "''')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbb9c940-40cd-4fe9-8278-ef064e43933e",
   "metadata": {},
   "source": [
    "You might note that there are duplicate names.  For our purposes here, we will treat it as entity resolution, meaning that we treat companies with the same name as belonging to the same overarching entity for semantic search. In a more rigorous setting, we would need to disambiguate with other EDGAR keys."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da6b2835-2dfe-42ff-854f-c639eee67672",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Show duplicates via HAS relationship\n",
    "gds.run_cypher('''\n",
    "MATCH(b:Company)\n",
    "RETURN count(b) AS totalCompanies, count(DISTINCT b.companyName) AS uniqueCompanyNames\n",
    "''')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64854d9c-1776-451b-a3cb-60414d8116e3",
   "metadata": {},
   "source": [
    "## View Embeddings as Clusters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4beb9e08-bbaf-494c-bf47-d9dc49944300",
   "metadata": {},
   "source": [
    "Vector embeddings generated by language models are nothing but numerical representation of words or sentences.  So, similar sentences will be located nearby.  The embeddings we generated earlier are higher dimensional ones.  To visualize them, we need to reduce the dimensionality.  Let's do that and visualize it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f690e19-1d53-47fe-974c-9fda3c5903ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "import altair as alt\n",
    "\n",
    "def generate_chart(df, xcol, ycol, lbl = 'on', color = 'basic', title = '', tooltips = ['documentId'], label = ''):\n",
    "  chart = alt.Chart(df).mark_circle(size=30).encode(\n",
    "    x = alt.X(xcol,\n",
    "        scale=alt.Scale(zero = False),\n",
    "        axis=alt.Axis(labels = False, ticks = False, domain = False)\n",
    "    ),\n",
    "    y = alt.Y(ycol,\n",
    "        scale=alt.Scale(zero = False),\n",
    "        axis=alt.Axis(labels = False, ticks = False, domain = False)\n",
    "    ),\n",
    "    color= alt.value('#333293') if color == 'basic' else color,\n",
    "    tooltip=tooltips\n",
    "    )\n",
    "\n",
    "  if lbl == 'on':\n",
    "    text = chart.mark_text(align = 'left', baseline = 'middle', dx = 7, size = 5, color = 'black').encode(text = label, color = alt.value('black'))\n",
    "  else:\n",
    "    text = chart.mark_text(align = 'left', baseline = 'middle', dx = 10).encode()\n",
    "\n",
    "  result = (chart + text).configure(background=\"#FDF7F0\"\n",
    "        ).properties(\n",
    "        width = 800,\n",
    "        height = 500,\n",
    "        title = title\n",
    "       ).configure_legend(\n",
    "  orient = 'bottom', titleFontSize = 18, labelFontSize = 18)\n",
    "        \n",
    "  return result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e875832-4e4c-47d3-a535-e0b5fb177284",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reduce dimensionality using PCA\n",
    "from sklearn.decomposition import PCA\n",
    "\n",
    "# Function to return the principal components\n",
    "def get_pc(arr, n):\n",
    "  pca = PCA(n_components = n)\n",
    "  embeds_transform = pca.fit_transform(arr)\n",
    "  return embeds_transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3aa76707-4150-4948-83b0-d454f086afe7",
   "metadata": {},
   "outputs": [],
   "source": [
    "emb_df = gds.run_cypher(\"MATCH (c:Company)-[:HAS]->(n:Document) RETURN c.companyName as companyName, n.documentId as documentId, n.text as text, n.textEmbedding as emb LIMIT 1000\")\n",
    "emb_df"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99731d73-3648-43f1-b99d-79300f7df67d",
   "metadata": {},
   "source": [
    "## K-Means Clustering on the Embeddings\n",
    "Let's run the K-Means Clustering algorithm and view similar document chunks. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e2335c8-34bc-42ee-b093-ec3c5f718a6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.cluster import KMeans\n",
    "\n",
    "embeds = np.array(emb_df['emb'].tolist())\n",
    "embeds_pc2 = get_pc(embeds, 2)\n",
    "\n",
    "df_clust = pd.concat([emb_df, pd.DataFrame(embeds_pc2)], axis = 1)\n",
    "n_clusters = 5\n",
    "\n",
    "kmeans_model = KMeans(n_clusters = n_clusters, n_init = 1, random_state = 0)\n",
    "classes = kmeans_model.fit_predict(embeds).tolist()\n",
    "df_clust['cluster'] = (list(map(str,classes)))\n",
    "\n",
    "df_clust.columns = df_clust.columns.astype(str)\n",
    "generate_chart(df_clust.iloc[:],'0', '1', lbl = 'off', color = 'cluster', title = 'K-Means Clustering with n Clusters', tooltips = ['documentId', 'text'], label = '')\n"
   ]
  }
 ],
 "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
}