Руководство по кодированию для MSEB от Google Research: написание звуковых энкодеров в соответствии с требованиями бенчмарка и их оценка для классификации, кластеризации, поиска и сегментации

В этом руководстве мы работаем с MSEB — масштабным бенчмарком звуковых эмбеддингов от Google Research — и рассматриваем его с точки зрения того, что на самом деле означает число в таблице лидеров: поверхности оценивания. Мы устанавливаем пакет и разбираем его на три уровня, а затем создаём два принципиально разных энкодера на основе собственного абстрактного базового класса фреймворка: один измеряет громкость во времени, другой — тембр. После этого мы кодируем небольшой синтетический корпус, сгенерированный прямо в ноутбуке, поэтому ничего скачивать не нужно. Мы запускаем оцениватели классификации, кластеризации, поиска и сегментации для этих эмбеддингов, вызываем функции метрик напрямую, чтобы увидеть, что именно оценивает каждая из них, и завершаем сборкой TaskMetadata, которую содержит настоящая отправка. В результате получается сравнение, в котором два энкодера меняются местами в зависимости от выбранного оценивателя, что и является численным аргументом в пользу многозадачного бенчмарка, а не просто утверждением в тексте.

import os
import sys
import json
import math
import traceback
import subprocess
import numpy as np
 
RESULTS = {}
BENCH = {}
 
 
def banner(title):
    print("\n" + "=" * 78)
    print(title)
    print("=" * 78)
 
 
def section(name):
    def wrap(fn):
        def run(*a, **kw):
            banner(name)
            try:
                out = fn(*a, **kw)
                RESULTS[name] = out if isinstance(out, str) else "ok"
                return out
            except Exception as e:
                RESULTS[name] = f"SKIPPED / FAILED -> {type(e).__name__}: {e}"
                print(f"\n[!] {name} did not complete: {type(e).__name__}: {e}")
                traceback.print_exc(limit=3)
                return None
        return run
    return wrap
 
 
banner("0. Install MSEB and map the three layers we will use")
subprocess.run([sys.executable, "-m", "pip", "install", "-q", "mseb==0.1.0"], check=True)
 
import mseb
from mseb import types, encoder as encoder_lib, evaluator as evaluator_lib, metrics
from mseb.evaluators import (
    classification_evaluator,
    clustering_evaluator,
    retrieval_evaluator,
    segmentation_evaluator,
)
 
print(f"  mseb {mseb.__version__}  |  Python {sys.version.split()[0]}  |  numpy {np.__version__}")
print("\n  MSEB is three layers, and a benchmark run walks down them:")
print("    types      -> Sound, SoundEmbedding, Score, TaskMetadata: the shapes every task speaks")
print("    encoder    -> MultiModalEncoder: the contract YOUR model implements")
print("    evaluators -> classification, clustering, retrieval, reranking, transcription, segmentation, ...")
print("\n  evaluator entry points we will drive:")
for module, cls in [(classification_evaluator, "ClassificationEvaluator"),
                    (clustering_evaluator, "ClusteringEvaluator"),
                    (retrieval_evaluator, "RetrievalEvaluator"),
                    (segmentation_evaluator, "SegmentationEvaluator")]:
    print(f"    {module.__name__.split('.')[-1]:28s} {cls}")
print("\n  Everything below runs on CPU with no dataset download: we synthesise the audio.")

Мы устанавливаем mseb и импортируем три уровня, через которые проходит запуск бенчмарка. Модуль types содержит структуры, используемые всеми задачами: Sound, SoundEmbedding, Score и TaskMetadata; модуль encoder содержит MultiModalEncoder — контракт, которому соответствует наша собственная модель; пакет evaluators содержит отдельный модуль для каждого семейства задач. Мы импортируем только четыре оценивателя, используемые в этом ноутбуке, поскольку модули классификации, кластеризации, поиска и сегментации зависят лишь от NumPy и scikit-learn. В отличие от них, оцениватели переранжирования и транскрибации загружают Whisper, а исполнитель задач использует TensorFlow и apache-beam. Поэтому всё ниже выполняется на бесплатном CPU-окружении без скачивания датасета и без ускорителя.

SR = 16000
 
 
@section("1. The type contract: Sound, SoundEmbedding, Score")
def type_contract():
    t = np.arange(SR) / SR
    waveform = (0.5 * np.sin(2 * np.pi * 440 * t)).astype(np.float32)
    sound = types.Sound(
        waveform=waveform,
        context=types.SoundContextParams(id="demo_000", sample_rate=SR, length=len(waveform),
                                         language="en_us", text="a 440 Hz tone"),
    )
    print(f"  Sound          id={sound.context.id!r}  {sound.waveform.shape} @ {sound.context.sample_rate} Hz"
          f"  -> {sound.size_bytes:,} bytes")
 
    embedding = types.SoundEmbedding(
        embedding=np.zeros((1, 16), dtype=np.float32),            # (N, D): one utterance-level vector
        timestamps=np.array([[0.0, 1.0]], dtype=np.float32),      # (M, 2): [start, end] in seconds
        context=sound.context,
        encoding_stats=types.EncodingStats(input_size_bytes=sound.size_bytes, embedding_size_bytes=16 * 4),
    )
    print(f"  SoundEmbedding embedding{embedding.embedding.shape}  timestamps{embedding.timestamps.shape}"
          f"  -> {embedding.size_bytes} bytes")
    print(f"                 compression_ratio = {embedding.encoding_stats.compression_ratio:.5f}"
          f"  ({1 / embedding.encoding_stats.compression_ratio:,.0f}x smaller than the audio)")
    print("  N embeddings and M timestamps: M == N is frame-aligned, M == 1 is utterance-level.")
    print("  `embedding` may also hold N strings instead of vectors - step 8 uses exactly that.")
 
    score = types.Score(metric="Accuracy", description="Overall classification accuracy",
                        value=0.875, min=0.0, max=1.0)
    print(f"\n  Score          {score.metric}={score.value} in [{score.min}, {score.max}] :: {score.description}")
    for bad, why in [(dict(metric="", description="d", value=0.5, min=0.0, max=1.0), "empty metric name"),
                     (dict(metric="m", description="d", value=0.5, min=1.0, max=0.0), "min > max")]:
        try:
            types.Score(**bad)
        except Exception as e:
            print(f"  rejected at construction ({why}): {type(e).__name__}: {e}")
    return f"Sound {sound.size_bytes:,} B -> embedding {embedding.size_bytes} B"
 
 
type_contract()

Мы начинаем с контракта типов, поскольку все остальные уровни выражаются через него. Sound содержит аудиосигнал и SoundContextParams — идентификатор, частоту дискретизации, длину, язык и необязательную расшифровку, которые сопровождают аудио на всём протяжении конвейера. SoundEmbedding содержит массив из N эмбеддингов и массив из M пар временных меток; соотношение между N и M составляет словарь бенчмарка: M, равное N, означает один вектор на кадр, а M, равное единице, — один вектор на всё высказывание, именно такой вариант создают наши энкодеры. EncodingStats записывает размеры входных данных и эмбеддинга и предоставляет compression_ratio — здесь это тысячекратное уменьшение размера аудио до размера вектора. Score содержит название метрики, её значение и границы и проверяет себя во время создания, отклоняя пустое имя метрики или минимальное значение, превышающее максимальное, чтобы некорректное число не попало в таблицу лидеров. Поле embedding также принимает N строк вместо N векторов — именно этим воспользуется шаг 8.

class EnergyEnvelopeEncoder(encoder_lib.MultiModalEncoder):
    """Baseline: average energy in `n_bins` equal time slices. Loud/quiet, nothing about timbre."""
 
    def __init__(self, n_bins: int = 16):
        super().__init__()
        self.n_bins = n_bins
 
    def _setup(self):
        self._ready = True                                    # a real encoder loads weights here
 
    def _check_input_types(self, batch):
        for item in batch:
            if not isinstance(item, types.Sound):
                raise ValueError(f"{type(self).__name__} takes types.Sound, got {type(item).__name__}")
 
    def _encode(self, batch) -> list[types.SoundEmbedding]:
        out = []
        for sound in batch:
            slices = np.array_split(sound.waveform.astype(np.float32), self.n_bins)
            vec = np.array([[float(np.sqrt(np.mean(s ** 2) + 1e-12)) for s in slices]], dtype=np.float32)
            vec /= np.linalg.norm(vec) + 1e-9
            out.append(types.SoundEmbedding(
                embedding=vec,
                timestamps=np.array([[0.0, sound.context.length / sound.context.sample_rate]], dtype=np.float32),
                context=sound.context))
        return out
 
 
class SpectralProfileEncoder(encoder_lib.MultiModalEncoder):
    """Contender: mean log-magnitude spectrum pooled into `n_bands` bands. Describes timbre."""
 
    def __init__(self, n_bands: int = 16, frame: int = 512):
        super().__init__()
        self.n_bands, self.frame = n_bands, frame
 
    def _setup(self):
        self._window = np.hanning(self.frame).astype(np.float32)
 
    def _check_input_types(self, batch):
        for item in batch:
            if not isinstance(item, types.Sound):
                raise ValueError(f"{type(self).__name__} takes types.Sound, got {type(item).__name__}")
 
    def _encode(self, batch) -> list[types.SoundEmbedding]:
        out = []
        for sound in batch:
            w = sound.waveform.astype(np.float32)
            n_frames = max(1, len(w) // self.frame)
            spectra = [np.abs(np.fft.rfft(w[i * self.frame:(i + 1) * self.frame] * self._window))
                       for i in range(n_frames)]
            mean_spectrum = np.log1p(np.mean(spectra, axis=0))
            vec = np.array([[float(b.mean()) for b in np.array_split(mean_spectrum, self.n_bands)]],
                           dtype=np.float32)
            vec /= np.linalg.norm(vec) + 1e-9
            out.append(types.SoundEmbedding(
                embedding=vec,
                timestamps=np.array([[0.0, sound.context.length / sound.context.sample_rate]], dtype=np.float32),
                context=sound.context))
        return out

Мы создаём два энкодера, наследуя их от MultiModalEncoder, у которого ровно три абстрактных метода: _setup загружает всё необходимое модели, _check_input_types отклоняет объекты, не являющиеся Sound, а _encode превращает пакет в объекты SoundEmbedding. Фреймворк самостоятельно управляет setup и encode; именно encode добавляет EncodingStats к каждому результату, поэтому наш код не заполняет это поле вручную. EnergyEnvelopeEncoder усредняет энергию в шестнадцати равных временных интервалах и описывает только изменения громкости; SpectralProfileEncoder объединяет средний логарифм амплитудного спектра в шестнадцать полос и описывает тембр. Оба энкодера нормализуют результат по L2, поэтому скалярное произведение равно косинусному сходству. Кодирование затухающей ноты обоими энкодерами сразу показывает разницу: первый видит затухание, а второй — единственный пик на частоте 440 Гц.

CLASSES = ["tone", "chirp", "noise"]
N_PER_CLASS = 12
 
 
def synthesize(kind: str, index: int, take: int) -> types.Sound:
    """One second of audio. `take` 0 is the document, take 1 is a noisier recording of the SAME clip.
    Two cues are deliberately separated: the spectrum says which class it is, and the amplitude
    envelope - drawn per item, independent of class - says which item it is.
    """
    item = np.random.default_rng(1000 + CLASSES.index(kind) * 100 + index)
    control = 0.25 + 0.75 * item.random(8)
    envelope = np.interp(np.linspace(0, 7, SR), np.arange(8), control).astype(np.float32)
 
    t = np.arange(SR) / SR
    if kind == "tone":
        w = np.sin(2 * np.pi * (380 + 80 * item.random()) * t)
    elif kind == "chirp":
        f0, f1 = 200 + 50 * item.random(), 3200 + 400 * item.random()
        w = np.sin(2 * np.pi * (f0 * t + 0.5 * (f1 - f0) * t ** 2))
    else:
        w = item.standard_normal(SR)
    w /= np.sqrt(np.mean(w ** 2)) + 1e-9
    take_rng = np.random.default_rng(50_000 + take * 10_000 + CLASSES.index(kind) * 100 + index)
    w = (0.4 + 0.2 * take_rng.random()) * envelope * (w + 0.02 * take_rng.standard_normal(SR))
    return types.Sound(waveform=w.astype(np.float32), context=types.SoundContextParams(
        id=f"{kind}_{index:02d}" + ("" if take == 0 else "_take2"), sample_rate=SR,
        length=SR, language="en_us", text=kind))

Мы синтезируем корпус, в котором два признака намеренно разделены. Спектр определяет класс записи — тон, чирп или шум, — тогда как амплитудная огибающая генерируется отдельно для каждого объекта и не зависит от класса. Поэтому она идентифицирует конкретный фрагмент, но ничего не говорит о его содержании. Перед применением огибающей мы нормализуем каждую форму сигнала до единичного RMS, оставляя огибающую единственным признаком громкости. Каждый из тридцати шести объектов создаётся дважды: один раз как документ и второй раз как более шумная запись того же фрагмента. Обе версии кодируются обоими энкодерами в кэши эмбеддингов MSEB — обычные словари, сопоставляющие идентификатор звука с SoundEmbedding, которые использует каждый оцениватель. Напечатанные средние косинусные сходства внутри класса и между классами позволяют предсказать результат следующих трёх шагов: только спектральный энкодер вообще разделяет классы.

ClassificationEvaluator принимает таблицу эмбеддингов классов в качестве weights и функцию расстояния; мы строим weights как прототипы классов — средние единичные векторы каждого класса. Два его метода чётко разделены: compute_predictions возвращает для каждого кэшированного эмбеддинга исходную оценку каждого класса, а compute_metrics преобразует эти оценки вместе с метками ClassificationReference в список объектов Score, сохраняемый таблицей лидеров. Параметр top_k_value, равный двум, добавляет Top-2 Accuracy к accuracy, balanced accuracy, а также взвешенным precision, recall и F1. Спектральный энкодер идеально классифицирует корпус, тогда как энкодер огибающей работает значительно лучше случайного угадывания, но заметно хуже него. Это соответствует порядку, предсказанному разницей косинусных сходств.

ClusteringEvaluator задаёт более сложную версию того же вопроса, поскольку не видит меток во время кодирования: он запускает KMeans над кэшем. Кластеры оцениваются по меткам с помощью V-measure — гармонического среднего однородности и полноты. Разрыв между двумя энкодерами здесь резко увеличивается по сравнению с классификацией, поскольку контролируемое считывание через прототипы может использовать слабый признак, который неконтролируемая кластеризация самостоятельно обнаружить не способна. Есть важная практическая деталь, которую стоит воспроизводить в любом бенчмарке: оцениватель создаёт MiniBatchKMeans без random_state, поэтому используется глобальный генератор NumPy. Без фиксации этого генератора неструктурированное пространство эмбеддингов даёт от запуска к запуску значения примерно от 0,01 до 0,08.

RetrievalEvaluator отвечает на вопрос, отличный от двух предыдущих, и мы специально настраиваем задачу так, чтобы это различие было заметно. Каждый запрос — это более шумная вторая запись ровно одного документа, поэтому требуется определить идентичность, а не категорию. Эмбеддинги документов индексируются в BruteForceSearcher, затем для кэша запросов вычисляются предсказания, а для каждого запроса передаётся один RetrievalReferenceId с указанием единственного правильного документа. Оцениватель возвращает MRR, exact match, recall при выбранном top_k и NDCG на десяти позициях. Результат меняет порядок из двух предыдущих шагов: энкодер огибающей извлекает каждый фрагмент на первом месте, поскольку огибающая служит отпечатком объекта. Спектральный энкодер работает немного хуже, так как записи одного класса выглядят для него похожими. Напечатанные списки первых пяти результатов показывают механизм: одно соседство случайно относительно классов, другое состоит преимущественно из объектов одного класса.

Мы вызываем функции метрик напрямую, без оболочки оценивателя, поскольку именно этот уровень используется разными семействами задач. compute_word_errors и compute_character_errors принимают две строки и отдельно возвращают число ошибок и общее количество элементов, поэтому вызывающий код сам выбирает способ агрегации по корпусу. Метрики ранжирования принимают эталон и упорядоченный список идентификаторов; сравнение exact match, reciprocal rank и nDCG на одном и том же ранжировании показывает, как каждая метрика учитывает позицию. Важный нюанс: compute_ndcg_at_k предполагает один релевантный документ и сравнивает его по равенству, поэтому передача списка релевантных идентификаторов незаметно даёт нулевую оценку. MRR, напротив, принимает список и по-прежнему выглядит корректно. В завершение мы используем compute_lp_norm и compute_dynamic_time_warping_distance — расстояния в пространстве эмбеддингов, лежащие в основе задач реконструкции и стабильности.

SegmentationEvaluator оценивает отдельно, что было сказано и где именно это было сказано, используя строковое представление SoundEmbedding, упомянутое на первом шаге: массив embedding содержит по одному термину на сегмент, а timestamps — временные интервалы сегментов. Процесс состоит из двух этапов: сначала compute_scores обрабатывает предсказания и эталонные данные, затем compute_metrics агрегирует полученный результат. Мы сравниваем четыре варианта сегментации одной и той же фразы с эталоном, допуская ошибку в пятьдесят миллисекунд. Точное совпадение и вариант с отклонением на пятьдесят миллисекунд получают максимальную оценку — именно для этого и нужна допустимая погрешность. Две последние строки показывают главный вывод: правильные слова в неправильных местах дают единицу по эмбеддингам и ноль по временным меткам, а правильные места с неправильными словами дают обратный результат. Только объединённая метрика учитывает одновременную корректность обоих аспектов.

Мы собираем TaskMetadata, которое содержит настоящая отправка: имя, тип, категорию, основную метрику, путь и ревизию датасета, оцениваемые разбиения и языки, а также сами объекты Score. Как и Score, TaskMetadata проверяется во время создания и отклоняет пустой список оценок. Затем мы сводим все полученные результаты в одну таблицу: одна строка на энкодер и один столбец на семейство задач. Победитель меняется от столбца к столбцу: энкодер, который не способен назвать звук, всё же может его распознать, а энкодер, правильно называющий каждый звук, может перепутать записи, принадлежащие одной группе. Одна основная цифра полностью скрыла бы это, что и служит аргументом в пользу бенчмарка, масштабного не только по данным, но и по числу задач.

banner("SUMMARY")
for name, res in RESULTS.items():
    print(f"  {name:<74s} {res}")
print("""
Where to go next
 - Swap in a real encoder: mseb/encoders/ ships wav2vec, Whisper, CLAP, EnCodec and SoundStream
   wrappers, plus CascadeEncoder for speech-to-text-to-embedding chains. Only the three methods
   from step 2 change; every evaluator above keeps working.
 - Run a published task: mseb.runner drives mseb.task over a real dataset with apache-beam; the
   task families live in mseb/tasks/ (classification, retrieval, reranking, transcription,
   segmentation, clustering, reasoning, brain_encoding, stability).
 - Compare against the leaderboard: https://huggingface.co/spaces/google/mseb-leaderboard
 - Read the contract you implemented: mseb/encoder.py and mseb/evaluator.py are ~500 lines total.
""")

В сводке печатается однострочный результат, возвращённый каждым разделом, а затем предлагаются три направления для продолжения работы: замена игрушечных моделей одним из реальных энкодеров, поставляемых вместе с пакетом — wav2vec, Whisper, CLAP, EnCodec, SoundStream или каскадной оболочкой, — при этом изменятся только три метода из шага 2, а все оцениватели продолжат работать; запуск опубликованной задачи через mseb.runner на настоящем датасете; и сравнение результата с общедоступной таблицей лидеров.

В заключение мы рассмотрели MSEB как то, чем он является на самом деле: контракт и набор оценивателей, которые можно полностью запустить без скачивания датасета и без ускорителя. Реализации трёх методов достаточно, чтобы сделать собственный код полноправной частью бенчмарка, после чего фреймворк сам обрабатывает пакетирование, статистику и проверку данных. Оцениватели задают действительно разные вопросы одним и тем же эмбеддингам: классификация и кластеризация спрашивают, чем является звук, поиск — какой именно это звук, а сегментация — что было сказано и где, причём эти аспекты оцениваются отдельно, чтобы ошибка во времени и ошибка распознавания не скрывались внутри одного среднего значения. Наши два энкодера менялись местами в зависимости от поставленного вопроса, и именно этот результат мы сохраняем, поскольку одно число не способно ранжировать звуковой эмбеддинг. Следующий шаг — заменить учебные энкодеры настоящей моделью и повторно запустить те же оцениватели: при улучшении эмбеддингов код вычисления оценок останется неизменным.

Посмотрите репозиторий GitHub с полным кодом. Все права и благодарности принадлежат исследователю, выполнившему этот проект. Также подписывайтесь на нас в Twitter и не забудьте присоединиться к нашему ML-сообществу на Reddit с более чем 150 тысячами участников и подписаться на нашу рассылку. Стоп! Вы есть в Telegram? теперь к нам можно присоединиться и в Telegram.

Хотите сотрудничать с нами для продвижения своего репозитория GitHub, страницы Hugging Face, релиза продукта, вебинара и т. д.? Свяжитесь с нами