PandasAI—让AI做数据分析

lichunericli2024-03-09 11:46

安装

pip install pandasai

!pip install --upgrade pandas pandasai

导入依赖项

import pandas as pd

from pandasai import PandasAI

from pandasai.llm.openai import OpenAI

使用pandas创建一个数据框

df = pd.DataFrame({
    "country": ["United States", "United Kingdom", "France", "Germany", "Italy", "Spain", "Canada", "Australia", "Japan", "China"],
    "gdp": [21400000, 2940000, 2830000, 3870000, 2160000, 1350000, 1780000, 1320000, 516000, 14000000],
    "happiness_index": [7.3, 7.2, 6.5, 7.0, 6.0, 6.3, 7.3, 7.3, 5.9, 5.0]
})

实例化LLM

OPENAI_API_KEY = "YOUR API KEY"
llm = OpenAI(api_token=OPENAI_API_KEY)

使用提供的大语言模型实例化PandasAI，并运行

pandas_ai = PandasAI(llm, conversational=False)
pandas_ai.run(df, prompt='What are the 5 happiest countries')

图表展示

pandas_ai.run(df, "Plot the histogram of countries showing for each the gpd, using different colors for each bar")

询问涉及多个数据框的问题

employees_data = {
    'EmployeeID': [1, 2, 3, 4, 5],
    'Name': ['John', 'Emma', 'Liam', 'Olivia', 'William'],
    'Department': ['HR', 'Sales', 'IT', 'Marketing', 'Finance']
}

salaries_data = {
    'EmployeeID': [1, 2, 3, 4, 5],
    'Salary': [5000, 6000, 4500, 7000, 5500]
}

employees_df = pd.DataFrame(employees_data)
salaries_df = pd.DataFrame(salaries_data)

pandas_ai([employees_df, salaries_df], "Who gets paid the most?")