08_EntityExtraction/entitiy_extraction.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "46fa37fd-e892-4504-ad32-edabb4760596",
   "metadata": {},
   "source": [
    "# Entity Extraction with Claude\n",
    "\n",
    "> *This notebook should work well with the **`Data Science 3.0`** kernel in SageMaker Studio*\n",
    "\n",
    "### Context\n",
    "Entity extraction is an NLP technique that allows us to automatically extract specific data from naturally written text, such as news, emails, books, etc.\n",
    "That data can then later be saved to a database, used for lookup or any other type of processing.\n",
    "\n",
    "Classic entity extraction programs usually limit you to pre-defined classes, such as name, address, price, etc. or require you to provide many examples of types of entities you are interested in.\n",
    "By using a LLM for entity extraction in most cases you are only required to specify what you need to extract in natural language. This gives you flexibility and accuracy in your queries while saving time by removing necessity of data labeling.\n",
    "\n",
    "In addition, LLM entity extraction can be used to help you assemble a dataset to later creat a customised solution for your use case, such as [Amazon Comprehend custom entity](https://docs.aws.amazon.com/comprehend/latest/dg/custom-entity-recognition.html) recognition."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "373675b6-cdc4-437e-83b5-7d897516b8fc",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Before running the rest of this notebook, you'll need to run the cells below to (ensure necessary libraries are installed and) connect to Bedrock.\n",
    "\n",
    "For more details on how the setup works and ⚠️ **whether you might need to make any changes**, refer to the [Bedrock boto3 setup notebook](../00_Intro/bedrock_boto3_setup.ipynb) notebook.\n",
    "\n",
    "In this notebook, we'll also need some extra dependencies:\n",
    "\n",
    "- [Beautiful Soup](https://www.crummy.com/software/BeautifulSoup/), to easily extract data from XML tags in Claude prompts and outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5e1cbda6-1a86-4183-bba2-576aa1ab9d1a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "source": [
    "%pip install --no-build-isolation --force-reinstall \\\n",
    "    \"boto3>=1.28.57\" \\\n",
    "    \"awscli>=1.29.57\" \\\n",
    "    \"botocore>=1.31.57\" \\\n",
    "    langchain==0.0.309 beautifulsoup4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b19a1c78",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d8ef0441-b424-403e-9394-d81b64e8332b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Create new client\n",
      "  Using region: us-east-1\n",
      "boto3 Bedrock client successfully created!\n",
      "bedrock(https://bedrock.us-east-1.amazonaws.com)\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "module_path = \"..\"\n",
    "sys.path.append(os.path.abspath(module_path))\n",
    "from utils import bedrock\n",
    "\n",
    "\n",
    "# ---- ⚠️ Un-comment and edit the below lines as needed for your AWS setup ⚠️ ----\n",
    "\n",
    "# os.environ[\"AWS_DEFAULT_REGION\"] = \"<REGION_NAME>\"  # E.g. \"us-east-1\"\n",
    "# os.environ[\"AWS_PROFILE\"] = \"<YOUR_PROFILE>\"\n",
    "# os.environ[\"BEDROCK_ASSUME_ROLE\"] = \"<YOUR_ROLE_ARN>\"  # E.g. \"arn:aws:...\"\n",
    "\n",
    "\n",
    "boto3_bedrock = bedrock.get_bedrock_client(\n",
    "    assumed_role=os.environ.get(\"BEDROCK_ASSUME_ROLE\", None),\n",
    "    region=os.environ.get(\"AWS_DEFAULT_REGION\", None),\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1fb9074b-d72e-4419-9165-421414d28f4b",
   "metadata": {},
   "source": [
    "## Configure langchain\n",
    "\n",
    "We begin with instantiating the LLM. Here we are using Anthropic Claude v2 for text generation.\n",
    "\n",
    "Note: It is possible to choose other models available with Bedrock. You can replace the `model_id` as follows to change the model.\n",
    "\n",
    "`llm = Bedrock(model_id=\"amazon.titan-tg1-large\")`\n",
    "\n",
    "Available model IDs include:\n",
    "\n",
    "- `amazon.titan-tg1-large`\n",
    "- `ai21.j2-grande-instruct`\n",
    "- `ai21.j2-jumbo-instruct`\n",
    "- `anthropic.claude-instant-v1`\n",
    "- `anthropic.claude-v1`\n",
    "- `anthropic.claude-v2`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "621790c0-332a-4bab-bf81-967a63cb52fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.llms.bedrock import Bedrock\n",
    "\n",
    "# - create the Anthropic Model\n",
    "llm = Bedrock(\n",
    "    model_id=\"anthropic.claude-v2\",\n",
    "    client=boto3_bedrock,\n",
    "    model_kwargs={\n",
    "        \"max_tokens_to_sample\": 200,\n",
    "        \"temperature\": 0, # Using 0 to get reproducible results\n",
    "        \"stop_sequences\": [\"\\n\\nHuman:\"]\n",
    "    }\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5f85e961-3530-4bf4-ac28-12611965d408",
   "metadata": {},
   "source": [
    "## Entity Extraction\n",
    "Now that we have our LLM initialised, we can start extracting entities.\n",
    "\n",
    "For this exercise we will pretend to be an online bookstore that receives questions and orders by email.\n",
    "Our task would be to extract relevant information from the email to process the order.\n",
    "\n",
    "Let's begin by taking a look at the sample email:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "b958f4c7-0ded-4537-9939-d1623337317f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dear Sir or Madam,\n",
      "\n",
      "I would like to order Treasure Island, do you have it in stock?\n",
      "\n",
      "Also, is it possible to pay by cheque?\n",
      "\n",
      "Yours sincerely,\n",
      "John Smith\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "emails_dir = Path(\".\") / \"emails\"\n",
    "with open(emails_dir / \"00_treasure_island.txt\") as f:\n",
    "    book_question_email = f.read()\n",
    "\n",
    "print(book_question_email)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "59f62564-cd46-4bff-bda3-c0f29a47dd9d",
   "metadata": {},
   "source": [
    "### Basic approach\n",
    "\n",
    "For basic cases we can directly ask the model to return the result.\n",
    "Let's try extracting the name of the book."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "efdc9062-64e9-4634-855c-d06ccb5efb50",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = f\"\"\"\n",
    "\n",
    "Human: Given the email inside triple-backticks, please read it and analyse the contents.\n",
    "If a name of a book is mentioned, return it, otherwise return nothing.\n",
    "\n",
    "Email: ```\n",
    "{book_question_email}\n",
    "```\n",
    "\n",
    "Assistant:\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4742618e-25e9-441e-a6f8-b47330a0bd05",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Treasure Island\n"
     ]
    }
   ],
   "source": [
    "result = llm(query)\n",
    "print(result.strip())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e31a3407-caca-445a-bb1a-d62d40ddccd2",
   "metadata": {},
   "source": [
    "### Model specific prompts\n",
    "\n",
    "While basic approach works, to achieve best results we recommend to customise your prompts for the particular model you will be using.\n",
    "In this example we are using `anthropic.claude-v2`, [prompt guide for which can be found here](https://docs.anthropic.com/claude/docs/introduction-to-prompt-design).\n",
    "\n",
    "Here is the a more optimised prompt for Claude v2."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a5a461a9-4bad-4634-b568-a07769b1d349",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"\n",
    "\n",
    "Human: Given the email provided, please read it and analyse the contents.\n",
    "If a name of a book is mentioned, return it.\n",
    "If no name is mentioned, return empty string.\n",
    "The email will be given between <email></email> XML tags.\n",
    "\n",
    "<email>\n",
    "{email}\n",
    "</email>\n",
    "\n",
    "Return the name of the book between <book></book> XML tags.\n",
    "\n",
    "Assistant:\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5142cff9-c2d3-451e-8d21-06ce8538adb5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<book>Treasure Island</book>\n"
     ]
    }
   ],
   "source": [
    "query = prompt.format(email=book_question_email)\n",
    "result = llm(query)\n",
    "print(result.strip())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e87ee5ab-33d9-4def-a462-8e5992032bd0",
   "metadata": {},
   "source": [
    "To extract results easier, we can use a helper function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bfa9d2d0-2bc8-465c-b89b-73b2fd76d4b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bs4 import BeautifulSoup\n",
    "\n",
    "def extract_by_tag(response: str, tag: str, extract_all=False) -> str | list[str] | None:\n",
    "    soup = BeautifulSoup(response)\n",
    "    results = soup.find_all(tag)\n",
    "    if not results:\n",
    "        return\n",
    "        \n",
    "    texts = [res.get_text() for res in results]\n",
    "    if extract_all:\n",
    "        return texts\n",
    "    return texts[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "fef7b280-71be-41ad-9f21-8c87d09226ae",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'Treasure Island'"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "extract_by_tag(result, \"book\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f19454e6-22cd-41ee-888c-39801cc72c74",
   "metadata": {},
   "source": [
    "We can check that our model doesn't return arbitrary results when no appropriate information is given (also know as 'hallucination'), by running our prompt on other emails."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "35fd1343-9b4b-4efd-846f-1312af18e15c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I didn't like the last book I ordered and would like to return it.\n"
     ]
    }
   ],
   "source": [
    "with open(emails_dir / \"01_return.txt\") as f:\n",
    "    return_email = f.read()\n",
    "\n",
    "print(return_email)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "4ec8a1b2-beb4-4ddf-9935-1fc7a3b08729",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<book></book>\n"
     ]
    }
   ],
   "source": [
    "query = prompt.format(email=return_email)\n",
    "result = llm(query)\n",
    "print(result.strip())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d154c270-41dc-4e58-bca2-f9fe5d021223",
   "metadata": {},
   "source": [
    "Using tags also allows us to extract multiple pieces of information at the same time and makes extraction much easier.\n",
    "In the following prompt we will extract not just the book name, but any questions, requests and customer name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "ea5b5c9b-b0c0-427d-a7fb-005253e9bbb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"\"\"\n",
    "\n",
    "Human: Given email provided , please read it and analyse the contents.\n",
    "\n",
    "Please extract the following information from the email:\n",
    "- Any questions the customer is asking, return it inside <questions></questions> XML tags.\n",
    "- The customer full name, return it inside <name></name> XML tags.\n",
    "- Any book names the customer mentions, return it inside <books></books> XML tags.\n",
    "\n",
    "If a particular bit of information is not present, return an empty string.\n",
    "Make sure that each question can be understoon by itself, incorporate context if requred.\n",
    "Each returned question should be concise, remove extra information if possible.\n",
    "The email will be given between <email></email> XML tags.\n",
    "\n",
    "<email>\n",
    "{email}\n",
    "</email>\n",
    "\n",
    "Return each question inside <question></question> XML tags.\n",
    "Return the name of each book inside <book></book> XML tags.\n",
    "\n",
    "Assistant:\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ac605eb5-2483-46ed-a205-6932051c8d2b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<name>John Smith</name>\n",
      "<books>\n",
      "<book>Treasure Island</book> \n",
      "</books>\n",
      "<questions>\n",
      "<question>Do you have Treasure Island in stock?</question>\n",
      "<question>Is it possible to pay by cheque?</question>\n",
      "</questions>\n"
     ]
    }
   ],
   "source": [
    "query = prompt.format(email=book_question_email)\n",
    "result = llm(query)\n",
    "print(result.strip())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "8a7cc2cb-8036-44a5-9fb6-db2172f9b601",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Do you have Treasure Island in stock?', 'Is it possible to pay by cheque?']"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "extract_by_tag(result, \"question\", extract_all=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "f5617e0d-0923-45b6-8e91-03748ad76d31",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'John Smith'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "extract_by_tag(result, \"name\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "66852eb2-97a2-4041-a76f-3fb03e1aaef5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Treasure Island']"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "extract_by_tag(result, \"book\", extract_all=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "d830e149-5c89-4f50-9833-b499ee70f3f3",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "Entity extraction is a powerful technique using which you can extract arbitrary data using plain text descriptions.\n",
    "\n",
    "This is particularly useful when you need to extract specific data which doesn't have clear structure. In such cases regex and other traditional extraction techniques can be very difficult to implement.\n",
    "\n",
    "### Take aways\n",
    "- Adapt this notebook to experiment with different models available through Amazon Bedrock such as Amazon Titan and AI21 Labs Jurassic models.\n",
    "- Change the prompts to your specific usecase and evaluate the output of different models.\n",
    "- Apply different prompt engineering principles to get better outputs. Refer to the prompt guide for your chosen model for recommendations, e.g. [here is the prompt guide for Claude](https://docs.anthropic.com/claude/docs/introduction-to-prompt-design)."
   ]
  }
 ],
 "metadata": {
  "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.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}