Всем привет, меня зовут Евгений Мунин. Я Senior ML Engineer в Ad Tech в платформе ставок для Web рекламы и автор ТГ канала ML Advertising. Сегодня расскажу, как мы применяем ML модели в рекламных платформах с бэк‑ендом на JVM.

Как работают платформы ставок?

Все начинается с момента когда пользователь заходит на сайт, на котором содержатся рекламные слоты. В этот момент издатель, который владеет сайтом логирует вход пользователя и отправляет запрос в аггрегатор т.н. Prebid. Он в свою очередь контактирует с Supply Side платформой, которая проводит аукцион на продажу рекламного слота. Участники аукциона оценивают запрос, чтобы понять, соответствует ли он критериям рекламных кампаний клиентов, и если да, то какова его ценность для рекламодателя. После этого участники отправляют ответ с ставкой обратно. Победитель аукциона выставляет креатив своего клиента на купленное рекламное мести для показа.

Платформы ставок обрабатывают запросы пользователя, фильтруют их и делают рекомендации близко к реальному времени. Зачастую, весь процесс от запроса до размещения креатива на сайте должен занимать не более 150 мс, т.е. то время за которое у пользователя с хорошим интернетом загружается web-страница. Чтобы отвечать этому требованию бэк платформ, как правило, пишется на языках под JVM (встречал на своей практике на Java, Scala) или еще на Rust.

Так, а в чем же здесь проблема, спросите вы меня. А проблема кроется в том, платформа, обрабатывая входящий запрос от пользователя должна применять на нем ML модели, которые, например, фильтруют fraud, отсекает низкокачественный трафик или подкручивают монетизацию. Эти ML модели, которые предполагается использовать на платформах ставок, в большинстве случаев пишутся на Python фреймворках: Sklearn, PyTorch, etc. (кроме разве-что SparkML, у которого есть обертки и на Python, и на Scala). Библиотеки не особо часто запариваются над вопросами совместимости с другими языками.

Чтобы решить этот вопрос был разработан единый формат Open Neural Network Exchange (сокращенно ONNX), в который можно записать ML модели с разных библиотек и сделать их доступными для использования на платформах, в том числе под JVM.

Сегодня мы рассмотрим на примере простой модели логистической регрессии на Sklearn, которая предсказывает вероятность наличия ставки на рекламной платформе. Для того, чтобы обернуть ее в ONNX формат мы воспользуемся библиотекой ONNXRuntime и ее Java байндингом, чтобы запустить модель на JVM.

Напишем модель логистической регрессии

Для начала, нам нужно прописать сам пайплайн модели. Здесь мы воспользуемся sklearn.pipeline, который будет включать два этапа:

• FeatureHasher , чтобы хешировать входящие признаки следующим образом: hash_feat = MurMurHash3(feat) % hash_size

• LogisticRegression модель, которая предсказывает вероятность ставки на рекламный запрос

class GammaPipeline(Pipeline):
  def __init__(self):
    super().__init__([
        ("hasher", FeatureHasher(n_features=2**5, input_type="string", dtype=np.float32)),
        ("logreg",
            LogisticRegression(
                max_iter=10,
                random_state=42,
                penalty="l2",
                solver="lbfgs",
                C=1.0,
                tol=1e-4,
            ),
        ),
    ], verbose=True)

Когда мы задали пайплайн, создадим выборку каких-нибудь данных.

data = {
    "timestamp": [datetime.datetime.now() for _ in range(3)],
    "feature1": ["val10", "val11", "val12"],
    "feature2": ["val20", "val21", "val22"],
    "target": [True, False, False],
}

df = pd.DataFrame(data)

x = df.drop(columns=["target"])
y = df["target"]

Зафитим пайплайн и отобразим его структуру

pipeline = GammaPipeline()
pipeline.fit(x.values, y)

Сериализуем ONNX предиктор

Теперь, когда пайплайн обучен, будем кастовать его в ONNX формат. Здесь нам понадобится определить initial_type фичей и их количество в выборке. Также стоит отключить zipmap для меток классов логистической регрессии, поскольку мы ожидаем выход модель, ни как список словарей с метками классов, а просто, как вектор вероятностнй.

initial_type = [
    ('input', StringTensorType([None, len(x[0])])),
]

options = {LogisticRegression: {'zipmap': False}}

onnx_model = convert_sklearn(
    pipeline,
    initial_types=initial_type,
    options=options
)

with open(path_data + "models/onnx_log_reg.onnx", "wb") as f:
  f.write(onnx_model.SerializeToString())

После того, как модель записана в ONNX, можем ее визуализировать с помощью netron.app.

Читаем ONNX предиктор и запускаем на Java

Теперь, когда модель записана в формат ONNX, можем ее десериализовать на JVM. В данном примере, я использовал Java 17.

Для начала создадим OrtEnvironment, откроем сессию для предсказания с помощью OrtSession класса и передадим путь до .onnx предиктора. Когда предиктор прочитан, запускаем инференс с помощью session.run. Поскольку модель ожидает на входе структуру в виде Map<String, OnnxTensor>, ее ключи колонок соответствовать названием входных колонок в пайплайн, а сами колонки должны иметь тип OnnxTensor .

public class OnnxModelRunner {
    private OrtSession session;
    private OrtEnvironment environment;

    public OnnxModelRunner(String modelPath) throws OrtException {
        environment = OrtEnvironment.getEnvironment();
        OrtSession.SessionOptions options = new OrtSession.SessionOptions();
        session = environment.createSession(modelPath, options);
    }

    public OrtSession.Result runModel(String[][] bidders) throws OrtException {
        OnnxTensor inputTensor = OnnxTensor.createTensor(OrtEnvironment.getEnvironment(), bidders);
        return session.run(Collections.singletonMap("input", inputTensor));
    }
}

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

public class ApplicationOnnx {

    public static void main(String[] args) throws OrtException {

        String[][] bidders = {
            {"val11", "val21"},
            {"val12", "val22"},
            {"val13", "val23"}
        };

        OnnxModelRunner modelRunner = new OnnxModelRunner("onnx_log_reg_v1_2.onnx");

        OrtSession.Result results = modelRunner.runModel(bidders);

        StreamSupport.stream(results.spliterator(), false)
                .filter(onnxItem -> Objects.equals(onnxItem.getKey(), "probabilities"))
                .forEach(onnxItem -> {
                    OnnxValue onnxValue = onnxItem.getValue();
                    OnnxTensor tensor = (OnnxTensor) onnxValue;
                    try {
                        float[][] probas = (float[][]) tensor.getValue();
                        System.out.println(
                                "    tensor.getValue(): " + tensor.getValue() +
                                        "\n    probas: " + Arrays.deepToString(probas)
                        );
                    } catch (OrtException e) {
                        throw new RuntimeException(e);
                    }
                });
    }
}

В OrtSession.Result у нас Map из двух элементов:

Отфильтруем мапу по ключу probabilities, выбрав только вероятности. Значения имеет тип OnnxValue, который мы сначана скастуем в OnnxTensor а потом в массив чисел float[][]. Расперсенный выход модели выглядит следующим образом:

Output: probabilities: OnnxTensor(info=TensorInfo(javaType=FLOAT,onnxType=ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT,shape=[3, 2]))
    tensor.getValue(): [[F@6438a396
    probas: [[0.29334846, 0.70665157], [0.853327, 0.14667302], [0.853327, 0.14667302]]

В заключение

Мы рассмотрели простой пример, как портировать Sklearn пайплайн в ONNX, потом прочитать его и запустить предсказания на Java. Естественно, в зависимости от усложнения задач будут подниматься новые вопросы, например:

• Как записывать в ONNX пайплайн с кастомными преобразованиями, и как задавать shape_calculator и converter

• Что делать, если требуется сохранить тип разных входных колонок, вместо то, чтобы их переводить все в string, float?

• Как сделать, чтобы ONNX модель могла запускаться на предсказании на GPU

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

Если вас интересует тема ML моделей в рекламных платформах, и как катить модели в прод, то заходите на мой ТГ канал ML Advertising!

Спасибо что дочитали о конца!