openai-cookbook/examples/Entity_extraction_for_long_documents.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Long Document Content Extraction\n",
    "\n",
    "We often need GPT-3 to help us extract key figures, dates or other bits of important content from documents that are too big to fit into the context window. One approach for solving this is to chunk the document up and process each chunk separately, before combining into one list of answers. \n",
    "\n",
    "In this notebook we'll run through this approach:\n",
    "- We'll load in a long PDF and pull the text out\n",
    "- We'll create a prompt to be used to extract key bits of information\n",
    "- We'll chunk up our document and process each chunk to pull any answers out\n",
    "- We'll then combine them at the end\n",
    "- This simple approach will then be extended to three more difficult questions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We'll take our PDF, a Formula 1 Financial Regulation document on Power Units, and extract the text from it for entity extraction. We'll use this to try to extract answers that are buried in the content."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install textract"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import textract\n",
    "import os\n",
    "import openai\n",
    "\n",
    "# Extract the raw text from each PDF using textract\n",
    "text = textract.process('data/fia_f1_power_unit_financial_regulations_issue_1_-_2022-08-16.pdf', method='pdfminer').decode('utf-8')\n",
    "clean_text = text.replace(\"  \", \" \").replace(\"\\n\", \"; \").replace(';',' ')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple Entity Extraction\n",
    "\n",
    "We'll start by extracting key bits of information from chunks of a document.\n",
    "\n",
    "We'll accomplish this by:\n",
    "- Creating a template prompt with our questions and an example of the format it expects\n",
    "- Create a function to take a chunk of text as input, combine with the prompt and get a response\n",
    "- Run a script to chunk the text, extract answers and output them for parsing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extract key pieces of information from this regulation document. \n",
      "If a particular piece of information is not present, output \"Not specified\". \n",
      "When you extract a key piece of information, include the closest page number.\n",
      "Use the following format:\n",
      "0. Author\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR\n",
      "2. What is the value of External Manufacturing Costs in USD\n",
      "3. What is the Capital Expenditure Limit in USD\n",
      "\n",
      "Document: \"\"\"<document>\"\"\"\n",
      "\n",
      "0. Author: Tom Anderson (Page 1)\n",
      "1.\n"
     ]
    }
   ],
   "source": [
    "# Example prompt - \n",
    "document = '<document>'\n",
    "template_prompt=f'''Extract key pieces of information from this regulation document. \n",
    "If a particular piece of information is not present, output \\\"Not specified\\\". \n",
    "When you extract a key piece of information, include the closest page number.\n",
    "Use the following format:\\n0. Author\\n1. What is the amount of the power unit cost cap in USD, GBP and EUR\\n2. What is the value of External Manufacturing Costs in USD\\n3. What is the Capital Expenditure Limit in USD\\n\\nDocument: \\\"\\\"\\\"{document}\\\"\\\"\\\"\\n\\n0. Author: Tom Anderson (Page 1)\\n1.'''\n",
    "print(template_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import openai\n",
    "\n",
    "\n",
    "def extract_chunk(document,template_prompt):\n",
    "    \n",
    "    prompt=template_prompt.replace('<document>',document)\n",
    "\n",
    "    response = openai.Completion.create(\n",
    "    model='text-davinci-003', \n",
    "    prompt=prompt,\n",
    "    temperature=0,\n",
    "    max_tokens=1500,\n",
    "    top_p=1,\n",
    "    frequency_penalty=0,\n",
    "    presence_penalty=0\n",
    "    )\n",
    "    return \"1.\" + response['choices'][0]['text']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: USD 95,000,000 (Page 3) GBP 76,459,000 (Page 4) EUR 90,210,000 (Page 4)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 1)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 1)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 1)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 8)\n",
      "2. What is the value of External Manufacturing Costs in USD: US Dollars 20,000,000 in respect of each of the Full Year Reporting Periods ending on 31 December 2023, 31 December 2024 and 31 December 2025, adjusted for Indexation (Page 10)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 10)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 1)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 1)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 1)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 15)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 15)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 15)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 18)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 18)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 18)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 1)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 1)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 1)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 25)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 25)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 25)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 29)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 29)\n",
      "3. What is the Capital Expenditure Limit in USD: US Dollars 30,000,000 (Page 32)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 34)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 35)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 35)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 36)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 36)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 36)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 41)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 41)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 41)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 45)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 45)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 45)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 48)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 48)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 49)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 1)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 1)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 56)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified (Page 1)\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified (Page 1)\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified (Page 1)\n",
      "1. What is the amount of the power unit cost cap in USD, GBP and EUR: Not specified\n",
      "2. What is the value of External Manufacturing Costs in USD: Not specified\n",
      "3. What is the Capital Expenditure Limit in USD: Not specified\n"
     ]
    }
   ],
   "source": [
    "chunks = []\n",
    "results = []\n",
    "\n",
    "# chunk up the contract into consecutive 10k character chunks\n",
    "for i in range(0, len(clean_text), 10000):\n",
    "    #print(clean_text[i:i+10000])\n",
    "    chunks.append(clean_text[i:i+10000])\n",
    "\n",
    "for chunk in chunks:\n",
    "    results.append(extract_chunk(chunk,template_prompt))\n",
    "    print(results[-1])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['1. What is the amount of the power unit cost cap in USD, GBP and EUR: USD 95,000,000 (Page 3) GBP 76,459,000 (Page 4) EUR 90,210,000 (Page 4)',\n",
       " '2. What is the value of External Manufacturing Costs in USD: US Dollars 20,000,000 in respect of each of the Full Year Reporting Periods ending on 31 December 2023, 31 December 2024 and 31 December 2025, adjusted for Indexation (Page 10)',\n",
       " '3. What is the Capital Expenditure Limit in USD: US Dollars 30,000,000 (Page 32)']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "groups = [r.split('\\n') for r in results]\n",
    "\n",
    "# zip the groups together\n",
    "zipped = list(zip(*groups))\n",
    "zipped = [x for y in zipped for x in y if \"Not specified\" not in x and \"__\" not in x]\n",
    "zipped"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Complex Entity Extraction\n",
    "\n",
    "Now we'll try to answer more difficult questions that need tougher logic to work out.\n",
    "\n",
    "We'll update the prompt with three more questions and see whether our model can work them out."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Extract key pieces of information from this regulation document. \n",
      "If a particular piece of information is not present, output \"Not specified\". \n",
      "When you extract a key piece of information, include the closest page number.\n",
      "Use the following format:\n",
      "0. Author\n",
      "1. How is a minor overspend breach calculated\n",
      "2. How is a major overspend breach calculated\n",
      "3.Which years do these financial regulations apply to\n",
      "\n",
      "Document: \"\"\"<document>\"\"\"\n",
      "\n",
      "0. Author: Tom Anderson (Page 1)\n",
      "1.\n"
     ]
    }
   ],
   "source": [
    "# Example prompt - \n",
    "template_prompt=f'''Extract key pieces of information from this regulation document. \n",
    "If a particular piece of information is not present, output \\\"Not specified\\\". \n",
    "When you extract a key piece of information, include the closest page number.\n",
    "Use the following format:\\n0. Author\\n1. How is a minor overspend breach calculated\\n2. How is a major overspend breach calculated\\n3.Which years do these financial regulations apply to\\n\\nDocument: \\\"\\\"\\\"{document}\\\"\\\"\\\"\\n\\n0. Author: Tom Anderson (Page 1)\\n1.'''\n",
    "print(template_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['1. How is a minor overspend breach calculated: Relevant Costs reported exceed the Power Unit Cost Cap by less than 5% (Page 24)',\n",
       " '2. How is a major overspend breach calculated: Relevant Costs reported exceed the Power Unit Cost Cap by more than 5% (Page 24)',\n",
       " '3. Which years do these financial regulations apply to: 2023-2026 (Page 2)',\n",
       " '3. Which years do these financial regulations apply to: 2022 (Page 4)',\n",
       " '3. Which years do these financial regulations apply to: 2022 (Page 8)',\n",
       " '3. Which years do these financial regulations apply to: 2022 (Page 1)',\n",
       " '3. Which years do these financial regulations apply to: 2022 (Page 1)',\n",
       " '3.Which years do these financial regulations apply to: 2022 (Page 1)',\n",
       " '3. Which years do these financial regulations apply to: 2023, 2024, 2025, 2026 onwards (Page 34)',\n",
       " '3.Which years do these financial regulations apply to: 2022 (Page 1)',\n",
       " '3.Which years do these financial regulations apply to: 2026 onwards (Page 41)',\n",
       " '3.Which years do these financial regulations apply to: 2026 to 2030 seasons (inclusive) (Page 46)',\n",
       " '3. Which years do these financial regulations apply to: 2022 (Page 1)',\n",
       " '3. Which years do these financial regulations apply to: 2022 (Page 57)',\n",
       " '3. Which years do these financial regulations apply to: 2022 (Page 61)']"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chunks = []\n",
    "results = []\n",
    "\n",
    "for i in range(0, len(clean_text), 10000):\n",
    "    chunks.append(clean_text[i:i+10000])\n",
    "\n",
    "for chunk in chunks:\n",
    "    results.append(extract_chunk(chunk,template_prompt))\n",
    "    \n",
    "groups = [r.split('\\n') for r in results]\n",
    "\n",
    "# zip the groups together\n",
    "zipped = list(zip(*groups))\n",
    "zipped = [x for y in zipped for x in y if \"Not specified\" not in x and \"__\" not in x]\n",
    "zipped"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Consolidation\n",
    "\n",
    "We've been able to extract the first two answers safely, while the third was confounded by the date that appeared on every page, though the correct answer is in there as well.\n",
    "\n",
    "To tune this further you can consider experimenting with:\n",
    "- A more descriptive or specific prompt\n",
    "- Negative prompting, telling the model to ignore a particular thing i.e. the footnote date on every page\n",
    "- If you have sufficient training data, fine-tuning a model to find a set of outputs very well\n",
    "- The way you chunk your data - we have gone for 10,000 characters with no overlap, but more intelligent chunking that breaks info into sections, cuts by tokens or similar may get better results\n",
    "\n",
    "However, with minimal tuning we have now answered 6 questions of varying difficulty using the contents of a long document, and have a reusable approach that we can apply to any long document requiring entity extraction. Look forward to seeing what you can do with this!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "embed_retrieve",
   "language": "python",
   "name": "embed_retrieve"
  },
  "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.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "5997d090960a54cd76552f75eca12ec3b416cf9d01a1a5af08ae48cf90878791"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}