Пишем сервис инференса ML-модели на go, на примере BERT-а

Привет, на связи команда аналитиков Х5 Tech. В статье пишем сервис инференса ML-NLP модели на go. Допустим, вам нужно внедрить ML-модель (разработанную/обученную на Рython-фреймворке) в сервис в вашей инфраструктуре. По какой-то причине (не важно какой) этот сервис должен быть на golang-е. Здесь покажем, как это можно сделать, используя ONNX.

Если вы это читаете, то, вероятно, или вы знакомы с обучением ML-моделей на Рython, библиотекой моделей huggingface, языковыми моделями BERT, или вы являетесь бэкенд разработчиком на golang.

В качестве примера будем использовать модель из библиотеки huggingface seara/rubert-tiny2-russian-sentiment, которая классифицирует сантимент текста.

Пара фраз про ONNX

Вот Гугл перевод с официального сайта.

Open Neural Network Exchange — это открытый формат, созданный для представления моделей машинного обучения. ONNX определяет общий набор операторов — строительных блоков моделей машинного и глубокого обучения — и общий формат файлов, позволяющий разработчикам ИИ использовать модели с различными платформами, инструментами, средами выполнения и компиляторами.

С ONNX на стадии деплоя пропадает зависимость от фреймворка модели. Сервис, выдающий предсказания, должен теперь уметь работать с сетками, сохранёнными в формате ONNX, какой при этом модель может быть как Яндексовый Catboost или sbert из библиотеки хагингфейса или какое-нибудь каффе2.

Постановка задачи

Есть модель классификации из текстов rubert-tiny2-russian-sentiment на 3 класса:

0: neutral 1: positive 2: negative

Есть чудесный код на Python для получения предсказаний принадлежности комментария к классу:

from transformers import AutoTokenizer, AutoModelForSequenceClassification  model_name = "seara/rubert-tiny2-russian-sentiment"  tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForSequenceClassification.from_pretrained(model_name)  input_text = ["Привет, ты мне нравишься!", "Ах ты черт мерзкий"] tokenized_text = tokenizer( input_text, return_tensors="pt", padding=True, truncation=True, add_special_tokens=True, ) outputs = model(**tokenized_text) predicted = outputs.logits.softmax(-1) print(predicted)    #[[0.0571, 0.9399, 0.0030], # [0.1483, 0.1615, 0.6902]]

Хотим научиться писать не менее чудесный код получения предсказаний на go при помощи ONNX.

Что будем использовать

— github.com/yalue/onnxruntime_go
обёртка для microsoft/onnxruntime

— github.com/daulet/tokenizers
токенайзер будет нам кодировать текст в вектора

Ещё есть github.com/knights-analytics/hugot – реализация пайплайнов Huggingface на go. В README репы уже есть решение нашей задачи.

Пользоваться этим мы, конечно, не будем.

А всё напишем сами, чтобы разобраться.

Подготовка среды

— Для daulet/tokenizers нужно сбилдить бинарник libtokenizers.a и добавить путь до него в переменную окружения CGO_LDFLAGS или разместить её в /usr/lib/libtokenizers.a

# кладём бинарь для токенайзера /usr/lib/libtokenizers.a  RUN wget https://github.com/daulet/tokenizers/releases/download/v1.20.2/libtokenizers.linux-amd64.tar.gz \ && tar -C /usr/lib -xzf libtokenizers.linux-amd64.tar.gz 

Для onnxruntime_go необходимо наличие libonnxruntime.so. Её скачаем отсюда https://github.com/microsoft/onnxruntime/releases/ и положим в путь /usr/lib/libonnxruntime.so

Всё вышесказанное в виде докерфайла:

ARG BUILD_PLATFORM=linux/amd64  FROM --platform=${BUILD_PLATFORM} golang:1.23.3-bullseye as runtime  WORKDIR /tmp # кладём бинарь для onnxruntime в /usr/lib/libonnxruntime.so RUN wget https://github.com/microsoft/onnxruntime/releases/download/v1.20.0/onnxruntime-linux-x64-1.20.0.tgz \ && tar -xzf onnxruntime-linux-x64-1.20.0.tgz \ && mv ./onnxruntime-linux-x64-1.20.0/lib/libonnxruntime.so.1.20.0 /usr/lib/libonnxruntime.so   # кладём бинарь для токенайзера /usr/lib/libtokenizers.a RUN wget https://github.com/daulet/tokenizers/releases/download/v1.20.2/libtokenizers.linux-amd64.tar.gz \ && tar -C /usr/lib -xzf libtokenizers.linux-amd64.tar.gz 

Экспорт модели в ONNX

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

Нужна будет библиотека для экспорта:

pip install 'optimum[exporters]'

Далее загрузим модель в память. Сохраним в папку модель и токенайзер:

from optimum.onnxruntime import ORTModelForSequenceClassification from transformers import AutoTokenizer  model_checkpoint = "seara/rubert-tiny2-russian-sentiment"  save_directory = "./data/rubert-tiny2-russian-sentiment-onnx"  # Load a model from transformers and export it to ONNX ort_model = ORTModelForSequenceClassification.from_pretrained( model_checkpoint, export=True ) tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)  # Save the onnx model and tokenizer ort_model.save_pretrained(save_directory) tokenizer.save_pretrained(save_directory) 

Получаем что-то вроде такого:

ls ./data/rubert-tiny2-russian-sentiment-onnx -rw-r--r--  1 dmitrij  staff   922B Nov 23 00:03 config.json -rw-r--r--  1 dmitrij  staff   111M Nov 23 00:03 model.onnx -rw-r--r--  1 dmitrij  staff   695B Nov 23 00:03 special_tokens_map.json -rw-r--r--  1 dmitrij  staff   2.3M Nov 23 00:03 tokenizer.json -rw-r--r--  1 dmitrij  staff   1.3K Nov 23 00:03 tokenizer_config.json -rw-r--r--  1 dmitrij  staff   1.0M Nov 23 00:03 vocab.txt 

— model.onnx – бинарник модели. Внезапно самый большой файл в папке. Путь до него будет указываться как путь до модели в go.
— tokenizer.json – словарь токенов, содержит также конфигурационную информацию (паддинги, специальные токены и прочее). Путь до него будем сообщать токенайзеру.

Параметры модели вход/выход

Как узнать из Питона

Далее небольшой кусочек кода, чтобы посмотреть названия входных слоёв и выходных:

from onnx import load, shape_inference  model = load("./data/rubert-tiny2-russian-sentiment-onnx/model.onnx")  inferred_model = shape_inference.infer_shapes(model) print(inferred_model.graph.output)  output =[node.name for node in model.graph.output]  input_all = [node.name for node in model.graph.input] input_initializer =  [node.name for node in model.graph.initializer] net_feed_input = list(set(input_all)  - set(input_initializer))  print('Inputs: ', net_feed_input) #Inputs:  ['attention_mask', 'input_ids', 'token_type_ids']  print('Outputs: ', output) # Outputs:  ['logits']  print(model.graph.output[0]) ''' name: "logits" type {   tensor_type { elem_type: 1 shape {   dim {     dim_param: "batch_size"   }   dim {     dim_value: 3   } }   } } ''' 

Фиксируем, что:
— входные слои 3 штуки: input_ids, attention_mask, token_type_ids
— выходной слой logits размером 3

Метод в Go

Ту же самую информацию можно получить в go, используя GetInputOutputInfo. Здесь мы узнаем ещё и тип данных, который ожидается на вход. В данном примере – INT64.

import ( "fmt" ort "github.com/yalue/onnxruntime_go" )      ...  input, output, err := ort.GetInputOutputInfo("/data/rubert-tiny2-russian-sentiment-onnx/model.onnx") fmt.Println(input) // [{input_ids ONNX_TYPE_TENSOR [-1 -1] ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64} {attention_mask ONNX_TYPE_TENSOR [-1 -1] ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64} {token_type_ids ONNX_TYPE_TENSOR [-1 -1] ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64}]      fmt.Println(output) // [{logits ONNX_TYPE_TENSOR [-1 3] ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT}] 

Эта информация нам понадобится при настройке входных/выходных параметров в модели в go.

Пишем вызов токенайзера

Проинициализируем объект токенайзера, указав путь до tokenizer.json, который ранее получили при экспорте:

package main  import ( "fmt" "github.com/daulet/tokenizers" )  func main() {  text := "Привет, ты мне нравишься!" tk, err := tokenizers.FromFile("/data/rubert-tiny2-russian-sentiment-onnx/tokenizer.json") if err != nil {     panic(err) } defer tk.Close() 

Далее настроим токенайзер. Наша модель ждёт три слоя:  input_ids, attention_mask, token_type_ids. В EncodeOption укажем, что нужно вернуть attention_mask и token_type_ids. Также выставим addSpecialTokens=true аналогично с кодом на Питоне.

encodeOptions := []tokenizers.EncodeOption{     tokenizers.WithReturnTypeIDs(),     tokenizers.WithReturnAttentionMask(), } // здесь 2ой параметр addSpecialTokens=true encodingResponse := tk.EncodeWithOptions(text, true, encodeOptions...)  fmt.Printf("IDs=%v\n", encodingResponse.IDs) fmt.Printf("AttentionMask=%v\n", encodingResponse.AttentionMask) fmt.Printf("TypeIDs=%v\n", encodingResponse.TypeIDs) // IDs=[2 51343 16 23101 20284 30100 64940 5 3] // AttentionMask=[1 1 1 1 1 1 1 1 1] // TypeIDs=[0 0 0 0 0 0 0 0 0] 

Получили токены для нашего текста и сравним их с тем, что выдаёт аналогичный код на Питоне.

In []: : input_text = "Привет, ты мне нравишься!" ...: tokenized_text = tokenizer( ...: input_text, ...: return_tensors="pt", ...: padding=True, ...: truncation=True, ...: add_special_tokens=True, ...: ) ...: print(tokenized_text) {'input_ids': tensor([[2, 51343,16, 23101, 20284, 30100, 64940, 5, 3]]), 'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1]])} 

Ура! Значения совпадают. Едем дальше.

Инициализация модели и предикт

По дефолту libonnxruntime.so ищется в /usr/lib/, но на всякий случай явно укажем:

ort.SetSharedLibraryPath("/usr/lib/libonnxruntime.so")  err = ort.InitializeEnvironment() if err != nil {     panic(err) } defer ort.DestroyEnvironment() 

Далее создадим сессию.

func NewDynamicAdvancedSession(onnxFilePath string, inputNames, outputNames []string, options *SessionOptions) (*DynamicAdvancedSession, error)

— onnxFilePath – путь до нашей модели с расширением .onnx;
— inputNames, outputNames – названия входных и выходных слоев соответственно.

Получается примерно такое:

session, err := ort.NewDynamicAdvancedSession(     "/data/rubert-tiny2-russian-sentiment-onnx/model.onnx",     []string{"input_ids", "token_type_ids", "attention_mask"},     []string{"logits"},     nil, ) if err != nil {     panic(err) } defer session.Destroy() 

Далее для получение прогноза нужно передать в сессию наши токены. Для этого нам понадобится небольшая функция, которая приведёт массив, полученный из токенайзера, к нужному типу данных для модели []uint32 -> ort.Tensor[int64]

func makeTensor(x []uint32) *ort.Tensor[int64] { inputShape := ort.NewShape(1, int64(len(x))) inputData := make([]int64, len(x)) for i, v := range x {     inputData[i] = int64(v) }  inputTensor, err := ort.NewTensor(inputShape, inputData) if err != nil {     panic(err) } return inputTensor }

Также нам потребуется создать выходной тензор для получения прогноза модели, куда будут записаны значения из слоя logits c размерностью (n x 3).

outputTensor, err := ort.NewEmptyTensor[float32](ort.NewShape(1, 3)) defer outputTensor.Destroy() if err != nil {     panic(err) } 

Получения предсказания модели выглядят так:

err = session.Run([]ort.Value{     makeTensor(encodingResponse.IDs),     makeTensor(encodingResponse.TypeIDs),     makeTensor(encodingResponse.AttentionMask), }, []ort.Value{outputTensor}) if err != nil {     panic(err) }  outputData := outputTensor.GetData() fmt.Println(outputData) // [-0.010913536 2.7902415 -2.9452484] 

Сравним с тем, что выдаёт нам аналогичный код на Питоне:

outputs = model(**tokenized_text) print(outputs.logits) # tensor([[-0.0109,  2.7902, -2.9452]], grad_fn=<AddmmBackward0>)

Аутпуты совпадают. Ура!

Полную версию скриптов можно глянуть в репе:
https://github.com/dsyubaev/onnx_bert/tree/main

Итого

Мы научились на примере языковой модели BERT:

— делать экспорт модели в ONNX;
— смотреть мета-информацию модели о её входах и выходах;
— написали программу на go получения предсказаний модели.

Теперь вы можете писать сервисы на golang, которые работают с ML-моделями, поддерживающими импорт/экспорт в ONNX формате.