"""GEC pair schema, dataset construction, augmentation, and ablation selection. A *pair* is one ``raw_asr -> gold_text`` example carrying the retrieved NEs and provenance. Real pairs come from running Gipformer over ViMedCSS audio (paper §3.1); synthetic pairs come from running Gipformer over voice-cloned TTS audio (paper §4.1 Step 3). ``augment_training_pairs`` merges them (paper §5: ``nsyn = n``), and ``select_variant_rows`` carves out the four ablation configurations (paper §5 "Comparison Methods"). Pure-python (validation, jsonl IO, augmentation, variant selection) imports with no ML deps; ``build_real_pairs`` / ``build_synthetic_pairs`` lazily import ``datasets`` and the ASR service. """ from __future__ import annotations import json import tempfile from dataclasses import dataclass from pathlib import Path from typing import Any, Iterable from gec.config import variant_uses_retrieval from gec.metrics import split_terms REAL_SOURCE = "vimedcss_real" SYNTHETIC_SOURCE = "darag_synthetic_tts" REAL_SOURCES = {REAL_SOURCE} VALID_SOURCES = {REAL_SOURCE, SYNTHETIC_SOURCE} GEC_PAIR_REQUIRED_FIELDS = { "split", "source_kind", "audio_id", "raw_asr", "gold_text", "gold_terms", "retrieved_terms", "asr_model", } SYNTHETIC_REQUIRED_FIELDS = { "source_kind", "synthetic_id", "clean_text", "topic", "seed_example_ids", "model", } @dataclass(frozen=True) class ValidationResult: ok: bool errors: list[str] def validate_gec_pair(row: dict[str, Any]) -> ValidationResult: errors = _missing(row, GEC_PAIR_REQUIRED_FIELDS) if row.get("source_kind") not in VALID_SOURCES: errors.append(f"source_kind must be one of {sorted(VALID_SOURCES)}") if not str(row.get("raw_asr", "")).strip(): errors.append("raw_asr must be non-empty") if not str(row.get("gold_text", "")).strip(): errors.append("gold_text must be non-empty") if not isinstance(row.get("gold_terms"), list): errors.append("gold_terms must be a list") if not isinstance(row.get("retrieved_terms"), list): errors.append("retrieved_terms must be a list") return ValidationResult(ok=not errors, errors=errors) def validate_synthetic_transcript(row: dict[str, Any]) -> ValidationResult: errors = _missing(row, SYNTHETIC_REQUIRED_FIELDS) clean = str(row.get("clean_text", "")).strip() if not clean: errors.append("clean_text must be non-empty") elif len(clean.split()) < 5: errors.append("clean_text is too short for ASR/GEC training") if row.get("source_kind") != "darag_synthetic_clean": errors.append("source_kind must be darag_synthetic_clean") return ValidationResult(ok=not errors, errors=errors) # --------------------------------------------------------------------------- IO def read_jsonl(path: Path) -> list[dict[str, Any]]: if not path.exists(): return [] return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line] def write_jsonl(path: Path, rows: Iterable[dict[str, Any]]) -> int: path.parent.mkdir(parents=True, exist_ok=True) count = 0 with path.open("w", encoding="utf-8") as handle: for row in rows: handle.write(json.dumps(row, ensure_ascii=False) + "\n") count += 1 return count # ------------------------------------------------------------- augment / select def augment_training_pairs( real_pairs: list[dict[str, Any]], synthetic_pairs: list[dict[str, Any]], nsyn_factor: float = 1.0, ) -> list[dict[str, Any]]: """Merge real ``train`` pairs with up to ``nsyn_factor * n`` synthetic pairs. Only the real ``train`` split is augmented; validation/test/hard stay frozen and real (paper keeps evaluation splits untouched). Synthetic pairs are forced onto the ``train`` split. """ real_train = [r for r in real_pairs if r.get("split") == "train"] held_out = [r for r in real_pairs if r.get("split") != "train"] budget = int(round(nsyn_factor * len(real_train))) if real_train else len(synthetic_pairs) chosen = [{**r, "split": "train"} for r in synthetic_pairs[:budget]] return real_train + chosen + held_out def select_variant_rows( augmented_pairs: list[dict[str, Any]], variant: str, ) -> tuple[list[dict[str, Any]], bool]: """Return (training rows, use_retrieval) for a DARAG ablation variant. * ``full`` — real + synthetic train rows, NEs in prompt. * ``wo_rac`` — same rows, NEs stripped from the prompt (use_retrieval=False). * ``wo_aug`` — real train rows only (drop synthetic), NEs in prompt. * ``only_synth`` — synthetic train rows only, NEs in prompt. Validation rows (real ``validation`` split) are returned separately by ``eval_split_rows`` and are never altered by the variant. """ use_retrieval = variant_uses_retrieval(variant) train_rows = [r for r in augmented_pairs if r.get("split") == "train"] if variant == "wo_aug": train_rows = [r for r in train_rows if r.get("source_kind") in REAL_SOURCES] elif variant == "only_synth": train_rows = [r for r in train_rows if r.get("source_kind") == SYNTHETIC_SOURCE] return train_rows, use_retrieval def eval_split_rows(pairs: list[dict[str, Any]], split: str) -> list[dict[str, Any]]: return [r for r in pairs if r.get("split") == split] # ----------------------------------------------------- build pairs (lazy/heavy) def build_real_pairs( dataset: str, splits: list[str], retriever, asr, limit_per_split: int | None = None, completed_ids: set[str] | None = None, n_best: int = 1, seed: int = 13, ): """Yield real GEC pairs by running ``asr`` over each ``dataset`` audio clip. ``retriever`` and ``asr`` are injected so this stays testable; the CLI wires in the gec NE retriever and ``carepath.services.asr``. ``n_best > 1`` adds the paper's other-hypotheses via ``nbest.other_hypotheses`` (perturbation decodes), computed while the temp wav still exists. """ import soundfile as sf # type: ignore from datasets import Audio, load_dataset # type: ignore from gec import nbest completed_ids = completed_ids or set() for split in splits: data = load_dataset(dataset, split=split).cast_column("audio", Audio(decode=True)) if limit_per_split: data = data.select(range(min(limit_per_split, len(data)))) for row in data: audio_id = f"{split}:{row['segment_id']}" if audio_id in completed_ids: continue audio = row["audio"] with tempfile.TemporaryDirectory(prefix="carepath_gec_") as temp_dir: wav_path = Path(temp_dir) / f"{row['segment_id']}.wav" sf.write(str(wav_path), audio["array"], int(audio["sampling_rate"])) result = asr.transcribe(wav_path) # Compute N-best before the temp dir (and wav) is removed. others = nbest.other_hypotheses(asr, wav_path, result.text, n_best, seed=seed) gold_terms = split_terms(row.get("cs_terms_list", "")) retrieval_text = " ".join(p for p in (result.text, row["segment_text"]) if p) retrieved = [item.term for item in retriever.retrieve(retrieval_text)] yield { "split": split, "source_kind": REAL_SOURCE, "audio_id": row["segment_id"], "raw_asr": result.text, "other_hypotheses": others, "gold_text": row["segment_text"], "gold_terms": gold_terms, "retrieved_terms": retrieved, "cs_terms_list": row.get("cs_terms_list", ""), "topic": row.get("topic"), "duration_seconds": row.get("duration_seconds"), "asr_model": result.model, "asr_metadata": result.metadata, } def build_synthetic_pairs( manifest_rows: list[dict[str, Any]], retriever, asr, completed_ids: set[str] | None = None, n_best: int = 1, seed: int = 13, ): """Yield synthetic GEC pairs by running ``asr`` over voice-cloned TTS audio.""" from gec import nbest completed_ids = completed_ids or set() for row in manifest_rows: synthetic_id = row["synthetic_id"] if synthetic_id in completed_ids: continue audio_path = Path(row["audio_path"]) result = asr.transcribe(audio_path) others = nbest.other_hypotheses(asr, audio_path, result.text, n_best, seed=seed) retrieval_text = f"{result.text} {row['clean_text']}" retrieved = [item.term for item in retriever.retrieve(retrieval_text)] yield { "split": row.get("split", "train"), "source_kind": SYNTHETIC_SOURCE, "audio_id": synthetic_id, "synthetic_id": synthetic_id, "raw_asr": result.text, "other_hypotheses": others, "gold_text": row["clean_text"], "gold_terms": row.get("intended_terms", []), "retrieved_terms": retrieved, "topic": row.get("topic"), "duration_seconds": result.metadata.get("duration_seconds"), "asr_model": result.model, "asr_metadata": result.metadata, "tts_model": row.get("tts_model"), "tts_provider": row.get("tts_provider"), "speaker_reference": row.get("speaker_reference"), "audio_path": row.get("audio_path"), } def _hash(text: str) -> str: import hashlib return hashlib.sha1(text.encode("utf-8")).hexdigest()[:12] def _missing(row: dict[str, Any], required: set[str]) -> list[str]: return [f"missing required field: {f}" for f in sorted(required) if f not in row]