Open In Colab

Pydantic Extractor#

Here we test out the capabilities of our PydanticProgramExtractor - being able to extract out an entire Pydantic object using an LLM (either a standard text completion LLM or a function calling LLM).

The advantage of this over using a “single” metadata extractor is that we can extract multiple entities with a single LLM call.

Setup#

import nest_asyncio

nest_asyncio.apply()

import os
import openai
os.environ["OPENAI_API_KEY"] = "YOUR_API_KEY"
openai.api_key = os.getenv("OPENAI_API_KEY")

Setup the Pydantic Model#

Here we define a basic structured schema that we want to extract. It contains:

  • entities: unique entities in a text chunk

  • summary: a concise summary of the text chunk

  • contains_number: whether the chunk contains numbers

This is obviously a toy schema. We’d encourage you to be creative about the type of metadata you’d want to extract!

from pydantic import BaseModel, Field
from typing import List
class NodeMetadata(BaseModel):
    """Node metadata."""

    entities: List[str] = Field(
        ..., description="Unique entities in this text chunk."
    )
    summary: str = Field(
        ..., description="A concise summary of this text chunk."
    )
    contains_number: bool = Field(
        ...,
        description=(
            "Whether the text chunk contains any numbers (ints, floats, etc.)"
        ),
    )

Setup the Extractor#

Here we setup the metadata extractor. Note that we provide the prompt template for visibility into what’s going on.

from llama_index.program.openai_program import OpenAIPydanticProgram
from llama_index.extractors import PydanticProgramExtractor

EXTRACT_TEMPLATE_STR = """\
Here is the content of the section:
----------------
{context_str}
----------------
Given the contextual information, extract out a {class_name} object.\
"""

openai_program = OpenAIPydanticProgram.from_defaults(
    output_cls=NodeMetadata,
    prompt_template_str="{input}",
    # extract_template_str=EXTRACT_TEMPLATE_STR
)

program_extractor = PydanticProgramExtractor(
    program=openai_program, input_key="input", show_progress=True
)

Load in Data#

We load in Eugene’s essay (https://eugeneyan.com/writing/llm-patterns/) using our LlamaHub SimpleWebPageReader.

# load in blog

from llama_hub.web.simple_web.base import SimpleWebPageReader
from llama_index.node_parser import SentenceSplitter

reader = SimpleWebPageReader(html_to_text=True)
docs = reader.load_data(urls=["https://eugeneyan.com/writing/llm-patterns/"])
from llama_index.ingestion import IngestionPipeline

node_parser = SentenceSplitter(chunk_size=1024)

pipeline = IngestionPipeline(transformations=[node_parser, program_extractor])

orig_nodes = pipeline.run(documents=docs)
orig_nodes

Extract Metadata#

Now that we’ve setup the metadata extractor and the data, we’re ready to extract some metadata!

We see that the pydantic feature extractor is able to extract all metadata from a given chunk in a single LLM call.

sample_entry = program_extractor.extract(orig_nodes[0:1])[0]
display(sample_entry)
{'entities': ['eugeneyan', 'HackerNews', 'Karpathy'],
 'summary': 'This section discusses practical patterns for integrating large language models (LLMs) into systems & products. It introduces seven key patterns and provides information on evaluations and benchmarks in the field of language modeling.',
 'contains_number': True}
new_nodes = program_extractor.process_nodes(orig_nodes)
display(new_nodes[5:7])