Классификация и встраивание текста: BERT, DistilBERT

ONNX Runtime Web (ORT Web) предоставляет возможность запускать модели машинного обучения, экспортированные в формат ONNX, непосредственно в браузере или в среде Node.js с использованием JavaScript. Библиотека оптимизирована для работы с WebAssembly (WASM) и WebGL, что обеспечивает высокую производительность даже на устройствах с ограниченными ресурсами. Основные компоненты ORT Web включают:

  • InferenceSession — основной объект для загрузки и выполнения модели.
  • Tensor — структура данных для представления входных и выходных массивов.
  • Execution Providers — интерфейсы для использования различных движков вычислений (WASM, WebGL, WebGPU).

Загрузка модели BERT и DistilBERT

Модели типа BERT и DistilBERT предназначены для задач обработки естественного языка, таких как классификация текста, выделение эмбеддингов и анализ сентимента. Для работы с ними в ORT Web необходимо:

  1. Экспортировать модель в формат ONNX, сохранив все веса и конфигурацию токенизатора.
  2. Создать экземпляр InferenceSession с указанием URL или пути к файлу модели.
import * as ort from 'onnxruntime-web';

const session = await ort.InferenceSession.create('bert_model.onnx', {
    executionProviders: ['wasm', 'webgl']
});

Подготовка входных данных

BERT и DistilBERT требуют специфической подготовки текста:

  • Токенизация — преобразование текста в последовательность токенов с использованием заранее обученного токенизатора.
  • Создание attention masks — бинарная маска для обозначения значимых токенов.
  • Паддинг — выравнивание последовательностей до одинаковой длины.

Пример подготовки входных данных:

import { BertTokenizer } from '@huggingface/tokenizers';

const tokenizer = await BertTokenizer.fromOptions({ vocabFile: 'vocab.txt' });

const text = "Пример текста для классификации";
const encoded = tokenizer.encode(text);
const inputIds = new ort.Tensor('int64', encoded.ids, [1, encoded.ids.length]);
const attentionMask = new ort.Tensor('int64', encoded.attentionMask, [1, encoded.ids.length]);

Выполнение инференса

После подготовки входов можно выполнить инференс, используя метод run:

const feeds = {
    input_ids: inputIds,
    attention_mask: attentionMask
};

const results = await session.run(feeds);
const logits = results.logits.data;

Ключевым моментом является соответствие имен входов и выходов модели тем, которые использовались при экспорте в ONNX. Для BERT и DistilBERT это обычно input_ids, attention_mask и token_type_ids (для BERT).

Получение эмбеддингов

Для задач встраивания текста (text embedding) DistilBERT чаще используется для генерации векторных представлений. В этом случае интерес представляет выход последнего скрытого слоя (last_hidden_state):

const embeddings = results.last_hidden_state.data;
// Вектор размерности [batch_size, sequence_length, hidden_size]

Эти эмбеддинги можно использовать для:

  • Кластеризации текста.
  • Поиска по семантическому сходству.
  • Фильтрации и ранжирования документов.

Оптимизация работы в браузере

Для больших моделей BERT и DistilBERT критично снизить задержки:

  1. WebAssembly (WASM) обеспечивает кросс-браузерную совместимость и стабильное потребление памяти.
  2. WebGL позволяет использовать GPU для ускорения операций с матрицами.
  3. WebGPU (при поддержке) даёт ещё более высокую производительность на современных устройствах.

Выбор движка осуществляется через executionProviders при создании InferenceSession.

Асинхронная обработка и многопоточность

ORT Web поддерживает асинхронное выполнение инференса, что позволяет не блокировать основной поток браузера. Использование Web Worker позволяет вынести вычисления на отдельный поток:

const worker = new Worker('inferenceWorker.js');
worker.postMessage({ text: "Пример текста" });
worker.onmess age = (e) => {
    const embeddings = e.data;
    console.log(embeddings);
};

Преобразование и сохранение результатов

Эмбеддинги и логиты, полученные с помощью ORT Web, обычно хранятся в виде массивов Float32Array или Int64Array. Их можно:

  • Сохранять локально в IndexedDB для кэширования.
  • Отправлять на сервер для дальнейшей обработки.
  • Использовать напрямую в браузере для поиска и рекомендаций.

Поддержка различных архитектур

ORT Web поддерживает как стандартные BERT, так и облегчённые DistilBERT модели. Отличие заключается в количестве слоёв и размере скрытых состояний, что влияет на производительность и точность:

  • BERT: полная модель с большим количеством слоёв (12 для base).
  • DistilBERT: уменьшенная версия, оптимизированная для быстрого инференса, но с минимальной потерей качества.

Использование DistilBERT оправдано для приложений, где критична скорость отклика и ограничены ресурсы клиента.

Интеграция с приложениями

ORT Web позволяет интегрировать BERT и DistilBERT в:

  • Поисковые системы на стороне клиента.
  • Системы рекомендаций и фильтрации контента.
  • Чат-боты и интерактивные NLP интерфейсы.

Главная особенность — возможность выполнения модели полностью на клиенте без обращения к серверу, что повышает скорость отклика и защищает конфиденциальность данных.

Оглавление