%matplotlib inline
import os, sys, time, warnings
os.environ.setdefault("OMP_NUM_THREADS", "4")   # before numpy/torch: keep CPU threads sane
sys.path.insert(0, "..")
warnings.filterwarnings("ignore")

from pathlib import Path
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import librosa
import torch
import torch.nn as nn
import torch.nn.functional as F

from src.utils import load, play, show_wave, show_spec, logmel, median_f0, SR, HOP
from src.probe import pool_features, probe_accuracy, probe_layers, layer_similarity, plot_layer_curves

np.random.seed(0)
torch.manual_seed(0)
torch.set_num_threads(min(4, os.cpu_count() or 1))

DATA = Path("../data")
OUT = Path("../outputs/day5"); OUT.mkdir(parents=True, exist_ok=True)
CACHE = Path.home() / ".cache"           # pretrained models (Section 5) are cached here

probe_meta = pd.read_csv(DATA / "meta/probe_clips.csv")   # 48 clips = 8 speakers x 6 sentences
demo_meta = pd.read_csv(DATA / "meta/demo_clips.csv")     # 20 clips, 10 speakers
print(f"probe clips: {len(probe_meta)}  speakers: {probe_meta.speaker.nunique()}  sentences: {probe_meta.sentence_id.nunique()}")
print(f"demo clips:  {len(demo_meta)}   total audio for the mini model: {probe_meta.duration.sum() + demo_meta.duration.sum():.0f} s")
print("torch", torch.__version__, "| threads", torch.get_num_threads())

probe clips: 48  speakers: 8  sentences: 6
demo clips:  20   total audio for the mini model: 272 s
torch 2.14.1+cu130 | threads 2

class ConvFrontEnd(nn.Module):
    '''The wav2vec 2.0 / HuBERT / WavLM feature encoder (random weights here; only the geometry matters).'''
    def __init__(self, kernels=(10, 3, 3, 3, 3, 2, 2), strides=(5, 2, 2, 2, 2, 2, 2), ch=512):
        super().__init__()
        layers, cin = [], 1
        for k, s in zip(kernels, strides):
            layers += [nn.Conv1d(cin, ch, k, stride=s, bias=False), nn.GELU()]   # real models add GroupNorm / LayerNorm
            cin = ch
        self.net = nn.Sequential(*layers)
        self.kernels, self.strides = kernels, strides

    def forward(self, wav):              # wav: (batch, samples)
        return self.net(wav[:, None, :]).transpose(1, 2)   # -> (batch, frames, 512)

    def geometry(self):
        stride, rf = 1, 1
        for k, s in zip(self.kernels, self.strides):
            rf += (k - 1) * stride       # each new layer widens the window by (k-1) input-strides
            stride *= s
        return stride, rf

fe = ConvFrontEnd()
stride, rf = fe.geometry()
print(f"total stride = {stride} samples = {1000*stride/SR:.0f} ms   ->  {SR/stride:.0f} frames/s")
print(f"receptive field = {rf} samples = {1000*rf/SR:.0f} ms")

x = load(DATA / "audio/khmer_studio.wav")[: 3 * SR]         # a 3 s clip
with torch.no_grad():
    feats = fe(torch.from_numpy(x)[None])
n_stft = 1 + len(x) // HOP                                   # librosa centre-pads: frames = 1 + samples/hop
print(f"3 s clip: CNN front-end gives {feats.shape[1]} frames x {feats.shape[2]} dims; STFT/log-mel gives {n_stft} frames (hop 10 ms)")
print(f"parameters in the front-end: {sum(p.numel() for p in fe.parameters())/1e6:.2f} M")

total stride = 320 samples = 20 ms   ->  50 frames/s
receptive field = 400 samples = 25 ms

3 s clip: CNN front-end gives 149 frames x 512 dims; STFT/log-mel gives 301 frames (hop 10 ms)
parameters in the front-end: 4.20 M

def load_clip(utt_id, folder):
    return load(DATA / f"audio/{folder}/{utt_id}.wav")

t0 = time.time()
probe_wavs = [load_clip(u, "probe") for u in probe_meta.utt_id]
demo_wavs = [load_clip(u, "demo") for u in demo_meta.utt_id]
all_wavs = probe_wavs + demo_wavs
mels = [logmel(w).T.astype(np.float32) for w in all_wavs]        # each (T, 80), dB, loudest bin = 0
n_frames = sum(m.shape[0] for m in mels)
print(f"{len(mels)} clips -> {n_frames} log-mel frames in {time.time()-t0:.1f} s  (e.g. clip 0: {mels[0].shape})")

# standardise per mel band so that k-means / the Transformer do not just see loudness
mel_mean = np.concatenate(mels).mean(0); mel_std = np.concatenate(mels).std(0) + 1e-5
mels_z = [(m - mel_mean) / mel_std for m in mels]

68 clips -> 27270 log-mel frames in 0.5 s  (e.g. clip 0: (317, 80))

from sklearn.cluster import MiniBatchKMeans

K = 32
t0 = time.time()
km = MiniBatchKMeans(n_clusters=K, random_state=0, batch_size=2048, n_init=3).fit(np.concatenate(mels_z))
labels = [km.predict(m).astype(np.int64) for m in mels_z]          # pseudo-label per frame
print(f"k-means: {K} clusters fitted on {n_frames} frames in {time.time()-t0:.1f} s")

# which clusters are silence / background noise? -> those whose frames are quiet on the real dB scale.
# Phone recordings have a noisy floor (around -65 dB) that is NOT the same cluster as clean studio silence (-78 dB),
# so "silence" is a small SET of quiet clusters, not one id.
all_ids = np.concatenate(labels); counts = np.bincount(all_ids, minlength=K)
cluster_db = np.array([np.concatenate(mels)[all_ids == k].mean() for k in range(K)])   # mean dB of each cluster's frames
SIL = set(int(k) for k in np.where(cluster_db < -62)[0])                              # quiet clusters = silence / noise floor
is_sil = lambda ids: np.isin(ids, list(SIL))
print(f"quiet (silence/noise) clusters: {sorted(SIL)}  mean dB {np.round(cluster_db[sorted(SIL)], 0)}  "
      f"-> {100*counts[list(SIL)].sum()/n_frames:.0f} % of all frames")
print("loudest clusters:", np.argsort(-cluster_db)[:5], "mean dB", np.round(np.sort(cluster_db)[::-1][:5], 0))
print("most common clusters:", np.argsort(-counts)[:5], "with shares", np.round(np.sort(counts)[::-1][:5] / n_frames, 2))

k-means: 32 clusters fitted on 27270 frames in 2.1 s
quiet (silence/noise) clusters: [1, 4, 8, 9, 18, 22, 30, 31]  mean dB [-62. -63. -70. -78. -74. -64. -64. -69.]  -> 29 % of all frames
loudest clusters: [ 6 13 19  2 21] mean dB [-32. -38. -38. -42. -43.]
most common clusters: [30 10 27 31 22] with shares [0.05 0.05 0.05 0.05 0.05]

# plot: spectrogram of one probe clip with its cluster-id sequence underneath
i = 0
x, ids = probe_wavs[i], labels[i]
text = probe_meta.text[i]; print(probe_meta.utt_id[i], probe_meta.speaker[i], "|", text)

fig, ax = plt.subplots(2, 1, figsize=(12, 5), sharex=True, gridspec_kw={"height_ratios": [3, 1]})
show_spec(x, ax=ax[0], mel=True, title=f"log-mel of {probe_meta.utt_id[i]}  ({probe_meta.sentence_id[i]})")
t = np.arange(len(ids)) * HOP / SR
ax[1].imshow(ids[None, :], aspect="auto", cmap="tab20", extent=[0, t[-1] + HOP / SR, 0, 1], interpolation="nearest")
ax[1].set(yticks=[], xlabel="time (s)", title=f"k-means cluster id per 10 ms frame (k={K}; white = quiet clusters {sorted(SIL)})")
# mark the quiet clusters in white so you can line them up with the dark gaps above
ax[1].imshow(np.where(is_sil(ids), 1.0, np.nan)[None, :], aspect="auto", cmap="gray", vmin=0, vmax=1,
             extent=[0, t[-1] + HOP / SR, 0, 1], interpolation="nearest")
plt.tight_layout(); plt.savefig(OUT / "kmeans_clusters.png", dpi=110); plt.show()
print("cluster ids, first 40 frames:", ids[:40])

iany_046e9e11_cc152fc8 iany-s-046e9e11 | ពួកយើងលេងបាល់ទាត់នៅទីលានក្រោយម៉ោងរៀន

cluster ids, first 40 frames: [31 22 31 30 22 22 30 30 30 30  4 15 15 15 22  4 22 22 22 22 30 22 22 22
 22 22 22 22 30 22 30 22 15 22 22 22 22 22 22 22]

play(x)

class MiniHuBERT(nn.Module):
    '''Masked prediction over log-mel frames. Returns per-layer hidden states for probing.'''
    def __init__(self, n_in=80, d=128, n_layers=2, n_heads=4, n_classes=K, max_len=4096):
        super().__init__()
        self.proj = nn.Linear(n_in, d)                       # "layer 0" = projected input frames
        self.mask_emb = nn.Parameter(torch.randn(d) * 0.1)   # the [MASK] vector that replaces hidden frames
        self.pos = nn.Embedding(max_len, d)                  # learned absolute positions (real models use a conv pos-encoder)
        layer = nn.TransformerEncoderLayer(d, n_heads, dim_feedforward=4 * d, dropout=0.1, batch_first=True, norm_first=True)
        self.layers = nn.ModuleList([nn.TransformerEncoderLayer(d, n_heads, 4 * d, 0.1, batch_first=True, norm_first=True)
                                     for _ in range(n_layers)])
        self.norm = nn.LayerNorm(d)
        self.head = nn.Linear(d, n_classes)                  # predicts the cluster id

    def forward(self, x, mask=None):                         # x: (B, T, 80); mask: (B, T) bool, True = hidden
        h = self.proj(x)
        if mask is not None:
            h = torch.where(mask[..., None], self.mask_emb.expand_as(h), h)
        h = h + self.pos(torch.arange(x.shape[1], device=x.device))[None]
        hidden = [h]
        for layer in self.layers:
            h = layer(h); hidden.append(h)
        return self.head(self.norm(h)), hidden               # logits (B, T, K), list of (B, T, d) per layer


def span_mask(B, T, p=0.30, span=8, rng=np.random):
    '''Boolean (B, T): choose span starts so that about p of all frames end up masked.'''
    m = np.zeros((B, T), bool)
    n_starts = int(p * T / span * 1.15)                      # 15 % extra because spans overlap
    for b in range(B):
        for s in rng.choice(T - span, n_starts, replace=False):
            m[b, s:s + span] = True
    return torch.from_numpy(m)


CHUNK = 200                                                   # 2 s windows
def make_chunks(feats, labs, chunk=CHUNK):
    X, Y = [], []
    for f, l in zip(feats, labs):
        for s in range(0, len(f) - chunk + 1, chunk):        # non-overlapping 2 s windows
            X.append(f[s:s + chunk]); Y.append(l[s:s + chunk])
    return torch.from_numpy(np.stack(X)), torch.from_numpy(np.stack(Y))

X_train, Y_train = make_chunks(mels_z, labels)
print(f"training windows: {X_train.shape}  (windows x frames x mels)")
share_top = counts.max() / n_frames
print(f"chance loss ln({K}) = {np.log(K):.2f} nats; 'always the most common cluster' baseline = {-np.log(share_top):.2f} nats")
print("(the clusters are nearly equally common, so that baseline is only a little below chance)")

model = MiniHuBERT()
print(f"mini model parameters: {sum(p.numel() for p in model.parameters())/1e6:.2f} M")

training windows: torch.Size([100, 200, 80])  (windows x frames x mels)
chance loss ln(32) = 3.47 nats; 'always the most common cluster' baseline = 2.95 nats
(the clusters are nearly equally common, so that baseline is only a little below chance)
mini model parameters: 0.94 M

EPOCHS, BATCH = 16, 16
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=2e-3, total_steps=EPOCHS * ((len(X_train) + BATCH - 1) // BATCH))
rng = np.random.RandomState(0)
hist = []
t0 = time.time()
for ep in range(EPOCHS):
    model.train(); tot, n = 0.0, 0
    for idx in np.array_split(rng.permutation(len(X_train)), (len(X_train) + BATCH - 1) // BATCH):
        xb, yb = X_train[idx], Y_train[idx]
        mask = span_mask(len(idx), CHUNK, rng=rng)
        logits, _ = model(xb, mask)
        loss = F.cross_entropy(logits[mask], yb[mask])        # loss ONLY at masked frames (the HuBERT rule)
        opt.zero_grad(); loss.backward(); opt.step(); sched.step()
        tot += loss.item() * mask.sum().item(); n += mask.sum().item()
    hist.append(tot / n)
    if ep % 2 == 0 or ep == EPOCHS - 1:
        print(f"epoch {ep:2d}  masked-frame loss {hist[-1]:.3f} nats   ({time.time()-t0:.0f} s)")
print(f"masked fraction in the last batch: {mask.float().mean():.2f}")
torch.save(model.state_dict(), OUT / "mini_hubert.pt")

# accuracy at masked frames vs the 'copy the most common id' baseline
model.eval()
with torch.no_grad():
    mask = span_mask(len(X_train), CHUNK, rng=np.random.RandomState(1))
    logits, _ = model(X_train, mask)
    acc_masked = (logits.argmax(-1)[mask] == Y_train[mask]).float().mean().item()
print(f"accuracy at masked frames: {100*acc_masked:.1f} %   (most-common-cluster baseline {100*share_top:.1f} %, chance {100/K:.1f} %)")

plt.figure(figsize=(12, 3)); plt.plot(hist, marker="o")
plt.axhline(np.log(K), ls=":", color="gray"); plt.text(0, np.log(K) + 0.05, "chance ln 32", color="gray")
plt.axhline(-np.log(share_top), ls="--", color="gray"); plt.text(0, -np.log(share_top) + 0.05, "most-common-cluster baseline", color="gray")
plt.title("mini-HuBERT: cross-entropy at masked frames"); plt.xlabel("epoch"); plt.ylabel("loss (nats)"); plt.grid(alpha=0.3)
plt.savefig(OUT / "mini_hubert_loss.png", dpi=110); plt.show()

epoch  0  masked-frame loss 3.621 nats   (2 s)

epoch  2  masked-frame loss 3.389 nats   (5 s)

epoch  4  masked-frame loss 3.133 nats   (8 s)

epoch  6  masked-frame loss 2.996 nats   (12 s)

epoch  8  masked-frame loss 2.887 nats   (16 s)

epoch 10  masked-frame loss 2.752 nats   (20 s)

epoch 12  masked-frame loss 2.734 nats   (26 s)

epoch 14  masked-frame loss 2.678 nats   (29 s)

epoch 15  masked-frame loss 2.679 nats   (31 s)
masked fraction in the last batch: 0.29

accuracy at masked frames: 22.7 %   (most-common-cluster baseline 5.2 %, chance 3.1 %)

@torch.no_grad()
def mini_layer_features(feats):
    '''Per-layer frame features (no masking) for a list of (T, 80) standardised log-mel clips.'''
    model.eval()
    layers = [[] for _ in range(len(model.layers) + 1)]
    for f in feats:
        _, hidden = model(torch.from_numpy(f)[None])
        for L, h in enumerate(hidden):
            layers[L].append(h[0].numpy())
    return layers

mini_layers = mini_layer_features(mels_z[:48])                 # probe clips only
layer_frames = [mels_z[:48]] + mini_layers                     # [log-mel, layer0, layer1, layer2]
layer_names = ["log-mel", "L0 (proj. input)", "L1", "L2"]

spk, sent = probe_meta.speaker.values, probe_meta.sentence_id.values
spk_m, spk_s = probe_layers(layer_frames, labels=spk, groups=sent)     # speaker probe, sentences held out
con_m, con_s = probe_layers(layer_frames, labels=sent, groups=spk)     # content probe, speakers held out

mini_results = pd.DataFrame({"layer": layer_names, "speaker_acc": spk_m.round(3), "speaker_std": spk_s.round(3),
                             "sentence_acc": con_m.round(3), "sentence_std": con_s.round(3)})
mini_results.to_csv(OUT / "mini_probe_results.csv", index=False)
print(mini_results.to_string(index=False))
print(f"chance: speaker {1/8:.3f}, sentence {1/6:.3f}")

           layer  speaker_acc  speaker_std  sentence_acc  sentence_std
         log-mel        0.979        0.047         0.271         0.116
L0 (proj. input)        0.979        0.047         0.208         0.072
              L1        0.979        0.047         0.188         0.055
              L2        0.979        0.047         0.229         0.116
chance: speaker 0.125, sentence 0.167

fig, ax = plt.subplots(1, 2, figsize=(12, 3.5))
w = 0.38; x = np.arange(len(layer_names))
ax[0].bar(x - w/2, spk_m, w, yerr=spk_s, capsize=3, label="speaker (8 classes)")
ax[0].bar(x + w/2, con_m, w, yerr=con_s, capsize=3, label="sentence (6 classes)")
ax[0].axhline(1/8, ls=":", color="C0"); ax[0].axhline(1/6, ls=":", color="C1")
ax[0].set(xticks=x, xticklabels=layer_names, ylim=(0, 1.05), ylabel="held-out accuracy", title="mini model: linear-probe accuracy per layer")
ax[0].legend(loc="upper right"); ax[0].grid(alpha=0.3, axis="y")

S = layer_similarity([pool_features(f) for f in layer_frames])
im = ax[1].imshow(S, vmin=0, vmax=1, cmap="magma")
ax[1].set(xticks=x, yticks=x, xticklabels=layer_names, yticklabels=layer_names, title="linear CKA between layers (pooled)")
for i in range(len(x)):
    for j in range(len(x)):
        ax[1].text(j, i, f"{S[i,j]:.2f}", ha="center", va="center", color="w" if S[i,j] < 0.6 else "k", fontsize=8)
plt.colorbar(im, ax=ax[1]); plt.tight_layout(); plt.savefig(OUT / "mini_probe_layers.png", dpi=110); plt.show()

# needs internet (first run only) -- microsoft/wavlm-base-plus, ~380 MB, cached under ~/.cache/huggingface
SSL_MODELS = {}                       # name -> loaded model (only the ones that could be downloaded)
def try_load_ssl(hf_name, short):
    try:
        from transformers import AutoModel, AutoFeatureExtractor
        fx = AutoFeatureExtractor.from_pretrained(hf_name, cache_dir=CACHE / "huggingface")
        m = AutoModel.from_pretrained(hf_name, cache_dir=CACHE / "huggingface").eval()
        SSL_MODELS[short] = (fx, m)
        print(f"loaded {short} ({hf_name}): {sum(p.numel() for p in m.parameters())/1e6:.0f} M params")
    except Exception as e:
        print(f"[offline] could not load {hf_name}: {type(e).__name__}: {str(e)[:120]}")
        print(f"          -> skipping {short}; run this cell on a machine with internet to fill in the real curves.")

try_load_ssl("microsoft/wavlm-base-plus", "WavLM-base+")

[offline] could not load microsoft/wavlm-base-plus: OSError: Can't load feature extractor for 'microsoft/wavlm-base-plus'. If you were trying to load it from 'https://huggingface.co
          -> skipping WavLM-base+; run this cell on a machine with internet to fill in the real curves.

@torch.no_grad()
def ssl_layer_features(short, wavs, max_clips=None):
    '''Run one SSL model on a list of waveforms. Returns list over layers (13) of lists over clips of (T, 768).'''
    fx, m = SSL_MODELS[short]
    layers = None
    for w in (wavs if max_clips is None else wavs[:max_clips]):
        inp = fx(w, sampling_rate=SR, return_tensors="pt")
        hs = m(**inp, output_hidden_states=True).hidden_states          # tuple of 13 x (1, T, 768)
        if layers is None:
            layers = [[] for _ in hs]
        for L, h in enumerate(hs):
            layers[L].append(h[0].numpy())
    return layers

ssl_probe = {}                    # short -> DataFrame of per-layer accuracies
if "WavLM-base+" in SSL_MODELS:
    t0 = time.time()
    wavlm_layers = ssl_layer_features("WavLM-base+", probe_wavs)
    print(f"WavLM features for 48 clips in {time.time()-t0:.0f} s; layer 6 of clip 0: {wavlm_layers[6][0].shape}")
    m_spk, s_spk = probe_layers(wavlm_layers, labels=spk, groups=sent)
    m_con, s_con = probe_layers(wavlm_layers, labels=sent, groups=spk)
    ssl_probe["WavLM-base+"] = pd.DataFrame({"layer": range(len(wavlm_layers)), "speaker_acc": m_spk, "speaker_std": s_spk,
                                             "sentence_acc": m_con, "sentence_std": s_con})
    ssl_probe["WavLM-base+"].to_csv(OUT / "wavlm_probe_results.csv", index=False)
    print(ssl_probe["WavLM-base+"].round(3).to_string(index=False))

    fig, ax = plt.subplots(1, 2, figsize=(12, 3.8), gridspec_kw={"width_ratios": [2, 1]})
    plot_layer_curves({"speaker (8)": (m_spk, s_spk), "sentence (6)": (m_con, s_con)},
                      chance={"spk": 1/8, "sent": 1/6}, title="WavLM-base+: linear-probe accuracy per layer", ax=ax[0])
    S = layer_similarity([pool_features(f) for f in wavlm_layers])
    im = ax[1].imshow(S, vmin=0, vmax=1, cmap="magma"); ax[1].set(title="linear CKA between layers", xlabel="layer", ylabel="layer")
    plt.colorbar(im, ax=ax[1]); plt.tight_layout(); plt.savefig(OUT / "wavlm_probe_layers.png", dpi=110); plt.show()
else:
    ref = pd.read_csv(DATA / "meta/ssl_probe_reference.csv")
    print("WavLM not available -> showing the shipped *qualitative* reference (no numbers; re-run on your laptop):")
    print(ref[ref.model == "microsoft/wavlm-base-plus"][["layer", "speaker_trend_expected", "content_trend_expected"]].to_string(index=False))
    print("\nnote:", ref.note.iloc[0])

WavLM not available -> showing the shipped *qualitative* reference (no numbers; re-run on your laptop):
 layer speaker_trend_expected content_trend_expected
     0                 rising                 rising
     1                 rising                 rising
     2       high (near peak)                 rising
     3       high (near peak)                 rising
     4       high (near peak)                 rising
     5       high (near peak)                 rising
     6       high (near peak)                 rising
     7                falling                 rising
     8                falling       high (near peak)
     9                falling       high (near peak)
    10                falling       high (near peak)
    11                falling       high (near peak)
    12                falling                falling

note: qualitative only; SUPERB layer-weight analysis (Chen et al. 2022) puts speaker tasks early, phonetic/ASR tasks late. Re-run Section 5 on your laptop to fill the numbers.

# needs internet (first run only) -- facebook/hubert-base-ls960 (~380 MB) and utter-project/mHuBERT-147 (~380 MB)
try_load_ssl("facebook/hubert-base-ls960", "HuBERT-base")
try_load_ssl("utter-project/mHuBERT-147", "mHuBERT-147")

for short in ["HuBERT-base", "mHuBERT-147"]:
    if short in SSL_MODELS:
        lay = ssl_layer_features(short, probe_wavs)
        m_spk, s_spk = probe_layers(lay, labels=spk, groups=sent)
        m_con, s_con = probe_layers(lay, labels=sent, groups=spk)
        ssl_probe[short] = pd.DataFrame({"layer": range(len(lay)), "speaker_acc": m_spk, "speaker_std": s_spk,
                                         "sentence_acc": m_con, "sentence_std": s_con})
        ssl_probe[short].to_csv(OUT / f"{short.lower().replace('-', '_')}_probe_results.csv", index=False)

if ssl_probe:
    fig, ax = plt.subplots(1, 2, figsize=(12, 3.8), sharey=True)
    for short, df in ssl_probe.items():
        ax[0].errorbar(df.layer, df.speaker_acc, yerr=df.speaker_std, marker="o", capsize=2, label=short)
        ax[1].errorbar(df.layer, df.sentence_acc, yerr=df.sentence_std, marker="o", capsize=2, label=short)
    ax[0].axhline(1/8, ls=":", color="gray"); ax[1].axhline(1/6, ls=":", color="gray")
    ax[0].set(title="speaker probe (8 classes) per layer", xlabel="layer", ylabel="held-out accuracy", ylim=(0, 1.05))
    ax[1].set(title="sentence probe (6 classes) per layer", xlabel="layer")
    for a in ax: a.grid(alpha=0.3); a.legend()
    plt.tight_layout(); plt.savefig(OUT / "ssl_probe_overlay.png", dpi=110); plt.show()
else:
    print("No pretrained encoder available in this session. The plotting code above is what to run on your laptop;")
    print("expected shapes are described in the reading guide, and data/meta/ssl_probe_reference.csv lists the qualitative trends.")

[offline] could not load facebook/hubert-base-ls960: OSError: Can't load feature extractor for 'facebook/hubert-base-ls960'. If you were trying to load it from 'https://huggingface.c
          -> skipping HuBERT-base; run this cell on a machine with internet to fill in the real curves.

[offline] could not load utter-project/mHuBERT-147: OSError: Can't load feature extractor for 'utter-project/mHuBERT-147'. If you were trying to load it from 'https://huggingface.co
          -> skipping mHuBERT-147; run this cell on a machine with internet to fill in the real curves.
No pretrained encoder available in this session. The plotting code above is what to run on your laptop;
expected shapes are described in the reading guide, and data/meta/ssl_probe_reference.csv lists the qualitative trends.

from sklearn.decomposition import PCA

# pick 4 speakers (2 f, 2 m) x 2 sentences so both colourings have something to show
sel = probe_meta[probe_meta.sentence_id.isin(["daily-04", "news-03"])]
sel = sel[sel.speaker.isin(sorted(sel.speaker.unique())[:4])]
idx = sel.index.values
if "WavLM-base+" in SSL_MODELS:
    frames = [wavlm_layers[6][i] for i in idx]; src_name = "WavLM-base+ layer 6"
    # WavLM runs at 50 fps, our k-means ids at 100 fps -> take every second id
    sound_ids = [labels[i][1::2][: len(fr)] for i, fr in zip(idx, frames)]
else:
    frames = [mini_layers[-1][i] for i in idx]; src_name = "mini model layer 2 (stand-in; WavLM not available)"
    sound_ids = [labels[i][: len(fr)] for i, fr in zip(idx, frames)]

Xf = np.concatenate(frames); Xf = (Xf - Xf.mean(0)) / (Xf.std(0) + 1e-5)
spk_of_frame = np.concatenate([[s] * len(fr) for s, fr in zip(sel.speaker, frames)])
snd_of_frame = np.concatenate(sound_ids)
keep = ~is_sil(snd_of_frame)                                  # drop silence/noise frames: they would form one big blob
P = PCA(2, random_state=0).fit_transform(Xf[keep])
spk_codes = pd.factorize(spk_of_frame[keep])[0]

fig, ax = plt.subplots(1, 2, figsize=(12, 4.5))
ax[0].scatter(P[:, 0], P[:, 1], c=spk_codes, cmap="tab10", s=4, alpha=0.6)
ax[0].set(title=f"{src_name}: frames coloured by SPEAKER (4 speakers)", xlabel="PC1", ylabel="PC2")
ax[1].scatter(P[:, 0], P[:, 1], c=snd_of_frame[keep], cmap="tab20", s=4, alpha=0.6)
ax[1].set(title="same frames coloured by k-means SOUND cluster", xlabel="PC1")
plt.tight_layout(); plt.savefig(OUT / "layer_pca_speaker_vs_sound.png", dpi=110); plt.show()

# a number instead of eyeballing: how well does a nearest-neighbour vote predict each label from the 2-D map?
from sklearn.neighbors import KNeighborsClassifier
from sklearn.model_selection import cross_val_score
for name, y in [("speaker", spk_codes), ("sound cluster", snd_of_frame[keep])]:
    acc = cross_val_score(KNeighborsClassifier(5), Xf[keep], y, cv=5).mean()
    print(f"5-NN accuracy in the full feature space, predicting {name:13s}: {100*acc:.0f} %  (chance {100/len(np.unique(y)):.0f} %)")
print(f"({keep.sum()} non-silent frames from {len(idx)} clips)")

5-NN accuracy in the full feature space, predicting speaker      : 100 %  (chance 25 %)
5-NN accuracy in the full feature space, predicting sound cluster: 44 %  (chance 4 %)
(2631 non-silent frames from 8 clips)

from sklearn.neighbors import NearestNeighbors

def pool_frames(folder):
    fr = [(logmel(load(p)).T - mel_mean) / mel_std for p in sorted((DATA / f"audio/{folder}").glob("*.wav"))]
    return np.concatenate(fr).astype(np.float32), len(fr)

pool_kh, n_kh = pool_frames("pool_khmer"); ids_kh = km.predict(pool_kh)
pool_en, n_en = pool_frames("pool_english"); ids_en = km.predict(pool_en)
print(f"Khmer pool: {n_kh} clips, {len(pool_kh)} frames ({len(pool_kh)/100:.0f} s), {100*is_sil(ids_kh).mean():.0f} % quiet")
print(f"English pool: {n_en} clips, {len(pool_en)} frames ({len(pool_en)/100:.0f} s), {100*is_sil(ids_en).mean():.0f} % quiet")

# drop quiet frames from the pools and from the probe clips
pool_kh_v, ids_kh_v = pool_kh[~is_sil(ids_kh)], ids_kh[~is_sil(ids_kh)]
pool_en_v, ids_en_v = pool_en[~is_sil(ids_en)], ids_en[~is_sil(ids_en)]
nn_kh = NearestNeighbors(n_neighbors=4).fit(pool_kh_v)
nn_en = NearestNeighbors(n_neighbors=4).fit(pool_en_v)

d_kh, d_en, same_kh, same_en, per_clip = [], [], [], [], []
for i in range(48):
    v = ~is_sil(labels[i]); f = mels_z[i][v]; l = labels[i][v]
    dk, jk = nn_kh.kneighbors(f); de, je = nn_en.kneighbors(f)
    d_kh.append(dk[:, 0]); d_en.append(de[:, 0])
    same_kh.append((ids_kh_v[jk] == l[:, None]).mean(1))      # do the 4 neighbours have the SAME sound cluster?
    same_en.append((ids_en_v[je] == l[:, None]).mean(1))
    per_clip.append((probe_meta.utt_id[i], probe_meta.gender[i], dk[:, 0].mean(), de[:, 0].mean(), same_kh[-1].mean(), same_en[-1].mean()))
d_kh, d_en, same_kh, same_en = map(np.concatenate, (d_kh, d_en, same_kh, same_en))
far = np.percentile(d_kh, 95)                                 # "far" = farther than 95 % of Khmer-pool matches
cov = pd.DataFrame(per_clip, columns=["utt_id", "gender", "mean_d_khmer", "mean_d_english", "same_sound_khmer", "same_sound_english"])
cov.to_csv(OUT / "pool_coverage.csv", index=False)
print(f"\nmedian NN distance:        Khmer pool {np.median(d_kh):.2f}   English pool {np.median(d_en):.2f}")
print(f"'far' frames (d > {far:.2f}):   Khmer pool {100*(d_kh>far).mean():.1f} %   English pool {100*(d_en>far).mean():.1f} %")
print(f"neighbours keep the sound: Khmer pool {100*same_kh.mean():.1f} %   English pool {100*same_en.mean():.1f} %  (4-NN same-cluster rate)")
print(cov.groupby("gender")[["mean_d_khmer", "mean_d_english", "same_sound_khmer", "same_sound_english"]].mean().round(3))

Khmer pool: 22 clips, 18208 frames (182 s), 26 % quiet
English pool: 22 clips, 19869 frames (199 s), 33 % quiet

median NN distance:        Khmer pool 4.28   English pool 4.31
'far' frames (d > 6.24):   Khmer pool 5.0 %   English pool 4.1 %
neighbours keep the sound: Khmer pool 64.0 %   English pool 63.1 %  (4-NN same-cluster rate)
        mean_d_khmer  mean_d_english  same_sound_khmer  same_sound_english
gender                                                                    
female         4.361           4.415             0.617               0.612
male           4.524           4.544             0.668               0.655

# sound-cluster coverage: share of each (non-quiet) cluster in the probe clips vs in each pool
probe_ids = np.concatenate(labels[:48]); probe_ids = probe_ids[~is_sil(probe_ids)]
share = lambda ids: np.bincount(ids, minlength=K) / len(ids)
occ = pd.DataFrame({"probe": share(probe_ids), "khmer_pool": share(ids_kh_v), "english_pool": share(ids_en_v), "mean_dB": cluster_db.round(0)})
occ = occ[~occ.index.isin(SIL)].sort_values("probe", ascending=False)
occ["english/probe"] = (occ.english_pool / occ.probe).round(2); occ["khmer/probe"] = (occ.khmer_pool / occ.probe).round(2)
occ.to_csv(OUT / "pool_cluster_coverage.csv")
under = lambda col: occ[occ[col] < 0.25]                     # pool has < 1/4 of the share the probe needs
print(f"sounds the ENGLISH pool under-covers: clusters {list(under('english/probe').index)} = {100*under('english/probe').probe.sum():.0f} % of probe frames")
print(f"sounds the KHMER pool under-covers:   clusters {list(under('khmer/probe').index)} = {100*under('khmer/probe').probe.sum():.0f} % of probe frames")
print(occ.head(8).round(3).to_string())

fig, ax = plt.subplots(1, 2, figsize=(12, 3.8), gridspec_kw={"width_ratios": [1, 1.6]})
bins = np.linspace(0, max(d_en.max(), d_kh.max()), 60)
ax[0].hist(d_kh, bins, alpha=0.6, label="Khmer pool (2 speakers)"); ax[0].hist(d_en, bins, alpha=0.6, label="English pool (1 speaker)")
ax[0].axvline(far, ls="--", color="k"); ax[0].set(title="nearest pool frame: distance", xlabel="distance (standardised log-mel units)", ylabel="probe frames"); ax[0].legend(fontsize=8)
x = np.arange(len(occ)); w = 0.28
ax[1].bar(x - w, occ.probe, w, label="Khmer probe clips (what is needed)")
ax[1].bar(x, occ.khmer_pool, w, label="Khmer pool"); ax[1].bar(x + w, occ.english_pool, w, label="English pool")
ax[1].set(xticks=x, xticklabels=occ.index, xlabel="sound cluster (sorted by how often the probe clips use it)", ylabel="share of non-quiet frames",
          title="sound-cluster coverage: needed vs available in each pool")
ax[1].legend(fontsize=8); plt.tight_layout(); plt.savefig(OUT / "pool_coverage.png", dpi=110); plt.show()

sounds the ENGLISH pool under-covers: clusters [10, 3, 23, 5, 24] = 21 % of probe frames
sounds the KHMER pool under-covers:   clusters [6, 3, 5, 26, 21, 0] = 18 % of probe frames
    probe  khmer_pool  english_pool  mean_dB  english/probe  khmer/probe
27  0.082       0.039         0.027    -49.0           0.33         0.48
2   0.067       0.096         0.034    -42.0           0.51         1.43
15  0.067       0.167         0.025    -59.0           0.38         2.50
10  0.065       0.054         0.008    -52.0           0.13         0.82
13  0.060       0.027         0.105    -38.0           1.75         0.46
20  0.055       0.087         0.046    -53.0           0.83         1.57
28  0.047       0.023         0.017    -57.0           0.37         0.49
11  0.044       0.014         0.040    -55.0           0.92         0.32

from src.probe import pool_features, probe_accuracy, probe_layers, layer_similarity, plot_layer_curves
m, s = probe_layers(layer_frames, labels=speaker_ids, groups=sentence_ids)   # list over layers of list over clips of (T, D)
plot_layer_curves({"speaker": (m, s)}, chance={"spk": 1/8})

spk_m2, spk_s2 = probe_layers(layer_frames, labels=spk, groups=sent, stats=True)
con_m2, con_s2 = probe_layers(layer_frames, labels=sent, groups=spk, stats=True)
cmp = pd.DataFrame({"layer": layer_names, "speaker_mean_pool": spk_m.round(3), "speaker_mean+std_pool": spk_m2.round(3),
                    "sentence_mean_pool": con_m.round(3), "sentence_mean+std_pool": con_m2.round(3)})
print(cmp.to_string(index=False))
cmp.to_csv(OUT / "mini_probe_pooling_comparison.csv", index=False)

plot_layer_curves({"speaker, mean pool": (spk_m, spk_s), "speaker, mean+std": (spk_m2, spk_s2),
                   "sentence, mean pool": (con_m, con_s), "sentence, mean+std": (con_m2, con_s2)},
                  chance={"spk": 1/8, "sent": 1/6}, title="src.probe on the mini model: pooling choice", layer_names=layer_names)
plt.tight_layout(); plt.savefig(OUT / "mini_probe_tool.png", dpi=110); plt.show()
print("saved outputs:", sorted(p.name for p in OUT.iterdir()))

           layer  speaker_mean_pool  speaker_mean+std_pool  sentence_mean_pool  sentence_mean+std_pool
         log-mel              0.979                  0.979               0.271                   0.396
L0 (proj. input)              0.979                  0.979               0.208                   0.333
              L1              0.979                  0.979               0.188                   0.333
              L2              0.979                  0.958               0.229                   0.271

saved outputs: ['kmeans_clusters.png', 'layer_pca_speaker_vs_sound.png', 'mini_hubert.pt', 'mini_hubert_loss.png', 'mini_probe_layers.png', 'mini_probe_pooling_comparison.csv', 'mini_probe_results.csv', 'mini_probe_tool.png', 'pool_cluster_coverage.csv', 'pool_coverage.csv', 'pool_coverage.png']

You see	What it is	Do you need it?
Purple heading	a new section and its explanation	yes: read the bold lines and the "Why it matters" line
Visualization Version	pictures (one is animated) that explain the section	yes: the first image of each section is the fastest way in
Additional reading	what the next code cell does and what its output means, plus extra notes	optional: read just the In short lines if you are busy
code cell + output	the real computation on 68 clips and two target pools of 22 clips each	run it, look at the result

model	pretext target	what is certain	what to check in the paper
wav2vec 2.0 (Baevski et al., 2020)	contrastive: pick the true quantized latent among distractors; the quantizer is learned jointly (Gumbel softmax, 2 codebooks x 320 codes)	CNN stride 320, masked spans, Transformer	the exact masking probability (p=0.065, span 10) and the diversity loss weight
HuBERT (Hsu et al., 2021)	cross-entropy on k-means cluster ids of MFCC (iteration 1) then of its own layer features (iteration 2+)	targets are offline cluster ids, the loss is only on masked frames	k = 100 then 500; which layer is clustered in iteration 2 (layer 6 of the base model)
WavLM (Chen et al., 2022)	HuBERT loss, but the input is mixed with noise or another speaker and the model must predict the clean speaker's labels	gated relative position bias in attention; better speaker tasks	the mixing ratio and the 94k-hour training set details
mHuBERT-147 (Boito et al., 2024)	HuBERT loss on 147 languages, faiss-based clustering	multilingual, small (95M)	whether Khmer is in the 147 languages (we believe yes, via VoxLingua107/other sources; verify in the language table)

real HuBERT	our mini version
CNN front-end on raw wave, 50 frames/s	log-mel frames (80-d, 100 frames/s) from Day 1
k-means on MFCC, k = 100 (iteration 1)	k-means on log-mel, k = 32
12-layer Transformer, d = 768, 960 h LibriSpeech	2-layer Transformer, d = 128, 100 windows of 2 s = 200 s of Khmer (from 272 s of clips)

model	HF name	download
WavLM-base-plus	`microsoft/wavlm-base-plus`	~380 MB
HuBERT-base	`facebook/hubert-base-ls960`	~380 MB
mHuBERT-147	`utter-project/mHuBERT-147`	~380 MB
wav2vec 2.0 base	`facebook/wav2vec2-base`	~380 MB

Day 5 · Self-supervised speech encoders (wav2vec 2.0, HuBERT, WavLM): what each layer knows about who vs what¶

How to read this notebook (for the team)¶

1. The big idea and the three pretext tasks¶

Visualization Version¶

Same skeleton for all three models. They differ in what must be guessed.¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

2. The CNN front-end by hand¶

Visualization Version¶

How the receptive field and the stride grow layer by layer, and what a 3 s clip becomes.¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

3. Masked prediction from scratch (a miniature HuBERT)¶

Visualization Version¶

The real k-means labels of one clip, a sketch of the masking, and the real training loss.¶

Additional reading: you can read all of it or just read the "In short"¶

Step (a): features¶

3a. Log-mel for 68 clips¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

Additional reading: you can read all of it or just read the "In short"¶

Step (b): pseudo-labels¶

3b. k-means pseudo-labels¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

Additional Reading¶

How to read this plot¶

Additional reading: you can read all of it or just read the "In short"¶

Step (c): masked prediction¶

3c. Mask spans, predict the cluster id¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

Additional Reading¶

How to read this plot¶

Additional reading: you can read all of it or just read the "In short"¶

What the model must learn¶

3d. The trick¶

4. Probing "who vs what" on the mini model¶

Visualization Version¶

Real probe accuracy of the mini model per layer, and how similar the layers are.¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

Additional Reading¶

How to read this plot¶

5. The real encoders (needs internet the first time)¶

Visualization Version¶

The red box says: not run here. The coloured table is the shipped qualitative reference, not a measurement.¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

Additional reading: you can read all of it or just read the "In short"¶

What the real curves should look like (not shown here)¶

Additional reading: you can read all of it or just read the "In short"¶

What this cell does¶

What the output means¶

6. Which layer for voice conversion?¶

Visualization Version¶

The red box says: this is the mini model's layer 2 (stand-in), not WavLM.¶

Additional reading: you can read all of it or just read the "In short"¶

Build: a layer-probing tool (`src/probe.py`)¶