Пришвидшуємо роботу з даними за допомогою PandasAI

💡 Усі статті, обговорення, новини про AI — в одному місці. Приєднуйтесь до AI спільноти!

Привіт, я Ігор Крохін, працюю Software Engineer у HOSTiQ понад 3 роки, паралельно роблю свій внесок у деякі OpenSource-продукти.

На сьогодні неможливо не помітити зацікавленість суспільства в напрямі штучного інтелекту та інструментів на його основі.

Тож я пропоную звернути увагу на один з таких інструментів, у розробці якого я брав участь. Ця стаття буде корисна будь-яким розробникам, які пов’язані з Data Analysis та використовують бібліотеку pandas. А також людям з досвідом, які шукають нові інструменти для розширення свого тулз-сету, і початківцям — для швидшого занурення в аналіз даних з pandas та в різноманітні AI-сервіси.

Короткий огляд

PandasAI ― це бібліотека, реалізована винятково на Python, яка має пришвидшити аналіз даних за допомогою pandas та сторонніх text2text AI-сервісів.

Що саме робить PandasAI? Якщо простими словами: бібліотека дозволяє проводити обробку даних вербальними запитами. Припустимо, у вас є дата-сет у csv-форматі, або excel/googlesheets, або ж DataFrame-об’єкт у контексті Python-модуля.

Бібліотека формує prompt, з вашого вербального запиту (зазвичай питання) приймає на вхід дані у вигляді DataFrame, аналізуючи його та додаючи до prompt. Модель (LLM) генерує та повертає python-код, який PandasAI виконує та надає йому в оточення клієнтський DataFrame. PandasAI забирає отриманий результат. Користувач натомість отримує новий DataFrame, який є відповіддю на його запитання.

Детальніше про використання йтиметься нижче в прикладах.

Встановлення

Вам знадобиться Python 3.9+. Пакет є на PyPI, встановлюємо через pip:

$ pip install pandasai

Приклади

Визначаємо найвищий хмарочос

Для базового прикладу я буду використовувати OpenAI у якості LLM-провайдеру.

Вам знадобиться API_TOKEN від OpenAI, його можна отримати наступним способом:

1. Перейти на platform.openai.com і зареєструватися, використовуючи адресу електронної пошти, або підключити обліковий запис Google.
2. Відкрити View API keys, натиснувши на іконку профілю праворуч.
3. Обрати «Створити новий секретний ключ».
4. Зберегти згенерований секрет.

Перед запуском додамо токен до оточення (опціонально):

$ export OPENAI_API_KEY=<your_openai_token>

(або можете вказувати одразу в коді).

Далі використовуємо PandasAI ось так:

>>> import os

>>> 

>>> import pandas as pd

>>> 

>>> from pandasai import PandasAI, SmartDataframe

>>> 

>>> from pandasai.llm.openai import OpenAI

>>> 

>>> source_dataframe = pd.DataFrame({

...     {

...         "name": [

...             "Burj Khalifa",

...             "Merdeka 118",

...             "Abraj Al-Bait Clock Tower",

...             "Ping An International Finance Centre",

...             "Lotte World Tower",

...             "One World Trade Center",

...             "Guangzhou CTF Finance Centre",

...             "Tianjin CTF Finance Centre",

...             "Taipei 101",

...             "Shanghai World Financial Center",

...         ],

...         "height": [

...             828.0,

...             678.9,

...             632.0,

...             601.0,

...             599.1,

...             554.5,

...             541.3,

...             530.0,

...             508.0,

...             492.0,

...         ],

...         "year": [2010, 2023, 2015, 2012, 2017, 2017, 2014, 2016, 2004, 2008],

...     }

... )

>>> llm = OpenAI(api_token=os.environ.get('OPENAI_API_KEY'))

>>> 

>>> smart_df = SmartDataframe(df=source_dataframe, config={"llm": llm})

>>> response = smart_df.chat("Which of buildings are in top-3 by height?")

>>> print(response

                        name  height  year

0               Burj Khalifa   828.0  2010

1                Merdeka 118   678.9  2023

2  Abraj Al-Bait Clock Tower   632.0  2015

Що ми зробили зверху? Покроково:

1. Створили DataFrame 10×3 з даними по хмарочосах (їх назви, висоту та рік забудови).
2. Інстаціювали SmartDataframe — об’єкт, через який будемо робити запит до зовнішнього провайдера.
3. Передали питання: «Які з хмарочосів входять у трійку найвищих?».

Як результат, користувач отримує DataFrame 3×3, з трійкою найвисотніших будівель. До речі, оскільки сам response є DataFrame, він буде (частково, або повністю) реалізувати методи абстрактного DataFrame.

Будуємо графік

Побудуємо графік хмарочосів з попереднього прикладу.

Для побудови графіків PandasAI використовує matplotlib. Щоб побудувати діаграму стовпців для хмарочосів з попереднього прикладу, просто змінимо запит:

>>> smart_df.chat("Plot all buildings by height (meters) in bar chart. Make labels to not overlap each other")

Унаслідок цього в новому вікні відображається графік, у моєму випадку він виглядає так:

На графіку бачимо, що лейбли (назви будов) перетинаються між собою, тож трохи доповнимо запит:

>>> smart_df.chat("Plot 5 tallest buildings by height (meters) in bar chart. Make the labels to not overlap each other. Use large canvas")

Результат:

Який саме код побудував графік вище?

Код надає LLM-провайдер, у поточному прикладі це OpenAI. У моєму випадку OpenAI повернув наступні інструкці (нагадую, df у контексті нижче — це клієнтський DataFrame):

# TODO import all the dependencies required

import pandas as pd

import matplotlib.pyplot as plt

# Analyze the data

# 1. Prepare: Preprocessing and cleaning data if necessary

# 2. Process: Manipulating data for analysis (grouping, filtering, aggregating, etc.)

# 3. Analyze: Conducting the actual analysis (if generating a plot, create a figure and axes using plt.subplots() and save it to an image in exports/charts/temp_chart.png and do not show the chart.)

# 4. Output: return a dictionary of:

# - type (possible values "text", "number", "dataframe", "plot")

# - value (can be a string, a dataframe or the path of the plot, NOT a dictionary)

def analyze_data(dfs: list[pd.DataFrame]) -> dict:

    # Combine all dataframes into one

    df = pd.concat(dfs)

    # Sort the dataframe by height in descending order

    df_sorted = df.sort_values(by='height', ascending=False)

    # Create a bar chart of buildings by height

    fig, ax = plt.subplots(figsize=(10, 6))

    ax.bar(df_sorted['name'], df_sorted['height'])

    ax.set_xlabel('Building')

    ax.set_ylabel('Height (meters)')

    ax.set_title('Buildings by Height')

    # Rotate x-axis labels to avoid overlapping

    plt.xticks(rotation=45, ha='right')

    # Save the chart to an image file

    chart_path = 'exports/charts/temp_chart.png'

    plt.savefig(chart_path)

    # Close the plot

    plt.close()

    # Return the path of the chart image

    return {'type': 'plot', 'value': chart_path}

# Declare a result variable

result = analyze_data(dfs)

Troubleshooting

Що робити, коли не вдається отримати очікуваного результату?

Звісно, перш за все для відладки вкрай рекомендовано перевіряти логи, які можна отримати через property logs в об’єкта SmartDataframe:

>>> from pprint import pprint

>>> pprint(smart_df.logs)

Наприклад, у цьому топіку користувач зауважив, що графік в нього не побудувався. На відміну від прикладів, він використовував інший LLM — StarCoder з хабу HuggingFace.

І ні, справа тут не в тому, що він неправильно сформулював prompt, або надав битий DataFrame. HF API повернув наступний код:

df.groupby('PRODUCTLINE').sum()['SALES'].plot(kind='pie')

plt.show()

plt.close('all')

Оскільки в інструкціях вище немає import statement для matplotlib, авжеж, під час виконання виникне NameError через невідомий аліас plt. Тобто модель просто забула додати import matplotlib.pyplot as plt на перших рядках.

На той момент я запропонував розв’язувати проблему через використання кастомного посередника (Middlewares — про них ще поговоримо детальніше нижче).

Посередник виглядав так:

from pandasai.middlewares.base import Middleware

class AddMissingImportsMiddleware(Middleware):

    def run(self, code: str) -> str:

        return f"from matplotlib import pyplot as plt\n{code}"

Описаний баг помічений у версії 0.8.3, у поточному релізі такої помилки статися не може (було додано те, що plt за дефолтом пролягає до контексту виконання), проте можуть бути інші помилки з імпортами — залежить від запиту.

Якщо посередники не допомогли вам з відладкою, спробуйте інший LLM. Якщо все ж таки рішення не знайдено — ласкаво просимо до issues.

Доступні мовні моделі:

Реліз 1.0.1 підтримує 5 зовнішніх LLM-провайдерів:

• OpenAI — потужна модель від компанії-розробника ChatGPT. Реалізована через клас OpenAI.
• StarCoder — достатньо нова модель, але стрімко розвивається. Підтримує HuggingFace. Реалізована через клас Starcoder.
• Falcon — відкрита й добре натренована модель. Реалізована через клас Falcon.
• GooglePalm — модель від Google. Реалізована через клас GooglePalm.
• Azure OpenAI — фактично є проксі до OpenAI API, через Azure instances. Реалізована через клас AzureOpenAI.

Деякі фічі

Кеш

PandasAI зберігає останні запити та результати для них на файловій системі користувача через shelve. Проте може некоректно працювати під Windows платформою (на цю тему я створив окремий issue, поки що не вирішений).

Кеш можна вимкнути, передавши config={"enable_cache": False} при інстанціюванні SmartDataframe, за замовченням True.

Колбеки

Інколи дуже корисна річ, яка розв’язує руки для того, щоб обробляти повернутий код. Існують декілька callbacks «з коробки». Наприклад, запис коду напряму в новий модуль:

>>> from pandasai import SmartDataframe from pandasai.callbacks import FileCallback

>>> llm = OpenAI(api_token=os.environ.get('OPENAI_API_KEY'))

>>> smart_df = SmartDataframe(df=source_dataframe, config={"llm": llm, "callback": FileCallback("code_generated_by_openai.py")})

Посередники

Як уже згадувалося вище, посередники дозволяють на власний смак модифікувати код, згенерований LLM. Кастомний посередник має реалізовувати метод run(), що приймає рядок з кодом аргументом. Наприклад:

class DelayExecMiddleware(Middleware):

    def run(self, code: str) -> str:

        return f"import time; time.sleep(3);\n{code}" 

Чи заощадить PandasAI час при роботі з даними?

Якщо ви працюєте Data Analyst з динамічними запитами (наприклад, постійно змінюються умови для формування вибірки) — так, PandasAI обов’язково буде вам в пригоді.

Більш того, бібліотека має потенціал для інтегрування у якості text2code-адаптеру до будь-якої бізнес-моделі, компоненти якої використовують pandas на DAL-рівні.

Корисні посилання:

• Офіційна документація
• Репозиторій проєкту