Написание запросов для решения аналитических задач — это основное занятие тех, кто работает с данными Pinterest. Но подбор подходящих данных и преобразование описания проблемы в корректный и эффективный SQL‑код могут оказаться непростыми делами. Ведь речь идёт о среде, которая быстро меняется, и о значительных объёмах данных, разбросанных по разным местам.

Мы увидели в росте доступности больших языковых моделей (Large Language Model, LLM) возможность исследовать один важный вопрос. Можно ли помочь тем, кто использует наши данные для решения вышеописанных задач, разработав систему преобразования текста в SQL (Text‑to‑SQL, Text2SQL), которая способна «переводить» вопросы аналитиков на язык кода?

Как в Pinterest работает система преобразования текста в SQL

Большинство задач анализа данных в Pinterest решается с привлечением Querybook — опенсорсного инструмента собственной разработки, ориентированного на работу с большими данными посредством выполнения SQL‑запросов. Этот инструмент представляется нам естественным местом для разработки и развёртывания систем, призванных помочь пользователям наших данных. К таким системам относится и Text‑to‑SQL.

Реализация Text-to-SQL

Исходная версия: реализация Text-to-SQL с использованием LLM

Первая версия системы включала в себя незамысловатый механизм преобразования текста в SQL с использованием большой языковой модели. Взглянем архитектуру этой системы.

https://miro.medium.com/v2/resize:fit:700/1*wDKR6-ToiX5UgsUYN41JiQ.png — Первая версия системы, преобразующей текст в SQL

Пользователь задаёт системе аналитические вопросы, выбирая таблицы, которые должны быть использованы.

Из хранилища метаданных таблиц извлекаются подходящие схемы таблиц.
Вопрос, сведения о выбранном диалекте SQL и схемы таблиц компилируются в Text-to-SQL-промпт.
Промпт передаётся LLM.
Генерируется ответ, который, в поточном режиме, выдаётся пользователю.

Схема таблицы

Схема таблицы, получаемая из хранилища метаданных, включает в себя следующие сведения:

Имя таблицы
Описание таблицы
Столбцы
Имена столбцов
Типы столбцов
Описания столбцов

Столбцы с низким кардинальным числом

Определённые аналитические запросы могут генерировать SQL‑запросы, которые не соответствуют реальным значениям, хранящимся в базе данных. Например — «сколько имеется активных пользователей на платформе «web». Ошибки могут происходить в том случае, если генерирование таких запросов выполняется без особых изысков. Например, предложение WHERE в ответе может выглядеть как WHERE platform=’web’. А правильный вариант этого запроса выглядел бы как WHERE platform=’WEB’. Для решения подобных проблем в схему таблицы внедряются уникальные значения, описывающие столбцы с низким кардинальным числом, которые планируется часто использовать для фильтрации данных. В результате LLM может воспользоваться этой информацией для генерирования точных SQL‑запросов.

Ограниченный размер контекстного окна LLM

Очень большие схемы таблиц могут не помещаться в типичное контекстное окно LLM. Для решения этой проблемы мы прибегли к нескольким подходам:

Уменьшенная версия схемы таблицы: в неё входят только самые важные элементы — такие, как имя таблицы, имена и типы столбцов.
Устранение лишних столбцов: в хранилище метаданных столбцам назначаются теги. Определённые столбцы убирают из схемы таблицы на основании их тегов.

Потоковая передача ответа модели

Получение полного ответа от LLM может занять десятки секунд. Для того чтобы не заставлять пользователей ждать, мы применили технологию WebSocket для организации потоковой передачи ответа модели. Учитывая то, что, помимо сгенерированного SQL‑кода, модель должна возвращать и другую информацию, для нас очень важно формирование ответов, структурированных так, как нам нужно. Обычный текст легко передавать в потоковом режиме, а вот потоковая передача данных в формате JSON может оказаться более сложной задачей. Мы задействовали систему для парсинга JSON‑кода, который ещё полностью не сформирован, взяв её из фреймворка LangChain. Её мы используем в подсистеме потоковой передачи данных сервера. JSON‑код, после парсинга, мы передаём клиенту через WebSocket.

Промпт

Вот как выглядит наш текущий промпт, используемый в Text2SQL:

https://miro.medium.com/v2/resize:fit:636/1*4tX9Fq-hsl_GpPStTarCVg.png

Оценка исходной версии системы и выводы о её работе

Первые оценки Text‑to‑SQL были, в основном, проведены, чтобы убедиться в том, что наша реализация этой системы обладает эффективностью работы, результаты которой сравнимы с теми, что описаны в различных публикациях. Причём, тут мы учитываем то, что в нашей реализации, по большей части, используются стандартные подходы. На наборе данных Spider мы выявили результаты, сравнимыми с теми, о которых пишут. Правда, мы заметили, что задачи, описываемые в различных источниках, гораздо легче тех, которые вынуждены решать наши пользователи. В частности, в публикациях речь идёт о малом числе заранее заданных таблиц с небольшим количеством хорошо промаркированных столбцов.

После того, как наша система начала использоваться в продакшне, мы смогли понаблюдать и за тем, как пользователи с ней взаимодействуют. По мере того, как реализация Text‑to‑SQL улучшалась, и по мере того, как пользователи лучше с ней знакомились, доля сгенерированного ей SQL‑кода, принятого с первого раза, выросла с 20% до 40%. На практике большинство генерируемых запросов требовали, для доведения их до полной готовности, проведения множества сеансов взаимодействия человека и ИИ. Для того чтобы оценить то, как система Text‑to‑SQL повлияла на продуктивность пользователей данных, мы прибегли к самому надёжному, в данном случае, методу — к эксперименту. Предыдущее исследование, в котором использовался подобный метод, показало, что помощь ИИ повысила скорость завершения задач на 50%. На реальных данных (мы, что важно, не учитывали различий данных в разных задачах) мы обнаружили 35%‑е улучшение в скорости завершения задач при написании SQL‑запросов с использованием помощи ИИ.

Вторая версия: включение в систему RAG для выбора таблиц

Первая версия системы работала достойно — учитывая то, что пользователи знают о том, какие таблицы им нужны для выполнения запроса. Но поиск подходящих таблиц среди сотен тысяч сущностей, имеющихся в нашем хранилище, представляет собой сложную, серьёзную задачу. Для того чтобы улучшить ситуацию, мы интегрировали в Text‑to‑SQL подсистему генерации с дополненной выборкой (Retrieval Augmented Generation, RAG). Её цель заключается в том, чтобы помогать пользователям выбирать подходящие таблицы для решения их задач. Вот обзор улучшенной инфраструктуры нашей системы преобразования текста в SQL, в состав которой включена подсистема RAG.

https://miro.medium.com/v2/resize:fit:700/1*8lgsIGY1LVd1HqPAXNsddw.png — Новая версия системы, в состав которой входит RAG

Для генерирования векторного индекса сводных сведений по таблицам и по истории запросов к этим таблицам используется оффлайновое задание.
Если пользователь не дал никаких указаний о таблицах, то его вопросы преобразуются в эмбеддинги, после чего в векторном индексе проводится поиск по сходству, который позволяет выявить первые N наиболее подходящих таблиц.
Первые N таблиц, вместе со схемами таблиц и с аналитическим вопросом, компилируют в промпт для LLM, что позволяет выбрать первые K наиболее уместных таблиц.
Первые K таблиц возвращаются пользователю, который может их проверить или изменить выбор таблиц.
Стандартный процесс работы Text‑to‑SQL возобновляется с использованием таблиц, применение которых подтвердил пользователь.

Оффлайновое создание векторного индекса

В векторном индексе существует два типа эмбеддингов документов:

Обобщение таблицы.
Обобщение запроса.

Обобщение таблиц

В Pinterest идёт непрерывный процесс по стандартизации таблиц, направленный на добавление к таблицам сведений об их уровне. Мы индексируем лишь таблицы самого высокого уровня, подталкивая пользователей к применению высококачественных наборов данных, хранящихся в таких таблицах. Процесс генерирования обобщения таблицы включает в себя следующие шаги:

Получение схемы таблицы из хранилища метаданных таблицы.
Сбор наиболее свежих примеров запросов, в которых используется эта таблица.
В промпт, с помощью которого будут обобщаться данные, включается схема таблицы и как можно большее количество примеров запросов. Мы, формируя этот промпт, ориентируемся на размер контекстного окна.
Передача промпта LLM для создания обобщения.
Генерирование эмбеддингов и их сохранение в хранилище векторных данных.

Обобщение таблицы включает в себя описание таблицы, данные, которые она содержит, а так же — потенциальные сценарии использования таблицы. Вот — промпт, который мы используем в настоящий момент для формирования обобщений таблиц:

https://miro.medium.com/v2/resize:fit:643/1*qK0r2iGh0NOGQIAK1nFlkg.png

Обобщение запросов

Образцы запросов, связанных с таблицами, помимо их роли в формировании обобщений таблиц, сами подвергаются обобщению. В частности, речь идёт о сведениях о том, какова цель того или иного запроса, и о том, к каким таблицам он обращается. Вот — промпт, который мы для этого используем:

https://miro.medium.com/v2/resize:fit:639/1*O1Xih_9w-HNvTY1dbOWWGg.png

NLP-поиск по таблицам

Когда пользователь задаёт аналитический вопрос, мы преобразуем его в эмбединнги, используя ту же самую модель эмбеддингов. Затем мы выполняем поиск и по векторному индексу таблиц, и по векторному индексу запросов. В качестве хранилища векторных данных мы применяем OpenSearch и используем его встроенную систему поиска, похожую на эту.

Учитывая то, что с одним запросом может быть связано несколько таблиц, одна и та же таблица при поиске по сходству может появиться в результатах поиска несколько раз. Сейчас мы применяем простую стратегию для агрегирования и оценки подобных результатов. А именно — обобщения таблиц обладают большим весом, чем обобщения запросов. В будущем эта система оценки результатов поиска может быть изменена.

Эта система поиска, основанная на NLP, применяется не только в Text‑to‑SQL. В Querybook эта система используется для выполнения обычных операций поиска по таблицам.

Уточнение выбора таблиц

После возвращения первых N наиболее подходящих таблиц из векторного индекса мы задействуем LLM для выбора наиболее уместных таблиц, оценивая вопрос вместе с обобщениями таблиц. В зависимости от контекстного окна мы включаем в промпт настолько много таблиц, насколько это возможно. Вот — промпт, который мы используем для уточнения выбора таблиц:

https://miro.medium.com/v2/resize:fit:647/1*-k1BOyHaauorCuxV_vxhWQ.png

После того, как выбор таблиц был уточнён, результаты возвращаются пользователю для проверки, а уже после этого передаются туда, где осуществляется создание SQL‑запросов.

Оценка второй версии системы и выводы о её работе

Мы оценивали компонент выбора таблиц нашей системы преобразования текста в SQL, используя оффлайновые данные из предыдущих сеансов поиска таблиц. Этих данных нам было недостаточно по одной важной причине: они отражали поведение пользователей, которое было характерно для них до того, как они узнали о возможности пользоваться NLP‑поиском. В результате эти данные были использованы, в основном, чтобы убедиться в том, что поиск таблиц, основанный на эмбеддингах, работает не хуже, чем существующий текстовый поиск. Мы не пытались оценить улучшение, которое даёт новая система поиска. Мы использовали это испытание для выбора метода и для установки весов для эмбеддингов, используемых при поиске таблиц. Этот подход показал нам, что метаданные таблиц, генерируемые в результате усилий, прилагаемых нами к управлению данными, вносят огромный вклад в общую эффективность работы системы. Процент попаданий поиска, проводимого без наличия документации к таблицам в эмбеддингах, составил 40%. Но при использовании весов, касающихся документации к таблицам, наблюдался линейный рост эффективности работы до уровня в 90%.

Дальнейшее развитие Text-to-SQL

Несмотря на то, что существующая в данный момент система Text‑to‑SQL значительно улучшила продуктивность наших аналитиков, этой системе ещё есть куда расти. Вот некоторые сферы, в которых её можно улучшить.

NLP-поиск таблиц

Улучшение метаданных

Сейчас векторный индекс связан лишь с обобщениями таблиц. Одно из возможных улучшений системы может означать включение в индекс дополнительных метаданных. Это могут быть сведения об уровнях таблиц, теги, данные о предметных областях и прочее подобное. Это позволит добиться более точного фильтрования данных при извлечении из хранилища похожих таблиц.

Запланированное обновление индекса или обновление индекса в реальном времени

Сейчас векторный индекс генерируется вручную. Реализация обновления индекса, выполняемого по расписанию, или даже в реальном времени (при создании новых таблиц или при выполнении к ним запросов) может значительно повысить эффективность работы системы.

Пересмотр механизма поиска по сходству и стратегии назначения оценок

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

Валидация запросов

В настоящий момент SQL‑запрос, сгенерированный LLM, возвращается сразу пользователю, не проходя какой‑либо валидации. Это означает потенциальный риск того, что запрос может выполняться не так, как ожидается. Реализация системы валидации запросов, возможно — с использованием поиска по алгоритму «Constrained Beam», могла бы обеспечить дополнительные гарантии корректности работы Text‑to‑SQL.

Отзывы пользователей

Если в систему будет добавлен интерфейс для эффективного сбора отзывов пользователей, касающихся поиска таблиц и генерирования запросов, это может дать разработчикам ценные подсказки относительно улучшения системы. Подобные отзывы могут обрабатываться и добавляться в векторный индекс, или в хранилище метаданных таблиц, что, в итоге, способно повысить эффективность работы системы.

Оценка работы систем, подобных нашей

Работая на этим проектом, мы поняли, что эффективность работы Text‑to‑SQL в реальных условиях сильно отличается от результатов существующих бенчмарков, в которых обычно используется небольшое количество хорошо нормализованных таблиц (они, к тому же, заданы заранее). Прикладным исследователям было бы полезно создавать более реалистичные бенчмарки, при работе которых используется большое количество денормализованных таблиц. Исследователям, кроме того, стоит рассматривать поиск таблиц, подходящих для формирования запросов, как один из центральных аспектов исследуемых ими проблем.

О, а приходите к нам работать? 🤗 💰

Мы в wunderfund.io занимаемся высокочастотной алготорговлей с 2014 года. Высокочастотная торговля — это непрерывное соревнование лучших программистов и математиков всего мира. Присоединившись к нам, вы станете частью этой увлекательной схватки.

Мы предлагаем интересные и сложные задачи по анализу данных и low latency разработке для увлеченных исследователей и программистов. Гибкий график и никакой бюрократии, решения быстро принимаются и воплощаются в жизнь.

Сейчас мы ищем плюсовиков, питонистов, дата-инженеров и мл-рисерчеров.

Присоединяйтесь к нашей команде

О создании системы, преобразующей текст в SQL для аналитиков Pinterest