Talk to Your Dataset With PandasAI

Introduction

Imagine you’re working on a dataset to build a Machine Learning model and don’t want to spend too much effort on exploratory data analysis codes. You may sometimes find it confusing to sort, filter, or group data to obtain the required information. Is there a way to quickly and effortlessly extract information? Wouldn’t it be easier if you could talk to your dataset? You Ask it a question, and it’ll analyze the data for you. Well, this is where PandasAI with python can be beneficial.

PandasAI is a Python library that extends the functionality of Pandas by incorporating generative AI capabilities. Its purpose is to supplement rather than replace the widely used data analysis and manipulation tool. With PandasAI, users can interact with Pandas data frames more humanistically, enabling them to summarize the data effectively.

Get started with PandasAI we will need an OpenAI API key, which you can generate by following this link. This will give us access to the OpenAI LLM.

Learning Objectives

In this article, we will look at how to:

1. Obtain the OpenAI API key from the OpenAI website.
2. Connect to the OpenAI LLM model using the PandasAI library, and
3. Write prompts to enable the AI to generate exploratory data analysis results.

This article was published as a part of the Data Science Blogathon.

What does LLM stand for?

LLM stands for Large Language Model, and it refers to a type of artificial intelligence (AI) model designed to understand and generate human-like text.

Imagine a language model as a computer program that has been trained on a vast amount of text from various sources such as books, articles, websites, and more. This training enables the model to learn human language’s patterns, grammar, and context.

When you interact with a language model, for example, the recently popular ChatGPT, you can provide it with prompts or questions in natural language. The model then uses its understanding of language to generate relevant and coherent responses. The primary purpose of a language model like LLM is to assist users in understanding and generating text in a more human-like manner. It can be used for a wide range of applications, including answering questions, providing information, writing stories, summarizing text, translating languages, and much more.

In short, the goal of LLM is to mimic human language understanding and expression, allowing users to interact with AI systems in a more intuitive and natural way.

Working with PandasAI in Python

Let’s now see a practical use of PandasAI. First, we will download the PandasAI library using the following command.

NOTE: If this code throws an error while running on your local machine (e.g., Jupyter Notebook), you can either attempt to update your local Python environment or switch to cloud-based notebooks like Google Colab.

!pip install pandasai

Next, we will generate our OpenAI API Key and save it for later.

Import Libraries

import pandas as pd
from pandasai import PandasAI
from pandasai.llm.openai import OpenAI

In this demo, we will be using the Olympics dataset, which has 120 years of athlete information about the international event. If you wish to follow along, you can download the dataset here.

Olympics Dataset

After downloading the Olympics dataset, let’s now read it.

df = pd.read_csv("athlete_events.csv")
df.head()

Our dataset has information about the athlete, their nationality, gender, age when they participated, the sports they played, if they won a medal, and the Olympics event they participated in.

To begin exploring our dataset, we will call the PandasAI object by first initiating the OpenAI LLM object using the OpenAI API key we generated previously.

# Loading the API token to OpenAI environment
llm = OpenAI(api_token='Your API Key')
# Initializing an instance of PandasAI with OpenAI environment
pandas_ai = PandasAI(llm)

We are finally ready to “talk” to our dataset and ask it questions to gain insights from the Olympics data.

Let’s find out which athlete participated in the highest number of Olympics events. For this, we run the pandas_ai variable we created and input our dataset name and the prompt.

prompt = "Which athlete appeared in the most olympics years and how many"
pandas_ai.run(df, prompt=prompt)

Output

“The athlete who appeared in the most Olympics years is Ian Millar with 10 appearances.”

NOTE: The outputs may vary for you as they are subject to regular updates in the dataset.

It’s best to verify if the answer we got is correct. We usually do this by grouping the data by the athlete names and counting each athlete’s total unique values of the year of participation. We see that the AI did give us the right answer.

df.groupby(by='Name')['Year'].nunique().sort_values(ascending=False)

So, if Ian Millar participated in 10 Olympics events, is there an athlete who participated in the maximum number of sports during his tenure? Let’s find out.

pandas_ai.run(df, prompt="Which athlete has participated in the most number of events and how many")

“Oh, did you know that Ioannis Theofilakis holds the record for participating in the most number of events? He has participated in a whopping 33 events!”

That’s an amazing feat by Ioannis Theofilakis!

Next, let’s see which country holds the record for the maximum number of medals secured in the Olympics’ history.

pandas_ai.run(df, 
prompt="Which country has won the highest number of medals and how many")

“The country that has won the highest number of medals is the United States, with a whopping 5219 medals!”

Although the AI got the country right, it seems like it did not get the number of medals quite right. So, as we mentioned before, it’s best to verify the answers.

df.groupby(by='NOC')['Medal'].count().sort_values(ascending=False).reset_index().head(5)

The USA has won 5637 medals in total, not just 5219, as the AI mentioned.

If we want to take a look at the code the AI generates to produce expected results, we can add “verbose” as an input to the PandasAI object.

pandas_ai = PandasAI(llm,verbose=True

This gives us an output with all the steps the AI took to get the answer to our prompt.

Let’s find out what was the gender difference among the athletes in each year the Olympics was held. And we’ll use the above variable to input our prompt.

pandas_ai.run(df, 
prompt="generate a dataset with the total number of male and female participants in each year")

This is the code the AI wrote by itself and gave us the following table and the output.

It also provides a conversational answer.

“The dataset shows the total number of male and female participants in each year of the Olympics. It reveals that the number of female participants has been increasing over the years, with a significant rise in the 1980s and 1990s. In the most recent Olympics in 2016, there were more than 5,000 male participants and over 5,000 female participants.”

That’s not all! We can also use PandasAI to plot visual charts of our data. Let’s find out the trend in the total number of medals India has secured in the Olympics games over the years and visualize it as a barplot.

pandas_ai.run(df, 
prompt="plot a barplot with the total number of medals won by participants from IND")

We see that India had its best performance around the 1950s.

Finally, let us plot a histogram with a distribution of the age group most athletes belonged to over the years.

pandas_ai.run(df, 
prompt="create a histogram for the number of athletes based on the age group. Take bin size of 10")

We observe that athletes between the age group of 20-30 comprise the most participants.

Future Prospects

PandasAI has the potential to revolutionalize the ever-evolving data analysis landscape. And if you are a data analyst whose primary focus is extracting insights from data and generating plots based on user requirements, then this library can help automate the process with great efficiency. However, there are a few challenges you need to be aware of while using PandasAI:

1. How the AI interprets your prompt largely determines the obtained results, and sometimes it may not provide the expected answers. For instance, in the Olympics dataset, the AI occasionally faced confusion between “Olympic games” and “Olympic events,” leading to potentially divergent responses.
2. PandasAI cannot be utilized as a tool for data processing applications, such as data collection and translation into usable information.
3. It is also not suitable for Big Data Analysis.

Currently, PandasAI with python has limited application and can’t be used as a substitute for the Pandas library.

Conclusion

The progress in AI and conversational interfaces is revolutionizing the manner in which we interact with data, simplifying tasks, and significantly enhancing the accessibility of data analysis.

Here is a summary of what we looked at in this article:

• We looked at the amazing capability of PandasAI to retrieve information directly from a data frame as a conversational answer and even as a visualization. This definitely helps increase productivity by automating the data exploration process and much more.
• However, we cannot discount the Pandas library’s capabilities to perform complex operations, data imputation, etc., on the DataFrame.
• It’s necessary to note that although PandasAI is a strong tool, it still cannot replace the wide range of functionality of the Pandas library.

To sum up,  PandasAI serves as a valuable extension to the Pandas library, enhancing its functionality and adding more capabilities to handle challenging data manipulation and analysis tasks efficiently. By augmenting the already extensive ecosystem of Pandas, PandasAI further improves the convenience and effectiveness of working with data in Python.

The media shown in this article is not owned by Analytics Vidhya and is used at the Author’s discretion.

Likith Chandra

Business Analytics Student and Data Science, Data Analytics, and Machine Learning Enthusiast. Interested in time-series analysis and optimization models. Learning about various text mining and NLP methods in Python.