yfyang86 / pandas-ai

PandasAI is a Python library that integrates generative artificial intelligence capabilities into Pandas, making dataframes conversational
https://pandas-ai.readthedocs.io/en/latest/
MIT License
1 stars 1 forks source link

PandasAI 🐼

Release CI CD Coverage Documentation Status Discord Downloads License: MIT Open in Colab

PandasAI is a Python library that adds Generative AI capabilities to pandas, the popular data analysis and manipulation tool. It is designed to be used in conjunction with pandas, and is not a replacement for it.

PandasAI

πŸ”§ Quick install

pip install pandasai

πŸ” Demo

Try out PandasAI in your browser:

Open in Colab

πŸ“– Documentation

The documentation for PandasAI can be found here.

πŸ’» Usage

Disclaimer: GDP data was collected from this source, published by World Development Indicators - World Bank (2022.05.26) and collected at National accounts data - World Bank / OECD. It relates to the year of 2020. Happiness indexes were extracted from the World Happiness Report. Another useful link.

PandasAI is designed to be used in conjunction with pandas. It makes pandas conversational, allowing you to ask questions to your data in natural language.

Queries

For example, you can ask PandasAI to find all the rows in a DataFrame where the value of a column is greater than 5, and it will return a DataFrame containing only those rows:

import pandas as pd
from pandasai import SmartDataframe

# 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 import OpenAI
llm = OpenAI(api_token="YOUR_API_TOKEN")

df = SmartDataframe(df, config={"llm": llm})
df.chat('Which are the 5 happiest countries?')

The above code will return the following:

6            Canada
7         Australia
1    United Kingdom
3           Germany
0     United States
Name: country, dtype: object

Of course, you can also ask PandasAI to perform more complex queries. For example, you can ask PandasAI to find the sum of the GDPs of the 2 unhappiest countries:

df.chat('What is the sum of the GDPs of the 2 unhappiest countries?')

The above code will return the following:

19012600725504

Charts

You can also ask PandasAI to draw a graph:

df.chat(
    "Plot the histogram of countries showing for each the gdp, using different colors for each bar",
)

Chart

You can save any charts generated by PandasAI by setting the save_charts parameter to True in the PandasAI constructor. For example, PandasAI(llm, save_charts=True). Charts are saved in ./pandasai/exports/charts .

Multiple DataFrames

Additionally, you can also pass in multiple dataframes to PandasAI and ask questions relating them.

import pandas as pd
from pandasai import SmartDatalake
from pandasai.llm import OpenAI

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)

llm = OpenAI()
dl = SmartDatalake([employees_df, salaries_df], config={"llm": llm})
dl.chat("Who gets paid the most?")

The above code will return the following:

Oh, Olivia gets paid the most.

You can find more examples in the examples directory.

⚑️ Shortcuts

PandasAI also provides a number of shortcuts (beta) to make it easier to ask questions to your data. For example, you can ask PandasAI to clean_data, impute_missing_values, generate_features, plot_histogram, and many many more.

# Clean data
df.clean_data()

# Impute missing values
df.impute_missing_values()

# Generate features
df.generate_features()

# Plot histogram
df.plot_histogram(column="gdp")

Learn more about the shortcuts here.

πŸ”’ Privacy & Security

In order to generate the Python code to run, we take the dataframe head, we randomize it (using random generation for sensitive data and shuffling for non-sensitive data) and send just the head.

Also, if you want to enforce further your privacy you can instantiate PandasAI with enforce_privacy = True which will not send the head (but just column names) to the LLM.

🀝 Contributing

Contributions are welcome! Please check out the todos below, and feel free to open a pull request. For more information, please see the contributing guidelines.

After installing the virtual environment, please remember to install pre-commit to be compliant with our standards:

pre-commit install

Contributors

Contributors

πŸ“œ License

PandasAI is licensed under the MIT License. See the LICENSE file for more details.

Acknowledgements