prosatts-work / scripts /pttrain14 /train_aligner.py

Checkpoints, banco de risos e scripts

63a4356 verified about 14 hours ago

8.6 kB

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

	"""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 [jHOP, (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()