Привет, на связи команда аналитиков Х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 формате.

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