download_fleurs.py

import hashlib
import argparse
import math
import numpy as np
import os
from datasets import Dataset, load_from_disk, get_dataset_config_names, load_dataset
from tqdm.auto import tqdm

FLEURS_CACHE = "data/fleurs_subset"
FLEURS_CACHE_NO_VOX = "data/fleurs_subset_no_vox"
SR = 16000

# VoxAngeles uses ISO-639-3 directory names; FLEURS uses locale-like config names.
# This maps VoxAngeles languages/close varieties that overlap FLEURS configs.
VOXANGELES_TO_FLEURS_CONFIGS = {
    "aeb": "ar_eg",   # Tunisian Arabic -> Arabic
    "afr": "af_za",
    "ajp": "ar_eg",   # South Levantine Arabic -> Arabic
    "apc": "ar_eg",   # North Levantine Arabic -> Arabic
    "asm": "as_in",
    "azb": "az_az",   # South Azerbaijani -> Azerbaijani
    "ben": "bn_in",
    "ces": "cs_cz",
    "dan": "da_dk",
    "ell": "el_gr",
    "ffm": "ff_sn",   # Maasina Fulfulde -> Fula
    "fin": "fi_fi",
    "fub": "ff_sn",   # Adamawa Fulfulde -> Fula
    "guj": "gu_in",
    "hau": "ha_ng",
    "heb": "he_il",
    "hin": "hi_in",
    "hrv": "hr_hr",
    "hun": "hu_hu",
    "hye": "hy_am",
    "ibo": "ig_ng",
    "isl": "is_is",
    "kan": "kn_in",
    "kea": "kea_cv",
    "khm": "km_kh",
    "lav": "lv_lv",
    "lit": "lt_lt",
    "lug": "lg_ug",
    "mal": "ml_in",
    "mlt": "mt_mt",
    "mya": "my_mm",
    "nld": "nl_nl",
    "pes": "fa_ir",
    "prs": "fa_ir",   # Dari -> Persian
    "yue": "yue_hant_hk",
}

def audio_key(audio):
    """Content-hashed, collision-proof key. FLEURS `id` is NOT unique (multiple speakers
    read the same sentence with the same id), so hashing the waveform is the only stable key."""
    return f"fleurs/{hashlib.sha1(np.asarray(audio, dtype=np.float32).tobytes()).hexdigest()[:16]}"

def discover_voxangeles_languages(data_dir="data"):
    vox_dir = os.path.join(data_dir, "voxangeles", "data", "audited_aligned")
    if not os.path.exists(vox_dir):
        return set()
    return {
        name
        for name in os.listdir(vox_dir)
        if os.path.isdir(os.path.join(vox_dir, name))
    }

def fleurs_configs_overlapping_voxangeles(data_dir="data"):
    vox_langs = discover_voxangeles_languages(data_dir)
    return {
        VOXANGELES_TO_FLEURS_CONFIGS[lang]
        for lang in vox_langs
        if lang in VOXANGELES_TO_FLEURS_CONFIGS
    }

def get_fleurs_subset(total_utterances=4000, max_dur=20.0, min_dur=0.3, cache_dir=FLEURS_CACHE,
                      exclude_voxangeles=True, data_dir="data"):
    """Small quota from every FLEURS language, persisted as an HF dataset so a crash
    never re-streams or re-decodes. Feats stay pickle-cached via content-hash keys."""
    if os.path.exists(cache_dir):
        print(f"{cache_dir} already exists. Doing nothing.")
        return None
    langs = [c for c in get_dataset_config_names("google/fleurs") if c != "all"]
    if exclude_voxangeles:
        excluded = fleurs_configs_overlapping_voxangeles(data_dir)
        langs = [lang for lang in langs if lang not in excluded]
        print(f"Excluding {len(excluded)} FLEURS configs overlapping VoxAngeles: {sorted(excluded)}")
    per_lang = math.ceil(total_utterances / len(langs))
    print(f"Targeting {total_utterances} utterances across {len(langs)} languages ({per_lang}/language)")
    keys, audios, seen = [], [], set()
    for lang in tqdm(langs, desc="FLEURS langs"):
        stream = load_dataset("google/fleurs", lang, split="train", streaming=True)
        n = 0
        for ex in stream:
            if len(keys) >= total_utterances: break
            if n >= per_lang: break
            a = ex["audio"]; dur = len(a["array"]) / a["sampling_rate"]
            if not (min_dur < dur <= max_dur): continue
            assert a["sampling_rate"] == SR, f"{lang}: {a['sampling_rate']}"
            audio = np.asarray(a["array"], dtype=np.float32)
            key = audio_key(audio)
            if key in seen: continue           # skip exact-duplicate waveforms
            seen.add(key)
            keys.append(key); audios.append(audio); n += 1
    assert len(set(keys)) == len(keys), "duplicate keys after content hashing"
    ds = Dataset.from_dict({"key": keys, "audio": audios})
    ds.save_to_disk(cache_dir)
    print(f"Saved {len(ds)} clips from {len(langs)} languages -> {cache_dir}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--exclude-voxangeles",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Exclude FLEURS configs that overlap local VoxAngeles languages.",
    )
    parser.add_argument(
        "--total-utterances",
        type=int,
        default=4000,
        help="Total utterance target. Per-language quota is derived from this.",
    )
    args = parser.parse_args()
    cache_dir = FLEURS_CACHE_NO_VOX if args.exclude_voxangeles else FLEURS_CACHE
    get_fleurs_subset(
        total_utterances=args.total_utterances,
        cache_dir=cache_dir,
        exclude_voxangeles=args.exclude_voxangeles,
    )

if __name__ == "__main__":
    main()