basic-data-extraction-llms/structured-data-extraction-with-guardrails-and-llms.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1cc7b68d-f469-4b2d-bdab-60597a2d42e8",
   "metadata": {},
   "source": [
    "# Structured Data Extraction using LLMs and Guardrails AI\n",
    "\n",
    "In this post, we're going to be extracting structured data from a podcast transcript. We've all seen the ability for generative models to be effective *generating* text.\n",
    "\n",
    "**But what about extracting data?**\n",
    "\n",
    "Data extraction is a far more common use case (today) than generation, in particular for businesses. Businesses have to process all kinds of documents, exchange them and so on.\n",
    "\n",
    "**Emerging Use Case: LLMs for data extraction**\n",
    "\n",
    "If you think about it, summarization is effectively data extraction. We've seen LLMs perform well at summarization, but did you know that they can extract structured data quite well?\n",
    "\n",
    "## What you'll learn from this post\n",
    "\n",
    "In this post, you'll learn how to extract structured data from LLMs into Pydantic objects using the [Guardrails](https://www.guardrailsai.com/) library.\n",
    "\n",
    "## Here's us setting up on environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "bfc5b8de-c56b-4104-95df-dc9246883705",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Python 3.9.18\n"
     ]
    }
   ],
   "source": [
    "!python --version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e9c40f04-4407-45f8-9985-b159c13ad7e2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "guardrails-ai==0.2.4\n"
     ]
    }
   ],
   "source": [
    "!pip freeze | grep guardrails-ai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "cfb484c3-41ae-477b-b53f-8cda908428f0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.10.9\n"
     ]
    }
   ],
   "source": [
    "import pydantic\n",
    "print(pydantic.__version__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "3b965611-0a5d-4004-b62d-0af1221a03d3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7cb680b8-5ca0-4190-98a6-0b5b9d1d6556",
   "metadata": {},
   "source": [
    "## Background on Guardrails AI\n",
    "\n",
    "[Guardrails](https://www.guardrailsai.com/) is a project that appears to be backed by a company. The project, aptly named Guardrails, is centered around the following concept:\n",
    "\n",
    "> What is Guardrails?\n",
    "> \n",
    "> Guardrails AI is an open-source library designed to ensure reliable interactions with Large Language Models (LLMs). It provides:\n",
    "> \n",
    "> ✅ A framework for creating custom validators\n",
    "> ✅ Orchestration of prompting, verification, and re-prompting\n",
    "> ✅ A library of commonly used validators for various use cases\n",
    "> ✅ A specification language for communicating requirements to LLM\n",
    "\n",
    "Guardrails shares a similar philosophy with Instructor and Marvin. \n",
    "\n",
    "> Guardrails AI enables you to define and enforce assurance for AI applications, from structuring output to quality controls. It achieves this by creating a 'Guard', a firewall-like bounding box around the LLM application, which contains a set of validators. A Guard can include validators from our library or a custom validator that enforces your application's intended function.\n",
    "> - from [the blog](https://www.guardrailsai.com/blog/0.2-release)\n",
    "\n",
    "![guardrails approach](https://images.learnbybuilding.ai/guardrails-architecture.webp)\n",
    "\n",
    "Currently, Guardrails is in beta and has not had a formal release, so it may not be suitable for production use cases.\n",
    "\n",
    "However, Guardrails has a broader mission. They aim to create a \"bounding box\" around LLM apps to validate and ensure quality. They plan to achieve this by introducing a `.RAIL` file type, a dialect of XML. This ambitious approach extends beyond the simple concept of \"using AI to extract data\".\n",
    "\n",
    "Here's their explanation of `RAIL`.\n",
    "\n",
    "> 🤖 What is `RAIL`?\n",
    "> \n",
    "> `.RAIL` is a dialect of XML, standing for \"**R**eliable **AI** markup **L**anguage\". It can be used to define:\n",
    "> \n",
    "> 1. The structure of the expected outcome of the LLM. (E.g. JSON)\n",
    "> 2. The type of each field in the expected outcome. (E.g. string, integer, list, object)\n",
    "> 3. The quality criteria for the expected outcome to be considered valid. (E.g. generated text should be bias-free, generated code should be bug-free)\n",
    "> 4. The corrective action to take in case the quality criteria is not met. (E.g. reask the question, filter the LLM, programmatically fix, etc.)\n",
    "\n",
    "**Warning...**\n",
    "\n",
    "Another challenge I had was the dependencies, Guardrails requires `pydantic version 1.10.9`.  [Pydantic](https://docs.pydantic.dev/latest/)'s latest version is [v.2.4.1](https://github.com/pydantic/pydantic/releases/tag/v2.4.1). This caused a number of issues for me during installation and I had to setup a dedicated environment for it. A challenge is that[ pydantic went through a ground up rewrite from version 1.10](https://github.com/pydantic/pydantic#pydantic-v110-vs-v2), so it may simply be challenging for them to rewrite everything.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b85a7dc3-15f8-46b6-870f-8f87ae56c2cc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<p><em>Thanks to the </em><em>over 11,000 people</em><em> who joined us for the first AI Engineer Su\n"
     ]
    }
   ],
   "source": [
    "import feedparser\n",
    "\n",
    "podcast_atom_link = \"https://api.substack.com/feed/podcast/1084089.rss\" # latent space podcastbbbbb\n",
    "parsed = feedparser.parse(podcast_atom_link)\n",
    "episode = [ep for ep in parsed.entries if ep['title'] == \"Why AI Agents Don't Work (yet) - with Kanjun Qiu of Imbue\"][0]\n",
    "episode_summary = episode['summary']\n",
    "print(episode_summary[:100])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c2af4fc2-6085-45b4-935f-e67f71f24450",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "First line of the transcript: 58\n"
     ]
    }
   ],
   "source": [
    "from unstructured.partition.html import partition_html\n",
    "\n",
    "parsed_summary = partition_html(text=''.join(episode_summary)) \n",
    "start_of_transcript = [x.text for x in parsed_summary].index(\"Transcript\") + 1\n",
    "print(f\"First line of the transcript: {start_of_transcript}\")\n",
    "text = '\\n'.join(t.text for t in parsed_summary[start_of_transcript:])\n",
    "text = text[:3508] # shortening the transcript for speed & cost"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9f8c993-f77f-4f5e-86cc-550665950f2f",
   "metadata": {},
   "source": [
    "## Using Guardrails\n",
    "\n",
    "Guardrails is similar to [other libraries for structured data extraction](https://learnbybuilding.ai/vs/marvin-ai-vs-guardrails-vs-instructor) that I've used. We define our pydantic models and then we specify an output class and a prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "cb8c469a-b16a-451f-9818-db889d98166c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import guardrails as gd\n",
    "\n",
    "from pydantic import BaseModel\n",
    "from typing import Optional, List\n",
    "from pydantic import Field\n",
    "\n",
    "class Person(BaseModel):\n",
    "    name: str\n",
    "    school: Optional[str] = Field(..., description=\"The school this person attended\")\n",
    "    company: Optional[str] = Field(..., description=\"The company this person works for\")\n",
    "\n",
    "class People(BaseModel):\n",
    "    people: List[Person]\n",
    "\n",
    "guard = gd.Guard.from_pydantic(output_class=People, prompt=\"Get the following objects from the text:\\n\\n ${text}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b679a4b9-e02d-428f-93d3-ca49a8eaa9b2",
   "metadata": {},
   "source": [
    "The important thing to note, like we mentioned above, is that Guardrails is translating this into a \"RAIL\" an underlying abstraction for representing these extractions.\n",
    "\n",
    "Now a key difference is that this abstraction might allow for structured data extraction from ANY model, not just OpenAI's function calling interface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6f02712c-3bdc-449f-84b6-e60aa2954f85",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'people': [{'name': 'Alessio', 'school': 'Residence at Decibel Partners', 'company': 'CTO'}, {'name': 'Swyx', 'school': 'Smol.ai', 'company': 'founder'}, {'name': 'Kanjun', 'school': 'Imbue', 'company': 'founder'}]}\n"
     ]
    }
   ],
   "source": [
    "import openai\n",
    "import os\n",
    "\n",
    "raw_llm_output, validated_output = guard(\n",
    "    openai.ChatCompletion.create,\n",
    "    prompt_params={\"text\": text},\n",
    ")\n",
    "\n",
    "print(validated_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "939da155-d095-46ef-a001-1bcf2128e08f",
   "metadata": {},
   "source": [
    "It did a fine job getting the people, although you can see that it got the company and schools wrong - especially compared to the other libraries we've looked at, this is a lower quality result.\n",
    "\n",
    "let's see how it does when we ask for more objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9a0ef84d-7dd3-4220-af91-4f1f8826654b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/williamchambers/miniconda3/envs/guardrails-extract/lib/python3.9/site-packages/guardrails/prompt/instructions.py:32: UserWarning: Instructions do not have any variables, if you are migrating follow the new variable convention documented here: https://docs.getguardrails.ai/0-2-migration/\n",
      "  warn(\n",
      "/Users/williamchambers/miniconda3/envs/guardrails-extract/lib/python3.9/site-packages/guardrails/prompt/prompt.py:23: UserWarning: Prompt does not have any variables, if you are migrating follow the new variable convention documented here: https://docs.getguardrails.ai/0-2-migration/\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "incorrect_value={'people': [{'name': 'Alessio', 'school': 'Decibel Partners', 'company': 'CTO'}, {'name': 'Swyx', 'school': 'Smol.ai', 'company': 'founder'}, {'name': 'Kanjun', 'school': 'Imbue', 'company': 'founder'}], 'companies': [{'name': 'Decibel Partners'}, {'name': 'Residence'}, {'name': 'Smol.ai'}, {'name': 'Imbue'}]} fail_results=[FailResult(outcome='fail', metadata=None, error_message='JSON does not match schema', fix_value=None)]\n"
     ]
    }
   ],
   "source": [
    "class Company(BaseModel):\n",
    "    name:str\n",
    "\n",
    "class ResearchPaper(BaseModel):\n",
    "    paper_name:str = Field(..., description=\"an academic paper reference discussed\")\n",
    "    \n",
    "class ExtractedInfo(BaseModel):\n",
    "    people: List[Person]\n",
    "    companies: List[Company]\n",
    "    research_papers: Optional[List[ResearchPaper]]\n",
    "\n",
    "guard = gd.Guard.from_pydantic(output_class=ExtractedInfo, prompt=\"Get the following objects from the text:\\n\\n ${text}\")\n",
    "raw_llm_output, validated_output = guard(\n",
    "    openai.ChatCompletion.create,\n",
    "    prompt_params={\"text\": text},\n",
    ")\n",
    "\n",
    "print(validated_output)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c192865a-7fce-4ceb-b659-90f393a2d39a",
   "metadata": {},
   "source": [
    "Unfortunately, the result failed to parse, while the other two [structured data extraction libraries](https://learnbybuilding.ai/vs/marvin-ai-vs-guardrails-vs-instructor) succeeded just fine. It might be in the implementation or might be a bug in this version, but it didn't work out of the box.\n",
    "\n",
    "While the example didn't work, you can see how well this *should work*. These models and libraries grab a whole bunch of structured data for us with basically no work. Right now we're just working on a excerpt of this data, but with basically zero NLP specific work, we're able to get some pretty powerful results.\n",
    "\n",
    "\n",
    "## Wrapping it all up\n",
    "\n",
    "The sky is really the limit here. There's so much structured data that's been locked up in unstructured text. I'm bullish on this space and can't wait to see what you build with this tool!\n",
    "\n",
    "If you found this interesting, you might want to see similar posts on:\n",
    "1. [Comparing 3 Data Extraction Libraries: Marvin, Instructor, and Guardrails](https://learnbybuilding.ai/vs/marvin-ai-vs-guardrails-vs-instructor)\n",
    "2. [Structured Data Extraction using LLMs and Marvin](https://learnbybuilding.ai/tutorials/structured-data-extraction-with-marvin-ai-and-llms)\n",
    "3. [Structured Data Extraction using LLMs and Instructor](https://learnbybuilding.ai/tutorials/structured-data-extraction-with-instructor-and-llms)"
   ]
  }
 ],
 "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.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}