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


简介
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",
)

赞叹之余,不免好奇这样集成是如何实现的呢?随手查阅一下仓库,再度惊喜,没想到核心代码量仅寥寥百余行。既然如此,不如让我们一同学习一下。
源码结构

作为支持 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 代码)四步,最终将返回符合要求的可执行代码。

结合三种指令与 LLM 类的交互,PandasAI 的整个响应流程就清晰浮现出来了——
1.
当用户手持一个 DataFrame(假设名为 df)向 PandasAI 提出问题(假设为 prompt)时,PandasAI 首先以 df 对应的行列数与表头若干行数据、加上日期等构造一个「任务指令」,带上 prompt 发送给 LLM;
2.
LLM 接收到输入,调用具体模型按照规定的格式生成回答,再解析提取其中的可执行代码返回给 PandasAI;
3.
PandasAI 在 df 上尝试执行这些代码,得到运行结果 answer;
4.
PandasAI 将 prompt 和 answer 结合构造「回答指令」发给 LLM,使其以对话形式输出。

尽管 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 :)