| |
| """Deterministic LAION-DISCO audio tranche ingestion.""" |
| from __future__ import annotations |
| import argparse, array, hashlib, heapq, json, math, os, re, shutil, subprocess, sys, time, unicodedata, wave |
| from collections import Counter |
| from concurrent.futures import ThreadPoolExecutor |
| from pathlib import Path |
| from typing import Any |
| ROOT=Path("/home/ubuntu/minimax-laion-corpus") |
| REVISION="6e7bf3758a77301e46a715af894fefd79bb1da53" |
| CANDIDATES=ROOT/"manifests/candidates.jsonl"; CANDIDATE_SUMMARY=ROOT/"manifests/candidates.summary.json" |
| KNOWN_HASHES=ROOT/"manifests/known_user22_hashes.json"; RESULTS=ROOT/"manifests/ingest-results.jsonl" |
| CANONICAL=ROOT/"canonical"; TMP=ROOT/"tmp"; YTDLP=ROOT/".venv/bin/yt-dlp" |
| FFMPEG=shutil.which("ffmpeg") or "/usr/bin/ffmpeg"; FFPROBE=shutil.which("ffprobe") or "/usr/bin/ffprobe" |
| ID_RE=re.compile(r"^[A-Za-z0-9_-]{11}$") |
| def sha256_file(path:Path,chunk:int=8<<20)->str: |
| h=hashlib.sha256() |
| with path.open("rb") as f: |
| while b:=f.read(chunk):h.update(b) |
| return h.hexdigest() |
| def norm(value:Any)->str: |
| if value is None:return "" |
| return " ".join(unicodedata.normalize("NFKC",str(value)).casefold().split()) |
| def stable_rank(song_id:str)->str:return hashlib.sha256((REVISION+"\0"+song_id).encode()).hexdigest() |
| def artist_keys(row:dict[str,Any])->list[str]: |
| ids=sorted({"id:"+norm(x) for x in (row.get("artist_ids") or []) if norm(x)}) |
| if ids:return ids |
| names=sorted({"name:"+norm(x) for x in (row.get("artist_names") or []) if norm(x)}) |
| return names or ["song:"+row["song_id"]] |
| def album_key(row:dict[str,Any])->str: |
| aid=norm(row.get("album_id")) |
| if aid:return "id:"+aid |
| name=norm(row.get("album_name")); artists=artist_keys(row) |
| if name:return "name:"+name+"|artist:"+(artists[0] if artists else row["song_id"]) |
| return "song:"+row["song_id"] |
| def atomic_text(path:Path,text:str)->None: |
| path.parent.mkdir(parents=True,exist_ok=True);tmp=path.with_suffix(path.suffix+".tmp");tmp.write_text(text,encoding="utf-8");os.replace(tmp,path) |
| def write_json(path:Path,obj:Any)->None:atomic_text(path,json.dumps(obj,sort_keys=True,indent=2,ensure_ascii=False)+"\n") |
| def read_jsonl(path:Path)->list[dict[str,Any]]: |
| with path.open(encoding="utf-8") as f:return [json.loads(x) for x in f if x.strip()] |
| def build_candidates(pool_size:int,target:int)->dict[str,Any]: |
| import pyarrow.parquet as pq |
| if pool_size<target:raise ValueError("pool_size must be >= target") |
| shards=sorted((ROOT/"metadata").glob("train-*-of-00005.parquet")) |
| if len(shards)!=5:raise RuntimeError(f"expected five shards, got {len(shards)}") |
| heap=[];seq=0;totals=Counter() |
| cols=["song_id","title","artist_names","artist_ids","album_name","album_id","isExplicit","views","duration"] |
| for shard_no,shard in enumerate(shards): |
| pf=pq.ParquetFile(shard);totals["metadata_rows"]+=pf.metadata.num_rows |
| for rg in range(pf.metadata.num_row_groups): |
| for raw in pf.read_row_group(rg,columns=cols).to_pylist(): |
| totals["rows_scanned"]+=1;sid=str(raw.get("song_id") or "").strip() |
| try:duration=int(raw.get("duration")) |
| except (TypeError,ValueError):totals["invalid_duration"]+=1;continue |
| if duration<45 or duration>360:totals["duration_filtered"]+=1;continue |
| if not ID_RE.fullmatch(sid):totals["invalid_song_id"]+=1;continue |
| rank=stable_rank(sid) |
| rec={"rank_sha256":rank,"song_id":sid,"url":"https://www.youtube.com/watch?v="+sid,"title":raw.get("title") or "", |
| "artist_names":raw.get("artist_names") or [],"artist_ids":raw.get("artist_ids") or [], |
| "album_name":raw.get("album_name") or "","album_id":raw.get("album_id") or "", |
| "is_explicit":bool(raw.get("isExplicit")) if raw.get("isExplicit") is not None else None, |
| "views":raw.get("views") or "","metadata_duration_seconds":duration,"source_shard":shard_no} |
| item=(-int(rank,16),sid,seq,rec);seq+=1 |
| if len(heap)<pool_size:heapq.heappush(heap,item) |
| elif item[0]>heap[0][0]:heapq.heapreplace(heap,item) |
| ranked=sorted((x[3] for x in heap),key=lambda x:(x["rank_sha256"],x["song_id"])) |
| accepted=[];seen_song=set();seen_album=set();artist_count=Counter() |
| for rec in ranked: |
| sid=rec["song_id"] |
| if sid in seen_song:totals["duplicate_song_id"]+=1;continue |
| seen_song.add(sid);ak=album_key(rec);artists=artist_keys(rec) |
| if ak in seen_album:totals["album_cap_rejected"]+=1;continue |
| if any(artist_count[a]>=2 for a in artists):totals["artist_cap_rejected"]+=1;continue |
| rec["album_cap_key"]=ak;rec["artist_cap_keys"]=artists;rec["candidate_index"]=len(accepted);accepted.append(rec);seen_album.add(ak) |
| for a in artists:artist_count[a]+=1 |
| if len(accepted)==target:break |
| if len(accepted)!=target:raise RuntimeError(f"only {len(accepted)} candidates after caps; increase pool_size") |
| atomic_text(CANDIDATES,"".join(json.dumps(x,sort_keys=True,ensure_ascii=False,separators=(",",":"))+"\n" for x in accepted)) |
| summary={"schema_version":1,"dataset":"laion/LAION-DISCO-12M","revision":REVISION, |
| "selection":{"rank":"ascending sha256(revision + NUL + song_id)","metadata_duration_seconds_inclusive":[45,360], |
| "album_cap":1,"artist_cap":2,"cap_scope":"all listed artist IDs; normalized names when IDs absent", |
| "pool_size":pool_size,"candidate_count":target,"intended_success_count":2000}, |
| "counts":dict(totals),"candidate_manifest_sha256":sha256_file(CANDIDATES), |
| "first_rank":accepted[0]["rank_sha256"],"last_rank":accepted[-1]["rank_sha256"]} |
| write_json(CANDIDATE_SUMMARY,summary);return summary |
| def validate_candidates()->dict[str,Any]: |
| rows=read_jsonl(CANDIDATES) |
| if len(rows)<2000:raise AssertionError("candidate pool cannot support 2,000 successes even with zero failures") |
| if len({r["song_id"] for r in rows})!=len(rows):raise AssertionError("duplicate song_id") |
| ordered=[(r["rank_sha256"],r["song_id"]) for r in rows] |
| if ordered!=sorted(ordered):raise AssertionError("not rank sorted") |
| albums=Counter(r["album_cap_key"] for r in rows) |
| if max(albums.values(),default=0)>1:raise AssertionError("album cap violated") |
| artists=Counter(a for r in rows for a in r["artist_cap_keys"]) |
| if max(artists.values(),default=0)>2:raise AssertionError("artist cap violated") |
| for i,r in enumerate(rows): |
| if r["candidate_index"]!=i:raise AssertionError("candidate index mismatch") |
| if stable_rank(r["song_id"])!=r["rank_sha256"]:raise AssertionError("rank mismatch") |
| if not 45<=int(r["metadata_duration_seconds"])<=360:raise AssertionError("duration filter violated") |
| if r["url"]!="https://www.youtube.com/watch?v="+r["song_id"]:raise AssertionError("URL mismatch") |
| out={"status":"PASS","candidate_count":len(rows),"unique_albums":len(albums),"unique_artist_keys":len(artists), |
| "max_artist_uses":max(artists.values(),default=0),"manifest_sha256":sha256_file(CANDIDATES),"bytes":CANDIDATES.stat().st_size} |
| print(json.dumps(out,sort_keys=True));return out |
| def probe(path:Path)->dict[str,Any]: |
| cp=subprocess.run([FFPROBE,"-v","error","-select_streams","a:0","-show_entries", |
| "stream=codec_name,codec_long_name,sample_fmt,sample_rate,channels,channel_layout,bits_per_sample,duration:format=duration,size", |
| "-of","json",str(path)],check=True,capture_output=True,text=True,timeout=60) |
| return json.loads(cp.stdout) |
| def wav_stats(path:Path)->dict[str,Any]: |
| pcm_hash=hashlib.sha256();n=ss=sall=peak=clipped=diff=ls=rs=0 |
| with wave.open(str(path),"rb") as w: |
| channels,rate,width,fc,ctype=w.getnchannels(),w.getframerate(),w.getsampwidth(),w.getnframes(),w.getcomptype() |
| if (channels,rate,width,ctype)!=(2,44100,2,"NONE"):raise ValueError(f"not canonical PCM16 stereo 44.1k: {(channels,rate,width,ctype)}") |
| while data:=w.readframes(65536): |
| pcm_hash.update(data);vals=array.array("h");vals.frombytes(data) |
| if sys.byteorder!="little":vals.byteswap() |
| n+=len(vals) |
| for i in range(0,len(vals),2): |
| l,r=vals[i],vals[i+1];ls+=l;rs+=r;diff+=(l-r)*(l-r) |
| for v in vals: |
| av=abs(v);peak=max(peak,av);clipped+=int(av>=32767);sall+=v;ss+=v*v |
| if n==0:raise ValueError("empty canonical audio") |
| fc=n//2 |
| return {"pcm_sha256":pcm_hash.hexdigest(),"sample_rate":44100,"channels":2,"sample_width_bytes":2,"frame_count":fc, |
| "duration_seconds":fc/44100.0,"peak_linear":peak/32768.0,"rms_linear":math.sqrt(ss/n)/32768.0, |
| "dc_linear":(sall/n)/32768.0,"left_dc_linear":(ls/fc)/32768.0,"right_dc_linear":(rs/fc)/32768.0, |
| "stereo_difference_rms_linear":math.sqrt(diff/fc)/32768.0,"clipped_sample_fraction":clipped/n} |
| def canonicalize(source:Path,dest:Path)->tuple[dict[str,Any],dict[str,Any]]: |
| dest.parent.mkdir(parents=True,exist_ok=True) |
| cp=subprocess.run([FFMPEG,"-nostdin","-hide_banner","-loglevel","error","-i",str(source),"-map","0:a:0","-vn","-sn","-dn", |
| "-ac","2","-ar","44100","-c:a","pcm_s16le","-fflags","+bitexact","-flags:a","+bitexact","-map_metadata","-1","-y",str(dest)], |
| capture_output=True,text=True,timeout=600) |
| if cp.returncode:raise RuntimeError("ffmpeg: "+cp.stderr[-2000:]) |
| return probe(dest),wav_stats(dest) |
| def classify_download_error(error:str)->str: |
| text=(error or "").casefold() |
| systemic=("http error 429","too many requests","rate limit","sign in to confirm","not a bot","authentication required") |
| if any(x in text for x in systemic):return "SYSTEMIC_RATE_LIMIT_OR_AUTH" |
| unavailable=("video unavailable","not available","private video","private","removed","copyright removal","geo","country","http error 403","members-only") |
| if any(x in text for x in unavailable):return "SKIP_UNAVAILABLE" |
| return "DOWNLOAD_ERROR" |
| def ytdlp_download(rec:dict[str,Any],attempts:int=3)->dict[str,Any]: |
| idx,sid=rec["candidate_index"],rec["song_id"];work=TMP/f"{idx:05d}_{sid}" |
| if work.exists():shutil.rmtree(work) |
| work.mkdir(parents=True);output=work/"%(id)s.%(ext)s";error="" |
| for attempt in range(attempts): |
| cp=subprocess.run([str(YTDLP),"--no-playlist","--no-cache-dir","--no-mtime","--quiet","--no-warnings","--socket-timeout","30", |
| "--retries","3","--fragment-retries","3","-f","bestaudio/best","-o",str(output),rec["url"]], |
| capture_output=True,text=True,timeout=900) |
| if cp.returncode==0:break |
| error=(cp.stderr or cp.stdout)[-4000:] |
| if attempt+1<attempts:time.sleep(2**attempt) |
| else: |
| shutil.rmtree(work,ignore_errors=True);category=classify_download_error(error);return {"ok":False,"stage":"skip_unavailable" if category=="SKIP_UNAVAILABLE" else "download","error_category":category,"error":error} |
| files=[p for p in work.iterdir() if p.is_file() and not p.name.endswith((".part",".ytdl",".json"))] |
| if len(files)!=1: |
| names=[p.name for p in files];shutil.rmtree(work,ignore_errors=True);return {"ok":False,"stage":"download_output","error":f"expected 1 media file, got {names}"} |
| raw=files[0];wavp=work/f"{sid}.canonical.wav" |
| try: |
| raw_probe=probe(raw);raw_sha=sha256_file(raw);canonical_probe,stats=canonicalize(raw,wavp);canonical_sha=sha256_file(wavp) |
| if not 45.0<=stats["duration_seconds"]<=360.0:raise ValueError(f"actual duration {stats['duration_seconds']:.6f}s outside [45,360]") |
| if stats["rms_linear"]<=0:raise ValueError("silent canonical audio") |
| raw.unlink() |
| return {"ok":True,"work":str(work),"temp_canonical":str(wavp),"raw_filename":raw.name,"raw_sha256":raw_sha, |
| "canonical_sha256":canonical_sha,"raw_probe":raw_probe,"canonical_probe":canonical_probe,"audio_stats":stats} |
| except Exception as exc: |
| shutil.rmtree(work,ignore_errors=True);return {"ok":False,"stage":"canonicalize_validate","error":repr(exc)} |
| def load_known()->tuple[dict[str,str],dict[str,str]]: |
| if not KNOWN_HASHES.exists():return {},{} |
| obj=json.loads(KNOWN_HASHES.read_text(encoding="utf-8"));by_wav={};by_pcm={} |
| for r in obj["files"]:by_wav[r["canonical_sha256"]]=r["source_label"];by_pcm[r["audio_stats"]["pcm_sha256"]]=r["source_label"] |
| return by_wav,by_pcm |
| def append_results(rows:list[dict[str,Any]])->None: |
| RESULTS.parent.mkdir(parents=True,exist_ok=True) |
| with RESULTS.open("a",encoding="utf-8") as f: |
| for r in rows:f.write(json.dumps(r,sort_keys=True,ensure_ascii=False,separators=(",",":"))+"\n") |
| f.flush();os.fsync(f.fileno()) |
| def systemic_failure(row:dict[str,Any])->bool: |
| if row.get("accepted"):return False |
| text=(str(row.get("error") or "")+" "+str(row.get("stage") or "")).casefold() |
| terms=("http error 429","too many requests","rate limit","sign in to confirm","not a bot","authentication required") |
| return any(x in text for x in terms) |
| def ingest(limit:int,workers:int,smoke:bool)->dict[str,Any]: |
| if workers<1 or workers>4:raise ValueError("workers must be 1..4") |
| candidates=read_jsonl(CANDIDATES);known_wav,known_pcm=load_known() |
| history=read_jsonl(RESULTS) if RESULTS.exists() else [];attempted={x["song_id"] for x in history} |
| accepted=[x for x in history if x.get("accepted")];seen_wav=dict(known_wav);seen_pcm=dict(known_pcm) |
| for x in accepted: |
| final=Path(x["canonical_path"]) |
| if not final.is_file():raise RuntimeError(f"accepted file missing: {final}") |
| seen_wav[x["canonical_sha256"]]=x["song_id"];seen_pcm[x["audio_stats"]["pcm_sha256"]]=x["song_id"] |
| successes=len(accepted);cursor=0;start=time.monotonic();base_successes=successes;base_attempts=len(history) |
| checkpoints=[x for x in (250,500,1000,1500,2000) if x<=limit] |
| while successes<limit: |
| while cursor<len(candidates) and candidates[cursor]["song_id"] in attempted:cursor+=1 |
| if cursor>=len(candidates):raise RuntimeError(f"candidate list exhausted at {successes}/{limit}") |
| next_checkpoint=next((x for x in checkpoints if x>successes),limit) |
| remaining=min(limit-successes,next_checkpoint-successes);batch=[] |
| while cursor<len(candidates) and len(batch)<min(workers,remaining): |
| rec=candidates[cursor];cursor+=1 |
| if rec["song_id"] not in attempted:batch.append(rec);attempted.add(rec["song_id"]) |
| with ThreadPoolExecutor(max_workers=len(batch)) as ex:outcomes=list(ex.map(ytdlp_download,batch)) |
| rows=[] |
| for rec,out in zip(batch,outcomes): |
| base={"candidate_index":rec["candidate_index"],"rank_sha256":rec["rank_sha256"],"song_id":rec["song_id"],"url":rec["url"], |
| "metadata_duration_seconds":rec["metadata_duration_seconds"],"attempted_at_unix":int(time.time()), |
| "retention_policy":"canonical_only_after_validation","raw_retained":False} |
| if not out["ok"]:rows.append({**base,**out,"accepted":False});continue |
| dup=seen_wav.get(out["canonical_sha256"]) or seen_pcm.get(out["audio_stats"]["pcm_sha256"]) |
| temp=Path(out.pop("temp_canonical"));work=Path(out.pop("work")) |
| if dup: |
| shutil.rmtree(work,ignore_errors=True);rows.append({**base,**out,"accepted":False,"stage":"exact_dedup","duplicate_of":dup,"error":"exact canonical WAV or PCM duplicate"});continue |
| dest=CANONICAL/f"{rec['candidate_index']:05d}_{rec['song_id']}.wav";CANONICAL.mkdir(parents=True,exist_ok=True) |
| os.replace(temp,dest);shutil.rmtree(work,ignore_errors=True) |
| seen_wav[out["canonical_sha256"]]=rec["song_id"];seen_pcm[out["audio_stats"]["pcm_sha256"]]=rec["song_id"];successes+=1 |
| rows.append({**base,**out,"accepted":True,"canonical_path":str(dest.resolve()),"canonical_bytes":dest.stat().st_size}) |
| append_results(rows);history.extend(rows) |
| recent=history[-200:] |
| if len(recent)==200 and sum(bool(x.get("accepted")) for x in recent)<100: |
| stop={"status":"STOPPED","reason":"rolling_200_acceptance_below_50_percent","accepted_total":successes, |
| "attempted_total":len(history),"rolling_200_accepted":sum(bool(x.get("accepted")) for x in recent)} |
| write_json(ROOT/"manifests/tranche.stopped.json",stop);raise RuntimeError(json.dumps(stop,sort_keys=True)) |
| consecutive=0 |
| for x in reversed(history): |
| if systemic_failure(x):consecutive+=1 |
| else:break |
| tail=history[-20:];systemic_tail=sum(systemic_failure(x) for x in tail) |
| if consecutive>=8 or (len(tail)==20 and systemic_tail>=16 and not any(x.get("accepted") for x in tail)): |
| stop={"status":"STOPPED","reason":"systemic_rate_limit_or_auth_failure","accepted_total":successes, |
| "attempted_total":len(history),"consecutive_systemic_failures":consecutive,"systemic_failures_last_20":systemic_tail} |
| write_json(ROOT/"manifests/tranche.stopped.json",stop);raise RuntimeError(json.dumps(stop,sort_keys=True)) |
| if successes in checkpoints: |
| elapsed=max(time.monotonic()-start,1e-9);new_successes=successes-base_successes;rate_s=new_successes/elapsed |
| checkpoint={"event":"checkpoint","accepted_total":successes,"attempted_total":len(history), |
| "new_accepted":new_successes,"new_attempts":len(history)-base_attempts, |
| "acceptance_rate_total":successes/len(history),"accepted_per_minute":rate_s*60, |
| "eta_seconds":(limit-successes)/rate_s if rate_s>0 else None, |
| "failure_counts":dict(Counter(x.get("stage","unknown") for x in history if not x.get("accepted"))), |
| "canonical_bytes":sum(p.stat().st_size for p in CANONICAL.glob("*.wav")),"elapsed_seconds":elapsed} |
| with (ROOT/"logs/checkpoints.jsonl").open("a",encoding="utf-8") as f:f.write(json.dumps(checkpoint,sort_keys=True)+"\n") |
| print(json.dumps(checkpoint,sort_keys=True),flush=True) |
| allrows=read_jsonl(RESULTS) |
| summary={"status":"PASS","mode":"smoke" if smoke else "tranche","target_successes":limit, |
| "accepted_total":sum(bool(x.get("accepted")) for x in allrows),"attempted_total":len(allrows), |
| "failure_counts":dict(Counter(x.get("stage","accepted") for x in allrows if not x.get("accepted"))), |
| "results_sha256":sha256_file(RESULTS),"canonical_bytes":sum(p.stat().st_size for p in CANONICAL.glob("*.wav")), |
| "known_hash_count":len(known_pcm),"workers":workers} |
| write_json(ROOT/("manifests/smoke.summary.json" if smoke else "manifests/tranche.summary.json"),summary);return summary |
| def build_known(inputs_manifest:Path)->dict[str,Any]: |
| obj=json.loads(inputs_manifest.read_text(encoding="utf-8"));rows=[];work=TMP/"known_hashes" |
| if work.exists():shutil.rmtree(work) |
| work.mkdir(parents=True) |
| try: |
| for i,item in enumerate(obj["files"]): |
| source=Path(item["path"]) |
| if not source.is_file():raise FileNotFoundError(source) |
| dest=work/f"{i:03d}.wav";canonical_probe,stats=canonicalize(source,dest) |
| rows.append({"source_label":item["source_label"],"raw_sha256":sha256_file(source),"canonical_sha256":sha256_file(dest), |
| "canonical_probe":canonical_probe,"audio_stats":stats}) |
| finally:shutil.rmtree(work,ignore_errors=True) |
| if len({x["source_label"] for x in rows})!=len(rows):raise AssertionError("duplicate source label") |
| out={"schema_version":1,"pipeline":"ffmpeg -ac 2 -ar 44100 -c:a pcm_s16le; metadata stripped; bitexact flags","file_count":len(rows),"files":rows} |
| write_json(KNOWN_HASHES,out);out["manifest_sha256"]=sha256_file(KNOWN_HASHES);return out |
| def main()->None: |
| ap=argparse.ArgumentParser();sub=ap.add_subparsers(dest="cmd",required=True) |
| b=sub.add_parser("build-candidates");b.add_argument("--pool-size",type=int,default=250000);b.add_argument("--target",type=int,default=20000) |
| sub.add_parser("validate-candidates");k=sub.add_parser("build-known-hashes");k.add_argument("inputs_manifest",type=Path) |
| s=sub.add_parser("smoke");s.add_argument("--successes",type=int,default=3) |
| r=sub.add_parser("run");r.add_argument("--successes",type=int,default=2000);r.add_argument("--workers",type=int,default=4) |
| a=ap.parse_args() |
| if a.cmd=="build-candidates":print(json.dumps(build_candidates(a.pool_size,a.target),indent=2,sort_keys=True)) |
| elif a.cmd=="validate-candidates":validate_candidates() |
| elif a.cmd=="build-known-hashes":print(json.dumps(build_known(a.inputs_manifest),indent=2,sort_keys=True)) |
| elif a.cmd=="smoke":print(json.dumps(ingest(a.successes,min(a.successes,3),True),indent=2,sort_keys=True)) |
| elif a.cmd=="run":print(json.dumps(ingest(a.successes,a.workers,False),indent=2,sort_keys=True)) |
| if __name__=="__main__":main() |
|
|
|
|
