Update dune_extraction.py

dune_extraction.py

@@ -20,16 +20,7 @@ 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 datasets import Audio, Dataset, DatasetDict, Features, Sequence, Value, load_dataset, load_from_disk
 from tqdm.auto import tqdm
 
 warnings.filterwarnings("ignore")
@@ -38,20 +29,21 @@ warnings.filterwarnings("ignore")
 MODEL_ID = "Respair/dune_codec"
 
 
-# Primary dataset:
-# Contains the metadata rows to enrich and a unique bridge key.
+# contains the metadata rows to enrich and a unique bridge key.
 DATASET_SOURCE = "/home/ubuntu/data"
 DATASET_CONFIG = None
 DATASET_SPLIT = "train"
 
-# Secondary dataset:
-# Must contain the same bridge key and the expensive audio column.
+
+# must contain the same bridge key and the expensive audio column. if your metadata (aka DATASET_SOURCE) is not a subsection of this, you can just copy DATASET_SOURCE here.
 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"
@@ -63,41 +55,31 @@ 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
 
-# Processing options.
-TEST_RUN = 0
+TEST_RUN = 0 # you probably won't need it
 
-# Performance settings.
 BATCH_SIZE = 64
 NUM_WORKERS = 32
 TARGET_SAMPLE_RATE = 22050
 INFERENCE_DTYPE = torch.bfloat16
 
-# Data-driven duration bucketing. The complete metadata duration column is
-# used to optimize between 10 and 15 bucket ceilings for minimum padding.
+
+MAX_AUDIO_DURATION_SECONDS = 31.0 # tweak this for your own data
+
+
+
+# you can ignore these
 MIN_DURATION_BUCKETS = 10
 MAX_DURATION_BUCKETS = 15
 DURATION_HISTOGRAM_RESOLUTION_SECONDS = 0.05
 DURATION_BUCKET_ELBOW_TOLERANCE = 0.03
-MAX_AUDIO_DURATION_SECONDS = 31.0
 ARROW_DURATION_SCAN_BATCH_SIZE = 250_000
-
-# Number of duration rows scanned at once while assigning global bucket
-# membership from the Arrow-backed metadata column.
 DURATION_BUCKET_INDEX_SCAN_BATCH_SIZE = 250_000
-
-# Number of key rows read at once while constructing the key -> audio-row index.
 AUDIO_INDEX_BATCH_SIZE = 100_000
-
-# Pre-quantization latent expected dimensionality.
 EXPECTED_PREQUANT_DIM = 52
 
-# ============================================================
-# Hugging Face dataset loading
-# ============================================================
 def load_dataset_source(source, config, split):
 
      source_path = Path(source).expanduser()