AutoML для NLU без ручной настройки: делимся библиотекой OpenAutoNLU

Большинство существующих AutoML-библиотек либо не поддерживают обучение моделей для понимания естественного языка (Natural Language Understanding, или NLU) из коробки, либо не умеют обучать хорошие out of scope детекторы, либо неудобны и требуют расширенной экспертизы для использования.  

Для того чтобы решить эти проблемы, мы в MWS AI разработали OpenAutoNLU — опенсорс-библиотеку для NLU, включающую диагностику качества данных, гибко настраиваемый пайплайн обучения модуля фильтра запросов, которые не относятся ни к одному из известных текстовых классификаторов меток OOD, и функции LLM. Она уже дополнила инструментарий для работы с ИИ-моделями на платформе MWS AI Agents Platform, и теперь мы делимся ей на GitHub, а вот демо. 

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

Рассказывать буду я, Григорий Аршинов. А под «мы» подразумевается целая команда исследователей из MWS AI, ИТМО и MBZUAI: Александр Борискин, Сергей Сеничев, Аяз Зарипов, Дарья Галимзянова, Даниил Карпов и Леонид Саночкин, а также Дарья Самсонова и Антон Ненашев. По-научному нашу работу над библиотекой мы описали в этой статье.

Почему AutoML в NLP всё еще требует ручной настройки 

Классификация текстов и распознавание именованных сущностей (NER) давно стали базовыми задачами, на которых держится большая часть прикладных NLU-сценариев. Это и интенты в саппорте, и маршрутизация запросов, и извлечение сущностей из документов. На уровне сборки моделей в этом случае помогают предобученные трансформеры, готовые пайплайны и разнообразные решения AutoML. Однако, когда эту модель уже нужно довести до продакшена в краткие сроки, открытые решения теряют в своей полезности.

Задачи автоматического машинного обучения обычно создают приличный объем рутины для разработчиков: 

• Во-первых, чтобы выбрать нужный подход, специалисту требуется либо заранее знать, какие гиперпараметры подойдут лучше всего, либо прогнать все возможные комбинации или пресеты в поисках лучшей метрики. Это долго и дорого с точки зрения вычислений. 

• Во-вторых, в реальных задачах разметка редко бывает чистой: встречаются ошибки, неоднозначные примеры, дубли и перекосы по классам. Всё это влияет на итоговую модель, но большинство AutoML-решений не поддерживают проверку качества данных прямо в пайплайне. В итоге перед обучением приходится отдельно чистить датасет или уже после разбираться, почему модель ведет себя нестабильно.

• В-третьих, даже если модель показывает хорошие метрики на тесте, в реальной среде она сталкивается с запросами (OOD-сценариями), которых не было в обучении. В типичных решениях их не учитывают вообще или требуют заранее размечать такие примеры и добавлять их в обучение, что не всегда возможно, да и концептуально бессмысленно.

Мы хотели решить эти проблемы и собрать пайплайн, который снимал бы с разработчика повторяющуюся ручную работу и не требовал каждый раз заново принимать одни и те же архитектурные решения. Так мы пришли к OpenAutoNLU.

Чем OpenAutoNLU отличается от других решений

OpenAutoNLU — это библиотека на Python, в которой простой API позволяет прототипировать модели NLU и обучать их для продакшена. 

Фреймворк уникален тем, что система сама принимает решение о способе обучения, учитывая распределение классов и минимальное количество примеров на класс. Именно этот фактор чаще всего определяет, будет ли работать few-shot-подход или уже пора идти в полноценный fine-tuning.

В текущей версии фреймворка встроены три режима обучения: 

• При очень малом количестве данных используется AncSetFit — вариант few-shot-метода с якорными примерами и contrastive learning.

• При умеренном объеме данных применяется SetFit с sentence-transformer и простым классификатором сверху.

• При достаточном количестве примеров система переходит к полному fine-tuning трансформера с опциональной оптимизацией гиперпараметров.

Во фреймворк встроен модуль обучения «головы» с OOD. Есть три встроенных типа OOD-методов: Marginal Mahalanobis Distance, Centroid Mahalanobis Distance и MaxProb. Каждому методу есть применение в разных ситуациях распределения данных во входящем датасете. Например, Marginal Mahalanobis Distance лучше проявляет себя только тогда, когда в минимальном классе есть хотя бы 81 пример. Основная особенность нашего OOD-пайплайна в том, что мы научились генерировать неплохую синтетику в качестве примеров out of scope для нахождения порогов отделения OOD от in domain классов, поэтому фактически мы поддерживаем обучение OOD «головы» без учителя, хоть и даем возможность использовать примеры out of scope из входящего датасета.  

Кроме того, мы встроили в систему функционал на базе LLM, который помогает расширить выборку через генерацию дополнительных примеров. Также есть пайплайн, генерирующий искусственные тестовые данные на основе few-shot-примеров из пользовательского датасета. Такая синтетика позволяет получить оценку работы классификатора с надежностью до 95%. 

К тому же OpenAutoNLU упакован в минималистичный low-code API — пользователю нужно лишь предоставить данные, запустить обучение и получить модель. Выбор режима, анализ данных, балансировка, обучение и подготовка к инференсу проходят автономно в OpenAutoNLU.

Архитектура фреймворка

Архитектурно OpenAutoNLU — это несколько слоев, через которые проходит датасет. Главный модуль — auto_classes — открывает четыре публичных класса пайплайнов, через которые пользователь запускает обучение или инференс. 

Каждый пайплайн наследуется от AbstractTrainingPipeline или AbstractInferencePipeline. В них уже зафиксирована последовательность шагов: загрузка данных, диагностика качества (опционально), выбор метода, обучение, оценка и экспорт. Все они завязаны в единый сценарий.

Методы обучения и их роль

Все алгоритмы обучения собраны в модуле methods и реализованы как подклассы базового метода. Сейчас реализованы: 

• Finetuner/FinetunerWithOOD — полный fine-tuning трансформера с ранней остановкой и опциональной оптимизацией гиперпараметров через Optuna.

• SetFitMethod/SetFitOOD — обучение sentence-transformer через contrastive learning с последующей логистической регрессией.

• AncSetFitMethod/AncSetFitOOD — вариант SetFit с якорными метками для few-shot в режиме 2–5 примеров на класс.

• TokenClassificationFinetuner — NER в схеме BIO.

Все они вызываются через общий интерфейс, а выбор конкретного происходит внутри пайплайна, без изменения внешнего API. 

Работа с данными и LLM

Модуль data отвечает за то, как данные попадают в обучение, поэтому для классификации используется SimpleDataProvider, а для NER — SimpleNerDataProvider. Они нормализуют входной формат и подготавливают данные под конкретный тип задачи.

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

Особенности NER-пайплайна

NER-пайплайн в OpenAutoNLU принимает два типа аннотаций — prodigy-like json объекты или встроенный в текст формат под названием brackets. Внутри фреймворка оба варианта приводятся к единой BIO-схеме, чтобы дальше обучение и оценка шли в одном формате. Стоит заметить, что читать BIO напрямую фреймворк тоже сможет, но в следующем релизе, следите за апдейтами в репозитории.

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

Экспорт модели и инференс

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

Дальше менеджер инференса сам определяет доступное окружение — CUDA, CoreML или CPU — и запускает модель в подходящем режиме. При этом размер батча подбирается автоматически, чтобы не ловить OOM на реальных данных.

Пайплайн обучения классификации текстов

Отдельно расскажу о том, как проходит обучение классификации текстов. 

Диагностика качества данных

OpenAutoNLU умеет оценивать корректность разметки объектов тренировочной выборки — внутри системы есть опциональный слой диагностики, который позволяет увидеть, какие примеры стоит переразметить.  

Сначала DynamicTuner обучает модель и сохраняет логиты для каждого примера, чтобы замерить динамику изменения сигнала от модели. После этого включаются методы расчета корректности разметки:

• Retag ищет несовпадения между предсказанием и разметкой;

• Uncertainty фиксирует примеры с низкой уверенностью модели (softmax ниже порога);

• V-Information показывает, дает ли пример обучающий сигнал или просто шум;

• Dataset Cartography делит данные на легкообучаемые, пограничные и труднообучаемые по динамике уверенности.

Для NER используется Label Aggregation на базе Dawid-Skene. Модель делает несколько предсказаний с Monte Carlo dropout, после чего они агрегируются как разные аннотаторы. Это позволяет выявить расхождения еще на уровне токенов.

Подготовка выборки

После оценки распределения мощностей классов выборки пайплайн «дружит» данные с методом тренировки. Если в датасете слишком много классов с малым количеством примеров, система сначала пытается поднять их до рабочего уровня (Augmentex и — при необходимости и возможности — LLM).

Если после выбора метода оказывается, что для few-shot-режима некоторые классы, наоборот, слишком велики, выборка ограничивается сверху. Для SetFit максимальный размер класса составляет 80 примеров, для AncSetFit — 5. Это нужно, чтобы режим обучения не терял устойчивости из-за сильного дисбаланса.

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

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

Дальше обучение запускается уже в выбранной конфигурации. Если пайплайн ушел в полный fine-tuning, целевой метрикой становится macro-F1 на отложенной валидации. При включенной оптимизации гиперпараметров используется Optuna с семплером TPE, отбираются learning rate, batch size и weight decay. По умолчанию — в течение десяти итераций с критерием раннего останова в пять шагов без увеличения целевой метрики.

В few-shot-режимах явный поиск гиперпараметров не запускается. Там используются заранее подобранные настройки, в том числе число итераций contrastive learning и learning rate для энкодера. 

Генерация тестового набора на базе LLM

Если изначально нет тестового датасета, опционально подключается модуль генерации данных через LLM. Он получает обучающие данные и синтезирует реалистичные примеры, которые можно использовать как прокси-оценку качества. В экспериментах мы использовали GPT-4o-mini. На небольших выборках, до 80 примеров на класс, разница между сгенерированным и реальным тестом была не больше 5%.

При этом модуль не завязан на одну конкретную модель. OpenAutoNLU поддерживает любой openAI-совместимый API-эндпоинт, в том числе локально поднятые модели.

Эксперименты и результаты

Все эксперименты мы запускали на одной машине с Intel Xeon Gold 6448H на 64 ядра, NVIDIA H100 с 80 ГБ VRAM и 756 ГБ оперативной памяти. Время обучения считали по wall-clock time, то есть по реальному времени на этой конфигурации. Для оценки взяли четыре датасета классификации интентов: banking77, massive, hwu64 и snips, — так как они дают и бинарные, и многоклассовые сценарии. 

Сами прогоны проходили в трех режимах по объему данных: на совсем малых выборках (с 5–10 примерами на класс), средних (с 81–100 примерами) и на полном датасете. Все результаты усреднялись по трем рандомным значениям со стандартными пресетами фреймворков.

Мы сравнивали OpenAutoNLU с популярными фреймворками AutoIntent, AutoGluon, LightAutoML и H2O. В качестве основных метрик смотрели на F1-macro, F1-in-scope и F1-OOD. Чтобы сравнение не превращалось в спор о том, у кого удачнее эмбеддинги, там, где это возможно, использовали bert-base-uncased. А для пайплайнов, которым нужны представления в стиле E5, в том числе AutoIntent, брали intfloat/multilingual-e5-large-instruct.

На малых выборках OpenAutoNLU показал себя лучше, чем классические AutoML-фреймворки. Это тот сценарий, где ошибки слишком дорого обходятся: если сразу уйти в тяжелый fine-tuning, качество может быть нестабильным, если слишком долго держаться за few-shot, модель упирается в свой потолок.

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

Самая показательная часть экспериментов связана с OOD. Мы проверяли несколько сценариев:

• в OOD-aware примеры были в обучающем датасете;

• в OOD-in-test модель сталкивалась с ними только на тесте;

• отдельно смотрели режим OOD-unaware, который ближе всего к реальной эксплуатации: в обучающем датасете есть только in-domain данные, а в тестовом система получает запросы, которых в обучении не было.

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

В итоге OpenAutoNLU не только конкурентен на обычной классификации интентов, но и заметно сильнее в OOD-сценариях и на малых выборках. Для нас это был главный результат: мы хотели, чтобы пайплайн вел себя устойчиво в тех местах, где обычно и начинается ручная доработка.

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

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