import io import json import math import os import shutil import sys import time import warnings from bisect import bisect_left from concurrent.futures import ThreadPoolExecutor from pathlib import Path from codec.audio_processing.dune_codec import load_dune_audio_tokenizer import librosa import numpy as np import pyarrow as pa import pyarrow.compute as pc import torch import torch.nn.functional as F from datasets import Audio, Dataset, DatasetDict, Features, Sequence, Value, load_dataset, load_from_disk from tqdm.auto import tqdm warnings.filterwarnings("ignore") MODEL_ID = "Respair/dune_codec" # contains the metadata rows to enrich and a unique bridge key. DATASET_SOURCE = "/home/ubuntu/data" DATASET_CONFIG = None DATASET_SPLIT = "train" # if your metadata and audio dataset are the same, copy DATASET_SOURCE here, otherwise ensure you have a index column (like `key`) so these two can be joined AUDIO_DATASET_SOURCE = ["/home/ubuntu/data"] AUDIO_DATASET_CONFIG = None AUDIO_DATASET_SPLIT = "train" KEY_COLUMN = "key" AUDIO_COLUMN = "audio" # output columns DURATION_COLUMN = "duration" PREQUANT_COLUMN = "latents" DISCRETE_TOKENS_COLUMN = "codes" ERROR_COLUMN = "codec_error" CHECKPOINT_DIR = "/home/ubuntu/out" CHECKPOINT_INTERVAL_STEPS = 1024 * 1 HF_MAX_SHARD_SIZE = "2GB" # What to save? discrete speech tokens, FSQ pre-quant latents or both? ## for darya or any other flow matching models, FSQ pre-quant latents is enough SAVE_PREQUANT = True SAVE_DISCRETE = True TEST_RUN = 0 # you probably won't need it BATCH_SIZE = 64 NUM_WORKERS = 32 TARGET_SAMPLE_RATE = 22050 INFERENCE_DTYPE = torch.bfloat16 MAX_AUDIO_DURATION_SECONDS = 30.0 # if your clips are capped at another length, tweak it here. # you can ignore these MIN_DURATION_BUCKETS = 10 MAX_DURATION_BUCKETS = 15 DURATION_HISTOGRAM_RESOLUTION_SECONDS = 0.05 DURATION_BUCKET_ELBOW_TOLERANCE = 0.03 ARROW_DURATION_SCAN_BATCH_SIZE = 250_000 DURATION_BUCKET_INDEX_SCAN_BATCH_SIZE = 250_000 AUDIO_INDEX_BATCH_SIZE = 100_000 EXPECTED_PREQUANT_DIM = 52 def load_dataset_source(source, config, split): source_path = Path(source).expanduser() if source_path.exists(): dataset = load_from_disk(str(source_path)) if isinstance(dataset, DatasetDict): if split not in dataset: raise KeyError( f"Split '{split}' is not present in local dataset " f"'{source_path}'. Available splits: {list(dataset.keys())}" ) dataset = dataset[split] return dataset return load_dataset( source, config, split=split, streaming=False, ) def load_input_dataset(): dataset = load_dataset_source( source=DATASET_SOURCE, config=DATASET_CONFIG, split=DATASET_SPLIT, ) if KEY_COLUMN not in dataset.column_names: raise KeyError( f"Primary dataset does not contain the bridge column " f"'{KEY_COLUMN}'. Available columns: {dataset.column_names}" ) if AUDIO_COLUMN in dataset.column_names: dataset = dataset.remove_columns([AUDIO_COLUMN]) if TEST_RUN > 0: dataset = dataset.select(range(min(TEST_RUN, len(dataset)))) return dataset def load_audio_datasets(): sources = ( list(AUDIO_DATASET_SOURCE) if isinstance(AUDIO_DATASET_SOURCE, (list, tuple)) else [AUDIO_DATASET_SOURCE] ) audio_datasets = [] for source in sources: dataset = load_dataset_source( source=source, config=AUDIO_DATASET_CONFIG, split=AUDIO_DATASET_SPLIT, ) required_columns = {KEY_COLUMN, AUDIO_COLUMN} missing_columns = required_columns.difference(dataset.column_names) if missing_columns: raise KeyError( f"Audio dataset '{source}' is missing required columns: " f"{sorted(missing_columns)}. " f"Available columns: {dataset.column_names}" ) audio_datasets.append( dataset.cast_column( AUDIO_COLUMN, Audio(sampling_rate=TARGET_SAMPLE_RATE), ) ) return audio_datasets def normalize_bridge_key(key, dataset_role): if key is None: raise ValueError( f"{dataset_role} row has a null '{KEY_COLUMN}' value." ) if isinstance(key, str): return key if isinstance(key, (int, np.integer)) and not isinstance(key, (bool, np.bool_)): return str(int(key)) raise TypeError( f"{dataset_role} '{KEY_COLUMN}' value must be a string or integer, " f"received {type(key).__name__}: {key!r}" ) def build_audio_key_index(audio_datasets): key_to_position = {} for dataset_position, audio_dataset in enumerate(audio_datasets): key_dataset = audio_dataset.select_columns([KEY_COLUMN]) number_of_batches = math.ceil( len(key_dataset) / AUDIO_INDEX_BATCH_SIZE ) progress = tqdm( key_dataset.iter(batch_size=AUDIO_INDEX_BATCH_SIZE), total=number_of_batches, desc=f"Indexing audio keys [{dataset_position + 1}/{len(audio_datasets)}]", unit="batch", dynamic_ncols=True, ) audio_position = 0 for key_batch in progress: for raw_key in key_batch[KEY_COLUMN]: key = normalize_bridge_key(raw_key, "Audio dataset") if key in key_to_position: previous_dataset_position, previous_audio_position = ( key_to_position[key] ) raise ValueError( f"Duplicate normalized audio key {key!r}: found in " f"audio dataset {previous_dataset_position} at row " f"{previous_audio_position} and audio dataset " f"{dataset_position} at row {audio_position}." ) key_to_position[key] = (dataset_position, audio_position) audio_position += 1 return key_to_position # audio i/o and batchh def decode_audio(audio_value, target_sample_rate): if isinstance(audio_value, dict) and audio_value.get("array") is not None: wav = np.asarray(audio_value["array"], dtype=np.float32) sample_rate = int(audio_value["sampling_rate"]) elif hasattr(audio_value, "get_all_samples"): samples = audio_value.get_all_samples() wav = samples.data.cpu().numpy().astype(np.float32, copy=False) sample_rate = int(samples.sample_rate) else: source = audio_value if isinstance(audio_value, dict): if audio_value.get("bytes") is not None: source = io.BytesIO(audio_value["bytes"]) else: source = audio_value.get("path") elif isinstance(audio_value, (bytes, bytearray)): source = io.BytesIO(audio_value) if source is None: raise ValueError("Audio value has no array, bytes, or path.") wav, sample_rate = librosa.load( source, sr=None, mono=True, ) wav = np.asarray(wav, dtype=np.float32) if wav.ndim > 1: wav = wav.mean(axis=0) if sample_rate != target_sample_rate: wav = librosa.resample( wav, orig_sr=sample_rate, target_sr=target_sample_rate, ).astype(np.float32, copy=False) if len(wav) == 0: raise ValueError("Decoded audio is empty.") return wav def prepare_audio_item(item): position, row, audio_datasets, audio_key_index, sample_rate = item row = dict(row) raw_key = row.get(KEY_COLUMN) try: key = normalize_bridge_key(raw_key, "Primary") if key not in audio_key_index: raise KeyError( f"No matching audio row was found for key {key!r}." ) dataset_position, audio_position = audio_key_index[key] audio_value = audio_datasets[dataset_position][audio_position][AUDIO_COLUMN] wav = decode_audio(audio_value, sample_rate) return position, row, torch.from_numpy(wav).float(), None except Exception as exc: return ( position, row, None, f"ERROR: audio lookup or decoding failed - {exc}", ) def get_duration_bucket_index(num_samples, bucket_boundaries_samples): bucket_index = bisect_left( bucket_boundaries_samples, num_samples, ) if bucket_index == len(bucket_boundaries_samples): return None return bucket_index def build_metadata_duration_histogram( metadata_dataset, sample_rate, histogram_bin_samples, ): if DURATION_COLUMN not in metadata_dataset.column_names: raise KeyError( f"Primary dataset does not contain duration column " f"'{DURATION_COLUMN}'. Available columns: " f"{metadata_dataset.column_names}" ) total_rows = len(metadata_dataset) if total_rows == 0: raise ValueError( "Cannot calculate duration buckets from an empty dataset." ) duration_table = metadata_dataset.select_columns( [DURATION_COLUMN] ).data duration_column = duration_table.column(DURATION_COLUMN) if not isinstance( duration_column, (pa.Array, pa.ChunkedArray), ): raise TypeError( f"Expected an Arrow Array or ChunkedArray for " f"'{DURATION_COLUMN}', received {type(duration_column)}." ) maximum_histogram_samples = int( math.ceil(MAX_AUDIO_DURATION_SECONDS * sample_rate) ) histogram_size = ( (maximum_histogram_samples - 1) // histogram_bin_samples ) + 1 counts = np.zeros(histogram_size, dtype=np.int64) sample_sums = np.zeros(histogram_size, dtype=np.float64) bin_maxima = np.zeros(histogram_size, dtype=np.int64) observed_max_samples = 0 processed_rows = 0 progress = tqdm( total=total_rows, desc="Scanning Arrow durations", unit="rows", dynamic_ncols=True, ) for batch_start in range( 0, total_rows, ARROW_DURATION_SCAN_BATCH_SIZE, ): batch_length = min( ARROW_DURATION_SCAN_BATCH_SIZE, total_rows - batch_start, ) arrow_batch = duration_column.slice( batch_start, batch_length, ) if isinstance(arrow_batch, pa.ChunkedArray): arrow_batch = arrow_batch.combine_chunks() if arrow_batch.null_count: null_mask = pc.is_null(arrow_batch).to_numpy( zero_copy_only=False ) local_indices = np.flatnonzero(null_mask)[:3] global_indices = ( local_indices + batch_start ).tolist() raise ValueError( f"Duration column '{DURATION_COLUMN}' contains null " f"values at rows {global_indices}." ) try: float_batch = pc.cast( arrow_batch, pa.float64(), safe=False, ) except ( pa.ArrowInvalid, pa.ArrowNotImplementedError, TypeError, ) as exc: raise TypeError( f"Duration column '{DURATION_COLUMN}' cannot be cast " f"to float64 from Arrow type {arrow_batch.type}." ) from exc duration_seconds = float_batch.to_numpy( zero_copy_only=False ) invalid_mask = ( ~np.isfinite(duration_seconds) | (duration_seconds <= 0.0) ) if invalid_mask.any(): local_indices = np.flatnonzero(invalid_mask)[:3] global_indices = ( local_indices + batch_start ).tolist() invalid_values = duration_seconds[ local_indices ].tolist() raise ValueError( f"Duration column '{DURATION_COLUMN}' contains invalid " f"values at rows {global_indices}: {invalid_values}." ) batch_max_seconds = float(duration_seconds.max()) if batch_max_seconds > MAX_AUDIO_DURATION_SECONDS: local_index = int(np.argmax(duration_seconds)) raise ValueError( f"Metadata duration column contains " f"{batch_max_seconds:.3f}s at row " f"{batch_start + local_index}, which exceeds " f"MAX_AUDIO_DURATION_SECONDS=" f"{MAX_AUDIO_DURATION_SECONDS}." ) duration_samples = np.ceil( duration_seconds * sample_rate ).astype(np.int64, copy=False) batch_max_samples = int(duration_samples.max()) observed_max_samples = max( observed_max_samples, batch_max_samples, ) bin_ids = ( duration_samples - 1 ) // histogram_bin_samples counts += np.bincount( bin_ids, minlength=histogram_size, ).astype(np.int64, copy=False) sample_sums += np.bincount( bin_ids, weights=duration_samples, minlength=histogram_size, ).astype(np.float64, copy=False) np.maximum.at( bin_maxima, bin_ids, duration_samples, ) processed_rows += batch_length progress.update(batch_length) progress.set_postfix( max_s=f"{observed_max_samples / sample_rate:.3f}", refresh=False, ) progress.close() if processed_rows != total_rows: raise RuntimeError( f"Arrow duration scan processed {processed_rows} rows, " f"but the metadata dataset contains {total_rows} rows." ) occupied = counts > 0 return ( counts[occupied], sample_sums[occupied], bin_maxima[occupied], observed_max_samples, ) def optimize_bucket_boundaries_for_count( counts, sample_sums, bin_maxima, bucket_count, max_allowed_samples, ): occupied_bins = len(counts) if bucket_count > occupied_bins: raise ValueError( f"Cannot create {bucket_count} non-empty duration buckets from " f"only {occupied_bins} occupied duration bins." ) prefix_counts = np.concatenate( ([0], np.cumsum(counts, dtype=np.int64)) ) prefix_sums = np.concatenate( ([0.0], np.cumsum(sample_sums, dtype=np.float64)) ) costs = np.full( (bucket_count + 1, occupied_bins), np.inf, dtype=np.float64, ) backpointers = np.full( (bucket_count + 1, occupied_bins), -1, dtype=np.int32, ) for end_index in range(occupied_bins): ceiling = ( max_allowed_samples if bucket_count == 1 and end_index == occupied_bins - 1 else int(bin_maxima[end_index]) ) segment_count = prefix_counts[end_index + 1] segment_sum = prefix_sums[end_index + 1] costs[1, end_index] = ( ceiling * segment_count - segment_sum ) for used_buckets in range(2, bucket_count + 1): for end_index in range(used_buckets - 1, occupied_bins): starts = np.arange( used_buckets - 1, end_index + 1, dtype=np.int64, ) previous_costs = costs[ used_buckets - 1, starts - 1, ] ceiling = ( max_allowed_samples if ( used_buckets == bucket_count and end_index == occupied_bins - 1 ) else int(bin_maxima[end_index]) ) segment_counts = ( prefix_counts[end_index + 1] - prefix_counts[starts] ) segment_sums = ( prefix_sums[end_index + 1] - prefix_sums[starts] ) candidate_costs = previous_costs + ( ceiling * segment_counts - segment_sums ) best_offset = int(np.argmin(candidate_costs)) costs[used_buckets, end_index] = candidate_costs[ best_offset ] backpointers[used_buckets, end_index] = int( starts[best_offset] ) end_index = occupied_bins - 1 endpoints = [] for used_buckets in range(bucket_count, 0, -1): endpoints.append(end_index) if used_buckets == 1: break start_index = int( backpointers[used_buckets, end_index] ) end_index = start_index - 1 endpoints.reverse() boundaries = [ int(bin_maxima[index]) for index in endpoints ] boundaries[-1] = max_allowed_samples return tuple(boundaries), float( costs[bucket_count, occupied_bins - 1] ) def calculate_duration_bucket_boundaries( metadata_dataset, sample_rate, ): histogram_bin_samples = max( 1, int( round( DURATION_HISTOGRAM_RESOLUTION_SECONDS * sample_rate ) ), ) ( counts, sample_sums, bin_maxima, max_allowed_samples, ) = build_metadata_duration_histogram( metadata_dataset=metadata_dataset, sample_rate=sample_rate, histogram_bin_samples=histogram_bin_samples, ) occupied_bins = len(counts) maximum_candidate_buckets = min( MAX_DURATION_BUCKETS, occupied_bins, ) minimum_candidate_buckets = min( MIN_DURATION_BUCKETS, maximum_candidate_buckets, ) candidates = [] actual_sample_total = float(sample_sums.sum()) bucket_counts_to_test = range( minimum_candidate_buckets, maximum_candidate_buckets + 1, ) optimization_progress = tqdm( bucket_counts_to_test, total=( maximum_candidate_buckets - minimum_candidate_buckets + 1 ), desc="Optimizing duration buckets", unit="candidate", dynamic_ncols=True, ) for bucket_count in optimization_progress: boundaries, padding_cost = ( optimize_bucket_boundaries_for_count( counts=counts, sample_sums=sample_sums, bin_maxima=bin_maxima, bucket_count=bucket_count, max_allowed_samples=max_allowed_samples, ) ) padding_ratio = padding_cost / actual_sample_total candidates.append( ( bucket_count, boundaries, padding_cost, padding_ratio, ) ) best_padding_cost = min( candidate[2] for candidate in candidates ) allowed_padding_cost = best_padding_cost * ( 1.0 + DURATION_BUCKET_ELBOW_TOLERANCE ) selected = next( candidate for candidate in candidates if candidate[2] <= allowed_padding_cost ) print("Duration bucket optimization:") for bucket_count, _, _, padding_ratio in candidates: marker = " *" if bucket_count == selected[0] else "" print( f" {bucket_count:2d} buckets: " f"estimated padding {padding_ratio * 100:.3f}%{marker}" ) boundaries_samples = selected[1] boundaries_seconds = tuple( boundary / sample_rate for boundary in boundaries_samples ) histogram_bucket_assignments = np.searchsorted( np.asarray(boundaries_samples, dtype=np.int64), bin_maxima, side="left", ) bucket_counts = np.bincount( histogram_bucket_assignments, weights=counts, minlength=len(boundaries_samples), ).astype(np.int64, copy=False) print( f"Selected {len(boundaries_samples)} duration buckets from " f"all {int(counts.sum()):,} metadata durations." ) for index, (ceiling, count) in enumerate( zip(boundaries_seconds, bucket_counts), start=1, ): print( f" Bucket {index:02d}: <= {ceiling:.3f}s " f"({int(count):,} metadata rows)" ) return boundaries_samples, boundaries_seconds def get_metadata_rows_by_indices(metadata_dataset, row_indices): """Read a small random-access metadata batch and return row dictionaries.""" indices = [int(index) for index in row_indices] columnar_batch = metadata_dataset[indices] return [ { column_name: columnar_batch[column_name][position] for column_name in metadata_dataset.column_names } for position in range(len(indices)) ] def build_duration_bucket_indices( metadata_dataset, sample_rate, bucket_boundaries_samples, ): duration_column = metadata_dataset.select_columns( [DURATION_COLUMN] ).data.column(DURATION_COLUMN) total_rows = len(metadata_dataset) boundaries = np.asarray( bucket_boundaries_samples, dtype=np.int64, ) index_chunks_by_bucket = [ [] for _ in bucket_boundaries_samples ] bucket_counts = np.zeros( len(bucket_boundaries_samples), dtype=np.int64, ) progress = tqdm( total=total_rows, desc="Assigning global duration buckets", unit="rows", dynamic_ncols=True, ) for batch_start in range( 0, total_rows, DURATION_BUCKET_INDEX_SCAN_BATCH_SIZE, ): batch_length = min( DURATION_BUCKET_INDEX_SCAN_BATCH_SIZE, total_rows - batch_start, ) arrow_batch = duration_column.slice( batch_start, batch_length, ) if isinstance(arrow_batch, pa.ChunkedArray): arrow_batch = arrow_batch.combine_chunks() if arrow_batch.null_count: raise ValueError( f"Duration column '{DURATION_COLUMN}' contains null values " f"while assigning global buckets near row {batch_start}." ) duration_seconds = pc.cast( arrow_batch, pa.float64(), safe=False, ).to_numpy(zero_copy_only=False) duration_samples = np.ceil( duration_seconds * sample_rate ).astype(np.int64, copy=False) assignments = np.searchsorted( boundaries, duration_samples, side="left", ) invalid_mask = assignments >= len(boundaries) if invalid_mask.any(): local_index = int(np.flatnonzero(invalid_mask)[0]) raise ValueError( f"Duration at metadata row {batch_start + local_index} " f"exceeds the final optimized bucket ceiling." ) global_indices = np.arange( batch_start, batch_start + batch_length, dtype=np.int64, ) for bucket_index in np.unique(assignments): bucket_mask = assignments == bucket_index bucket_indices = global_indices[bucket_mask] index_chunks_by_bucket[int(bucket_index)].append( bucket_indices ) bucket_counts[int(bucket_index)] += len(bucket_indices) progress.update(batch_length) progress.set_postfix( assigned=batch_start + batch_length, refresh=False, ) progress.close() bucket_indices = [] for chunks in index_chunks_by_bucket: if chunks: bucket_indices.append(np.concatenate(chunks)) else: bucket_indices.append(np.empty(0, dtype=np.int64)) if int(bucket_counts.sum()) != total_rows: raise RuntimeError( f"Assigned {int(bucket_counts.sum())} rows to duration buckets, " f"but the metadata dataset contains {total_rows} rows." ) print("Global duration bucket membership:") for bucket_index, (boundary, count) in enumerate( zip(bucket_boundaries_samples, bucket_counts), start=1, ): print( f" Bucket {bucket_index:02d}: <= " f"{boundary / sample_rate:.3f}s " f"({int(count):,} rows, " f"{math.ceil(int(count) / BATCH_SIZE):,} batches)" ) return bucket_indices, bucket_counts def prepare_global_bucket_batch( rows, audio_pool, audio_datasets, audio_key_index, sample_rate, padded_length, bucket_upper_seconds, ): items = ( ( position, row, audio_datasets, audio_key_index, sample_rate, ) for position, row in enumerate(rows) ) prepared = list(audio_pool.map(prepare_audio_item, items)) batch_rows = [None] * len(prepared) errors = [None] * len(prepared) audio_by_position = {} for position, row, audio, error in prepared: batch_rows[position] = row errors[position] = error if audio is None: continue if len(audio) > padded_length: errors[position] = ( f"ERROR: decoded audio length is " f"{len(audio) / sample_rate:.3f}s, which exceeds its " f"metadata-assigned bucket ceiling of " f"{bucket_upper_seconds:.3f}s" ) continue audio_by_position[position] = audio valid_positions = sorted(audio_by_position) if not valid_positions: unique_errors = [] for error in errors: if error is not None and error not in unique_errors: unique_errors.append(error) if len(unique_errors) == 3: break details = " | ".join(unique_errors) or "No error detail was returned." raise RuntimeError( "Every row in the global duration-bucket batch failed before " f"codec extraction. First errors: {details}" ) audio_tensor, audio_lens = build_padded_audio_batch( audio_by_position=audio_by_position, positions=valid_positions, padded_length=padded_length, ) return { "rows": batch_rows, "errors": errors, "valid_positions": valid_positions, "audio_tensor": audio_tensor, "audio_lens": audio_lens, } def build_padded_audio_batch( audio_by_position, positions, padded_length, ): audio_list = [audio_by_position[position] for position in positions] original_lengths = [len(audio) for audio in audio_list] padded_batch = [] for audio in audio_list: pad_len = padded_length - len(audio) if pad_len < 0: raise ValueError( f"Audio length {len(audio)} exceeds its bucket padding " f"length {padded_length}." ) if pad_len > 0: audio = F.pad( audio, (0, pad_len), mode="constant", value=0, ) padded_batch.append(audio) return ( torch.stack(padded_batch, dim=0), torch.tensor(original_lengths, dtype=torch.long), ) @torch.no_grad() def encode_batch(model, pre_quantized, latent_len): encoded_tokens = model.codec.vector_quantizer.encode( inputs=pre_quantized, input_len=latent_len, ) trimmed_codes = [] for idx in range(pre_quantized.shape[0]): tlen = int(latent_len[idx].item()) trimmed_codes.append(encoded_tokens[:, idx, :tlen]) return trimmed_codes def get_audio_encoder(model): if hasattr(model, "audio_encoder"): return model.audio_encoder if hasattr(model, "codec") and hasattr(model.codec, "audio_encoder"): return model.codec.audio_encoder raise AttributeError("Could not find audio_encoder on codec model.") def parse_audio_encoder_output(enc_out, batch_size): pre_quantized = None latent_len = None if torch.is_tensor(enc_out): pre_quantized = enc_out elif isinstance(enc_out, tuple): pre_quantized = enc_out[0] for item in enc_out[1:]: if ( torch.is_tensor(item) and item.dim() == 1 and item.numel() == batch_size ): latent_len = item break elif isinstance(enc_out, dict): pre_quantized = ( enc_out.get("encoded") or enc_out.get("latents") or enc_out.get("encoder_out") or enc_out.get("pre_quantized") ) latent_len = ( enc_out.get("encoded_len") or enc_out.get("latents_len") or enc_out.get("encoder_out_len") or enc_out.get("lengths") or enc_out.get("audio_codes_len") ) else: raise ValueError( f"Unsupported audio_encoder output type: {type(enc_out)}" ) if pre_quantized is None: raise ValueError( "Could not locate pre_quantized latents in audio_encoder output." ) return pre_quantized, latent_len @torch.no_grad() def extract_prequant_batch(model, wav_batch, audio_lens, device): wav_batch = wav_batch.to( device=device, dtype=INFERENCE_DTYPE, non_blocking=True, ) audio_lens = audio_lens.to(device, non_blocking=True) encoder = get_audio_encoder(model) enc_out = encoder(audio=wav_batch, audio_len=audio_lens) pre_quantized, latent_len = parse_audio_encoder_output( enc_out, wav_batch.shape[0], ) if pre_quantized.dim() != 3: raise ValueError( f"Expected 3D pre_quantized tensor, got shape " f"{tuple(pre_quantized.shape)}" ) if latent_len is None: raise ValueError( "The codec audio encoder did not return latent lengths." ) time_major_latents = pre_quantized.transpose(1, 2) trimmed_latents = [] for idx in range(time_major_latents.shape[0]): tlen = int(latent_len[idx].item()) trimmed_latents.append(time_major_latents[idx, :tlen, :]) return trimmed_latents, pre_quantized, latent_len #checkpoint def build_checkpoint_features(metadata_dataset): generated_columns = {ERROR_COLUMN} if SAVE_PREQUANT: generated_columns.add(PREQUANT_COLUMN) if SAVE_DISCRETE: generated_columns.add(DISCRETE_TOKENS_COLUMN) conflicting_columns = generated_columns.intersection( metadata_dataset.column_names ) if conflicting_columns: raise KeyError( f"Primary dataset already contains generated output columns: " f"{sorted(conflicting_columns)}" ) features = dict(metadata_dataset.features) features.pop(AUDIO_COLUMN, None) if SAVE_PREQUANT: features[PREQUANT_COLUMN] = Sequence( Sequence(Value("float16")) ) if SAVE_DISCRETE: features[DISCRETE_TOKENS_COLUMN] = Sequence( Sequence(Value("int64")) ) features[ERROR_COLUMN] = Value("string") return Features(features) def find_saved_state(): checkpoint_root = Path(CHECKPOINT_DIR) checkpoint_root.mkdir(parents=True, exist_ok=True) shard_paths = sorted( path for path in checkpoint_root.glob("shard-*") if path.is_dir() ) processed_rows = 0 checkpoint_progress = tqdm( shard_paths, total=len(shard_paths), desc="Scanning checkpoint shards", unit="shard", dynamic_ncols=True, disable=not shard_paths, ) for path in checkpoint_progress: processed_rows += len(load_from_disk(str(path))) checkpoint_progress.set_postfix( rows=processed_rows, refresh=False, ) state_path = checkpoint_root / "state.json" if not state_path.exists(): if shard_paths: raise RuntimeError( "Checkpoint shards exist without state.json. Global bucket " "resume requires both; clear CHECKPOINT_DIR for a fresh run." ) return 0, 0, 0, None, 0, 0 with state_path.open("r", encoding="utf-8") as handle: state = json.load(handle) if state.get("processing_order") != "global_duration_buckets": raise RuntimeError( "The checkpoint state was not created by the global " "duration-bucket pipeline. Clear CHECKPOINT_DIR before running." ) saved_processed_rows = int(state["processed_rows"]) if saved_processed_rows != processed_rows: raise RuntimeError( f"state.json reports {saved_processed_rows} processed rows, but " f"checkpoint shards contain {processed_rows}." ) saved_bucket_boundaries = tuple( int(boundary) for boundary in state["duration_bucket_boundaries_samples"] ) return ( len(shard_paths), processed_rows, int(state["processed_steps"]), saved_bucket_boundaries, int(state["current_bucket_index"]), int(state["bucket_offset"]), ) def save_state( processed_rows, processed_steps, duration_bucket_boundaries_samples, current_bucket_index, bucket_offset, ): checkpoint_root = Path(CHECKPOINT_DIR) state_path = checkpoint_root / "state.json" temp_path = checkpoint_root / ".state.json.tmp" with temp_path.open("w", encoding="utf-8") as handle: json.dump( { "processing_order": "global_duration_buckets", "processed_rows": processed_rows, "processed_steps": processed_steps, "current_bucket_index": current_bucket_index, "bucket_offset": bucket_offset, "duration_bucket_boundaries_samples": list( duration_bucket_boundaries_samples ), "duration_bucket_source": f"metadata:{DURATION_COLUMN}", }, handle, ) os.replace(temp_path, state_path) def save_shard( rows, shard_index, processed_rows, processed_steps, checkpoint_features, duration_bucket_boundaries_samples, current_bucket_index, bucket_offset, ): if not rows: return shard_index checkpoint_root = Path(CHECKPOINT_DIR) final_path = checkpoint_root / f"shard-{shard_index:06d}" temp_path = checkpoint_root / f".shard-{shard_index:06d}.tmp" if final_path.exists(): raise FileExistsError( f"Checkpoint shard already exists: {final_path}" ) if temp_path.exists(): shutil.rmtree(temp_path) Dataset.from_list( rows, features=checkpoint_features, ).save_to_disk( str(temp_path), max_shard_size=HF_MAX_SHARD_SIZE, ) os.replace(temp_path, final_path) save_state( processed_rows=processed_rows, processed_steps=processed_steps, duration_bucket_boundaries_samples=( duration_bucket_boundaries_samples ), current_bucket_index=current_bucket_index, bucket_offset=bucket_offset, ) print(f"Saved {len(rows)} rows to {final_path}") return shard_index + 1 def process_global_bucket_batch( model, batch, device, bucket_upper_seconds, ): rows = batch["rows"] errors = list(batch["errors"]) valid_positions = batch["valid_positions"] discrete_values = [None] * len(rows) prequant_values = [None] * len(rows) dim_skipped = 0 try: with torch.autocast( device_type="cuda", dtype=INFERENCE_DTYPE, ): ( trimmed_latents, pre_quantized, latent_len, ) = extract_prequant_batch( model, batch["audio_tensor"], batch["audio_lens"], device, ) trimmed_codes = None if SAVE_DISCRETE: trimmed_codes = encode_batch( model, pre_quantized, latent_len, ) for model_position, row_position in enumerate(valid_positions): if SAVE_PREQUANT and trimmed_latents is not None: latent = trimmed_latents[model_position] if latent.shape[-1] != EXPECTED_PREQUANT_DIM: errors[row_position] = ( f"ERROR: prequant dim is {latent.shape[-1]}, " f"expected {EXPECTED_PREQUANT_DIM}" ) dim_skipped += 1 continue prequant_values[row_position] = ( latent.cpu() .to(torch.float16) .numpy() .tolist() ) if SAVE_DISCRETE and trimmed_codes is not None: discrete_values[row_position] = ( trimmed_codes[model_position] .cpu() .to(torch.int64) .numpy() .tolist() ) except Exception as exc: raise RuntimeError( f"Codec extraction failed for the " f"{bucket_upper_seconds:.3f}s global duration bucket: {exc}" ) from exc processed_rows = [] successful_rows = 0 error_rows = 0 for position, row in enumerate(rows): output_row = dict(row) output_row.pop(AUDIO_COLUMN, None) if SAVE_PREQUANT: output_row[PREQUANT_COLUMN] = prequant_values[position] if SAVE_DISCRETE: output_row[DISCRETE_TOKENS_COLUMN] = discrete_values[position] output_row[ERROR_COLUMN] = errors[position] if errors[position] is None: successful_rows += 1 else: error_rows += 1 processed_rows.append(output_row) return processed_rows, successful_rows, error_rows, dim_skipped def main(): print("=" * 60) print("Unified Codec Extraction Script") print(f" Metadata dataset: {DATASET_SOURCE}") print(f" Audio dataset: {AUDIO_DATASET_SOURCE}") print(f" Checkpoint dir: {CHECKPOINT_DIR}") print(f" Save discrete: {SAVE_DISCRETE}") print(f" Save prequant: {SAVE_PREQUANT}") print(f" Batch size: {BATCH_SIZE}") print( f" Duration buckets: automatic " f"({MIN_DURATION_BUCKETS}-{MAX_DURATION_BUCKETS})" ) print(f" Processing order: global bucket-by-bucket") print(f" Duration column: {DURATION_COLUMN}") print(f" Maximum duration: {MAX_AUDIO_DURATION_SECONDS:.1f}s") print(f" Workers: {NUM_WORKERS}") print(f" Inference dtype: {INFERENCE_DTYPE}") print("=" * 60) if not SAVE_DISCRETE and not SAVE_PREQUANT: print( "ERROR: Both SAVE_DISCRETE and SAVE_PREQUANT are False. " "Nothing to do." ) sys.exit(1) print("Loading metadata dataset...") metadata_dataset = load_input_dataset() print(f"Loaded {len(metadata_dataset):,} metadata rows.") print("Building checkpoint schema...") checkpoint_features = build_checkpoint_features(metadata_dataset) print("Checkpoint schema ready.") total_rows = len(metadata_dataset) ( shard_index, processed_count, processed_steps, saved_bucket_boundaries_samples, current_bucket_index, bucket_offset, ) = find_saved_state() print( f"Checkpoint state: {processed_count:,}/{total_rows:,} " f"rows completed across {shard_index:,} shards." ) if saved_bucket_boundaries_samples is None: print( "Calculating duration buckets from the Arrow-backed " f"'{DURATION_COLUMN}' column..." ) ( duration_bucket_boundaries_samples, duration_bucket_boundaries_seconds, ) = calculate_duration_bucket_boundaries( metadata_dataset=metadata_dataset, sample_rate=TARGET_SAMPLE_RATE, ) else: duration_bucket_boundaries_samples = ( saved_bucket_boundaries_samples ) duration_bucket_boundaries_seconds = tuple( boundary / TARGET_SAMPLE_RATE for boundary in duration_bucket_boundaries_samples ) print( "Reusing duration buckets saved in state.json: " f"{tuple(round(value, 3) for value in duration_bucket_boundaries_seconds)}" ) print("Building global metadata row indices for each duration bucket...") duration_bucket_indices, duration_bucket_counts = ( build_duration_bucket_indices( metadata_dataset=metadata_dataset, sample_rate=TARGET_SAMPLE_RATE, bucket_boundaries_samples=( duration_bucket_boundaries_samples ), ) ) if current_bucket_index > len(duration_bucket_indices): raise ValueError( f"Checkpoint current_bucket_index={current_bucket_index} " f"exceeds the {len(duration_bucket_indices)} buckets." ) if current_bucket_index < len(duration_bucket_indices): if bucket_offset > len( duration_bucket_indices[current_bucket_index] ): raise ValueError( f"Checkpoint bucket_offset={bucket_offset} exceeds " f"bucket {current_bucket_index + 1} size." ) elif bucket_offset != 0: raise ValueError( "A completed checkpoint must have bucket_offset=0." ) expected_processed_rows = sum( len(duration_bucket_indices[index]) for index in range(current_bucket_index) ) if current_bucket_index < len(duration_bucket_indices): expected_processed_rows += bucket_offset if expected_processed_rows != processed_count: raise RuntimeError( f"Checkpoint cursor implies {expected_processed_rows} processed " f"rows, but checkpoint shards contain {processed_count}." ) print(f"Loading audio datasets from: {AUDIO_DATASET_SOURCE}") audio_datasets = load_audio_datasets() total_audio_rows = sum(len(dataset) for dataset in audio_datasets) print( f"Building audio lookup from {total_audio_rows:,} audio rows " f"across {len(audio_datasets)} datasets." ) audio_key_index = build_audio_key_index(audio_datasets) print(f"Indexed {len(audio_key_index):,} unique audio keys.") if not torch.cuda.is_available(): raise RuntimeError("CUDA is required for bfloat16 codec inference.") if not torch.cuda.is_bf16_supported(): raise RuntimeError( "The active CUDA device does not support bfloat16 inference." ) device = "cuda" print(f"Loading Kanadec nano audio tokenizer on {device}...") codec_model = load_dune_audio_tokenizer( MODEL_ID, device=device, ) codec_model = codec_model.to( device=device, dtype=INFERENCE_DTYPE, ).eval() codec_model.codec.audio_encoder = torch.compile(codec_model.codec.audio_encoder, mode="max-autotune-no-cudagraphs") floating_dtypes = { tensor.dtype for tensor in ( *codec_model.parameters(), *codec_model.buffers(), ) if tensor.is_floating_point() } if floating_dtypes != {INFERENCE_DTYPE}: raise RuntimeError( f"Codec floating dtypes are {sorted(map(str, floating_dtypes))}, " f"expected only {INFERENCE_DTYPE}." ) print(f"Target Sample Rate: {TARGET_SAMPLE_RATE}") print(f"Confirmed model dtype: {INFERENCE_DTYPE}") total_progress_steps = sum( math.ceil(int(count) / BATCH_SIZE) for count in duration_bucket_counts ) progress = tqdm( total=total_progress_steps, initial=processed_steps, desc="Extracting codec features", unit="step", dynamic_ncols=True, ) pending_rows = [] steps_since_checkpoint = 0 total_successful = 0 total_errors = 0 total_dim_skipped = 0 started_at = time.time() try: with ThreadPoolExecutor(max_workers=NUM_WORKERS) as audio_pool: with torch.no_grad(): for bucket_index in range( current_bucket_index, len(duration_bucket_indices), ): row_indices = duration_bucket_indices[bucket_index] start_offset = ( bucket_offset if bucket_index == current_bucket_index else 0 ) bucket_upper_samples = ( duration_bucket_boundaries_samples[bucket_index] ) bucket_upper_seconds = ( duration_bucket_boundaries_seconds[bucket_index] ) print( f"Processing bucket {bucket_index + 1}/" f"{len(duration_bucket_indices)}: <= " f"{bucket_upper_seconds:.3f}s, " f"{len(row_indices):,} rows, starting at " f"offset {start_offset:,}." ) for batch_start in range( start_offset, len(row_indices), BATCH_SIZE, ): batch_indices = row_indices[ batch_start:batch_start + BATCH_SIZE ] rows = get_metadata_rows_by_indices( metadata_dataset, batch_indices, ) batch = prepare_global_bucket_batch( rows=rows, audio_pool=audio_pool, audio_datasets=audio_datasets, audio_key_index=audio_key_index, sample_rate=TARGET_SAMPLE_RATE, padded_length=bucket_upper_samples, bucket_upper_seconds=bucket_upper_seconds, ) ( processed_batch, successful_rows, error_rows, dim_skipped, ) = process_global_bucket_batch( model=codec_model, batch=batch, device=device, bucket_upper_seconds=bucket_upper_seconds, ) pending_rows.extend(processed_batch) processed_count += len(processed_batch) processed_steps += 1 steps_since_checkpoint += 1 total_successful += successful_rows total_errors += error_rows total_dim_skipped += dim_skipped next_offset = batch_start + len(batch_indices) if next_offset >= len(row_indices): current_bucket_index = bucket_index + 1 bucket_offset = 0 else: current_bucket_index = bucket_index bucket_offset = next_offset progress.update(1) progress.set_postfix( bucket=( f"{bucket_index + 1}/" f"{len(duration_bucket_indices)}" ), bucket_rows=( f"{min(next_offset, len(row_indices))}/" f"{len(row_indices)}" ), rows=processed_count, errors=total_errors, dim_skipped=total_dim_skipped, checkpoint_in=max( 0, CHECKPOINT_INTERVAL_STEPS - steps_since_checkpoint, ), refresh=False, ) if ( steps_since_checkpoint >= CHECKPOINT_INTERVAL_STEPS ): shard_index = save_shard( rows=pending_rows, shard_index=shard_index, processed_rows=processed_count, processed_steps=processed_steps, checkpoint_features=checkpoint_features, duration_bucket_boundaries_samples=( duration_bucket_boundaries_samples ), current_bucket_index=( current_bucket_index ), bucket_offset=bucket_offset, ) pending_rows = [] steps_since_checkpoint = 0 except KeyboardInterrupt: print("\nInterrupted; saving completed rows before exit.") finally: shard_index = save_shard( rows=pending_rows, shard_index=shard_index, processed_rows=processed_count, processed_steps=processed_steps, checkpoint_features=checkpoint_features, duration_bucket_boundaries_samples=( duration_bucket_boundaries_samples ), current_bucket_index=current_bucket_index, bucket_offset=bucket_offset, ) progress.close() elapsed = time.time() - started_at print("=" * 60) print("Done.") print(f" Successful: {total_successful}") print(f" Errors: {total_errors}") print(f" Dim-skipped: {total_dim_skipped}") print(f" Elapsed: {elapsed:.2f} seconds") print(f" Shards: {CHECKPOINT_DIR}") print("=" * 60) if __name__ == "__main__": main()