File size: 6,016 Bytes
256a3b6 | 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 | """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" # English is the most common partner in real Nigerian code-switching
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")
|