code/training/prep_data.py

"""Build ASR fine-tuning manifests from licensed, public-safe Nigerian speech already on dynacc-h100.
-> /scratch/asr/data/{train.jsonl, dev.jsonl, realcs_test.jsonl, stats.json}
Each row: id, src, lang, spk, text (ASR target, normalized), dur, pcm (24 kHz int16 file), off24, n24, licence.
Excluded (never touched): African Voices / WaZoBia, NaijaVoices, AfriSpeech (NC), Flash-AV. Held out: WAXAL test+validation,
FLEURS dev/test/val, NaijaS2ST dev, IgboSynCorp speakers that produced natural code-switched clips (-> realcs test).
Target normalization (matches duplex asr_score.norm): NFC, lowercase, punctuation -> space (apostrophes kept);
Igbo: tone marks (acute/grave/macron) stripped, dot-below and n-dot kept (WAXAL/standard Igbo orthography);
Yoruba/Hausa: kept as written by the source."""
import json, os, re, unicodedata, random, collections
H = "/home/dev"; OUT = "/scratch/asr/data"; os.makedirs(OUT, exist_ok=True)
random.seed(11)
LIC = {"waxal": "CC-BY-SA-4.0 (WAXAL Media Trust subsets)", "fleurs": "CC-BY-4.0 (Google FLEURS)",
       "slr70": "CC-BY-SA-4.0 (OpenSLR 70, Nigerian English)", "slr86": "CC-BY-SA-4.0 (OpenSLR 86, Yoruba)",
       "bible_yor": "CC-BY-SA-4.0 (Open.Bible / Biblica)", "bible_hau": "CC-BY-SA-4.0 (Open.Bible / Biblica)",
       "open_bible_igbo": "CC-BY-SA-4.0 (Open.Bible / Biblica)", "igboapi": "CC-BY-4.0 (IgboAPI)",
       "naijas2st_igbo": "CC-BY-4.0 (NaijaS2ST)", "igbosyncorp": "CC0-1.0 (IgboSynCorp)"}
# per-source sampling multiplier per epoch (Bible + dictionary reads down-weighted, Pidgin up-weighted)
W = {"waxal:ibo": 1.0, "waxal:yor": 1.0, "waxal:hau": 1.0, "waxal:pcm": 3.0, "fleurs:eng": 1.0, "fleurs:hau": 1.0,
     "fleurs:ibo": 1.0, "fleurs:yor": 1.0, "slr70:eng": 0.6, "slr86:yor": 1.0, "bible_yor:yor": 0.35, "bible_hau:hau": 0.35,
     "open_bible_igbo:ibo": 0.15, "igboapi:ibo": 0.25, "naijas2st_igbo:ibo": 0.5, "igbosyncorp:ibo": 1.5}
TONE = {"\u0300", "\u0301", "\u0304", "\u0302", "\u030c"}

def norm(s, lang):
    s = re.sub(r"<[^>]+>", " ", s); s = re.sub(r"\[\d+\]", " ", s)
    s = unicodedata.normalize("NFC", s.lower())
    if lang == "ibo":
        s = unicodedata.normalize("NFC", "".join(c for c in unicodedata.normalize("NFD", s) if c not in TONE))
    s = s.replace("\u2019", "'").replace("`", "'")
    s = re.sub(r"[^\w\s']", " ", s); s = re.sub(r"(^|\s)'+|'+(\s|$)", " ", s)
    return re.sub(r"\s+", " ", s).strip()

def load(d):
    return [json.loads(l) for l in open(f"{d}/meta.jsonl")]

rows = []
for ds in ["waxal", "fleurs", "slr70", "slr86", "bible_yor", "bible_hau"]:
    d = f"{H}/data3/{ds}"
    for r in load(d):
        rows.append(dict(id=r["id"], src=ds, lang=r["lang"], spk=r.get("spk"), split=r.get("split"), raw=r["text"], dur=r["dur"],
                         pcm=f"{d}/pcm24k.i16", off24=r["off24"], n24=r["n24"]))
for ds in ["open_bible_igbo", "igboapi", "naijas2st_igbo", "igbosyncorp"]:
    d = f"{H}/data_p3/p3set/{ds}"
    for r in load(d):
        rows.append(dict(id=r["id"], src=ds, lang=r["lang"], spk=r.get("speaker"), split=r.get("split"), raw=r.get("text_plain") or r.get("text_human"),
                         dur=r.get("dur24") or r["dur"], pcm=f"{d}/pcm24k.i16", off24=r["off24"], n24=r["n24"]))
for r in rows:
    r["text"] = norm(r["raw"] or "", r["lang"]); r["licence"] = LIC[r["src"]]
print("all rows", len(rows), collections.Counter((r["src"], r["split"]) for r in rows))

# English lexicon (for detecting real English insertions in held-out Igbo/Yoruba/Hausa speech)
ec, nc = collections.Counter(), collections.Counter()
for r in rows:
    if r["split"] != "train": continue
    (ec if r["lang"] == "eng" else nc if r["lang"] in ("ibo", "yor", "hau") else collections.Counter()).update(r["text"].split())
ENG = {w for w, c in ec.items() if c >= 2 and len(w) >= 3 and w.isascii() and w.isalpha() and nc[w] <= max(1, c // 50)}
def eng_words(t):
    return [w for w in t.split() if w in ENG]

# real code-switch test: held-out ibo/yor/hau clips with >=1 English word (insertional CS, real speech) + IgboSynCorp natural CS
realcs, hold_spk = [], set()
for r in rows:
    if r["lang"] not in ("ibo", "yor", "hau") or not r["text"]: continue
    ew = eng_words(r["text"]); nw = len(r["text"].split())
    if r["src"] == "igbosyncorp" and len(ew) >= 2 and nw >= 3:
        hold_spk.add(r["spk"])
    if r["split"] in ("test", "dev") and r["src"] in ("waxal", "naijas2st_igbo") and ew and nw >= 4 and len(ew) < nw:
        realcs.append(dict(r, cs_eng=ew, cs_kind="insertional"))
for r in rows:
    if r["src"] == "igbosyncorp" and r["spk"] in hold_spk and eng_words(r["text"]) and len(r["text"].split()) >= 3:
        realcs.append(dict(r, cs_eng=eng_words(r["text"]), cs_kind="natural_conversational"))
print("realcs", len(realcs), collections.Counter((r["src"], r["lang"]) for r in realcs), "held-out igbosyncorp speakers", len(hold_spk))

# dev: WAXAL validation 40/lang + FLEURS eng dev/test 40 (FLEURS val = the duplex eval English set, untouched here)
dev = []
for L in ["ibo", "yor", "hau", "pcm"]:
    c = [r for r in rows if r["src"] == "waxal" and r["lang"] == L and r["split"] in ("validation", "val") and 1.5 < r["dur"] < 15]
    random.shuffle(c); dev += c[:40]
c = [r for r in rows if r["src"] == "slr70" and r["split"] != "train" and 1.5 < r["dur"] < 15]
if not c:
    sp = sorted({r["spk"] for r in rows if r["src"] == "slr70"}); hold70 = set(sp[:4])
    c = [r for r in rows if r["src"] == "slr70" and r["spk"] in hold70]
else:
    hold70 = set()
random.shuffle(c); dev += c[:40]

train, drop = [], collections.Counter()
for r in rows:
    k = f"{r['src']}:{r['lang']}"
    if r["split"] != "train": drop["not_train"] += 1; continue
    if r["src"] == "slr70" and r["spk"] in hold70: drop["slr70_dev_spk"] += 1; continue
    if r["src"] == "igbosyncorp" and r["spk"] in hold_spk: drop["realcs_spk"] += 1; continue
    if k not in W: drop["no_weight:" + k] += 1; continue
    if not (0.8 <= r["dur"] <= 20) or len(r["text"]) < 2: drop["dur_or_empty"] += 1; continue
    if len(r["text"]) > r["dur"] * 25: drop["too_many_chars"] += 1; continue
    r["w"] = W[k]; train.append(r)
print("dropped", drop)
def dump(name, R):
    with open(f"{OUT}/{name}.jsonl", "w") as f:
        for r in R:
            f.write(json.dumps({k: v for k, v in r.items() if k != "raw"} | {"ref_raw": r["raw"]}, ensure_ascii=False) + "\n")
dump("train", train); dump("dev", dev); dump("realcs_test", realcs)
st = collections.defaultdict(lambda: [0, 0.0, 0.0])
for r in train:
    s = st[f"{r['src']}:{r['lang']}"]; s[0] += 1; s[1] += r["dur"] / 3600; s[2] += r["dur"] * r["w"] / 3600
lang_h = collections.Counter()
for k, v in st.items(): lang_h[k.split(":")[1]] += v[2]
out = dict(sources={k: dict(n=v[0], hours=round(v[1], 2), hours_per_epoch=round(v[2], 2), licence=LIC[k.split(':')[0]]) for k, v in sorted(st.items())},
           lang_hours_per_epoch={k: round(v, 1) for k, v in lang_h.items()}, total_hours=round(sum(v[1] for v in st.values()), 1),
           epoch_hours=round(sum(v[2] for v in st.values()), 1), dev=len(dev), realcs=len(realcs), eng_lexicon=len(ENG))
json.dump(out, open(f"{OUT}/stats.json", "w"), indent=1); print(json.dumps(out, indent=1))
json.dump(sorted(ENG), open(f"{OUT}/eng_lexicon.json", "w"))
print("PREP_DONE")