Создание собственного сервиса встраивания Bert с помощью TorchServe

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

К счастью, TorchServe существует, и мы можем запустить эту модель в производство и, таким образом, использовать ее в нашем новом приложении, но отсутствие документации может превратить простой процесс во что-то довольно сложное.
Это было то, что случилось со мной с моей моделью Берта, поэтому я буду использовать этот пост, чтобы объяснить шаги по ее использованию и созданию службы встраивания слов.

Первые шаги

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

Я начну с разбора того, что мы добавили. Функция инициализации проверит спецификацию системы и загрузит все файлы, используемые для загрузки модели. preprocess() отвечает за этап обработки данных в нашем конвейере, в данном случае за токенизацию. В inference() мы применим наши предварительно обработанные данные к модели и выполним объединение средних значений с результатами. Этот шаг необходим, потому что модель выводит вложения Берта для отдельных слов, а не вложений предложений. Наконец, постобработка проверит, успешно ли были получены вложения, и преобразует ответ в объект JSON для упрощения использования.

После определения нашего обработчика пришло время создать архив модели с помощью команды:

torch-model-archiver --model-name "bert" --serialized-file ./bert_model/pytorch_model.bin --extra-files "./bert_model/config.json, ./bert_model/vocab.txt" --handler "./BertHandler.py"

Где сериализованный файл — это наша модель .bin Pytorch, а дополнительные файлы — это дополнительные файлы, необходимые для запуска нашей модели. В данном случае это файл config.json и словарь для нашего токенизатора. После запуска команды bert.mar появится в каталоге, который затем необходимо поместить в каталог нашего хранилища моделей.
Мы можем создать файл config.properties, в котором указаны общие определения, такие как указание URL-адреса и количества рабочих процессов, которые нам нужны. Я использовал следующую конфигурацию:

default_workers_per_model=3
default_response_timeout=300
unregister_model_timeout=300
inference_address=http://0.0.0.0:8443
management_address=http://0.0.0.0:8444
metrics_address=http://0.0.0.0:8445

Запуск службы

Затем мы можем запустить нашу службу локально, используя следующую команду:

torchserve --start --model-store model-store --models bert=bert.mar

Где model-store — это имя папки, в которой находится архив ou bert.mar, а models — это имя и файл модели, которую мы будем использовать. После запуска мы увидим несколько журналов, указывающих, что конечные точки будут использоваться и что модель была загружена. TorchServe предоставляет несколько конечных точек, которые можно найти здесь, но сейчас мы будем обращаться к конечной точке ‹URL›:8443/prediction/‹model_name› для получения вложений.

Через несколько секунд, в зависимости от размера модели, сервис запустится, и мы сможем начать делать запросы. В этом случае я буду использовать простой скрипт с POST-запросом к указанной конечной точке. Согласно моим спецификациям, весь текст, отправляемый в службу, должен быть закодирован в UTF-8 , но вы можете изменить это, изменив функцию предварительной обработки в обработчике.

Размещение TorchServe в Docker

Как мы видим, все работает нормально. Но запуск локально — не лучшее решение, поэтому мы будем использовать докер для создания более доступного сервиса.
Для этого нам нужно создать образ и запустить его в контейнере с. Для создания docker-образа нам нужно определить конкретный файл (пример файла показан здесь) со всеми настройками. Используемый dockerfile следующий:

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

docker build --tag torchserve:1.0 .
docker run --publish 8443:8443 --publish 8444:8444 --publish 8445:8445 --detach --name bb torchserve:1.0

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