Back to Timeline

大模型赋能数据分析：PandasAI

LLM 代码生成的局部应用

AI原生分享
自 ChatGPT 展示其能力以来，人们便热衷于将其接入各种不同的工具和场景中去，以实现既有流程的智能化改造。其中不少都贴近打工人的应用场景，一定程度上解放生产力，比如能帮你读 PDF 的 pdfGPT，除了 PDF 还能读 Markdown 和 EPUB 的 ChatFIles，以及直接嵌入 Slack 为你总结网页和视频的 myGPTReader。
不过对于常年跟数据打交道的分析师而言，数据清洗处理的繁琐往复似乎还没能乘上此股东风。一方面 ChatExcel 之流更像是玩具，有限的灵活性使其难当生产力；另一方面常年开着 Jupyter Notebook 和 Pandas 文档，边写边查敲代码处理数据的体验又属实有提升的空间。于是，PandasAI 来了，以 AI 赋能 Pandas，秒变功夫熊猫，短短一周就在 GitHub 收获了数千星。
Image
Image

简介

PandasAI 将大语言模型的能力无缝地接入数据分析工作流，如同有一个随时待命的 AI 助手，对 Python 和 Pandas 文档烂熟于心，并了解当前你手上的 DataFrame，告诉他分析需求即可自动处理，若是生成了新的 DataFrame 还可以继续迭代，体验堪称丝滑。
比如我们通过 Pandas 读取了不同国家 GDP 与居民幸福指数的数据集，得到了一个 DataFrame，便可以向 PandasAI 提问：哪五个国家的居民最开心？
import pandas as pd
from pandasai import PandasAI

# Sample DataFrame
df = pd.DataFrame({
    "country": ["United States", "United Kingdom", "France", "Germany", "Italy", "Spain", "Canada", "Australia", "Japan", "China"],
    "gdp": [19294482071552, 2891615567872, 2411255037952, 3435817336832, 1745433788416, 1181205135360, 1607402389504, 1490967855104, 4380756541440, 14631844184064],
    "happiness_index": [6.94, 7.16, 6.66, 7.07, 6.38, 6.4, 7.23, 7.22, 5.87, 5.12]
})

# Instantiate a LLM
from pandasai.llm.openai import OpenAI
llm = OpenAI()

pandas_ai = PandasAI(llm)
pandas_ai.run(df, prompt='Which are the 5 happiest countries?')
随即便能得到答案，并且是以 DataFrame 形式返回的：
6            Canada
7         Australia
1    United Kingdom
3           Germany
0     United States
Name: country, dtype: object
不仅如此，还可以交由 PandasAI 帮你算数、画图：
>>> pandas_ai.run(df, prompt='What is the sum of the GDPs of the 2 unhappiest countries?')
>>> 19012600725504

>>> pandas_ai.run(
    df,
    "Plot the histogram of countries showing for each the gpd, using different colors for each bar",
)
Image
赞叹之余，不免好奇这样集成是如何实现的呢？随手查阅一下仓库，再度惊喜，没想到核心代码量仅寥寥百余行。既然如此，不如让我们一同学习一下。

源码结构

Image
作为支持 pip install 安装的包，PandasAI 结构标准清晰，在主要的 pandasai 包内，有三组代码：
•
默认类 PandasAI 位于顶层，包含类构建 __init__.py 、常量 constants.py 和异常 exceptions.py ，其中类构建是核心代码，后做展开分析
•
大语言模型接入相关代码，位于 llm 子包内，包含 base.py 基类及 openai.py 等不同模型的衍生类，后以基类简单看 LLM 做了什么
•
辅助功能代码，位于 helpers 子包内，包含匿名化功能 anonymizer.py 和 Jupyter Notebook 支持功能 notebook.py

核心类

我们知道 PandasAI 是借助大语言模型的能力来实现智能的，但在 Python 解释器或 Notebook 中，DataFrame 应该在内存中才是，且格式繁杂、大小不一，如何将其传递给 LLM 吗？PandasAI 的思路可谓大巧不工，通过 df.head() 把 DataFrame 表头打印出来作为 LLM 输入的一部份，就像是甩给了 ChatGPT 一个小表格令其理解，本质上仍是巧妙的 Prompting。
这种传给 LLM 的指令，在 PandasAI 中被分为三类：
•
任务指令 _task_instruction，提供 DataFrame 的表头及行列数，要求其以固定格式返回用于完成某问题的代码
Today is {today_date}.
You are provided with a pandas dataframe (df) with {num_rows} rows and {num_columns} columns.
This is the result of `print(df.head({rows_to_display}))`:
{df_head}.

Return the python code (do not import anything) and make sure to prefix the requested python code with {START_CODE_TAG} exactly and suffix the code with {END_CODE_TAG} exactly to get the answer to the following question:
•
回答指令 _response_instruction，面向用户的问题，将结果以对话形式重写，目的是使代码运行结果与提问关联，从而使回答更连贯
Question: {question}
Answer: {answer}

Rewrite the answer to the question in a conversational way.
•
纠错指令 _error_correct_instruction，如果 LLM 给出的代码执行报错，通过纠错指令令其自查自纠
Today is {today_date}.
You are provided with a pandas dataframe (df) with {num_rows} rows and {num_columns} columns.
This is the result of `print(df.head({rows_to_display}))`:
{df_head}.

The user asked the following question:
{question}

You generated this python code:
{code}

It fails with the following error:
{error_returned}

Correct the python code and return a new python code (do not import anything) that fixes the above mentioned error. Do not generate the same code again.
Make sure to prefix the requested python code with {START_CODE_TAG} exactly and suffix the code with {END_CODE_TAG} exactly.
这些指令会经由 llm 子包传递给大模型，其基类代码（base.py）简单清晰，核心是对 LLM 所生成代码的提取（_extract_code），包括生成（call 调用具体 LLM）、解析（因格式已被要求固定，使用正则匹配）、完善（_polish_code 去除 import 等）、验证（_is_python_code 判断是否为正确的 Python 代码）四步，最终将返回符合要求的可执行代码。
Image
结合三种指令与 LLM 类的交互，PandasAI 的整个响应流程就清晰浮现出来了——
1.
当用户手持一个 DataFrame（假设名为 df）向 PandasAI 提出问题（假设为 prompt）时，PandasAI 首先以 df 对应的行列数与表头若干行数据、加上日期等构造一个「任务指令」，带上 prompt 发送给 LLM；
2.
LLM 接收到输入，调用具体模型按照规定的格式生成回答，再解析提取其中的可执行代码返回给 PandasAI；
3.
PandasAI 在 df 上尝试执行这些代码，得到运行结果 answer；
4.
PandasAI 将 prompt 和 answer 结合构造「回答指令」发给 LLM，使其以对话形式输出。
Image
尽管 ChatGPT 的代码能力有目共睹，但考虑到手上 DataFame 的复杂性及指令的多样性，直接运行其生成的代码还是存在不少风险，对此开发者也做了许多工作来避免危险的事情发生。其一是通过维护一个简单的白名单包列表，以阻止不安全的包被导入，目前白名单内只有 numpy 和 matplotlib 两个包，若有需求自行添加；其二是错误重试机制，在上述流程第 3 步中，如果代码执行出错，PandasAI 会将报错信息与 df 、prompt 等上下文构造「纠错指令」返回给 LLM 责令其重试，从而实现了托管式自动纠错。
此外，为防止 DataFrame 中可能存在的隐私被发送给外部模型，PandasAI 还提供了匿名化功能。具体而言，如果在构造指令的过程中，发现即将向模型输入的 df.head() 中可能包含 Email、电话号码、信用卡号等信息，PandasAI 便会自动将其改换为随机生成的对应字段值，以避免隐私泄漏。

结语

PandasAI 以优雅而简洁的方式将大语言模型的能力成功嫁接给 Pandas，尤其是代码能力，因为可以直接在 DataFrame 上执行以获取结果，从而使数据处理与分析过程的人机交互往智能化迈进了一大步，至少对于文档依赖的减轻能够大大降低 Pandas 的使用门槛，以后大概人均可玩转 Notebook 了。
但是等下，到那时为何还要写代码做分析呢？🤔️
代码地址：https://github.com/gventuri/pandas-ai 主要作者 Gabriele 是一位现居德国的意大利裔软件工程师，和我一样，他也在学习如何精调 Stable Diffusion :)