AI 日报hiw3c.com

A Coding Guide to Google Research’s MSEB: Writing Sound Encoders to the Benchmark Contract and Scoring Them Across Classification, Clustering, Retrieval and Segmentation

MarkTechPost www.marktechpost.com RSS 全文
正文为英文,可一键机器翻译(仅首次需要等待)

In this tutorial, we work with MSEB, the Massive Sound Embedding Benchmark from Google Research, and approach it from the perspective of what a leaderboard number actually means: the evaluator surface. We install the package and map its three layers, then write two deliberately different encoders against the framework’s own abstract base class: one that measures loudness over time and one that measures timbre, and encode a small synthetic corpus we generate in the notebook so nothing has to be downloaded. We drive the classification, clustering, retrieval, and segmentation evaluators over those embeddings, call the metric functions directly to see what each one rewards, and finish by assembling the TaskMetadata a real submission carries. The result is a comparison in which the two encoders trade places depending on which evaluator is asked, which is the argument for a multi-task benchmark made in numbers rather than in prose.

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.")

We install mseb and import the three layers that a benchmark run walks down. The types module holds the shapes every task speaks, Sound, SoundEmbedding, Score and TaskMetadata; the encoder module holds MultiModalEncoder, the contract our own model implements; and the evaluators package holds one module per task family. We import only the four evaluators this notebook drives, because the classification, clustering, retrieval, and segmentation modules depend on nothing heavier than NumPy and scikit-learn. In contrast, the reranking and transcription evaluators pull in Whisper and the task runner pulls in TensorFlow and apache-beam. Everything below therefore runs on a free CPU runtime with no dataset download and no accelerator.

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()

We start with the type contract, because every other layer is expressed in it. A Sound carries a waveform, along with SoundContextParams, the identifier, sample rate, length, language, and optional transcript, which follow the audio through the whole pipeline. A SoundEmbedding carries an array of N embeddings and an array of M timestamp pairs, and the relation between N and M is the benchmark’s vocabulary: M equal to N means one vector per frame, while M equal to one means a single utterance-level vector, which is what our encoders produce. EncodingStats records the input and embedding sizes and exposes compression_ratio, here a thousandfold reduction from audio to vector. A Score is a metric name, a value and its bounds, and it validates itself at construction, rejecting an empty metric name or a minimum above its maximum, so a malformed number cannot reach a leaderboard. The embedding field also accepts N strings instead of N vectors, which is the door that step 8 walks through.

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
 
 
@section("2. The encoder contract: three methods, and the framework does the rest")
def encoder_contract():
    print("  MultiModalEncoder abstract methods a subclass must implement:")
    for name in sorted(encoder_lib.MultiModalEncoder.__abstractmethods__):
        print(f"    {name}")
    print("  final (framework-owned, do not override): setup(), encode()")
 
    t = np.arange(SR) / SR
    fade = np.exp(-2.5 * t).astype(np.float32)                # a decaying note, so the envelope is not flat
    sound = types.Sound(waveform=(0.5 * fade * np.sin(2 * np.pi * 440 * t)).astype(np.float32),
                        context=types.SoundContextParams(id="demo_000", sample_rate=SR, length=SR))
    for enc in (EnergyEnvelopeEncoder(), SpectralProfileEncoder()):
        enc.setup()
        emb = enc.encode([sound])[0]
        stats = emb.encoding_stats                            # attached by encode(), not by our code
        print(f"\n  {type(enc).__name__:24s} -> {emb.embedding.shape} {emb.embedding.dtype}"
              f"   output_type={enc.output_type().__name__}")
        print(f"  {'':24s}    EncodingStats(input={stats.input_size_bytes:,} B, "
              f"embedding={stats.embedding_size_bytes} B, flops={stats.flops})")
        print(f"  {'':24s}    first 6 dims: {np.round(emb.embedding[0][:6], 3)}")
    print("\n  The envelope encoder sees the note decay; the spectral encoder sees one peak at 440 Hz.")
 
    try:
        EnergyEnvelopeEncoder().encode(["not a Sound"])
    except ValueError as e:
        print(f"\n  wrong input type is caught by _check_input_types: {e}")
    return "two encoders satisfying MultiModalEncoder"
 
 
encoder_contract()

We write two encoders by subclassing MultiModalEncoder, whose abstract methods are exactly three: _setup loads whatever the model needs, _check_input_types rejects anything that is not a Sound, and _encode turns a batch into SoundEmbedding objects. The framework owns setup and encode, and encode is what attaches EncodingStats to every result, so our code never fills that in by hand. EnergyEnvelopeEncoder averages energy in sixteen equal time slices and therefore describes only how loudness moves; SpectralProfileEncoder pools the mean log-magnitude spectrum into sixteen bands and therefore describes timbre. Both L2-normalise their output so a dot product is a cosine. Encoding one decaying note through each shows the difference immediately: the envelope encoder sees the decay, and the spectral encoder sees a single peak at 440 Hz.

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                      # unit RMS: the envelope is the only loudness cue
 
    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))
 
 
@section("3. A synthetic corpus, encoded into MSEB embedding caches")
def build_corpus():
    corpus = [synthesize(k, i, 0) for k in CLASSES for i in range(N_PER_CLASS)]
    queries = [synthesize(k, i, 1) for k in CLASSES for i in range(N_PER_CLASS)]
    labels = {s.context.id: s.context.text for s in corpus + queries}
    print(f"  {len(corpus)} documents + {len(queries)} second takes of the same clips,"
          f" {len(CLASSES)} classes, 1.0s each @ {SR} Hz")
 
    caches, query_caches = {}, {}
    for enc in (EnergyEnvelopeEncoder(), SpectralProfileEncoder()):
        enc.setup()
        embeddings = enc.encode(corpus)                        # one batched call, like a real runner
        caches[type(enc).__name__] = {e.context.id: e for e in embeddings}
        query_caches[type(enc).__name__] = {e.context.id: e for e in enc.encode(queries)}
 
        matrix = np.vstack([e.embedding for e in embeddings])
        within, between = [], []
        for i in range(len(corpus)):
            for j in range(i + 1, len(corpus)):
                sim = float(matrix[i] @ matrix[j])
                (within if labels[corpus[i].context.id] == labels[corpus[j].context.id] else between).append(sim)
        print(f"  {type(enc).__name__:24s} cache of {len(embeddings)} embeddings, dim {matrix.shape[1]}"
              f"   mean cosine: same-class {np.mean(within):.3f} vs other-class {np.mean(between):.3f}"
              f"   (gap {np.mean(within) - np.mean(between):+.3f})")
 
    print("\n  Read that gap as a prediction: only the spectral encoder separates the classes at all.")
    print("  Steps 4-6 check whether the evaluators agree - and whether the gap is the whole story.")
    globals().update(CORPUS=corpus, QUERIES=queries, LABELS=labels, CACHES=caches, QCACHES=query_caches)
    return f"{len(corpus)} documents + {len(queries)} queries encoded by 2 encoders"
 
 
build_corpus()

We synthesize a corpus in which two cues are deliberately separated. The spectrum says which class a clip belongs to, a tone, a chirp or noise, while the amplitude envelope is drawn per item and is independent of class, so it identifies which clip it is without saying anything about what it is. We normalize every waveform to unit RMS before applying the envelope, leaving the envelope as the only loudness cue. We render each of the thirty-six items twice, once as the document and once as a noisier second take of the same clip, and encode both sets with both encoders into MSEB embedding caches, the plain dictionaries from sound id to SoundEmbedding that every evaluator consumes. The mean same-class and other-class cosine similarities printed here read as a prediction about the next three steps: only the spectral encoder separates the classes at all.