| """Phase 2 data prep (H100 node): every source -> MMS forced alignment -> sentence segmentation of long clips -> |
| punctuation-consistent pause normalisation -> 24 kHz int16 + meta. Output ~/data2/<ds>/{pcm24k.i16, meta.jsonl, stats.json} |
| |
| python -m da.prep2 waxal bible_hau bible_yor slr86 slr70 fleurs |
| """ |
| import os, sys, io, json, glob, csv, hashlib, time |
| import numpy as np, soundfile as sf, librosa |
| from multiprocessing import Pool |
| sys.path.insert(0, os.path.expanduser("~/da")) |
| from da.p2common import words_for_align, Aligner, normalise_pauses |
|
|
| RAW = "/scratch/raw" |
| OUT = os.path.expanduser(os.environ.get("P2OUT", "~/data2")) |
| OLD = os.path.expanduser("~/data") |
| SR = 24000 |
| BIBLE_HAU_CAP_H = 30.0 |
| FLOW_OK = {"waxal", "bible_yor", "bible_hau", "slr86", "slr70"} |
|
|
|
|
| def clean_text(t): |
| import unicodedata, re |
| t = unicodedata.normalize("NFC", t or "") |
| t = t.replace("\u2019", "'").replace("\u2018", "'").replace("\u201c", '"').replace("\u201d", '"') |
| return re.sub(r"\s+", " ", t).strip() |
|
|
|
|
| def hsplit(i): |
| return "val" if int(hashlib.md5(i.encode()).hexdigest(), 16) % 100 == 0 else "train" |
|
|
|
|
| |
| def items_waxal(): |
| import pyarrow.parquet as pq |
| for lang in ["yor", "ibo", "hau", "pcm"]: |
| files = sorted(glob.glob(f"{RAW}/waxal/data/TTS/{lang}/*.parquet")) |
| files = [f for f in files if not os.path.basename(f).startswith("yor_tts")] |
| for f in files: |
| bn = os.path.basename(f) |
| split = "test" if "test" in bn else ("val" if "validation" in bn else "train") |
| pf = pq.ParquetFile(f) |
| for b in pf.iter_batches(batch_size=8): |
| for r in b.to_pylist(): |
| yield dict(id=f"waxal_{lang}_{r['id']}", ds="waxal", lang=lang, |
| spk=f"waxal_{lang}_{r.get('speaker_id') or 'spk'}", gender=r.get("gender") or "", |
| split=split, text=r["text"]), ("bytes", r["audio"]["bytes"]) |
|
|
|
|
| FLEURS = {"hau": "ha_ng", "ibo": "ig_ng", "yor": "yo_ng", "eng": "en_us"} |
|
|
|
|
| def items_fleurs(): |
| """raw FLEURS train audio (16 kHz); speaker ids = Phase 1 ECAPA pseudo-speakers from ~/data/fleurs/meta.jsonl""" |
| for l in open(f"{OLD}/fleurs/meta.jsonl", encoding="utf-8"): |
| m = json.loads(l) |
| fn = m["id"].split("_", 2)[2] + ".wav" |
| p = f"{RAW}/fleurs/data/{FLEURS[m['lang']]}/audio/train/{fn}" |
| yield dict(id=m["id"], ds="fleurs", lang=m["lang"], spk=m["spk"], gender=m.get("gender", ""), |
| split=hsplit(m["id"]), text=m["text"]), ("path", p) |
|
|
|
|
| def items_slr(num, lang, prefix): |
| for idx in sorted(glob.glob(f"{RAW}/slr{num}/**/line_index.tsv", recursive=True)): |
| d = os.path.dirname(idx) |
| for line in open(idx, encoding="utf-8"): |
| parts = line.rstrip("\n").split("\t") |
| if len(parts) < 2: |
| continue |
| fid, text = parts[0].strip(), parts[-1] |
| p = os.path.join(d, fid + ".wav") |
| if os.path.exists(p): |
| spk = "_".join(fid.split("_")[:2]) |
| yield dict(id=f"{prefix}_{fid}", ds=prefix, lang=lang, spk=f"{prefix}_{spk}", |
| gender="female" if fid[2] == "f" else "male", split=hsplit(fid), text=text), ("path", p) |
|
|
|
|
| def items_bible(lang_dir, lang, ds, cap_h=None): |
| root = f"{RAW}/bible/{lang_dir}" |
| auds = sorted(glob.glob(f"{root}/**/*.flac", recursive=True) + glob.glob(f"{root}/**/*.wav", recursive=True)) |
| if cap_h: |
| auds.sort(key=lambda p: hashlib.md5(p.encode()).hexdigest()) |
| tot = 0.0 |
| for p in auds: |
| txt = os.path.splitext(p)[0] + ".txt" |
| if not os.path.exists(txt): |
| continue |
| split = "test" if "/test/" in p else ("val" if "/dev/" in p else "train") |
| if cap_h and split == "train": |
| if tot > cap_h * 3600: |
| continue |
| try: |
| tot += sf.info(p).duration |
| except Exception: |
| continue |
| yield dict(id=f"{ds}_{os.path.splitext(os.path.basename(p))[0]}", ds=ds, lang=lang, spk=ds, gender="", |
| split=split, text=open(txt, encoding="utf-8").read()), ("path", p) |
|
|
|
|
| def enum(ds): |
| return {"waxal": items_waxal, "fleurs": items_fleurs, |
| "slr86": lambda: items_slr(86, "yor", "slr86"), "slr70": lambda: items_slr(70, "eng", "slr70"), |
| "bible_yor": lambda: items_bible("yoruba", "yor", "bible_yor"), |
| "bible_hau": lambda: items_bible("hausa", "hau", "bible_hau", BIBLE_HAU_CAP_H)}[ds]() |
|
|
|
|
| |
| def load(arg): |
| meta, src = arg |
| try: |
| if src[0] == "bytes": |
| y, sr = sf.read(io.BytesIO(src[1]), dtype="float32", always_2d=True); y = y.mean(1) |
| elif src[0] == "path": |
| y, sr = sf.read(src[1], dtype="float32", always_2d=True); y = y.mean(1) |
| else: |
| pcm = np.memmap(src[1], dtype=np.int16, mode="r") |
| y, sr = pcm[src[2]:src[2] + src[3]].astype(np.float32) / 32768.0, 16000 |
| meta["text"] = clean_text(meta["text"]) |
| meta["sr_in"] = int(sr) |
| if not meta["text"] or len(y) < 0.5 * sr: |
| return meta, None, None, "empty" |
| y24 = librosa.resample(y, orig_sr=sr, target_sr=SR, res_type="soxr_hq") if sr != SR else y |
| y16 = librosa.resample(y, orig_sr=sr, target_sr=16000, res_type="soxr_hq") if sr != 16000 else y |
| return meta, y16, y24, None |
| except Exception as e: |
| return meta, None, None, f"err:{type(e).__name__}" |
|
|
|
|
| def plan_segments(words, ali, dur): |
| """word-index ranges [(a, b)] (inclusive). Short clips: one segment. Long clips: whole sentences grouped |
| into 3-14 s segments; an over-long sentence is cut at clause/comma marks.""" |
| n = len(words) |
| if dur <= 20.0: |
| return [(0, n - 1)] |
| sents, a = [], 0 |
| for i, w in enumerate(words): |
| if w["punct"] == "sent" or i == n - 1: |
| sents.append((a, i)); a = i + 1 |
| pieces = [] |
| for a, b in sents: |
| if ali[b]["end"] - ali[a]["start"] <= 16.0: |
| pieces.append((a, b)); continue |
| s = a |
| for i in range(a, b + 1): |
| if (words[i]["punct"] in ("clause", "comma") and ali[i]["end"] - ali[s]["start"] > 3.0) or i == b: |
| pieces.append((s, i)); s = i + 1 |
| segs, cur = [], None |
| for a, b in pieces: |
| if cur is None: |
| cur = [a, b]; continue |
| if ali[b]["end"] - ali[cur[0]]["start"] <= 14.0 and ali[cur[1]]["end"] - ali[cur[0]]["start"] < 4.0: |
| cur[1] = b |
| else: |
| segs.append(tuple(cur)); cur = [a, b] |
| if cur: |
| segs.append(tuple(cur)) |
| return segs |
|
|
|
|
| def post(arg): |
| meta, y24, words, ali = arg |
| out, why = [], {} |
| segs = plan_segments(words, ali, len(y24) / SR) |
| from da.p2common import silence_mask |
| sil = silence_mask(y24, SR) |
| for k, (a, b) in enumerate(segs): |
| sc = float(np.mean([ali[i]["score"] for i in range(a, b + 1)])) |
| maxw = max(ali[i]["end"] - ali[i]["start"] for i in range(a, b + 1)) |
| text = " ".join(words[i]["word"] for i in range(a, b + 1)) |
| if sc < 0.30: |
| why["low_score"] = why.get("low_score", 0) + 1; continue |
| if maxw > 2.5: |
| why["long_word"] = why.get("long_word", 0) + 1; continue |
| lo = ali[a - 1]["end"] + 0.02 if a > 0 else 0.0 |
| hi = ali[b + 1]["start"] - 0.02 if b + 1 < len(ali) else None |
| y, pauses = normalise_pauses(y24, SR, words, ali, a, b, sil=sil, lo_s=lo, hi_s=hi) |
| dur = len(y) / SR |
| cps = len(text) / dur |
| if not (0.8 <= dur <= 20.0): |
| why["dur"] = why.get("dur", 0) + 1; continue |
| if not (4.0 <= cps <= 25.0): |
| why["cps"] = why.get("cps", 0) + 1; continue |
| rms = np.sqrt((y ** 2).mean() + 1e-12) |
| y = y * (0.07 / rms) |
| pk = np.abs(y).max() |
| if pk > 0.99: |
| y = y * (0.99 / pk) |
| m = dict(meta) |
| if len(segs) > 1: |
| m["id"] = f"{meta['id']}_s{k:03d}"; m["src_id"] = meta["id"]; m["longform"] = True |
| m.update(text=text, dur=round(dur, 3), ascore=round(sc, 3), cps=round(cps, 2), |
| pauses=[round(p, 2) for p in pauses], |
| puncts=[words[i]["punct"] for i in range(a, b)], flow_ok=meta["ds"] in FLOW_OK) |
| out.append((m, np.clip(y * 32767, -32768, 32767).astype(np.int16))) |
| return out, why |
|
|
|
|
| def run(ds): |
| od = f"{OUT}/{ds}" |
| os.makedirs(od, exist_ok=True) |
| al = Aligner() |
| stats = dict(src_items=0, src_hours=0.0, kept=0, hours=0.0, rej={}, longform_src=0, longform_hours_in=0.0, |
| longform_segments=0) |
| t0 = time.time(); off = 0 |
| pending = [] |
| with open(f"{od}/pcm24k.i16", "wb") as pf, open(f"{od}/meta.jsonl", "w", encoding="utf-8") as jf, \ |
| Pool(20) as lp, Pool(6) as pp: |
| def drain(block=False): |
| nonlocal off |
| keep = [] |
| for r in pending: |
| if block or r.ready(): |
| segs, why = r.get() |
| for k, v in why.items(): |
| stats["rej"][k] = stats["rej"].get(k, 0) + v |
| for m, pcm in segs: |
| m["off24"] = off; m["n24"] = int(len(pcm)) |
| pf.write(pcm.tobytes()); off += len(pcm) |
| jf.write(json.dumps(m, ensure_ascii=False) + "\n") |
| stats["kept"] += 1; stats["hours"] += m["dur"] / 3600 |
| if m.get("longform"): |
| stats["longform_segments"] += 1 |
| else: |
| keep.append(r) |
| pending[:] = keep |
| import itertools |
| src = enum(ds) |
| if os.environ.get("LIMIT"): |
| src = itertools.islice(src, 0, None, int(os.environ.get("STRIDE", "1"))) |
| src = itertools.islice(src, int(os.environ["LIMIT"])) |
| for meta, y16, y24, err in lp.imap(load, src, chunksize=2): |
| stats["src_items"] += 1 |
| if err: |
| stats["rej"][err] = stats["rej"].get(err, 0) + 1; continue |
| d = len(y24) / SR |
| stats["src_hours"] += d / 3600 |
| if d > 20: |
| stats["longform_src"] += 1; stats["longform_hours_in"] += d / 3600 |
| if d > 2400: |
| stats["rej"]["too_long"] = stats["rej"].get("too_long", 0) + 1; continue |
| words = words_for_align(meta["text"]) |
| ali = al.align(y16, words) if words else None |
| if ali is None: |
| stats["rej"]["align_fail"] = stats["rej"].get("align_fail", 0) + 1; continue |
| pending.append(pp.apply_async(post, ((meta, y24, words, ali),))) |
| if len(pending) > 64: |
| drain(block=False) |
| while len(pending) > 128: |
| time.sleep(0.05); drain(block=False) |
| if stats["src_items"] % 2000 == 0: |
| print(ds, stats["src_items"], "kept", stats["kept"], "h", round(stats["hours"], 2), |
| "rej", stats["rej"], round(time.time() - t0), "s", flush=True) |
| drain(block=True) |
| stats = {k: (round(v, 3) if isinstance(v, float) else v) for k, v in stats.items()} |
| json.dump(stats, open(f"{od}/stats.json", "w"), indent=1) |
| print(ds, "PREP2_DONE", json.dumps(stats), round(time.time() - t0), "s", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| for ds in sys.argv[1:]: |
| run(ds) |
|
|
