Toxicity classifier

TensorFlow.js представляет собой библиотеку для разработки и запуска моделей машинного обучения в браузере или на Node.js. Одной из актуальных задач в области обработки естественного языка является классификация токсичности текста — определение, содержит ли текст оскорбления, ненавистнические высказывания или другие формы агрессивного контента. В основе этой задачи лежат нейронные сети, способные работать с последовательностями слов и извлекать семантические признаки.


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

Для классификатора токсичности необходимо подготовить корпус текстов с метками. Чаще всего используют бинарную классификацию: токсично / нетоксично, либо многоклассовую с различными уровнями агрессии. Формат данных может быть следующим:

[
  {"text": "Пример токсичного комментария", "label": 1},
  {"text": "Пример дружелюбного комментария", "label": 0}
]

Важный момент — токенизация текста. В TensorFlow.js это можно сделать с помощью готовых инструментов или собственных функций. Токенизация разбивает текст на слова или подсловные единицы, которые затем кодируются в числа для подачи на вход модели.

const tf = require('@tensorflow/tfjs');

function tokenize(text, wordIndex) {
  const tokens = text.toLowerCase().replace(/[^\w\s]/g, '').split(' ');
  return tokens.map(word => wordIndex[word] || 0);
}

Создание модели

Для классификации токсичности чаще всего используют модели на основе последовательностей: LSTM, GRU, или трансформеры. В TensorFlow.js модель можно строить по следующему принципу:

const model = tf.sequential();

// Входной слой
model.add(tf.layers.embedding({
  inputDim: vocabSize,
  outputDim: 128,
  inputLength: maxLen
}));

// Рекуррентный слой
model.add(tf.layers.lstm({
  units: 64,
  returnSequences: false
}));

// Полносвязный слой с активацией ReLU
model.add(tf.layers.dense({units: 64, activation: 'relu'}));

// Выходной слой с сигмоидной активацией для бинарной классификации
model.add(tf.layers.dense({units: 1, activation: 'sigmoid'}));

model.compile({
  optimizer: 'adam',
  loss: 'binaryCrossentropy',
  metrics: ['accuracy']
});

Ключевые моменты:

  • embedding слой превращает числовые индексы слов в векторы фиксированной размерности, что позволяет модели работать с семантикой слов.
  • LSTM слой учитывает порядок слов и контекст внутри последовательности.
  • Выходной слой с sigmoid дает вероятность токсичности текста.

Подготовка данных для обучения

Перед подачей на вход модели данные должны быть преобразованы в тензоры одинаковой длины:

function padSequences(sequences, maxLen) {
  return sequences.map(seq => {
    if (seq.length > maxLen) {
      return seq.slice(0, maxLen);
    } else if (seq.length < maxLen) {
      return Array(maxLen - seq.length).fill(0).concat(seq);
    }
    return seq;
  });
}

const xTrain = tf.tensor2d(padSequences(trainSequences, maxLen));
const yTrain = tf.tensor2d(trainLabels, [trainLabels.length, 1]);

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


Обучение модели

Обучение модели в TensorFlow.js проходит через метод fit, который поддерживает как браузерный, так и Node.js режим:

await model.fit(xTrain, yTrain, {
  epochs: 10,
  batchSize: 32,
  validationSplit: 0.2,
  callbacks: tf.callbacks.earlyStopping({monitor: 'val_loss', patience: 2})
});

Особенности:

  • validationSplit позволяет контролировать переобучение.
  • earlyStopping прекращает обучение, если метрика валидации перестала улучшаться.
  • Подбор batchSize влияет на стабильность градиента и скорость обучения.

Оценка и использование модели

После обучения важно проверить качество классификатора на отдельном тестовом наборе:

const xTest = tf.tensor2d(padSequences(testSequences, maxLen));
const yTest = tf.tensor2d(testLabels, [testLabels.length, 1]);

const result = model.evaluate(xTest, yTest);
result[1].print(); // Вывод accuracy

Для прогнозирования токсичности конкретного текста:

const inputSeq = padSequences([tokenize("Пример текста", wordIndex)], maxLen);
const inputTensor = tf.tensor2d(inputSeq);
const prediction = model.predict(inputTensor);
prediction.print(); // Вероятность токсичности

Оптимизация и улучшение модели

  1. Регуляризация: добавление dropout слоев между LSTM и Dense слоями снижает переобучение.
  2. Словарь и токенизация: использование подсловных токенов (subword) уменьшает проблему редких слов.
  3. Предобученные эмбеддинги: можно интегрировать Word2Vec или GloVe, преобразовав их в tfjs-формат.
  4. Аугментация данных: генерация синонимов или перефразирование текстов помогает модели лучше обобщать.

Расширенные подходы

Для более точной классификации токсичности применяют трансформеры, такие как BERT или DistilBERT, перенесенные в TensorFlow.js через конвертацию моделей:

import * as use from '@tensorflow-models/universal-sentence-encoder';

const model = await use.load();
const embeddings = await model.embed(["Пример текста"]);

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


TensorFlow.js предоставляет полный стек инструментов для построения токсик-классификаторов, начиная с предобработки текста и заканчивая сложными трансформерными архитектурами, способными работать в реальном времени прямо в браузере. Это делает его мощным инструментом для анализа текстового контента и фильтрации токсических сообщений.