| """Phase 2: S3 speech tokens (speech_tokenizer_v2.onnx, CUDA EP) + CAM++ speaker embeddings for ~/data2/<ds>, |
| plus code-switched training pairs (ds "cs_pairs"). |
| |
| python -m da.s3_2 tok waxal fleurs ... -> s3tok.i32, s3tok.jsonl, campplus.npy |
| python -m da.s3_2 cs 4000 -> ~/data2/cs_pairs/{pcm24k.i16, meta.jsonl} (+ run tok on it) |
| """ |
| import os, sys, json, random, hashlib |
| import numpy as np, librosa, torch |
| from multiprocessing import Pool |
|
|
| D2 = os.path.expanduser(os.environ.get("P2OUT", "~/data2")) |
| MD = os.path.expanduser("~/models/CosyVoice2-0.5B") |
| _S = {} |
| _FB = torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=128)).float() |
| _WIN = torch.hann_window(400) |
|
|
|
|
| def log_mel_128(a): |
| st = torch.stft(a, 400, 160, window=_WIN, return_complex=True) |
| mag = st[..., :-1].abs() ** 2 |
| lm = torch.clamp(_FB @ mag, min=1e-10).log10() |
| lm = torch.maximum(lm, lm.max() - 8.0) |
| return (lm + 4.0) / 4.0 |
|
|
|
|
| def sess(name, gpu): |
| if name not in _S: |
| import onnxruntime as ort |
| torch.set_num_threads(1) |
| o = ort.SessionOptions(); o.intra_op_num_threads = 2; o.inter_op_num_threads = 1 |
| prov = ["CUDAExecutionProvider", "CPUExecutionProvider"] if gpu else ["CPUExecutionProvider"] |
| _S[name] = ort.InferenceSession(f"{MD}/{name}", sess_options=o, providers=prov) |
| return _S[name] |
|
|
|
|
| def work(args): |
| path, off, n = args |
| import torchaudio.compliance.kaldi as kaldi |
| pcm = np.memmap(path, dtype=np.int16, mode="r") |
| y24 = pcm[off:off + n].astype(np.float32) / 32768.0 |
| y16 = librosa.resample(y24, orig_sr=24000, target_sr=16000, res_type="soxr_hq") |
| a = torch.from_numpy(y16) |
| f = log_mel_128(a)[None] |
| s = sess("speech_tokenizer_v2.onnx", True) |
| ins = s.get_inputs() |
| tok = s.run(None, {ins[0].name: f.numpy(), ins[1].name: np.array([f.shape[2]], dtype=np.int32)})[0] |
| fb = kaldi.fbank(a[None], num_mel_bins=80, dither=0, sample_frequency=16000) |
| fb = fb - fb.mean(0, keepdim=True) |
| c = sess("campplus.onnx", False) |
| emb = c.run(None, {c.get_inputs()[0].name: fb[None].numpy()})[0].reshape(-1).astype(np.float32) |
| return np.asarray(tok).reshape(-1).astype(np.int32), emb |
|
|
|
|
| def tok(ds): |
| od = f"{D2}/{ds}" |
| meta = [json.loads(l) for l in open(f"{od}/meta.jsonl", encoding="utf-8")] |
| jobs = [(f"{od}/pcm24k.i16", m["off24"], m["n24"]) for m in meta] |
| toff, embs = 0, [] |
| with open(f"{od}/s3tok.i32", "wb") as tf, open(f"{od}/s3tok.jsonl", "w") as jf, Pool(12) as pool: |
| for i, (m, (t, e)) in enumerate(zip(meta, pool.imap(work, jobs, chunksize=4))): |
| tf.write(t.tobytes()); embs.append(e) |
| jf.write(json.dumps({"id": m["id"], "toff": toff, "ntok": int(len(t))}) + "\n") |
| toff += len(t) |
| if i % 5000 == 0: |
| print(ds, i, len(meta), "tok/s", round(len(t) / max(m["dur"], 1e-3), 1), flush=True) |
| np.save(f"{od}/campplus.npy", np.stack(embs)) |
| print(ds, "S3_DONE", toff, flush=True) |
|
|
|
|
| def build_cs(n_pairs=4000, seed=0): |
| """Join two train utterances in different languages (prefer same gender) into one code-switched item. |
| Junction silence follows the first clip's final punctuation (sentence 0.45 s, comma 0.25 s, else 0.20 s).""" |
| rng = random.Random(seed) |
| pools = {} |
| for ds in ["waxal", "slr70", "slr86", "fleurs"]: |
| p = f"{D2}/{ds}/meta.jsonl" |
| if not os.path.exists(p): |
| continue |
| for l in open(p, encoding="utf-8"): |
| m = json.loads(l) |
| if m["split"] != "train" or not (1.0 <= m["dur"] <= 6.5): |
| continue |
| if ds == "fleurs" and m["lang"] != "eng": |
| continue |
| pools.setdefault(m["lang"], []).append(m) |
| langs = sorted(pools) |
| print("cs pools", {l: len(v) for l, v in pools.items()}) |
| od = f"{D2}/cs_pairs"; os.makedirs(od, exist_ok=True) |
| mm = {} |
| off = 0 |
| with open(f"{od}/pcm24k.i16", "wb") as pf, open(f"{od}/meta.jsonl", "w", encoding="utf-8") as jf: |
| for k in range(n_pairs): |
| l1, l2 = rng.sample(langs, 2) |
| if rng.random() < 0.5 and "eng" in langs and "eng" not in (l1, l2): |
| l2 = "eng" |
| a = rng.choice(pools[l1]) |
| g = (a.get("gender") or "").lower()[:1] |
| cand = [b for b in rng.sample(pools[l2], min(50, len(pools[l2]))) if (b.get("gender") or "").lower()[:1] == g] |
| b = cand[0] if cand else rng.choice(pools[l2]) |
| ys = [] |
| for m in (a, b): |
| if m["ds"] not in mm: |
| mm[m["ds"]] = np.memmap(f"{D2}/{m['ds']}/pcm24k.i16", dtype=np.int16, mode="r") |
| ys.append(np.asarray(mm[m["ds"]][m["off24"]:m["off24"] + m["n24"]])) |
| last = a["text"].rstrip()[-1:] |
| junc = 0.45 if last in ".?!" else (0.25 if last in ",;:" else 0.20) |
| have = 0.18 + 0.12 |
| y1 = ys[0] |
| if junc < have: |
| y1 = y1[: len(y1) - int((have - junc) * 24000)] |
| gap = np.zeros(0, np.int16) |
| else: |
| gap = np.zeros(int((junc - have) * 24000), np.int16) |
| y = np.concatenate([y1, gap, ys[1]]) |
| text = f"<{l1}>{a['text']}</{l1}> <{l2}>{b['text']}</{l2}>" |
| m = dict(id=f"cs_{k:05d}", ds="cs_pairs", lang="cs", langs=[l1, l2], spk=a["spk"], gender=a.get("gender", ""), |
| split="train", text=text, dur=round(len(y) / 24000, 3), off24=off, n24=int(len(y)), |
| src=[a["id"], b["id"]], flow_ok=False) |
| pf.write(y.tobytes()); off += len(y) |
| jf.write(json.dumps(m, ensure_ascii=False) + "\n") |
| print("cs pairs", n_pairs, "hours", round(off / 24000 / 3600, 2), flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| if sys.argv[1] == "tok": |
| for ds in sys.argv[2:]: |
| tok(ds) |
| elif sys.argv[1] == "cs": |
| build_cs(int(sys.argv[2])) |
| tok("cs_pairs") |
|
|
