1-parsing-data.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "19ae6874-7585-4a0b-848c-965f639def41",
   "metadata": {},
   "source": [
    "# Parsing Data\n",
    "In this notebook, let's explore how to leverage generative AI to build and consume a knowledge graph in Neo4j.\n",
    "\n",
    "This notebook parses Form-13 data from SEC EDGAR. This is partially structured data, a mix of text and XML.  Instead of spending our time writing a bespoke parser to extract data from these files and load into Neo4j, we can prompt a Large Language Model (LLM) to do this for us automatically.  We will then also use the LLM to generate Cypher statements to load the extracted data into a Neo4j graph."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9f13ff0d-b2d4-4627-a95c-0c5078436c0f",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c76cc41-77ed-466d-94c8-27c14dc906c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --user graphdatascience"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5aa4433-a120-47bc-9e53-0eac953986b1",
   "metadata": {},
   "source": [
    "Now restart the kernel. That will allow the Python evironment to import the new packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "146d4f56-f83e-49d6-8d5a-6c4a971a44b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import IPython\n",
    "\n",
    "app = IPython.Application.instance()\n",
    "app.kernel.do_shutdown(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08a1a385-8d3a-4c69-9817-ed14ba8e6511",
   "metadata": {},
   "source": [
    "## Prompt Definition\n",
    "We will extract knowledge adhering to the same schema we used previously.  To teach the LLM about the schema, we will use a series of prompts.  Each prompt is focused on only one task, extracting a specific entity:\n",
    "\n",
    "1. Manager Information\n",
    "2. Filing Information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffeab081-6bc7-41a1-92a6-1f02b8adcf8f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "mgr_info_tpl = \"\"\"From the text below, extract the following as json. Do not miss any of these information.\n",
    "* The tags mentioned below may or may not namespaced. So extract accordingly. Eg: <ns1:tag> is equal to <tag>\n",
    "* \"managerName\" - The name from the <name> tag under <filingManager> tag\n",
    "* \"street1\" - The manager's street1 address from the <com:street1> tag under <address> tag\n",
    "* \"street2\" - The manager's street2 address from the <com:street2> tag under <address> tag\n",
    "* \"city\" - The manager's city address from the <com:city> tag under <address> tag\n",
    "* \"stateOrCounty\" - The manager's stateOrCounty address from the <com:stateOrCountry> tag under <address> tag\n",
    "* \"zipCode\" - The manager's zipCode from the <com:zipCode> tag under <address> tag\n",
    "* \"reportCalendarOrQuarter\" - The reportCalendarOrQuarter from the <reportCalendarOrQuarter> tag under <address> tag\n",
    "* Just return me the JSON enclosed by 3 backticks. No other text in the response\n",
    "\n",
    "Text:\n",
    "$ctext\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2483efd5-fedf-49b0-9f01-534bd7390968",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "filing_info_tpl = \"\"\"The text below contains a list of investments. Each instance of <infoTable> tag represents a unique investment. \n",
    "For each investment, please extract the below variables into json then combine into a list enclosed by 3 backticks. Please use the quoted names below while doing this\n",
    "* \"cusip\" - The cusip from the <cusip> tag under <infoTable> tag\n",
    "* \"companyName\" - The name under the <nameOfIssuer> tag.\n",
    "* \"value\" - The value from the <value> tag under <infoTable> tag. Return as a number. \n",
    "* \"shares\" - The sshPrnamt from the <sshPrnamt> tag under <infoTable> tag. Return as a number. \n",
    "* \"sshPrnamtType\" - The sshPrnamtType from the <sshPrnamtType> tag under <infoTable> tag\n",
    "* \"investmentDiscretion\" - The investmentDiscretion from the <investmentDiscretion> tag under <infoTable> tag\n",
    "* \"votingSole\" - The votingSole from the <votingSole> tag under <infoTable> tag\n",
    "* \"votingShared\" - The votingShared from the <votingShared> tag under <infoTable> tag\n",
    "* \"votingNone\" - The votingNone from the <votingNone> tag under <infoTable> tag\n",
    "\n",
    "Output format:\n",
    "* DO NOT output XML tags in the response. The output should be a valid JSON list enclosed by 3 backticks\n",
    "\n",
    "Text:\n",
    "$ctext\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "481717bb-6db2-4465-8667-b0c07da381d5",
   "metadata": {},
   "source": [
    "## Functions for Using LLMs\n",
    "Let's create some helper function to talk to the LLM with our prompt and text input. \n",
    "\n",
    "The [Bedrock documentation](https://aws.amazon.com/bedrock/developer-experience/) describes the available foundation models.  We will use Anthropic Claude v2 100k base model. In some cases, there may be a need to fine-tune LLM models for KG creation. [Bedrock provides an elegant way to fine-tune](https://docs.aws.amazon.com/bedrock/latest/userguide/fine-tuning.html) where the updated weights/model stay within your tenant and the base model is frozen."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e42db05-86be-4171-9ef3-538af376057e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "SERVICE_NAME = 'bedrock-runtime'\n",
    "REGION_NAME = 'us-east-1'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc39b09b-1ba9-487f-b1b0-5e6f9de26a92",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "import json\n",
    "bedrock = boto3.client(\n",
    " service_name=SERVICE_NAME,\n",
    " region_name=REGION_NAME,\n",
    " endpoint_url=f'https://{SERVICE_NAME}.{REGION_NAME}.amazonaws.com'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54780c17-da6e-4e31-a2ee-5d835abbabc9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import traceback\n",
    "def call_language_model(prompt_data):\n",
    "    try:\n",
    "        body = json.dumps({\"prompt\": f\"Human: {prompt_data} \\n Assistant:\",\n",
    "                           \"temperature\":0,\n",
    "                           \"top_k\":1, \"top_p\":0.88,\n",
    "                           \"anthropic_version\":\"bedrock-2023-05-31\",\n",
    "                           \"max_tokens_to_sample\": 2048\n",
    "                          })\n",
    "        modelId = 'anthropic.claude-v2' # change this to use a different version from the model provider\n",
    "        accept = 'application/json'\n",
    "        contentType = 'application/json'\n",
    "\n",
    "        response = bedrock.invoke_model(body=body, modelId=modelId, accept=accept, contentType=contentType)\n",
    "        response_body = json.loads(response.get('body').read())\n",
    "\n",
    "        return response_body.get('completion')\n",
    "    except Exception as e:\n",
    "        traceback.format_exc()\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79ca2aab-87a9-4628-9f17-6a8fd7d08153",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Wrapper for entity extraction and parsing\n",
    "def extract_entities_relationships(prompt):\n",
    "    try:\n",
    "        res = call_language_model(prompt)\n",
    "        return res\n",
    "    except Exception as e:\n",
    "        print(e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "726ef3e9-e66f-45ed-be59-d386ce1db4a2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import re\n",
    "import numpy as np\n",
    "\n",
    "# splitting function for chunking up filing information to avoid hitting LLM token limits\n",
    "def split_filing_info(s, chunk_size=10):\n",
    "    pattern = '(</(\\w+:)?infoTable>)'\n",
    "    splitter = re.findall(pattern, s)[0][0]\n",
    "    _parts = s.split(splitter)\n",
    "    if len(_parts) > chunk_size:\n",
    "        chunks_of_list = np.array_split(_parts, len(_parts)/chunk_size) # max 20 filings per part\n",
    "        chunks_of_str = map(lambda x: splitter.join(x)+splitter, chunks_of_list)\n",
    "        l = list(chunks_of_str)\n",
    "        if len(l) > 0:\n",
    "            l[len(l)-1] = re.sub(f'{splitter}$', '', l[len(l)-1])\n",
    "        return l\n",
    "    else:\n",
    "        return [s]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bb9dee8-c5dc-42ad-b05d-0f0d7e11d551",
   "metadata": {},
   "source": [
    "## Test Example for Parsing\n",
    "Let's start with one Form 13 file to see how we can parse it with Generative AI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57cc2dfa-4947-4afc-8b2a-6cd16338029a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "s3 = boto3.client('s3')\n",
    "obj = s3.get_object(Bucket='neo4j-dataset', Key='sec-demo/form13/raw_2023-05-15_archives_edgar_data_11544_0000899140-23-000649.txt')\n",
    "inp_text = obj['Body'].read().decode('utf-8')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d82c1ec-d49d-4444-b4fe-bc0be2777308",
   "metadata": {},
   "source": [
    "We can take a look at the file.  Note that it is an oddball mix of XML, delimeted and fixed spacing formatting that no standard parser could make sense of."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33cda3eb-25c6-47cf-8d11-e1981766aeae",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(inp_text[:1500])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b41717f-0c6e-4642-b6fd-ca327ce474cf",
   "metadata": {},
   "source": [
    "We can split data into manager and filing info pieces using `<XML>` tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11f23c86-3274-40fd-9c97-70337d8dd535",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "contents = inp_text.split('<XML>')\n",
    "manager_info = contents[1].split('</XML>')[0].strip()\n",
    "filing_info = contents[2].split('</XML>')[0].strip()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21286ffa-2ff2-40cd-8eee-8b4626f32c1f",
   "metadata": {},
   "source": [
    "## Parsing Manager Information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed35d7c5-ef70-419b-b9c3-69c34003a75e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from string import Template\n",
    "\n",
    "prompt = Template(mgr_info_tpl).substitute(ctext=manager_info)\n",
    "print(prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b13100f-de4b-49f9-89f6-58e01c19727e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# Use LLM to parse out manager information\n",
    "manager_data = json.loads(extract_entities_relationships(prompt).split('```')[1].strip('json'))\n",
    "manager_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "775186c0-d2bc-4711-88c4-0e61488ad2bd",
   "metadata": {},
   "source": [
    "## Parse Filing Information\n",
    "We will parse filing info in a similar manner to manager information. Because the filings include a list of many entries however, we will want to split the input into chunks so as not to exceed input or output token limits. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aa1690a-012b-4056-8fcc-6e04e4bb26f4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "filing_info_chunks = split_filing_info(filing_info)\n",
    "len(filing_info_chunks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "125ed4c0-22b9-4fc8-83f0-a1ccc5eaa112",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "prompt = Template(filing_info_tpl).substitute(ctext=filing_info_chunks[0][0:2])\n",
    "response = extract_entities_relationships(prompt)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62a2560b-2e6c-4b96-abc3-7c96196d0eb5",
   "metadata": {},
   "source": [
    "## Test Example\n",
    "\n",
    "Let's walk through the steps to do this with just the 1 form above first, then we can move on to parsing and ingesting multiple form13s"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c0a67de-b18f-4c65-a529-b890086663b8",
   "metadata": {},
   "source": [
    "To start we can run the LLM parsing over all the filing info from the form and then combine the resulting JSON into a list conducive for Neo4j loading.\n",
    "\n",
    "The code below may take around 12 minutes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "827621eb-2f7d-4ef4-9f00-cddbc5e2cfd2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "filings_list = []\n",
    "for filing_info_chunk in filing_info_chunks:\n",
    "    prompt = Template(filing_info_tpl).substitute(ctext=filing_info_chunk)\n",
    "    response = extract_entities_relationships(prompt)\n",
    "    if '```' in response:\n",
    "        response = response.split('```')[1].strip('json')\n",
    "    filings_list.extend(json.loads(response))\n",
    "\n",
    "for item in filings_list:\n",
    "    item['managerName'] = manager_data['managerName']\n",
    "    item['reportCalendarOrQuarter'] = manager_data['reportCalendarOrQuarter']\n",
    "filings_list[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09f6b661-7d08-42fa-9dda-6ea9e8dfcffc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "len(filings_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6b83ca6-c44b-4f43-9dcd-c88cf525a978",
   "metadata": {},
   "source": [
    "## Establish Neo4j Connection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93fffd14-e6bb-4843-90e9-d1d4e7cef3cc",
   "metadata": {
    "tags": []
   },
   "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://89d38941.databases.neo4j.io'\n",
    "NEO4J_PASSWORD = '65ye5i0QZfje9J6KU7CPLPXSghRhh3kXrNTJKqHaXd'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f38e2ef1-f5bd-4561-87e7-a014b8af2e62",
   "metadata": {
    "tags": []
   },
   "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": "97f6adf6-a299-47ca-a965-6b6cdec67f77",
   "metadata": {},
   "source": [
    "Before loading, we should create node key constraints for nodes.  This acts as a unique id and an index and is necessary for fast, efficient queries.  In general, if you notice ingestion is super slow (and getting slower) with Neo4j, double-check that you created indexes.  For this small sample, it won't matter, but it will undoubtedly impact as we ingest more data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3084e75-da30-42c1-8ca4-10cf55e7c3d2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "gds.run_cypher('CREATE CONSTRAINT unique_manager IF NOT EXISTS FOR (n:Manager) REQUIRE (n.managerName) IS NODE KEY')\n",
    "gds.run_cypher('CREATE CONSTRAINT unique_company_id IF NOT EXISTS FOR (n:Company) REQUIRE (n.cusip) IS NODE KEY')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "814834e0-54f6-4e6f-9a6e-67a526cf3bf6",
   "metadata": {},
   "source": [
    "To merge the data, we can use parameterized Cypher queries.  Basically, we will send filings in batches (in this sample case, just one batch) for each node and relationship type and insert them as parameters in the query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "feb25f22-39a0-4930-bf05-8adaa45c2f75",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Merge company nodes\n",
    "gds.run_cypher('''\n",
    "UNWIND $records AS record\n",
    "MERGE (c:Company {cusip: record.cusip})\n",
    "SET c.companyName = record.companyName\n",
    "RETURN count(c) AS company_node_merge_count\n",
    "''', params={'records':filings_list})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "635b9393-a2b6-4da6-94a4-cca6bcf515bd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Merge manager node\n",
    "gds.run_cypher('''\n",
    "MERGE (m:Manager {managerName: $name})\n",
    "RETURN count(m) AS manager_node_merge_count\n",
    "''', params={'name':manager_data['managerName']})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5394e270-45c3-44c9-96bb-6ea210dbb28a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Merge owns Relationship\n",
    "gds.run_cypher('''\n",
    "UNWIND $records AS record\n",
    "MATCH (m:Manager {managerName: record.managerName})\n",
    "MATCH (c:Company {cusip: record.cusip})\n",
    "MERGE(m)-[r:OWNS]->(c)\n",
    "SET r.reportCalendarOrQuarter = record.reportCalendarOrQuarter,\n",
    "    r.value = record.value,\n",
    "    r.shares = record.shares\n",
    "RETURN count(r) AS owns_relationship_merge_count\n",
    "''', params={'records':filings_list})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a7ebde6-7e39-40fc-8457-4d39bb45efd1",
   "metadata": {},
   "source": [
    "## Ingest Multiple Form 13 Files\n",
    "We will make a pipeline using the methods above.  In this case we will take a two-step approach, first parse all the data, then chunk that data and ingest into Neo4j."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70a0900b-0312-4213-9a3d-13e512226f93",
   "metadata": {},
   "source": [
    "For purposes of this lab we will just use a few form13 files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29c1d593-dd40-4f08-ac0d-9c8e9ff201d1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sample_file_names = [\n",
    "   'sec-demo/form13/raw_2023-05-15_archives_edgar_data_1004244_0001085146-23-002452.txt',\n",
    "   'sec-demo/form13/raw_2023-05-15_archives_edgar_data_1009232_0001009232-23-000007.txt',\n",
    "   'sec-demo/form13/raw_2023-05-15_archives_edgar_data_1009258_0000950159-23-000173.txt',\n",
    "   'sec-demo/form13/raw_2023-05-15_archives_edgar_data_1016150_0001140361-23-024548.txt',\n",
    "   'sec-demo/form13/raw_2023-05-15_archives_edgar_data_1026710_0001398344-23-009678.txt',\n",
    "   'sec-demo/form13/raw_2023-05-15_archives_edgar_data_1027451_0000919574-23-003245.txt', \n",
    "   'sec-demo/form13/raw_2023-05-15_archives_edgar_data_1027570_0001027570-23-000003.txt',\n",
    "   'sec-demo/form13/raw_2023-05-15_archives_edgar_data_104889_0000104889-23-000038.txt',\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "568c27a4-d6bf-4e12-af6e-478e20fe5ad3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Helper function for getting filing info\n",
    "def get_manager_and_filing_info(raw_txt):\n",
    "    contents = raw_txt.split('<XML>')\n",
    "    manager_info = contents[1].split('</XML>')[0].strip()\n",
    "    filing_info = contents[2].split('</XML>')[0].strip()\n",
    "    \n",
    "    return manager_info, filing_info"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a39f7ed5-d156-4191-bb1a-f82863ad14a7",
   "metadata": {},
   "source": [
    "This parsing will probably take around 45 mins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d46cc84-53b5-4088-a5aa-578f7eb5d0e4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "print(f'=== Parsing {len(sample_file_names)} Form 13 Files ===')\n",
    "\n",
    "filings_list = []\n",
    "manager_list = []\n",
    "\n",
    "for file_name in sample_file_names:\n",
    "    \n",
    "    print(f'--- parsing {file_name} ---')\n",
    "    try:\n",
    "        # Get raw form13 file\n",
    "        print('getting file text from S3....')\n",
    "        obj = s3.get_object(Bucket='neo4j-dataset', Key=file_name)\n",
    "        raw_text = obj['Body'].read().decode('utf-8')\n",
    "\n",
    "        # Get raw manager and filing info from file\n",
    "        print('getting file contents...')\n",
    "        manager_info, filing_info = get_manager_and_filing_info(raw_text)\n",
    "\n",
    "        # Parse manager info into dict using LLM\n",
    "        print('Parsing submission and manager info...')\n",
    "        mng_prompt = Template(mgr_info_tpl).substitute(ctext=manager_info)\n",
    "        mng_response = extract_entities_relationships(mng_prompt).split('```')[1].strip('json')\n",
    "        manager_data = json.loads(mng_response)\n",
    "        manager_list.append({'managerName': manager_data['managerName']})\n",
    "\n",
    "        # Parse filing info into list of dicts using LLM\n",
    "        print('Parsing filing info...')\n",
    "        tmp_filing_list = []\n",
    "        for filing_info_chunk in split_filing_info(filing_info):\n",
    "            filing_prompt = Template(filing_info_tpl).substitute(ctext=filing_info_chunk)\n",
    "            filing_response = extract_entities_relationships(filing_prompt)\n",
    "            if '```' in filing_response:\n",
    "                filing_response = filing_response.split('```')[1].strip('json')\n",
    "            tmp_filing_list.extend(json.loads(filing_response))\n",
    "        for item in tmp_filing_list: #Add information from manager_info to enable OWNS relationship loading\n",
    "            item['managerName'] = manager_data['managerName']\n",
    "            item['reportCalendarOrQuarter'] = manager_data['reportCalendarOrQuarter']\n",
    "        filings_list.extend(tmp_filing_list)\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        raise e\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b656917-8b98-4b94-a905-fcd4a4a5f8bd",
   "metadata": {},
   "source": [
    "Now we can merge the manager nodes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df9206bb-d619-4649-99cf-f5939461555a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Merge manager nodes\n",
    "gds.run_cypher('''\n",
    "UNWIND $records AS record\n",
    "MERGE (m:Manager {managerName: record.managerName})\n",
    "RETURN count(m) AS manager_node_merge_count\n",
    "''', params={'records':manager_list})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "064c27c4-4e48-4a08-8aba-26a3874e7f3e",
   "metadata": {},
   "source": [
    "For filings lets check ther length of the list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6c31f300-1c53-4795-ae1c-7228f240ddf9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "len(filings_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "197ede58-e686-42c8-90f7-fd8919a17f36",
   "metadata": {},
   "source": [
    "While we should not need chunking for this example, below is an example of how to chunk up a parameterized function for loading in case you need to scale up. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ceb716e1-c71c-4992-b654-016f9496f08b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# As the dataset gets bigger we will want to chunk up the filings we send to Neo4j\n",
    "def chunks(xs, n=10_000):\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": "e2216015-d8ca-475d-8240-a5717081c8d3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Merge company nodes\n",
    "for d in chunks(filings_list):\n",
    "    res = gds.run_cypher('''\n",
    "    UNWIND $records AS record\n",
    "    MERGE (c:Company {cusip: record.cusip})\n",
    "    SET c.companyName = record.companyName\n",
    "    RETURN count(c) AS company_node_merge_count\n",
    "    ''', params={'records':d})\n",
    "    print(res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "750215eb-ec16-4b42-909f-cf763b9955b5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Merge owns Relationships\n",
    "for d in chunks(filings_list):\n",
    "    res = gds.run_cypher('''\n",
    "    UNWIND $records AS record\n",
    "    MATCH (m:Manager {managerName: record.managerName})\n",
    "    MATCH (c:Company {cusip: record.cusip})\n",
    "    MERGE(m)-[r:OWNS]->(c)\n",
    "    SET r.reportCalendarOrQuarter = record.reportCalendarOrQuarter,\n",
    "        r.value = record.value,\n",
    "        r.shares = record.shares\n",
    "    RETURN count(r) AS owns_relationship_merge_count\n",
    "    ''', params={'records':d})\n",
    "    print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f88924cd-b43a-42ca-b975-87f1362987fc",
   "metadata": {},
   "source": [
    "This type of workflow can be applied to other unstructured data to parse entities and relationships with language models and load them into a Neo4j knowledge graph. "
   ]
  }
 ],
 "metadata": {
  "availableInstances": [
   {
    "_defaultOrder": 0,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.t3.medium",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 1,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.t3.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 2,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.t3.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 3,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.t3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 4,
    "_isFastLaunch": true,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 5,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 6,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 7,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 8,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 9,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 10,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 11,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 12,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.m5d.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 13,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.m5d.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 14,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.m5d.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 15,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.m5d.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 16,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.m5d.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 17,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.m5d.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 18,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.m5d.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 19,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.m5d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 20,
    "_isFastLaunch": false,
    "category": "General purpose",
    "gpuNum": 0,
    "hideHardwareSpecs": true,
    "memoryGiB": 0,
    "name": "ml.geospatial.interactive",
    "supportedImageNames": [
     "sagemaker-geospatial-v1-0"
    ],
    "vcpuNum": 0
   },
   {
    "_defaultOrder": 21,
    "_isFastLaunch": true,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 4,
    "name": "ml.c5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 22,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 8,
    "name": "ml.c5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 23,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.c5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 24,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.c5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 25,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 72,
    "name": "ml.c5.9xlarge",
    "vcpuNum": 36
   },
   {
    "_defaultOrder": 26,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 96,
    "name": "ml.c5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 27,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 144,
    "name": "ml.c5.18xlarge",
    "vcpuNum": 72
   },
   {
    "_defaultOrder": 28,
    "_isFastLaunch": false,
    "category": "Compute optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.c5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 29,
    "_isFastLaunch": true,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g4dn.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 30,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g4dn.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 31,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g4dn.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 32,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g4dn.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 33,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g4dn.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 34,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g4dn.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 35,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 61,
    "name": "ml.p3.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 36,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 244,
    "name": "ml.p3.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 37,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 488,
    "name": "ml.p3.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 38,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.p3dn.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 39,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.r5.large",
    "vcpuNum": 2
   },
   {
    "_defaultOrder": 40,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.r5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 41,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.r5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 42,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.r5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 43,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.r5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 44,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.r5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 45,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.r5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 46,
    "_isFastLaunch": false,
    "category": "Memory Optimized",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.r5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 47,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 16,
    "name": "ml.g5.xlarge",
    "vcpuNum": 4
   },
   {
    "_defaultOrder": 48,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.g5.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 49,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 64,
    "name": "ml.g5.4xlarge",
    "vcpuNum": 16
   },
   {
    "_defaultOrder": 50,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 128,
    "name": "ml.g5.8xlarge",
    "vcpuNum": 32
   },
   {
    "_defaultOrder": 51,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 1,
    "hideHardwareSpecs": false,
    "memoryGiB": 256,
    "name": "ml.g5.16xlarge",
    "vcpuNum": 64
   },
   {
    "_defaultOrder": 52,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 192,
    "name": "ml.g5.12xlarge",
    "vcpuNum": 48
   },
   {
    "_defaultOrder": 53,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 4,
    "hideHardwareSpecs": false,
    "memoryGiB": 384,
    "name": "ml.g5.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 54,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 768,
    "name": "ml.g5.48xlarge",
    "vcpuNum": 192
   },
   {
    "_defaultOrder": 55,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4d.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 56,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 8,
    "hideHardwareSpecs": false,
    "memoryGiB": 1152,
    "name": "ml.p4de.24xlarge",
    "vcpuNum": 96
   },
   {
    "_defaultOrder": 57,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 32,
    "name": "ml.trn1.2xlarge",
    "vcpuNum": 8
   },
   {
    "_defaultOrder": 58,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.trn1.32xlarge",
    "vcpuNum": 128
   },
   {
    "_defaultOrder": 59,
    "_isFastLaunch": false,
    "category": "Accelerated computing",
    "gpuNum": 0,
    "hideHardwareSpecs": false,
    "memoryGiB": 512,
    "name": "ml.trn1n.32xlarge",
    "vcpuNum": 128
   }
  ],
  "environment": {
   "kernel": "conda-root-py",
   "name": "workbench-notebooks.m111",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/workbench-notebooks:m111"
  },
  "instance_type": "ml.t3.medium",
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}