File size: 8,601 Bytes
63a4356 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | """A CTC phoneme aligner trained on Jenny's OWN audio, to replace the DTW duration transfer.
Why this exists. DTW-transferring the teacher's durations onto Jenny passed the energy gate
(+4.34 dB vowel-minus-quiet against a -0.18 dB reversed-label control) and still produced bad
labels: finetuning on them tripled free-running WER (0.041 -> 0.145), and -- the giveaway --
teacher-forcing those same durations was WORSE than free-running (0.258), which can only happen
when the labels disagree with the audio. Feeding the DTW durations to the ALREADY-GOOD shipped
model degraded it too (0.041 -> 0.079), so the durations, not the student, are the fault.
The energy gate is necessary but not sufficient: it only asks whether loud frames land on
vowels, which survives several frames of per-phoneme slop.
CTC needs only the phoneme SEQUENCE, not its timing, so this trains directly on Jenny's real
recordings -- no synthetic-to-real domain gap, unlike an aligner trained on teacher audio.
Framing is locked to the student's 40 Hz grid (hop 600 at 24 kHz) so a frame index means the
same thing here and in the trainer; converting rates between the two is where a previous
project in this line lost a run.
"""
from __future__ import annotations
import argparse, json, math, os, random, time
from pathlib import Path
import numpy as np, torch, torch.nn as nn, torch.nn.functional as F
import soundfile as sf
from torch.utils.data import Dataset, DataLoader
SR, HOP, NMELS = 24000, 600, 80
NVOCAB = 178 # Kokoro token ids 0..177
BLANK = NVOCAB # CTC blank gets its own class
def melspec(dev):
import torchaudio
return torchaudio.transforms.MelSpectrogram(
sample_rate=SR, n_fft=1200, hop_length=HOP, n_mels=NMELS, center=True, power=2.0).to(dev)
class Aligner(nn.Module):
def __init__(self, d=256):
super().__init__()
self.conv = nn.Sequential(
nn.Conv1d(NMELS, d, 5, padding=2), nn.GELU(), nn.BatchNorm1d(d),
nn.Conv1d(d, d, 5, padding=2), nn.GELU(), nn.BatchNorm1d(d))
self.rnn = nn.LSTM(d, d, num_layers=2, batch_first=True, bidirectional=True)
self.out = nn.Linear(2 * d, NVOCAB + 1)
def forward(self, mel): # mel [B, NMELS, T]
x = self.conv(mel).transpose(1, 2)
x, _ = self.rnn(x)
return self.out(x) # [B, T, NVOCAB+1]
class ClipSet(Dataset):
"""Jenny audio paired with the teacher's G2P phoneme ids for the same sentence."""
def __init__(self, rows, max_frames):
self.rows, self.max_frames = rows, max_frames
def __len__(self):
return len(self.rows)
def __getitem__(self, i):
r = self.rows[i]
y, _ = sf.read(r["wav"], dtype="float32")
if y.ndim > 1:
y = y.mean(axis=1)
n = min(len(y) // HOP, self.max_frames)
# frame j must cover samples [j*HOP, (j+1)*HOP) to match the student's grid; a
# centered STFT centres frame j at j*HOP, so drop half a hop first.
y = y[HOP // 2: HOP // 2 + n * HOP]
return {"y": torch.from_numpy(np.ascontiguousarray(y)), "ids": r["ids"], "n": n}
def collate(b):
T = max(x["n"] for x in b)
L = max(len(x["ids"]) for x in b)
y = torch.zeros(len(b), T * HOP)
tgt = torch.zeros(len(b), L, dtype=torch.long)
ylen = torch.tensor([x["n"] for x in b], dtype=torch.long)
tlen = torch.tensor([len(x["ids"]) for x in b], dtype=torch.long)
for i, x in enumerate(b):
y[i, : len(x["y"])] = x["y"]
tgt[i, : len(x["ids"])] = torch.tensor(x["ids"], dtype=torch.long)
return y, tgt, ylen, tlen
def load_rows(max_frames, manifest=None):
"""Rows of {wav, ids, text, n}.
A manifest that already carries phoneme ids (build_v7.py writes these) is used directly.
Otherwise fall back to joining Jenny's audio against the teacher's G2P by text, which is how
the English corpus was assembled."""
rows = []
if manifest:
for l in open(manifest, encoding="utf-8"):
d = json.loads(l)
p = [i for i in d["ids"] if i != 0]
n = d.get("frames") or int(d["duration"] * SR) // HOP
# DROP clips past the cap, never truncate them: the phoneme sequence describes the
# whole recording, so a clipped waveform paired with the full label set teaches the
# aligner to compress. v9 is long-form (up to 123 s), where this actually bites.
if not p or n > max_frames or n < max(len(p) + 2, 42):
continue
rows.append({"wav": d["wav"], "ids": p, "text": d["text"], "n": n})
return rows
ids = {}
for l in open("dist/index.jsonl", encoding="utf-8"):
d = json.loads(l)
ids[d["text"]] = [i for i in d["ids"] if i != 0] # strip the $ pad tokens
for l in open("data/jenny.jsonl", encoding="utf-8"):
d = json.loads(l)
t = d["text"].strip().replace("\n", " ")
p = ids.get(t)
if not p:
continue
n = min(int(d["duration"] * SR) // HOP, max_frames)
if n < len(p) + 2 or n < 42: # CTC needs at least one frame per token
continue
rows.append({"wav": d["audio"], "ids": p, "text": t, "n": n})
return rows
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--steps", type=int, default=12000)
ap.add_argument("--batch", type=int, default=16)
ap.add_argument("--lr", type=float, default=3e-4)
ap.add_argument("--max-frames", type=int, default=800) # 20 s
ap.add_argument("--manifest", default=None,
help="jsonl with wav/ids/frames already resolved")
ap.add_argument("--out", default="runs/aligner")
ap.add_argument("--holdout", type=int, default=200)
ap.add_argument("--resume", action="store_true")
a = ap.parse_args()
dev = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(42)
out = Path(a.out); out.mkdir(parents=True, exist_ok=True)
rows = load_rows(a.max_frames, a.manifest)
random.Random(1234).shuffle(rows)
tr, ev = rows[: -a.holdout], rows[-a.holdout:]
print(f"[*] aligner train {len(tr):,} | held out {len(ev):,}", flush=True)
mel = melspec(dev)
model = Aligner().to(dev)
print(f"[*] aligner {sum(p.numel() for p in model.parameters())/1e6:.2f}M", flush=True)
opt = torch.optim.AdamW(model.parameters(), lr=a.lr, weight_decay=1e-2)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=a.lr, total_steps=a.steps,
pct_start=0.05)
step = 0
ck = out / "last.pt"
if a.resume and ck.exists():
c = torch.load(ck, map_location="cpu", weights_only=False)
model.load_state_dict(c["model"]); opt.load_state_dict(c["opt"])
sched.load_state_dict(c["sched"]); step = c["step"]
print(f"[*] resumed at {step}", flush=True)
dl = DataLoader(ClipSet(tr, a.max_frames), batch_size=a.batch, shuffle=True,
collate_fn=collate, num_workers=4, drop_last=True, pin_memory=True)
t0 = time.time(); acc = 0.0; nacc = 0
model.train()
while step < a.steps:
for y, tgt, ylen, tlen in dl:
if step >= a.steps:
break
y, tgt = y.to(dev, non_blocking=True), tgt.to(dev)
with torch.no_grad():
m = torch.log(mel(y) + 1e-6)
m = (m - m.mean(dim=(1, 2), keepdim=True)) / (m.std(dim=(1, 2), keepdim=True) + 1e-5)
m = m[..., : int(ylen.max())]
lp = F.log_softmax(model(m), dim=-1).transpose(0, 1) # [T,B,C]
loss = F.ctc_loss(lp, tgt, ylen.to(dev), tlen.to(dev), blank=BLANK,
zero_infinity=True)
opt.zero_grad(set_to_none=True)
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), 5.0)
opt.step(); sched.step(); step += 1
acc += float(loss); nacc += 1
if step % 200 == 0:
print(f" step {step}/{a.steps} ctc {acc/nacc:.4f} "
f"{step*a.batch/(time.time()-t0):.0f} clip/s", flush=True)
acc = 0.0; nacc = 0
if step % 2000 == 0 or step == a.steps:
tmp = out / "last.pt.tmp"
torch.save({"model": model.state_dict(), "opt": opt.state_dict(),
"sched": sched.state_dict(), "step": step}, tmp)
os.replace(tmp, ck)
print(f"[+] aligner done at {step}", flush=True)
if __name__ == "__main__":
main()
|