PandasAI — кратно ускоряем работу аналитика данных в одну строчку

Привет, чемпион!

За прошедший год появилось много полезных AI инструментов для упрощения работы разработчиков, аналитиков данных и даже дизайнеров: Copilot допишет за тебя код, EverSQL оптимизирует SQl-запрос, а Kittl нарисует логотип. А сейчас поговорим о PandasAI -  аналоге классической библиотеки pandas на стероидах ChatGPT.

Как работает PandasAI?

По сути это LLM агент, которые имеет доступ к Pandas. Агенту необходимо передать вопрос или указание на естественном языке, а он на основе данных найдет, построит диаграмму или преобразует данные. Примеры запросов:

PandasAI vs просты запросы

Для установки используем команду pip install pandasai. Есть github библиотеки. Импортируем библиотеки, передаем ключ API OpenAI и создаем датасет с данными о продажах книг.

Пример кода и запрос в PandasAI

import os
import pandas as pd
from pandasai import Agent

os.environ["PANDASAI_API_KEY"] = "YOUR_API_KEY"

sales_by_book = pd.DataFrame({
  'book' : ['A', 'D', 'E', 'U', 'C', 'R', 'E'],
  "revenue": [1500, 3500, 1290, 4610, 7200, 9100, 2200]
})

agent = Agent(sales_by_book)
# Пишем вопрос “Какие топ-3 книги по продажам?”
agent.chat('Which are the top 3 book by sales?')
# PandasAI ответит нам R, C, U

Можно попробовать более сложные вопросы, например “Каковы общие продажи для топ-3 книг?”.

agent.chat("What is the total sales for the top 3 book by sales?")
# The total sales for the top 3 books is: 20910

Отлично! Получаем ожидаемо верный ответ. Вопрос для llm-ки оказался слишком простым. Зато все сделано всего в одну строчку. Давайте попробуем порисовать гистограммы?!

PandasAI vs графики

# Построй гистограмму книг по их продажаем. Используй разные цвета для каждой
agent.chat("Plot the histogram of books showing for each one. Use different colors for each bar"
)

И получаем вполне приличный график за пару секунд, который мы скорее всего строили бы дольше. 

Результат

Идеально отработало!

И что самое крутое, у нас есть код для построения этого графика. Конечно я бы строил через seaborn. Т.к он удобнее и чем matplotlib, но тут решение было такое:

Пример кода

import pandas as pd
import matplotlib.pyplot as plt

# Data for books and revenue
data = {'book': ['R', 'A', 'E'], 'revenue': [9100, 7200, 3500]}
books_df = pd.DataFrame(data)

# Plotting histogram with different colors for each bar
plt.figure(figsize=(10, 6))
colors = ['skyblue', 'salmon', 'lightgreen']
plt.bar(books_df['book'], books_df['revenue'], color=colors)
plt.xlabel('Book')
plt.ylabel('Revenue')
plt.title('Revenue by Book')
plt.show()

Также PandasAI имеет возможность создавать объекты SmartDatalake - хранилище нескольких датафреймов, которое позволяет выполнять запросы для объединенных данных. 

PandasAI vs мердж таблиц

Для этого нам потребоваться работать с несколькими фреймами данных одновременно. В таких случаях вместо использования SmartDataframe следует использовать SmartDatalake. Концепция очень похожа на SmartDataframe, но вместо того, чтобы принимать в качестве входных данных только 1 df, она может принимать несколько значений.

К примеру, имеем 2 датафрейма с данными о продаж книг и их цветах, которые связаны между собой по book_id сотрудника. Спросим PandasAI , какой цвето самой продаваемой книги?

Код

from pandasai import SmartDatalake

sales_by_book = pd.DataFrame({
    'book_id' :   ['A', 'D', 'E', ' U', 'C', 'R', 'E'],
    "revenue": [1500, 3500, 1290, 4610, 7200, 9100, 2200]
})

color_by_book = pd.DataFrame(
    {
        'book_id' :   ['A', 'D', 'E', ' U', 'C', 'R', 'E'],
        "color":   ['green',  'red',  'green',  'red', 'black', 'green', 'red'],
    }
)

lake = SmartDatalake([sales_by_book, color_by_book])
lake.chat("What's the color of book that was saled the most?")

# LLM-ка ответ нам:
# The color of the book that was sold the most is green.

Для выполнения такого запроса в Pandas пришлось бы мерджить два датафрейма между собой по id, а потом выполнять запрос поиска максимальных продаж, но ИИ-шка справилась на отлично. Такая классная фича сильно упрощает работу с несколькими источниками данных, а это все же одна из основных задач в работе аналитика данных.

Еще одним вариантом использования SmartDatalake является преобразование данных для дальнейшей работы с ними, например для построения моделей машинного обучения, потому что PandasAI понимает не только вопросы, но и запросы на изменения данных.

lake.chat("Merge this dataframes in one and rank it by the salary descending and delete the highest and the book with lowest sales")

Результат

Результат запроса. Верный!

PandasAI vs моя фирменная сложная задач

До этого были вполне себе простые вещи. Давайте попробуем справится с задачей, которая в свое время заняла у меня несколько часов экспериментов. Тогда я не смог сходу решил эту задачу, но посмотрим, смог бы PandasAI сохранить мне драгоценное время?!

Формулировка задания. Попробуйте решить сами без ChatGPT и PandasAI

Подаем таблицу в SmartDataframe и делаем запрос. В конце для удобства выводим сразу код генерации через атрибут last_code_generated.

Код и запрос

import os
import pandas as pd
from pandasai import SmartDataframe

os.environ["PANDASAI_API_KEY"] = 'API_KEY'

table_df = pd.DataFrame({
    'book' : ['A','A','A','A', 'D','D','D','D', 'E','E','E', 'E'],
     'page' : [15, 25, 54, 44, 75, 22, 14, 60, 18, 10, 20, 95],
     'prediction' : [0.1, 0.2, 0.26, 0.5, 0.1, 0.2, 0.4, 0.6, 0.1, 0.97, 0.2, 0.95]
})

lake = SmartDataframe(table_df)

lake.chat("Aggregate this dataframe and print top-3 pages for every book with the highest prediction, print only book and a list of pages. I'm expecting 'book' and 'pages' columns in anwer")
print('Предлагаемый код:')
print(lake.last_code_generated)

Результат

В точку! Прям как у меня в примере

А еще круто, мы можем брать в работу дальше код результа. Запускаем предложенный код. И он тоже работает! Вот он ниже. Можно брать и использовать!

Код из last_code_generated

table_df = (table_df
            .sort_values(by='prediction', ascending=False)
            .groupby('book')
            .head(3)
            .groupby('book')['page']
            .apply(list)
            .reset_index()
            )

Но что еще интереснее, PandasAI можно дообучить на своих данных, подробнее это описано в документации. Так же можно использовать другие LLM-ки. С другими вариантами использования PandasAI можно познакомиться здесь (ссылка на ноутбук).

Плюсы и минусы PandasAI

Плюсы

• Очевидно, самые большое преимущество - ускорение работы с данными, поскольку не приходится периодически подсматривать в документацию Pandas в поисках нужных методов.

• В некоторых случаях это позволит снизить количество аd-hoс задач на аналитика: продакт или бизнес-аналитик сможет самостоятельно провести EDA или проверить гипотезу.

Минусы

• Несмотря на универсальность, инструмент все еще кажется сырым, хотя его активно развивают

• Многие запросы по генерации кода можно сделать напрямую в ChatGPT (но не все)

• Для работы необходим ключ API OpenAI.

• PandasAI не получится использовать в рамках корпоративных правил работы с данными, так что если не хотите нарушать NDA - осторожнее с этим.

А вы бы использовали PandasAI в своей работе? Пишите в комменты!

А если вам нравится применение подобных AI инструментов в проекции на Data Scientist'ов - подписывайтесь на мой телеграм. Нас уже почти 10 000. Удачи!