03d_dialect.py

#!/usr/bin/env python3
"""
Step 03d — Swiss-German dialect classification.

Reads the base manifest (manifest.jsonl from step 02), classifies the dialect
of each segment's audio using a wav2vec2 phoneme extractor + sklearn pipeline,
and writes two output files:

  manifest_dialect.jsonl
      The full input manifest with four fields appended per row:
        dialect_segment        predicted label ID for this segment (e.g. "0")
        dialect_segment_name   human name for the segment label  (e.g. "Zürich")
        dialect_speaker        majority label ID for this (source_audio, speaker)
        dialect_speaker_name   human name for the speaker majority label

  manifest_dialect_speaker_map.jsonl   (one row per (source_audio, speaker) pair)
      dialect_speaker, dialect_speaker_name, vote_mode, num_segments,
      vote_counts, vote_counts_named, vote_scores

Dialect labels
--------------
  0 Zürich   1 Innerschweiz   2 Wallis      3 Graubünden   4 Ostschweiz
  5 Basel    6 Bern           7 Deutsch     8 Französisch  9 Italienisch
 10 Englisch

Model
-----
  dialect/model/best_dialect_model.joblib  (sklearn pipeline, bundled in repo)
  Feature: wav2vec2-xlsr-53-espeak-cv-ft phoneme string (see dialect/wav2vec2_phonemes.py)

Requirements
------------
  pip install torch transformers librosa joblib tqdm
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import time
from pathlib import Path

from dialect import (
    DEFAULT_MODEL_PATH,
    Wav2Vec2PhonemeExtractor,
    build_speaker_summary_rows,
    label_name,
    load_label_map,
    load_model,
    source_key_from_record,
    update_group_state,
)

DEFAULT_INPUT = Path(os.environ.get("MANIFEST", "manifest.jsonl"))
DEFAULT_OUTPUT = Path(os.environ.get("MANIFEST_DIALECT", "manifest_dialect.jsonl"))


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Classify Swiss-German dialect per segment and speaker.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    p.add_argument("--input", type=Path, default=DEFAULT_INPUT)
    p.add_argument("--output", type=Path, default=DEFAULT_OUTPUT)
    p.add_argument(
        "--speaker-map",
        type=Path,
        default=None,
        help="Output path for per-(source_audio, speaker) summary JSONL. "
             "Default: <output_stem>_speaker_map.jsonl next to --output.",
    )
    p.add_argument(
        "--model",
        type=Path,
        default=DEFAULT_MODEL_PATH,
        help="Path to trained dialect sklearn pipeline (joblib).",
    )
    p.add_argument(
        "--wav2vec2-model",
        default="facebook/wav2vec2-xlsr-53-espeak-cv-ft",
        help="HuggingFace wav2vec2 CTC phoneme model ID or local path.",
    )
    p.add_argument(
        "--device",
        choices=["auto", "cpu", "cuda"],
        default="auto",
        help="Inference device for wav2vec2. 'auto' picks CUDA if available.",
    )
    p.add_argument(
        "--vote-mode",
        choices=["prob_duration", "count"],
        default="prob_duration",
        help="prob_duration: sum(class_prob * segment_duration), then argmax. "
             "count: plain majority vote.",
    )
    p.add_argument(
        "--audio-field",
        default="audio_path",
        help="Manifest field that contains the per-segment audio file path.",
    )
    p.add_argument(
        "--speaker-field",
        default="speaker",
        help="Manifest field containing the speaker ID.",
    )
    p.add_argument(
        "--source-field",
        default="source_audio",
        help="Manifest field used as the grouping key for speaker-level voting.",
    )
    p.add_argument(
        "--overwrite",
        action="store_true",
        help="Re-run even if the output file already exists.",
    )
    p.add_argument(
        "--log-level",
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
    )
    return p.parse_args()


def _read_jsonl(path: Path) -> list[dict]:
    records = []
    with path.open(encoding="utf-8") as f:
        for i, line in enumerate(f, 1):
            line = line.strip()
            if not line:
                continue
            try:
                records.append(json.loads(line))
            except json.JSONDecodeError as exc:
                raise ValueError(f"Invalid JSON at {path}:{i}") from exc
    return records


def _write_jsonl(path: Path, rows: list[dict]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


def main() -> None:
    args = parse_args()
    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(asctime)s - %(levelname)s - %(message)s",
    )
    logger = logging.getLogger("dialect")

    if not args.input.exists():
        raise SystemExit(f"Input not found: {args.input}")
    if not args.model.exists():
        raise SystemExit(f"Model not found: {args.model}")
    if not args.overwrite and args.output.exists():
        logger.info("Output exists, skipping (--overwrite to force): %s", args.output)
        return

    speaker_map_path = args.speaker_map or args.output.with_name(
        args.output.stem + "_speaker_map.jsonl"
    )

    logger.info("Loading manifest: %s", args.input)
    records = _read_jsonl(args.input)
    if not records:
        raise SystemExit(f"Manifest is empty: {args.input}")
    logger.info("%d segments", len(records))

    logger.info("Loading dialect model: %s", args.model)
    model = load_model(args.model)
    label_map = load_label_map()

    logger.info("Loading wav2vec2 phoneme extractor: %s (device=%s)", args.wav2vec2_model, args.device)
    extractor = Wav2Vec2PhonemeExtractor(model_name=args.wav2vec2_model, device=args.device)

    try:
        from tqdm import tqdm
        iterator = tqdm(records, desc="Dialect classification", unit="seg")
    except ImportError:
        iterator = records

    t0 = time.perf_counter()
    state: dict = {}
    segment_preds: list[str] = []

    for rec in iterator:
        audio_path = rec.get(args.audio_field, "")
        if not isinstance(audio_path, str) or not audio_path.strip():
            logger.warning("Missing %s in record, skipping", args.audio_field)
            segment_preds.append("")
            continue
        if not Path(audio_path).exists():
            logger.warning("Audio file not found: %s", audio_path)
            segment_preds.append("")
            continue

        source_key = source_key_from_record(rec, args.source_field)
        speaker = str(rec.get(args.speaker_field, "")).strip() or "<unknown_speaker>"
        duration = float(rec.get("duration") or 0.0)

        phonemes = extractor.phonemize_audio_path(audio_path)
        pred = update_group_state(
            state,
            source_key=source_key,
            speaker=speaker,
            phonemes=phonemes,
            duration=duration,
            model=model,
            vote_mode=args.vote_mode,
        )
        segment_preds.append(pred)

    elapsed = time.perf_counter() - t0
    logger.info("Classification done in %.1f s (%.1f seg/s)", elapsed, len(records) / max(elapsed, 1e-6))

    summary_rows = build_speaker_summary_rows(state, args.vote_mode, label_map)
    majority_by_group: dict[tuple[str, str], str] = {
        (row["source_audio"], row["speaker"]): row["dialect_speaker"]
        for row in summary_rows
    }

    augmented: list[dict] = []
    for rec, pred in zip(records, segment_preds):
        source_key = source_key_from_record(rec, args.source_field)
        speaker = str(rec.get(args.speaker_field, "")).strip() or "<unknown_speaker>"
        majority = majority_by_group.get((source_key, speaker), "")
        out = dict(rec)
        out["dialect_segment"] = pred
        out["dialect_segment_name"] = label_name(pred, label_map) if pred else ""
        out["dialect_speaker"] = majority
        out["dialect_speaker_name"] = label_name(majority, label_map) if majority else ""
        augmented.append(out)

    _write_jsonl(args.output, augmented)
    _write_jsonl(speaker_map_path, summary_rows)

    logger.info("Wrote dialect manifest (%d rows): %s", len(augmented), args.output)
    logger.info("Wrote speaker map (%d entries): %s", len(summary_rows), speaker_map_path)

    from collections import Counter
    dist = Counter(r["dialect_speaker_name"] for r in summary_rows)
    logger.info("Speaker dialect distribution: %s", dict(dist.most_common()))


if __name__ == "__main__":
    main()