scripts/prep_data.py

#!/usr/bin/env python3
"""Extraction audio -> FLAC 16 kHz mono + manifests JSONL, pour WAXAL et AfriVoice.

Sorties par jeu : /scratch/prep/manifests/{name}.jsonl et /scratch/prep/audio/{name}/*.flac
Champs manifest : id, audio (chemin flac), duration, text, speaker, source.
En fin de run : rapport de recouvrement texte/locuteur AfriVoice <-> WAXAL val/test
(les deux corpus lin/sna viennent de Digital Umuganda : risque de fuite).
"""
import glob
import hashlib
import json
import os
import subprocess
import sys
import unicodedata
from concurrent.futures import ProcessPoolExecutor, as_completed

import numpy as np
import pyarrow.parquet as pq
import soundfile as sf

SR = 16000
PREP = "/scratch/prep"

JOBS = [
    # (name, glob parquets, colonne texte)
    ("waxal_lug_train", "/scratch/data/waxal/data/ASR/lug/lug-train-*.parquet", "transcription"),
    ("waxal_lug_validation", "/scratch/data/waxal/data/ASR/lug/lug-validation-*.parquet", "transcription"),
    ("waxal_lug_test", "/scratch/data/waxal/data/ASR/lug/lug-test-*.parquet", "transcription"),
    ("waxal_lin_train", "/scratch/data/waxal/data/ASR/lin/lin-train-*.parquet", "transcription"),
    ("waxal_lin_validation", "/scratch/data/waxal/data/ASR/lin/lin-validation-*.parquet", "transcription"),
    ("waxal_lin_test", "/scratch/data/waxal/data/ASR/lin/lin-test-*.parquet", "transcription"),
    ("waxal_sna_train", "/scratch/data/waxal/data/ASR/sna/sna-train-*.parquet", "transcription"),
    ("waxal_sna_validation", "/scratch/data/waxal/data/ASR/sna/sna-validation-*.parquet", "transcription"),
    ("waxal_sna_test", "/scratch/data/waxal/data/ASR/sna/sna-test-*.parquet", "transcription"),
    ("afrivoice_lin_train", "/scratch/data/afrivoice_ln/data/train-*.parquet", "text"),
    ("afrivoice_lin_validation", "/scratch/data/afrivoice_ln/data/validation-*.parquet", "text"),
    ("afrivoice_lin_test", "/scratch/data/afrivoice_ln/data/test-*.parquet", "text"),
    ("afrivoice_sna_train", "/scratch/data/afrivoice_sna/data/train-*.parquet", "transcription"),
]


def norm_text(t):
    if t is None:
        return ""
    t = unicodedata.normalize("NFC", str(t))
    return " ".join(t.split())


def decode_to_16k(raw):
    """Decode n'importe quel format audio (mp3 44.1/48/16 kHz...) -> float32 mono 16 kHz."""
    p = subprocess.run(
        ["ffmpeg", "-v", "error", "-i", "pipe:0", "-f", "f32le", "-ac", "1", "-ar", str(SR), "pipe:1"],
        input=raw, capture_output=True)
    if p.returncode != 0:
        raise RuntimeError("ffmpeg: " + p.stderr.decode(errors="replace")[:200])
    return np.frombuffer(p.stdout, dtype=np.float32)


def speaker_mapping(pf):
    """Mapping ClassLabel int -> nom (UID Firebase) via les metadonnees HF du parquet."""
    try:
        meta = pf.schema_arrow.metadata or {}
        info = json.loads(meta.get(b"huggingface", b"{}"))
        feat = info.get("info", {}).get("features", {}).get("speaker_id", {})
        names = feat.get("names") or (feat.get("class_label", {}) or {}).get("names")
        if isinstance(names, dict):
            return {int(k): v for k, v in names.items()}
        if isinstance(names, list):
            return dict(enumerate(names))
    except Exception:
        pass
    return None


def process_parquet(task):
    name, pf_path, text_col, shard_idx = task
    audio_dir = os.path.join(PREP, "audio", name)
    os.makedirs(audio_dir, exist_ok=True)
    rows, errors = [], 0
    pf = pq.ParquetFile(pf_path)
    spk_map = speaker_mapping(pf)
    for batch in pf.iter_batches(batch_size=16):
        for r in batch.to_pylist():
            try:
                audio = r.get("audio")
                raw = audio.get("bytes") if isinstance(audio, dict) else None
                if raw is None:
                    errors += 1
                    continue
                text = norm_text(r.get(text_col))
                rid = r.get("id") or (audio.get("path") if isinstance(audio, dict) else None)
                if not rid:
                    rid = hashlib.md5(raw[:4096]).hexdigest()[:16]
                rid = str(rid).replace("/", "_").replace(".mp3", "").replace(".wav", "")
                wav = decode_to_16k(raw)
                if len(wav) < int(0.1 * SR):
                    errors += 1
                    continue
                path = os.path.join(audio_dir, f"{rid}.flac")
                sf.write(path, wav, SR, format="FLAC")
                spk = r.get("speaker_id", "")
                if isinstance(spk, int) and spk_map:
                    spk = spk_map.get(spk, spk)
                rows.append({
                    "id": rid,
                    "audio": path,
                    "duration": round(len(wav) / SR, 3),
                    "text": text,
                    "speaker": str(spk),
                    "source": name,
                })
            except Exception:
                errors += 1
    return name, shard_idx, rows, errors


def main():
    os.makedirs(os.path.join(PREP, "manifests"), exist_ok=True)
    tasks = []
    for name, pattern, text_col in JOBS:
        files = sorted(glob.glob(pattern))
        if not files:
            print(f"!! aucun fichier pour {name} ({pattern})", flush=True)
            continue
        for i, f in enumerate(files):
            tasks.append((name, f, text_col, i))

    results = {}
    done = 0
    with ProcessPoolExecutor(max_workers=20) as ex:
        futs = {ex.submit(process_parquet, t): t for t in tasks}
        for fut in as_completed(futs):
            name, shard_idx, rows, errors = fut.result()
            results.setdefault(name, {"rows": [], "errors": 0})
            results[name]["rows"].extend(rows)
            results[name]["errors"] += errors
            done += 1
            print(f"[{done}/{len(tasks)}] {name} shard {shard_idx}: {len(rows)} ok, {errors} err", flush=True)

    for name, res in results.items():
        rows = sorted(res["rows"], key=lambda r: r["id"])
        with open(os.path.join(PREP, "manifests", f"{name}.jsonl"), "w", encoding="utf-8") as f:
            for r in rows:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
        hours = sum(r["duration"] for r in rows) / 3600
        print(f"=> {name}: {len(rows)} clips, {hours:.1f} h, {res['errors']} erreurs", flush=True)

    # ---- Rapport de fuite AfriVoice <-> WAXAL (texte exact normalise + locuteurs) ----
    def load(name):
        p = os.path.join(PREP, "manifests", f"{name}.jsonl")
        if not os.path.exists(p):
            return []
        return [json.loads(l) for l in open(p, encoding="utf-8")]

    def key(t):
        t = unicodedata.normalize("NFC", t).lower()
        t = "".join(c for c in t if c.isalnum() or c.isspace())
        return " ".join(t.split())

    print("\n===== RAPPORT DE FUITE =====", flush=True)
    for lang in ("lin", "sna"):
        av_names = [n for n in results if n.startswith(f"afrivoice_{lang}")]
        av = [r for n in av_names for r in load(n)]
        av_texts = {key(r["text"]) for r in av if r["text"]}
        av_speakers = {r["speaker"] for r in av}
        for split in ("validation", "test", "train"):
            wx = load(f"waxal_{lang}_{split}")
            n_text = sum(1 for r in wx if key(r["text"]) in av_texts)
            n_spk = sum(1 for r in wx if r["speaker"] in av_speakers)
            print(f"{lang} waxal-{split} vs afrivoice: {n_text}/{len(wx)} textes identiques, "
                  f"{n_spk}/{len(wx)} clips de locuteurs partages", flush=True)
    print("PREP_DONE", flush=True)


if __name__ == "__main__":
    main()