| |
| |
| import os, glob, subprocess, json, re, time, numpy as np, soundfile as sf, torch |
| from transformers import Qwen3ASRForConditionalGeneration, AutoProcessor |
| from opencc import OpenCC |
| s2twp=OpenCC("s2twp"); dev="cuda"; WIN=16.0 |
| YTDLP="/home/luigi/jetson-stt/.venv/bin/yt-dlp" |
| N_VIDEOS=int(os.environ.get("N_VIDEOS","150")) |
| os.makedirs("yt_scale", exist_ok=True) |
| OUT="kd_dataset.jsonl" |
| ids=[]; seen=set() |
| for l in open("yt_manifest.tsv",encoding="utf-8"): |
| fn=l.split("\t")[0].rsplit("/",1)[-1].rsplit(".",1)[0]; v="_".join(fn.split("_")[:-1]) |
| if v not in seen: seen.add(v); ids.append(v) |
| done=set() |
| if os.path.exists("scale_done.txt"): done=set(open("scale_done.txt").read().split()) |
| todo=[v for v in ids[15:] if v not in done][:N_VIDEOS] |
| print(f"scale: {len(todo)} new videos (target), {len(done)} already done",flush=True) |
| proc=AutoProcessor.from_pretrained("/tmp/qwen3-asr-hf") |
| model=Qwen3ASRForConditionalGeneration.from_pretrained("/tmp/qwen3-asr-hf", dtype=torch.float16).to(dev).eval() |
| def dl(vid,out): |
| subprocess.run([YTDLP,"--js-runtimes","deno","-f","140/bestaudio[ext=m4a]/bestaudio", |
| "--match-filter","!is_live & duration>120 & duration<14400","-x","--audio-format","wav", |
| "--postprocessor-args","ExtractAudio:-ar 16000 -ac 1","--no-playlist","-q","-o",out, |
| f"https://www.youtube.com/watch?v={vid}"], |
| env=dict(os.environ,PATH="/home/luigi/.local/bin:"+os.environ.get("PATH","")), |
| timeout=600,stdout=subprocess.DEVNULL,stderr=subprocess.DEVNULL) |
| @torch.no_grad() |
| def label(a): |
| conv=[{"role":"user","content":[{"type":"audio","audio":a},{"type":"text","text":""}]}] |
| inp=proc.apply_chat_template(conv,add_generation_prompt=True,tokenize=True,return_dict=True,return_tensors="pt",sampling_rate=16000) |
| inp={k:(v.to(dev).half() if v.dtype==torch.float32 else v.to(dev)) for k,v in inp.items()} |
| o=model.generate(**inp,max_new_tokens=96,do_sample=False) |
| return re.sub(r".*<asr_text>","",proc.batch_decode(o[:,inp["input_ids"].shape[1]:],skip_special_tokens=True)[0]).strip() |
| out=open(OUT,"a",encoding="utf-8"); dlog=open("scale_done.txt","a") |
| t0=time.time(); nseg=0; nh=0.0 |
| for vi,vid in enumerate(todo): |
| raw=f"yt_scale/_tmp_{vid}.wav" |
| try: |
| dl(vid,raw); a,_=sf.read(raw,dtype="float32"); a=a.mean(1) if a.ndim>1 else a |
| except Exception: |
| for p in glob.glob(f"yt_scale/_tmp_{vid}*"): |
| try: os.remove(p) |
| except: pass |
| continue |
| W=int(WIN*16000); k=0 |
| for s in range(0,len(a),W): |
| seg=a[s:s+W] |
| if len(seg)<16000 or float(np.abs(seg).max())<0.01: continue |
| raw_txt=label(seg) |
| cjk=sum(1 for c in raw_txt if '一'<=c<='鿿'); lat=sum(c.isalpha() and ord(c)<128 for c in raw_txt) |
| if cjk<2 and lat<4: continue |
| sp=f"yt_scale/{vid}_{k}.wav"; sf.write(sp,seg,16000) |
| out.write(json.dumps({"wav":os.path.abspath(sp),"text":s2twp.convert(raw_txt)},ensure_ascii=False)+"\n") |
| nseg+=1; nh+=len(seg)/16000/3600; k+=1 |
| try: os.remove(raw) |
| except: pass |
| out.flush(); dlog.write(vid+"\n"); dlog.flush() |
| print(f" +{vid} ({k} segs) -> {nseg} segs, {nh:.1f}h [{vi+1}/{len(todo)}] {(time.time()-t0)/60:.0f}min",flush=True) |
| print(f"SCALE-DONE: {nseg} segs, {nh:.1f}h -> {OUT}",flush=True) |
|
|
