00-Entity recognition.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/libbyseline/nicar24/blob/main/00-Entity%20recognition.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c4d92934-3169-472f-91c3-4cebe9a99e3d",
      "metadata": {
        "id": "c4d92934-3169-472f-91c3-4cebe9a99e3d"
      },
      "outputs": [],
      "source": [
        "!wget --quiet -O items.zip https://github.com/jsoma/nicar24-beyond-chatgpt/raw/main/items.zip\n",
        "!unzip -o -q items.zip"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "5e6a45ed-fe04-4e76-adfc-2b89f134b928",
      "metadata": {
        "id": "5e6a45ed-fe04-4e76-adfc-2b89f134b928"
      },
      "source": [
        "# Named entity recognition\n",
        "\n",
        "Finding names, dates, places, and more. If you want to do this without programming, [Google Pinpoint](https://journaliststudio.google.com/pinpoint/collections) will work great (until they [unceremoniously shut it down](https://killedbygoogle.com/))."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b5a2fccf-2082-4a42-b228-292fc02e2d76",
      "metadata": {
        "id": "b5a2fccf-2082-4a42-b228-292fc02e2d76"
      },
      "source": [
        "## spaCy\n",
        "\n",
        "After the dark days of NLTK but before the rise of ChatGPT, [spaCy](https://spacy.io/) was the champion of text analysis. Still is, in many ways! It has great documentation, is very fast and easy to use – everything you could ask for in a piece of software.\n",
        "\n",
        "### Installing spacy\n",
        "\n",
        "spaCy is a Python package, you install it the normal way with `pip`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a33e3cc2-f220-42cb-bf8b-01896644c320",
      "metadata": {
        "id": "a33e3cc2-f220-42cb-bf8b-01896644c320"
      },
      "outputs": [],
      "source": [
        "!pip install --quiet spacy"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "146ea24d-cfbb-460f-b483-47c969e7f7f1",
      "metadata": {
        "id": "146ea24d-cfbb-460f-b483-47c969e7f7f1"
      },
      "source": [
        "Before we use spaCy, we first need to download a *model*. spaCy itself doesn't understand language or text, it's just an interface! We'll end up using different models based on the languages we're analyzing and how important speed is.\n",
        "\n",
        "Down below we use `en_core_web_sm`, which is a small, quick model for the English language."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3ca86070-682a-4671-b450-ae3c176ed241",
      "metadata": {
        "id": "3ca86070-682a-4671-b450-ae3c176ed241"
      },
      "outputs": [],
      "source": [
        "!python -m spacy download en_core_web_sm"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "2f024199-b417-4f1f-a987-610f1b99607e",
      "metadata": {
        "id": "2f024199-b417-4f1f-a987-610f1b99607e"
      },
      "outputs": [],
      "source": [
        "import spacy\n",
        "\n",
        "nlp = spacy.load('en_core_web_sm')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "13b08e4b-de71-4ad6-80f3-435e0e616402",
      "metadata": {
        "id": "13b08e4b-de71-4ad6-80f3-435e0e616402"
      },
      "outputs": [],
      "source": [
        "text = \"\"\"\n",
        "Trump Seeks to Bump Haley From Race, Eyeing November Rematch With Biden\n",
        "\n",
        "It is a seemingly do-or-die moment for Nikki Haley’s campaign, as Donald Trump’s\n",
        "probable victory lines him up to face off with President Biden in the fall.\n",
        "\"\"\"\n",
        "\n",
        "nlp = spacy.load('en_core_web_sm')\n",
        "doc = nlp(text)\n",
        "\n",
        "for entity in doc.ents:\n",
        "  print(entity.label_, ' | ', entity.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "55bb4da4-8a71-4348-b923-68a1f3a7900a",
      "metadata": {
        "id": "55bb4da4-8a71-4348-b923-68a1f3a7900a"
      },
      "source": [
        "You can find [a list of models on the spaCy site](https://spacy.io/usage/models). We might not be not very impressed with the results above, so we'll try one that's bigger, better, fancier..! and slower: `en_core_web_trf`.\n",
        "\n",
        "We first need to download `en_core_web_trf` in the same way we downloaded the smaller model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "609a1be2-4713-4d88-bc62-29722c86660e",
      "metadata": {
        "id": "609a1be2-4713-4d88-bc62-29722c86660e"
      },
      "outputs": [],
      "source": [
        "!python -m spacy download en_core_web_trf"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "96803752-3d46-4a5d-b204-64d6e158c208",
      "metadata": {
        "id": "96803752-3d46-4a5d-b204-64d6e158c208"
      },
      "source": [
        "And we'll use it in the same way as the last one, only changing the line where we load it in through spaCy."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "e7bcf260-e28a-403f-afa9-eee295ebaaa4",
      "metadata": {
        "id": "e7bcf260-e28a-403f-afa9-eee295ebaaa4"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "\n",
        "text = \"\"\"\n",
        "Trump Seeks to Bump Haley From Race, Eyeing November Rematch With Biden\n",
        "\n",
        "It is a seemingly do-or-die moment for Nikki Haley’s campaign, as Donald Trump’s\n",
        "probable victory lines him up to face off with President Biden in the fall.\n",
        "\"\"\"\n",
        "\n",
        "nlp = spacy.load('en_core_web_trf')\n",
        "doc = nlp(text)\n",
        "\n",
        "for entity in doc.ents:\n",
        "  print(entity.label_, ' | ', entity.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ecb82e53-22da-4734-909b-a1963bce5d1d",
      "metadata": {
        "id": "ecb82e53-22da-4734-909b-a1963bce5d1d"
      },
      "source": [
        "Is that great? I don't know, *maybe*. If we poke around on [the packages link to all of the available English models](https://spacy.io/models/en) we see a couple more: `en_core_web_md` and `en_core_web_lg`. Might as well give the large one a try?"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3d96ed0d-9219-41eb-b0bc-215e99c9c053",
      "metadata": {
        "id": "3d96ed0d-9219-41eb-b0bc-215e99c9c053"
      },
      "outputs": [],
      "source": [
        "!python -m spacy download en_core_web_lg"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1213ec98-5558-4e88-91b6-3b91368e692c",
      "metadata": {
        "id": "1213ec98-5558-4e88-91b6-3b91368e692c"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "\n",
        "text = \"\"\"\n",
        "Trump Seeks to Bump Haley From Race, Eyeing November Rematch With Biden\n",
        "\n",
        "It is a seemingly do-or-die moment for Nikki Haley’s campaign, as Donald Trump’s\n",
        "probable victory lines him up to face off with President Biden in the fall.\n",
        "\"\"\"\n",
        "\n",
        "nlp = spacy.load('en_core_web_lg')\n",
        "doc = nlp(text)\n",
        "\n",
        "for entity in doc.ents:\n",
        "  print(entity.label_, ' | ', entity.text)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b943efed-0ea4-4a17-9bc9-beb07f15e246",
      "metadata": {
        "id": "b943efed-0ea4-4a17-9bc9-beb07f15e246"
      },
      "source": [
        "While we like to think of computers as always giving the One True Right Answer, text analysis is a little different: it's much closer to **the model's opinion** than anything else.\n",
        "\n",
        "People do spend time comparing models to humans, though! Take a look at the **Accuracy evaluation** tabs on the [models page](https://spacy.io/models/en) to see how each of the models compare on different tasks."
      ]
    },
    {
      "cell_type": "markdown",
      "id": "aeb5344c-a56b-4476-b05e-961611eefd9e",
      "metadata": {
        "id": "aeb5344c-a56b-4476-b05e-961611eefd9e"
      },
      "source": [
        "## Hugging Face Transformers\n",
        "\n",
        "The `_trf` in `en_core_web_trf` stands for *transformers*, which is a fancy approach to machine learning that showed up well after spaCy was created. The introduction of transformers is a big reason why we've seen such an explosion in AI and machine learning since 2017!\n",
        "\n",
        "While you *can* use spaCy with these fancy new models, it's not really what most people do. Instead, there's a much bigger ecosystem that you can find browsing around at [Hugging Face](https://huggingface.co/). You're no longer restricted to the four model options spaCy presented to you – everyone from Microsoft to random people upload new models every day.\n",
        "\n",
        "Hugging Face even has a package called `transformers` that helps you use these new models. Let's install it!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f5822928-5868-4b84-9b71-fdc94cfc7ff6",
      "metadata": {
        "id": "f5822928-5868-4b84-9b71-fdc94cfc7ff6"
      },
      "outputs": [],
      "source": [
        "!pip install transformers torch"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "9e378390-7d38-484f-a9ae-0c357d3edb01",
      "metadata": {
        "id": "9e378390-7d38-484f-a9ae-0c357d3edb01"
      },
      "source": [
        "Which model should we use? Well... [how about we just browse?](https://huggingface.co/models)\n",
        "\n",
        "You *immediately* see things are different when we're using Hugging Face as compared to spaCy. Even though spaCy wasn't the modern cutting-edge, it was at least easy to know what was going on! Now that we're in the Hugging Face world we're dropped in to selecting a model without a clue what's going on.\n",
        "\n",
        "The models page has models that do all sorts of things: generate images, categorize text, transcribe audio, and a million other options. Pulling names and places out of text is called **entity extraction** or **named entity recognition** (NER for short). It also falls under **token classification** problem, which is (kind of) the idea of putting words into categories.\n",
        "\n",
        "We can go [search for token classification models](https://huggingface.co/models?pipeline_tag=token-classification&sort=trending), but then we have to decide which one to pick. The hottest trending one? The most downloaded one? The most recent?\n",
        "\n",
        "Most of the time there's no way to know the right answer: **you just try them out and see what works best!**\n",
        "\n",
        "> That's kind of a lie, you can always look at [what's currently state-of-the-art on benchmarks](https://paperswithcode.com/sota/token-classification-on-conll2003).\n",
        "\n",
        "We can start with [dslim/bert-base-NER](https://huggingface.co/dslim/bert-base-NER) since it's (currently) on top of the list of trending models.\n",
        "\n",
        "We can see a suggestion on how to use the model in the column of text to the left, but I like to click the **Use in Transformers** button to see a little quick-start version. It's usually shorter!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "367d17bf-89a3-4591-ad77-c0ec4905c697",
      "metadata": {
        "id": "367d17bf-89a3-4591-ad77-c0ec4905c697"
      },
      "outputs": [],
      "source": [
        "# https://huggingface.co/dslim/bert-base-NER\n",
        "# I'm going to use \"ner\" instead of \"token-classification\"\n",
        "nlp = pipeline(\"ner\", model=\"dslim/bert-base-NER\")\n",
        "\n",
        "results = nlp(text)\n",
        "for entity in results:\n",
        "    print(entity)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4d72af71-3688-4082-b181-b759fdfbf7f0",
      "metadata": {
        "id": "4d72af71-3688-4082-b181-b759fdfbf7f0"
      },
      "source": [
        "It splits up the *pieces* of Biden for some reason, along with 'Nikki' and 'Haley' being split. If you scour the internet you can find that you need to also tell it the **aggregation strategy**, ways of combining entities that are next to each other."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f5361181-50ff-434c-9866-bb3b9fa5b3dd",
      "metadata": {
        "id": "f5361181-50ff-434c-9866-bb3b9fa5b3dd"
      },
      "outputs": [],
      "source": [
        "# https://huggingface.co/dslim/bert-base-NER\n",
        "nlp = pipeline(\"token-classification\",\n",
        "               model=\"dslim/bert-base-NER\",\n",
        "               aggregation_strategy='simple')\n",
        "\n",
        "results = nlp(text)\n",
        "for entity in results:\n",
        "    print(entity)\n",
        "    # print(entity['word'], entity['entity_group'])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "d5af45e9-888e-4610-9e5c-a4d490ad7705",
      "metadata": {
        "id": "d5af45e9-888e-4610-9e5c-a4d490ad7705"
      },
      "source": [
        "It also suggests a [larger model](https://huggingface.co/dslim/bert-large-NER), `bert-large` instead of `bert-base`. We can try that, I guess??"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "df2a058d-958c-4fbb-a131-fd1578ce2407",
      "metadata": {
        "id": "df2a058d-958c-4fbb-a131-fd1578ce2407"
      },
      "outputs": [],
      "source": [
        "# https://huggingface.co/dslim/bert-large-NER\n",
        "nlp = pipeline(\"ner\",\n",
        "               model=\"dslim/bert-large-NER\",\n",
        "               aggregation_strategy='simple')\n",
        "\n",
        "results = nlp(text)\n",
        "for entity in results:\n",
        "    print(entity['word'], '|', entity['entity_group'])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "68967d0a-8526-4d55-98e7-7faaabefb040",
      "metadata": {
        "id": "68967d0a-8526-4d55-98e7-7faaabefb040"
      },
      "source": [
        "Seems pretty good, yeah?\n",
        "\n",
        "If we go back to the [models page and sort by most downloads](https://huggingface.co/models?pipeline_tag=token-classification&sort=downloads), we see another model to try: [Jean-Baptiste/roberta-large-ner-english](https://huggingface.co/Jean-Baptiste/roberta-large-ner-english). Might as well give it a shot while we're here!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "433fc260-5d9d-43f6-a641-d4237936feb5",
      "metadata": {
        "id": "433fc260-5d9d-43f6-a641-d4237936feb5"
      },
      "outputs": [],
      "source": [
        "from transformers import pipeline\n",
        "\n",
        "# https://huggingface.co/Jean-Baptiste/roberta-large-ner-english\n",
        "nlp = pipeline(\"ner\",\n",
        "               model=\"Jean-Baptiste/roberta-large-ner-english\",\n",
        "               aggregation_strategy='simple')\n",
        "\n",
        "results = nlp(text)\n",
        "for entity in results:\n",
        "    print(entity['word'], '|', entity['entity_group'])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f12c6fdd-a53c-44f7-b699-72cfd0f6f6d2",
      "metadata": {
        "id": "f12c6fdd-a53c-44f7-b699-72cfd0f6f6d2"
      },
      "source": [
        "If we go back to the [bert-large-NER page](https://huggingface.co/dslim/bert-large-NER) hiding on the bottom right-hand side, we see a section labeled **Evaluation results**. If we click the **View on Papers With Code** link, we end up on [a leaderboard for NER tasks](https://paperswithcode.com/sota/token-classification-on-conll2003).\n",
        "\n",
        "Sitting pretty at the top of the list is `microsoft-deberta-v3-large_ner_conll2003` from Microsoft. If we click it we end up at [Gladiator/microsoft-deberta-v3-large_ner_conll2003](https://huggingface.co/Gladiator/microsoft-deberta-v3-large_ner_conll2003), where *some random guy who doesn't work at Microsoft* has tweaked a Microsoft model and turned it into something that performs well on entity recognition.\n",
        "\n",
        "Let's support his efforts by giving it a try!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a8ddf47d-1f9f-4f1c-a324-31cdef235bd6",
      "metadata": {
        "id": "a8ddf47d-1f9f-4f1c-a324-31cdef235bd6"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "\n",
        "# https://huggingface.co/Gladiator/microsoft-deberta-v3-large_ner_conll2003\n",
        "nlp = pipeline(\"ner\",\n",
        "               model=\"Gladiator/microsoft-deberta-v3-large_ner_conll2003\",\n",
        "               aggregation_strategy='simple')\n",
        "\n",
        "results = nlp(text)\n",
        "for entity in results:\n",
        "    print(entity['word'], '|', entity['entity_group'])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8da2d1e6-3071-4ddf-8b42-c91b9613a73e",
      "metadata": {
        "id": "8da2d1e6-3071-4ddf-8b42-c91b9613a73e"
      },
      "source": [
        "Now it's your turn: find [another model on the token classification model list](https://huggingface.co/models?pipeline_tag=token-classification&sort=downloads) that you can try out. Maybe it's in [French](https://huggingface.co/Jean-Baptiste/camembert-ner) or [Japanese](https://huggingface.co/tsmatz/xlm-roberta-ner-japanese), or instead of recognizing people and places it can [analyze medical records](https://huggingface.co/Clinical-AI-Apollo/Medical-NER)!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "757fa597-1f6f-456f-bba0-1030db00ad5f",
      "metadata": {
        "id": "757fa597-1f6f-456f-bba0-1030db00ad5f"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "\n",
        "text = \"\"\"\n",
        "This is some text\n",
        "\"\"\"\n",
        "\n",
        "nlp = pipeline(\"ner\",\n",
        "               model=\"Gladiator/microsoft-deberta-v3-large_ner_conll2003\",\n",
        "               aggregation_strategy='simple')\n",
        "\n",
        "results = nlp(text)\n",
        "for entity in results:\n",
        "    print(entity['word'], '|', entity['entity_group'])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "60bd80a2-03da-4a9d-87ff-34f6a439dd37",
      "metadata": {
        "id": "60bd80a2-03da-4a9d-87ff-34f6a439dd37"
      },
      "source": [
        "Want to make your own? Try [Hugging Face's AutoTrain](https://huggingface.co/autotrain)! [Docs are here](https://huggingface.co/docs/autotrain/en/index) and kiiinda hard to understand until you get things working. Maybe there's a good YouTube video??\n",
        "\n",
        "### Wrap-up\n",
        "\n",
        "So far we've seen that named entity recognition is fun, but also kind of tough to get a hold onto.\n",
        "\n",
        "**Different models perform at different levels,** and finding the best one (\"state of the art\") is tough. We saw [that one leaderboard](https://paperswithcode.com/sota/token-classification-on-conll2003), but how do we know that it's the best one? The Papers With Code website even has [a whole section of token classification tests and scores](https://paperswithcode.com/task/token-classification)!\n",
        "\n",
        "**The models can only tag what they know,** so you can't use a names-organizations-locations tagger on a dataset to find money measurements and dates. It isn't as customizable"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "0907605a-b73a-4c70-9a56-68cd4598f366",
      "metadata": {
        "id": "0907605a-b73a-4c70-9a56-68cd4598f366"
      },
      "source": [
        "## A new world: LLM models\n",
        "\n",
        "After the introduction of transformers, one of the biggest splashes in the world of AI was ChatGPT. You'd ask it for a poem about your cat and it would deliver! Amazing!\n",
        "\n",
        "Beyond cat poems, ChatGPT is impressive because it's just so *flexible* – it can do so much! Even if it doesn't always give you the right answer, hallucinates frequently, won't talk about things it determines are too sensitive, expresses all sorts of bias, and is owned by a big corporation, *it can still amaze you*.\n",
        "\n",
        "One of the wildest things about ChatGPT and other large language models is that they can do all sorts of traditional natural language processing tasks, [including named entity recognition](https://chat.openai.com/share/dfba5755-04ae-483d-9c87-cf5889341fca):\n",
        "\n",
        "> **ME:** List the entities in the text below. On each line, present one entity,\n",
        "a comma, and the category it belongs to (PERSON, ORG, DATE)\n",
        ">\n",
        "> TEXT:\n",
        ">\n",
        "> Trump Seeks to Bump Haley From Race, Eyeing November Rematch With Biden\n",
        ">\n",
        "> It is a seemingly do-or-die moment for Nikki Haley’s campaign, as Donald Trump’s\n",
        "probable victory lines him up to face off with President Biden in the fall.\n",
        ">\n",
        "> **ChatGPT:**\n",
        ">\n",
        "> - Trump, PERSON\n",
        "> - Haley, PERSON\n",
        "> - November, DATE\n",
        "> - Nikki Haley, PERSON\n",
        "> - Donald Trump, PERSON\n",
        "> - President Biden, PERSON\n",
        "\n",
        "Amazing, right? It would be a pain to do with a hundred thousand entries in a database, though, so instead we'll directly to OpenAI through Python. This involves using the **API**, a way for computers to talk to each other, and an **API key**, a long string that identifies me to the OpenAI systems. They'll use it to track my usage and charge me!\n",
        "\n",
        "We'll start by installing the package that allows us to talk to GPT. Note that we're using the official OpenAI package here, **not** the popular [LangChain](https://langchain.readthedocs.io/) package. LangChain is always updating and breaking and a lot of people find it overly complicated for what it does, so we're going to set it aside for now."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "78a08f06-a5b6-460f-bbee-118f08b546b7",
      "metadata": {
        "id": "78a08f06-a5b6-460f-bbee-118f08b546b7"
      },
      "outputs": [],
      "source": [
        "!pip install --quiet openai"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "a16c23ac-6b48-4edb-9e6a-4d685e2d1428",
      "metadata": {
        "id": "a16c23ac-6b48-4edb-9e6a-4d685e2d1428"
      },
      "source": [
        "Even though we're just writing Python code, we're still having a conversation. We give GPT a list of messages that pass back and forth between us and the system - `user` and the `assistant`. The first message is a `system` prompt that sets the stage for the interaction, [you can see the complete OpenAI one here](https://pastebin.com/vnxJ7kQk), we're just using the incredibly boring \"You are a helpful assistant.\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "1c2954f2-e667-4117-81fd-acc9ee8c93c1",
      "metadata": {
        "id": "1c2954f2-e667-4117-81fd-acc9ee8c93c1"
      },
      "outputs": [],
      "source": [
        "from openai import OpenAI\n",
        "\n",
        "client = OpenAI(api_key='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')\n",
        "\n",
        "prompt = \"\"\"\n",
        "List the entities in the text below. On each line, present one entity,\n",
        "a comma, and the category it belongs to (PERSON, ORG, DATE).\n",
        "\n",
        "TEXT:\n",
        "\n",
        "Trump Seeks to Bump Haley From Race, Eyeing November Rematch With Biden\n",
        "\n",
        "It is a seemingly do-or-die moment for Nikki Haley’s campaign, as Donald Trump’s\n",
        "probable victory lines him up to face off with President Biden in the fall.\n",
        "\"\"\"\n",
        "\n",
        "messages = [\n",
        "    { \"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
        "    { \"role\": \"user\", \"content\": prompt}\n",
        "]\n",
        "\n",
        "chat_completion = client.chat.completions.create(\n",
        "    messages=messages,\n",
        "    model=\"gpt-3.5-turbo\",\n",
        "    temperature=0\n",
        ")\n",
        "\n",
        "print(chat_completion.choices[0].message.content)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "97c1e725-ff34-4b42-96f4-0a15ef999384",
      "metadata": {
        "id": "97c1e725-ff34-4b42-96f4-0a15ef999384"
      },
      "source": [
        "If we want to create a reusable piece of code, it's best to transform it into a **function**. The code below will allow us to write `get_entities_llm(\"This is a story about Joe Biden\")` to pull the entities out of the text instead of writing the big long long long piece of code."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "450d1f44-6363-467f-924b-4718fe9ceea4",
      "metadata": {
        "id": "450d1f44-6363-467f-924b-4718fe9ceea4"
      },
      "outputs": [],
      "source": [
        "client = OpenAI(api_key='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')\n",
        "\n",
        "def get_entities_llm(text):\n",
        "    prompt_template = \"\"\"\n",
        "    List the entities in the text below. On each line, present one entity,\n",
        "    a comma, and the category it belongs to (PERSON, ORG, DATE).\n",
        "\n",
        "    TEXT:\n",
        "\n",
        "    {text}\n",
        "    \"\"\"\n",
        "\n",
        "    prompt = prompt_template.format(text=text)\n",
        "\n",
        "    messages = [\n",
        "        { \"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
        "        { \"role\": \"user\", \"content\": prompt}\n",
        "    ]\n",
        "\n",
        "    chat_completion = client.chat.completions.create(\n",
        "        messages=messages,\n",
        "        model=\"gpt-3.5-turbo\",\n",
        "        temperature=0\n",
        "    )\n",
        "\n",
        "    return chat_completion.choices[0].message.content"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f18d7d3f-a9e3-4cd8-a3e1-28bb16f63404",
      "metadata": {
        "id": "f18d7d3f-a9e3-4cd8-a3e1-28bb16f63404"
      },
      "source": [
        "Now way every time we use it, it'll be nice and short."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "03deb1ee-1213-4023-a9e4-1e9209fa3042",
      "metadata": {
        "id": "03deb1ee-1213-4023-a9e4-1e9209fa3042"
      },
      "outputs": [],
      "source": [
        "text = \"\"\"\n",
        "Trump Seeks to Bump Haley From Race, Eyeing November Rematch With Biden\n",
        "\n",
        "It is a seemingly do-or-die moment for Nikki Haley’s campaign, as Donald Trump’s\n",
        "probable victory lines him up to face off with President Biden in the fall.\n",
        "\"\"\"\n",
        "\n",
        "results = get_entities_llm(text)\n",
        "print(results)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fe00d8d0-a5d8-41a7-bda1-133a9f8ec355",
      "metadata": {
        "id": "fe00d8d0-a5d8-41a7-bda1-133a9f8ec355"
      },
      "source": [
        "There's a problem, though: if we use this again and again and again and again across tens of thousands of pieces of text, **sometimes it's going to give us incorrect output.**"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "34bdba5a-07fd-4458-9afd-ae7926da2a88",
      "metadata": {
        "id": "34bdba5a-07fd-4458-9afd-ae7926da2a88"
      },
      "source": [
        "### Guardrails\n",
        "\n",
        "[Guardrails](https://github.com/guardrails-ai/guardrails) is a Python library that magically convinces the large language model to format the response in the right way. Along with *asking* for it, it will also attempt to validate a response. If the validation fails, Guardrails will *demand that the LLM fix it!*\n",
        "\n",
        "It's kind of a pain to set up, but it's a must-have for larger projects."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c86bdce7-0e1d-473d-8189-5261f46b619c",
      "metadata": {
        "id": "c86bdce7-0e1d-473d-8189-5261f46b619c"
      },
      "outputs": [],
      "source": [
        "from guardrails import validators, Guard\n",
        "from rich import print\n",
        "from pydantic import BaseModel, Field\n",
        "import openai\n",
        "import json\n",
        "\n",
        "prompt = \"\"\"\n",
        "    Analyze the following text, extracting all of the entities.\n",
        "\n",
        "    ${text}\n",
        "\n",
        "    ## Detailed instructions\n",
        "\n",
        "    ${gr.complete_json_suffix_v2}\n",
        "\"\"\"\n",
        "\n",
        "from pydantic import BaseModel\n",
        "from typing import List\n",
        "\n",
        "class Entity(BaseModel):\n",
        "    text: str = Field(description=\"Text of entity\")\n",
        "    category: str = Field(\"Entity category\", validators=[\n",
        "        validators.ValidChoices(\n",
        "            choices=['PERSON', 'ORGANIZATION', 'LOCATION'],\n",
        "            on_fail='reask'\n",
        "        )\n",
        "    ])\n",
        "\n",
        "class EntityList(BaseModel):\n",
        "    entities: List[Entity]\n",
        "\n",
        "guard = Guard.from_pydantic(output_class=EntityList, prompt=prompt)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bbfb0f6f-1639-45c4-9c99-7631cd5e541c",
      "metadata": {
        "id": "bbfb0f6f-1639-45c4-9c99-7631cd5e541c"
      },
      "source": [
        "Now that we've set up our expectated results, let's try it out!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "99c9f135-443b-47bd-8715-ba48ed693c46",
      "metadata": {
        "id": "99c9f135-443b-47bd-8715-ba48ed693c46"
      },
      "outputs": [],
      "source": [
        "text = \"\"\"\n",
        "Trump Seeks to Bump Haley From Race, Eyeing November Rematch With Biden\n",
        "\n",
        "It is a seemingly do-or-die moment for Nikki Haley’s campaign, as Donald Trump’s\n",
        "probable victory lines him up to face off with President Biden in the fall.\n",
        "\"\"\"\n",
        "\n",
        "response = guard(\n",
        "    openai.chat.completions.create,\n",
        "    prompt_params={\"text\": text},\n",
        "    response_format={\"type\": \"json_object\"},\n",
        "    max_tokens=1024,\n",
        "    temperature=0.3,\n",
        ")\n",
        "\n",
        "# Print the validated output from the LLM\n",
        "print(json.dumps(response.validated_output, indent=2))"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "095d45e7-605f-4a0b-aad8-5c7b8a46d2bd",
      "metadata": {
        "id": "095d45e7-605f-4a0b-aad8-5c7b8a46d2bd"
      },
      "source": [
        "We can look at `guard.history.last.tree` to see the entire conversation - it's a lot more than just our prompt!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "a9facdd4-79d5-4f8b-98e4-a60df0c03f45",
      "metadata": {
        "id": "a9facdd4-79d5-4f8b-98e4-a60df0c03f45"
      },
      "outputs": [],
      "source": [
        "print(guard.history.last.tree)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "8f88f28b-0703-4eca-ba04-e0f9875a0c01",
      "metadata": {
        "id": "8f88f28b-0703-4eca-ba04-e0f9875a0c01"
      },
      "source": [
        "## Reflection\n",
        "\n",
        "Wow, doing entity recognition with LLMs is a lot more flexible and potentially accurate than traditional models! It costs more and is slower, but it might be worth it??"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "ce0f94a9-c1fc-46ba-962d-5cdd79201ba2",
      "metadata": {
        "id": "ce0f94a9-c1fc-46ba-962d-5cdd79201ba2"
      },
      "outputs": [],
      "source": []
    }
  ],
  "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.13"
    },
    "colab": {
      "provenance": [],
      "include_colab_link": true
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}