使用OpenAIPydanticProgram进行DataFrame结构化数据提取

本文链接：https://blog.csdn.net/qq_29929123/article/details/140981826

使用OpenAIPydanticProgram进行DataFrame结构化数据提取

介绍

在数据处理及分析中，将原始文本转换为结构化的数据形式（如DataFrame）是非常常见的需求。本文将向大家介绍如何使用OpenAIPydanticProgram从文本中提取表格化的DataFrame数据。

我们将使用OpenAIPydanticProgram，它是围绕OpenAI LLM的一个封装，支持函数调用，并返回结构化输出，例如Pydantic对象。

环境准备

在开始之前，请确保你已经安装了必要的库。以下是一些必要的安装命令：

%pip install llama-index-llms-openai
%pip install llama-index-program-openai
!pip install llama-index

示例代码

使用OpenAIPydanticProgram构建DF Extractor

下面是一个完整的示例代码，展示了如何使用OpenAIPydanticProgram从文本中提取DataFrame：

from llama_index.program.openai import OpenAIPydanticProgram
from llama_index.core.program import DFFullProgram, DataFrame, DataFrameRowsOnly
from llama_index.llms.openai import OpenAI

# 初始化OpenAIPydanticProgram对象
program = OpenAIPydanticProgram.from_defaults(
    output_cls=DataFrame,
    llm=OpenAI(temperature=0, model="gpt-4-0613"),
    prompt_template_str=(
        "Please extract the following query into a structured data according"
        " to: {input_str}.Please extract both the set of column names and a"
        " set of rows."
    ),
    verbose=True,
)

# 测试示例
response_obj = program(
    input_str="""My name is John and I am 25 years old. I live in 
        New York and I like to play basketball. His name is 
        Mike and he is 30 years old. He lives in San Francisco 
        and he likes to play baseball. Sarah is 20 years old 
        and she lives in Los Angeles. She likes to play tennis.
        Her name is Mary and she is 35 years old. 
        She lives in Chicago."""
)
print(response_obj)

注释 : //中转API

在上面的代码中，我们定义了一个Prompt模板，并使用OpenAIPydanticProgram解析文本，将其转换为结构化的DataFrame数据。

使用现有的DataFrame架构提取数据

接下来我们展示如何使用已有的DataFrame架构进行数据提取：

from llama_index.program.openai import OpenAIPydanticProgram
from llama_index.core.program import DFRowsProgram
import pandas as pd

# 初始化空的DataFrame
df = pd.DataFrame(
    {
        "Name": pd.Series(dtype="str"),
        "Age": pd.Series(dtype="int"),
        "City": pd.Series(dtype="str"),
        "Favorite Sport": pd.Series(dtype="str"),
    }
)

# 使用已存在的DataFrame架构初始化程序
df_rows_program = DFRowsProgram.from_defaults(
    pydantic_program_cls=OpenAIPydanticProgram, df=df
)

# 解析文本
result_obj = df_rows_program(
    input_str="""My name is John and I am 25 years old. I live in 
        New York and I like to play basketball. His name is 
        Mike and he is 30 years old. He lives in San Francisco 
        and he likes to play baseball. Sarah is 20 years old 
        and she lives in Los Angeles. She likes to play tennis.
        Her name is Mary and she is 35 years old. 
        She lives in Chicago."""
)

# 将结果转换为DataFrame
df_result = result_obj.to_df(existing_df=df)
print(df_result)