scripts/core/prepare_data.py

import argparse
import re
import os
from pathlib import Path
import sys
import time
import soundfile as sf
import librosa
import torch
import torchaudio
import warnings

# SUPPRESS WARNINGS (User Request)
warnings.filterwarnings("ignore")
try:
    # Ensure logs flush even when stdout is not a TTY (prevents silent crashes with buffered prints)
    sys.stdout.reconfigure(line_buffering=True)
    sys.stderr.reconfigure(line_buffering=True)
except Exception:
    pass

# --- Monkeypatch for DGX Spark Torchaudio (2.10+) ---
# Pyannote expects torchaudio.AudioMetaData, which might be moved/missing in bleeding edge versions
if not hasattr(torchaudio, "AudioMetaData"):
    from collections import namedtuple

    print("Monkeypatching torchaudio.AudioMetaData...")
    torchaudio.AudioMetaData = namedtuple(
        "AudioMetaData", ["sample_rate", "num_frames", "num_channels", "bits_per_sample", "encoding"]
    )
    # Also patch backend.common if it exists, as pyannote might look there
    if hasattr(torchaudio, "backend") and hasattr(torchaudio.backend, "common"):
        torchaudio.backend.common.AudioMetaData = torchaudio.AudioMetaData

if not hasattr(torchaudio, "list_audio_backends"):
    print("Monkeypatching torchaudio.list_audio_backends...")
    # Mock return value - soundfile/ffmpeg are standard
    torchaudio.list_audio_backends = lambda: ["ffmpeg", "soundfile"]

# Monkeypatch torch.load to default weights_only=False for PyTorch 2.6+ compatibility
original_torch_load = torch.load


def unsafe_torch_load(*args, **kwargs):
    # Force weights_only=False even if present
    kwargs["weights_only"] = False
    print(f"Intercepted torch.load, forced weights_only=False. Args: {args[1:] if len(args) > 1 else '?'}")
    return original_torch_load(*args, **kwargs)


torch.load = unsafe_torch_load
print("Monkeypatched torch.load for weights_only=False (FORCED)")

# Attempt to safe-list TorchVersion if possible
try:
    # torch.torch_version.TorchVersion is the class
    # We need to find where it is exposed.
    # Usually it's not public. But let's try to locate it via the instance.
    from torch.torch_version import TorchVersion

    torch.serialization.add_safe_globals([TorchVersion])
    print("Added TorchVersion to safe globals")
except Exception as e:
    print(f"Could not add safe globals (TorchVersion): {e}")

# --- Monkeypatching BEFORE Pyannote Imports ---
# --- Monkeypatching BEFORE Pyannote Imports ---
import semver

# 1. Nuclear Option: Patch semver.VersionInfo.parse
# PyTorch/Torchaudio versions on DGX Spark (e.g. 2.10.0a0+...) are not valid SemVer.
original_semver_parse = semver.VersionInfo.parse


def safe_semver_parse(version_str):
    try:
        return original_semver_parse(version_str)
    except ValueError:
        print(f"Warning: Bypassing invalid SemVer: {version_str}")
        # Return a dummy version that satisfies constraints (usually > 2.0.0)
        return semver.VersionInfo(3, 0, 0)  # Mock as 3.0.0


semver.VersionInfo.parse = safe_semver_parse
print("Monkeypatched semver.VersionInfo.parse (Nuclear Option)")

# 2. Try patching pyannote check_version too for good measure
try:
    from pyannote.audio.utils import version

    version.check_version = lambda library, mine, yours: None
    print("Monkeypatched pyannote.audio.utils.version.check_version")
except Exception as e:
    print(f"Could not patch pyannote check_version directly: {e}")

# Monkeypatch torchaudio.load to force soundfile backend (avoid torchcodec error)
# Robust replacement using soundfile directly


def robust_torchaudio_load(filepath, **kwargs):
    # Ignore backend arg if present
    # Directly use soundfile to load
    try:
        # soundfile.read returns (data, samplerate)
        # data is (frames, channels) if multichannel, or (frames,) if mono
        data, sr = sf.read(filepath)

        # Convert to torch tensor
        # Torchaudio expects (channels, time)
        if data.ndim == 1:
            # Mono
            # Must cast to float (float32) because soundfile returns float64 (Double)
            waveform = torch.from_numpy(data).float().unsqueeze(0)
        else:
            # Multichannel (time, channels) -> (channels, time)
            waveform = torch.from_numpy(data.T).float()

        return waveform, sr
    except Exception as e:
        print(f"Fallback load failed for {filepath}: {e}")
        raise e


torchaudio.load = robust_torchaudio_load
print("Monkeypatched torchaudio.load to use soundfile directly (ROBUST)")


# Monkeypatch torchaudio.info to use soundfile (MISSING API FIX)
class MockAudioInfo:
    def __init__(self, num_frames, sample_rate):
        self.num_frames = num_frames
        self.sample_rate = sample_rate


def robust_info(filepath, **kwargs):
    sinfo = sf.info(filepath)
    return MockAudioInfo(sinfo.frames, sinfo.samplerate)


torchaudio.info = robust_info
print("Monkeypatched torchaudio.info to use soundfile directly (ROBUST)")

# ... (Previous torchaudio hacks) ...

import numpy as np
import subprocess
from tqdm import tqdm
from transformers import AutoProcessor, pipeline

# Pyannote imports MUST happen AFTER patches
from pyannote.audio import Pipeline, Inference, Model
from scipy.spatial.distance import cosine
import json
from datasets import Dataset

# Re-enable TF32 (Pyannote disables it, but GB10 might need it or crash without it)
import torch

print("Re-enabling TF32/CuDNN benchmark to fix CUBLAS errors on GB10...")
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.benchmark = True  # Try ensuring optimal algo selection
print("TF32 Enabled. CuDNN Benchmark Enabled.")

# Robust environment check
try:
    import pytorch_metric_learning

    print(
        f"Successfully imported pytorch_metric_learning: {pytorch_metric_learning.__version__ if hasattr(pytorch_metric_learning, '__version__') else 'unknown'}"
    )
except ImportError as e:
    print(f"CRITICAL ERROR: pytorch_metric_learning failed to import: {e}")
    # Don't exit yet, let the pipeline try to proceed or crash loudly later

print("CUDA Architecture Check:")
if torch.cuda.is_available():
    print(f"  System Cuda Version: {torch.version.cuda}")
    print(f"  Device Name: {torch.cuda.get_device_name(0)}")
    print(f"  Device Arch: {torch.cuda.get_arch_list()}")
else:
    print("  CUDA NOT AVAILABLE")


def resolve_device(requested_device):
    if not requested_device:
        return "cuda" if torch.cuda.is_available() else "cpu"
    req = str(requested_device).lower()
    if req.startswith("cuda") or req == "gpu":
        if torch.cuda.is_available():
            return "cuda"
        print("Warning: CUDA requested but not available. Falling back to CPU.")
        return "cpu"
    return req


def env_flag(name, default=False):
    val = os.environ.get(name)
    if val is None:
        return default
    return val.strip().lower() in ("1", "true", "yes", "y", "on")


class DataPreparer:
    def __init__(self, output_dir, device=None, hf_token=None):
        # Safety: Check dependencies again
        pass
        self.device = resolve_device(device)
        self.device_index = 0 if self.device == "cuda" else -1
        self.hf_token = hf_token or os.environ.get("HF_TOKEN")
        print(f"Initializing DataPreparer on {self.device} (HF Token Present: {bool(self.hf_token)})")

        self.output_dir = Path(output_dir)
        self.wavs_dir = self.output_dir / "wavs"
        self.wavs_dir.mkdir(parents=True, exist_ok=True)
        self.metadata_path = self.output_dir / "metadata.csv"

        # Load Models
        print("Loading Whisper (Transformers)...")
        self.transcriber = self.load_whisper_pipeline()
        self.whisper_language = os.environ.get("WHISPER_LANGUAGE", "spanish")
        self.min_asr_words = int(os.environ.get("WHISPER_MIN_WORDS", "2"))
        self.min_asr_alpha_ratio = float(os.environ.get("WHISPER_MIN_ALPHA_RATIO", "0.5"))

        print("Loading Pyannote Diarization 3.1...")
        # Use configured device (GPU if available)
        self.diarization_pipeline = Pipeline.from_pretrained(
            "pyannote/speaker-diarization-3.1", use_auth_token=self.hf_token
        ).to(torch.device(self.device))

        print("Loading Pyannote Embedding Model for Verification...")
        self.embedding_model = Model.from_pretrained(
            "pyannote/wespeaker-voxceleb-resnet34-LM", use_auth_token=self.hf_token
        )
        # Use configured device
        self.device_embedding = torch.device(self.device)
        print(f"Moving Embedding Model to {self.device_embedding}")
        self.embedding_model.to(self.device_embedding)
        self.inference = Inference(self.embedding_model, window="whole", device=self.device_embedding)

        self.target_sr = 24000  # For F5-TTS

    def _text_quality_ok(self, text: str) -> bool:
        words = re.findall(r"[\wáéíóúñüÁÉÍÓÚÑÜ]+", text)
        if len(words) < self.min_asr_words:
            return False
        alpha = sum(1 for c in text if c.isalpha())
        ratio = alpha / max(1, len(text))
        return ratio >= self.min_asr_alpha_ratio

    def _normalize_embedding(self, emb):
        if isinstance(emb, torch.Tensor):
            if emb.ndim > 1:
                emb = emb.mean(dim=0)
            return emb.detach().cpu().numpy()
        if hasattr(emb, "ndim") and emb.ndim > 1:
            emb = emb.mean(axis=0)
        return np.asarray(emb)

    def _call_diarization_pipeline(self, wav_path):
        try:
            return self.diarization_pipeline(wav_path, batch_size=1, num_workers=0)
        except TypeError:
            return self.diarization_pipeline(wav_path)

    def diarize_file(self, wav_path, target_sr=16000):
        temp_path = None
        try:
            data, sr = sf.read(str(wav_path))
            if data.ndim == 1:
                data = data[None, :]
            else:
                data = data.T
            waveform = torch.from_numpy(data).float()

            if sr != target_sr:
                waveform = torchaudio.functional.resample(waveform, sr, target_sr)
                sr = target_sr

            if waveform.shape[0] > 1:
                waveform = torch.mean(waveform, dim=0, keepdim=True)

            # Trim to a multiple of 10ms to avoid padding shape mismatches
            frame = int(sr * 0.01)
            if frame > 0:
                trim = (waveform.shape[1] // frame) * frame
                if trim > 0:
                    waveform = waveform[:, :trim]

            temp_path = self.output_dir / "temp_diarization" / f"{wav_path.stem}_{int(time.time() * 1000)}.wav"
            temp_path.parent.mkdir(parents=True, exist_ok=True)
            sf.write(temp_path, waveform.squeeze(0).cpu().numpy(), sr)
            return self._call_diarization_pipeline(str(temp_path))
        finally:
            if temp_path and temp_path.exists():
                temp_path.unlink()

    def load_whisper_pipeline(self):
        whisper_device = os.environ.get("WHISPER_DEVICE")
        if whisper_device:
            whisper_device = resolve_device(whisper_device)
        else:
            whisper_device = self.device
        whisper_device_index = 0 if whisper_device == "cuda" else -1

        model_id = os.environ.get("WHISPER_MODEL")
        if not model_id:
            model_id = "openai/whisper-large-v3"
            if whisper_device == "cuda":
                try:
                    total_mem_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
                    if total_mem_gb < 16:
                        model_id = "openai/whisper-medium"
                        print(
                            f"  GPU memory {total_mem_gb:.1f}GB < 16GB; using {model_id} for stability. "
                            "Set WHISPER_MODEL to override."
                        )
                except Exception as e:
                    print(f"  Warning: could not read GPU memory ({e}); using {model_id}.")
        use_fast_env = os.environ.get("WHISPER_USE_FAST")
        processor_kwargs = {}
        if use_fast_env is not None:
            processor_kwargs["use_fast"] = use_fast_env.strip().lower() in ("1", "true", "yes", "y")

        dtype = torch.float16 if whisper_device == "cuda" else torch.float32
        model_kwargs = {
            "low_cpu_mem_usage": True,
            "use_safetensors": True,
        }

        use_device_map = whisper_device == "cuda"
        print(
            f"  Whisper model: {model_id} (device={whisper_device}, device_map={'auto' if use_device_map else 'none'})"
        )
        processor = AutoProcessor.from_pretrained(model_id, **processor_kwargs)

        pipeline_kwargs = {
            "model": model_id,
            "tokenizer": processor.tokenizer,
            "feature_extractor": processor.feature_extractor,
            "torch_dtype": dtype,
            "model_kwargs": model_kwargs,
        }
        if use_device_map:
            pipeline_kwargs["device_map"] = "auto"
        else:
            pipeline_kwargs["device"] = whisper_device_index

        return pipeline("automatic-speech-recognition", **pipeline_kwargs)

    def compute_embedding(self, wav_path):
        """Compute embedding for a wav file using Pyannote Inference"""
        # Pyannote inference handles loading/resampling internally usually,
        # but explicit loading is safer for ensuring device
        emb = self.inference(str(wav_path))
        return self._normalize_embedding(emb)

    def get_speaker_embeddings(self, audio_path, diarization, top_k=5):
        """
        Extracts embeddings for each speaker found in the diarization.
        Returns generic 'speaker_label' -> averaged embedding vector.
        """
        speaker_embeddings = {}

        # Group segments by speaker
        speaker_segments = {}
        for turn, _, speaker in diarization.itertracks(yield_label=True):
            if speaker not in speaker_segments:
                speaker_segments[speaker] = []
            speaker_segments[speaker].append(turn)

        # Compute embedding for longest segments of each speaker
        full_audio, sr = torchaudio.load(audio_path)

        for speaker, segments in speaker_segments.items():
            # Sort by duration, take top K longest
            segments.sort(key=lambda s: s.duration, reverse=True)
            top_segments = segments[:top_k]

            embeddings = []
            print(f"  Computing embedding for {speaker} using {len(top_segments)} segments...")

            for seg in top_segments:
                # Extract audio
                start_sample = int(seg.start * sr)
                end_sample = int(seg.end * sr)
                clip = full_audio[:, start_sample:end_sample]

                # Save temp to compute embedding (Pyannote Inference takes path or tensor, path is safer/standard api)
                temp_path = self.output_dir / f"temp_{speaker}_{start_sample}.wav"
                # Avoid torchaudio.save torchcodec dependency
                sf.write(temp_path, clip.squeeze().cpu().numpy(), sr)

                try:
                    emb = self.compute_embedding(temp_path)
                    embeddings.append(emb)
                finally:
                    if temp_path.exists():
                        temp_path.unlink()

            if embeddings:
                # Average them
                avg_emb = np.mean(np.stack(embeddings), axis=0)
                speaker_embeddings[speaker] = avg_emb

        return speaker_embeddings

    def _speaker_durations(self, diarization):
        durations = {}
        total = 0.0
        for turn, _, speaker in diarization.itertracks(yield_label=True):
            dur = float(turn.end - turn.start)
            durations[speaker] = durations.get(speaker, 0.0) + dur
            total += dur
        return durations, total

    def _select_target_speakers(
        self,
        speaker_embs,
        diarization,
        master_ref_emb,
        threshold,
        selection,
        min_margin,
        min_share,
    ):
        durations, total = self._speaker_durations(diarization)
        scored = []
        for spk, emb in speaker_embs.items():
            dist = float(cosine(emb, master_ref_emb))
            scored.append(
                {
                    "speaker": spk,
                    "dist": dist,
                    "duration": float(durations.get(spk, 0.0)),
                }
            )
        scored.sort(key=lambda x: x["dist"])
        if not scored:
            return [], scored, "no_speakers"

        if selection == "threshold_all":
            target = [s["speaker"] for s in scored if s["dist"] <= threshold]
            return target, scored, "threshold_all" if target else "no_match"

        # default: closest speaker only
        best = scored[0]
        if best["dist"] > threshold:
            return [], scored, "best_above_threshold"

        if len(scored) > 1:
            margin = scored[1]["dist"] - best["dist"]
            if margin < min_margin:
                return [], scored, f"ambiguous_margin_{margin:.4f}"

        share = (best["duration"] / total) if total > 0 else 0.0
        if min_share > 0.0 and share < min_share:
            return [], scored, f"low_share_{share:.3f}"

        return [best["speaker"]], scored, "closest"

    def chunk_large_file(self, wav_path, chunk_duration_min=10):
        """Splits a large wav file into smaller chunks using ffmpeg"""
        try:
            # Check duration first using our robust info
            info = torchaudio.info(str(wav_path))
            duration_s = info.num_frames / info.sample_rate

            if duration_s <= (chunk_duration_min * 60):
                return [wav_path]

            print(
                f"Splitting large file {wav_path.name} ({duration_s / 60:.2f} min) into {chunk_duration_min} min chunks..."
            )

            # Create temp dir for chunks
            chunk_dir = self.output_dir / "temp_chunks" / wav_path.stem
            chunk_dir.mkdir(parents=True, exist_ok=True)

            # Use ffmpeg to split
            # segment_time is compatible with most ffmpeg versions
            out_pattern = str(chunk_dir / f"{wav_path.stem}_%03d.wav")

            cmd = [
                "ffmpeg",
                "-y",
                "-i",
                str(wav_path),
                "-f",
                "segment",
                "-segment_time",
                str(chunk_duration_min * 60),
                "-c",
                "copy",
                out_pattern,
            ]

            subprocess.run(cmd, check=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

            chunks = list(chunk_dir.glob("*.wav"))
            print(f"  -> Generated {len(chunks)} chunks.")
            return sorted(chunks)

        except Exception as e:
            print(f"Error chunking file {wav_path}: {e}")
            return [wav_path]  # Fallback to whole file

    def run(
        self,
        trusted_dir,
        untrusted_dir=None,
        threshold=0.4,
        skip_trusted=False,
        speaker_selection="closest",
        speaker_margin=0.05,
        min_speaker_share=0.4,
        segment_verify=False,
        segment_threshold=None,
    ):  # Cosine distance threshold (lower is better)
        metadata_lines = []
        speaker_audit_path = self.output_dir / "speaker_audit.jsonl"
        segment_verify = bool(segment_verify)
        if segment_threshold is None:
            segment_threshold = threshold

        # 1. Build Reference Embedding from Trusted Data
        print("--- Phase 1: Processing Trusted Data (Building Reference) ---")
        trusted_files = list(Path(trusted_dir).rglob("*.wav"))
        if not trusted_files:
            raise ValueError(f"No trusted wav files found in {trusted_dir}")

        # DEFENSIVE: Filter out huge trusted files to prevent slowdowns/OOM
        safe_trusted = []
        MAX_TRUSTED_SIZE_BYTES = 150 * 1024 * 1024  # 150MB limit
        for f in trusted_files:
            size = f.stat().st_size
            if size > MAX_TRUSTED_SIZE_BYTES:
                print(
                    f"WARNING: Skipping trusted file {f.name} (Size: {size / 1024 / 1024:.2f} MB) - Exceeds safety limit."
                )
                continue
            safe_trusted.append(f)
        if not safe_trusted:
            raise ValueError(f"No trusted wav files under size limit in {trusted_dir}")

        ref_embeddings = []
        # Use first 5 or specific named files as anchors if available
        safe_trusted = sorted(safe_trusted, key=lambda p: p.stat().st_size)
        anchors = safe_trusted[: min(len(safe_trusted), 5)]

        for anchor in anchors:
            print(f"  Encoding Anchor: {anchor.name}")
            try:
                emb = self.compute_embedding(anchor)
                ref_embeddings.append(emb)
            except Exception as e:
                print(f"    Failed to encode anchor {anchor}: {e}")

        if not ref_embeddings:
            print("Critical: No reference embeddings created!")
            return

        # Average Reference for downstream untrusted filtering
        master_ref_emb = np.mean(np.stack(ref_embeddings), axis=0)
        try:
            np.save(self.output_dir / "ref_embedding.npy", master_ref_emb)
            with open(self.output_dir / "ref_anchors.json", "w", encoding="utf-8") as f:
                json.dump({"anchors": [str(p) for p in anchors]}, f)
        except Exception as e:
            print(f"Warning: failed to save reference embedding: {e}")

        if skip_trusted:
            print("--- Phase 1b: Transcribing Trusted - SKIPPED (User Request) ---")
        else:
            # Transcribing trusted is intentionally skipped to save time/avoid repetition
            print("--- Phase 1b: Transcribing Trusted - SKIPPED (Optimized) ---")

        # 2. Process Untrusted Data (Diarization -> Verify -> Transcribe)
        if untrusted_dir:
            print("--- Phase 2: Processing Untrusted Data (Diarization + Filtering) ---")
            untrusted_files = sorted(list(Path(untrusted_dir).rglob("*.wav")))

            print(f"Scanning {len(untrusted_files)} files in {untrusted_dir}...")

            # DEFENSIVE: Filter out huge files to prevent OOM
            safe_files = []
            MAX_SIZE_BYTES = 150 * 1024 * 1024  # 150MB Limit

            for f in untrusted_files:
                size = f.stat().st_size
                if size > MAX_SIZE_BYTES:
                    print(
                        f"WARNING: Skipping file {f.name} (Size: {size / 1024 / 1024:.2f} MB) - Exceeds safety limit."
                    )
                    continue
                safe_files.append(f)

            print(f"Processing {len(safe_files)} safe files (filtered from {len(untrusted_files)})...")

            audit_fh = open(speaker_audit_path, "a", encoding="utf-8")
            try:
                for f in safe_files:
                    print(f"Diarizing {f.name}...")
                    try:
                        # A. Run Diarization
                        diarization_start = time.time()
                        diarization = self.diarize_file(f)
                        diarization_elapsed = time.time() - diarization_start
                        print(f"  Diarization complete in {diarization_elapsed:.1f}s")

                        # B. Identify Target Speaker
                        # Get embeddings for all speakers found
                        emb_start = time.time()
                        speaker_embs = self.get_speaker_embeddings(f, diarization)
                        emb_elapsed = time.time() - emb_start
                        print(f"  Speaker embeddings complete in {emb_elapsed:.1f}s")

                        target_speakers, scored, reason = self._select_target_speakers(
                            speaker_embs,
                            diarization,
                            master_ref_emb,
                            threshold,
                            speaker_selection,
                            speaker_margin,
                            min_speaker_share,
                        )
                        for s in scored:
                            print(f"  Speaker {s['speaker']}: Distance {s['dist']:.4f}, Dur {s['duration']:.1f}s")
                        print(f"  Selection: {reason} -> {target_speakers}")

                        audit_fh.write(
                            json.dumps(
                                {
                                    "file": str(f),
                                    "selection": speaker_selection,
                                    "threshold": threshold,
                                    "min_margin": speaker_margin,
                                    "min_share": min_speaker_share,
                                    "reason": reason,
                                    "chosen": target_speakers,
                                    "scores": scored,
                                },
                                ensure_ascii=False,
                            )
                            + "\n"
                        )
                        audit_fh.flush()

                        if not target_speakers:
                            print(f"  Warning: No target speaker found in {f.name}!")
                            continue

                        # C. Extract Valid Segments & Transcribe
                        print("  Extracting and Transcribing valid segments...")

                        # Careful load for slicing
                        full_audio, sr = torchaudio.load(str(f))  # Using path str for my patched load

                        # Resampler for F5
                        resampler_f5 = None
                        if sr != self.target_sr:
                            resampler_f5 = torchaudio.transforms.Resample(sr, self.target_sr).to(full_audio.device)

                        # Optional segment-level verification
                        temp_verify_dir = None
                        if segment_verify:
                            temp_verify_dir = self.output_dir / "temp_verify"
                            temp_verify_dir.mkdir(parents=True, exist_ok=True)

                        # Iterate tracks
                        valid_segments_count = 0
                        rejected_segments = 0
                        rejected_by_similarity = 0
                        for turn, _, speaker in diarization.itertracks(yield_label=True):
                            if speaker not in target_speakers:
                                continue
                            if turn.duration < 1.5:
                                continue  # Skip short

                            # Extract Audio
                            start_s = int(turn.start * sr)
                            end_s = int(turn.end * sr)

                            # Boundary check
                            if end_s > full_audio.shape[1]:
                                end_s = full_audio.shape[1]

                            seg_audio = full_audio[:, start_s:end_s]

                            # Mix to mono
                            if seg_audio.shape[0] > 1:
                                seg_audio_mono = torch.mean(seg_audio, dim=0, keepdim=True)
                            else:
                                seg_audio_mono = seg_audio

                            # Transcribe
                            # Note: Transcribing short segments individually can be hallucination-prone.
                            # Preferable to transcribe whole file and align, BUT here we want to ensure we ONLY get target audio.
                            # So specific transcription is safer for data purity.

                            # Better: Transcribe ONLY this segment
                            try:
                                audio_input = seg_audio_mono.squeeze(0).cpu().numpy()
                                # Convert to 16k for Whisper to avoid sampling_rate incompatibility
                                if sr != 16000:
                                    audio_input = librosa.resample(audio_input, orig_sr=sr, target_sr=16000)
                                res = self.transcriber(
                                    audio_input,
                                    return_timestamps=False,
                                    generate_kwargs={"language": self.whisper_language, "task": "transcribe"},
                                )
                                text = res["text"].strip()

                                if len(text) < 2:
                                    rejected_segments += 1
                                    continue
                                if not self._text_quality_ok(text):
                                    rejected_segments += 1
                                    continue

                                if segment_verify:
                                    temp_verify_path = (
                                        temp_verify_dir / f"verify_{f.stem}_{speaker}_{int(turn.start * 1000)}.wav"
                                    )
                                    sf.write(temp_verify_path, seg_audio_mono.squeeze().cpu().numpy(), sr)
                                    try:
                                        seg_emb = self.compute_embedding(temp_verify_path)
                                        seg_dist = float(cosine(seg_emb, master_ref_emb))
                                    finally:
                                        if temp_verify_path.exists():
                                            temp_verify_path.unlink()
                                    if seg_dist > segment_threshold:
                                        rejected_by_similarity += 1
                                        continue

                                # Use original SR audio for saving to avoid double resampling quality loss?
                                # Actually we need target_sr for F5.

                                # Resample
                                if resampler_f5:
                                    seg_audio_f5 = resampler_f5(
                                        seg_audio
                                    )  # Re-use stereo/original channels or mono? F5 usually mono.
                                else:
                                    seg_audio_f5 = seg_audio

                                if seg_audio_f5.shape[0] > 1:
                                    seg_audio_f5 = torch.mean(seg_audio_f5, dim=0, keepdim=True)

                                seg_name = f"{f.stem}_{speaker}_{turn.start:.2f}.wav"
                                seg_path = self.wavs_dir / seg_name
                                sf.write(seg_path, seg_audio_f5.squeeze().cpu().numpy(), self.target_sr)

                                metadata_lines.append(f"{seg_path.absolute()}|{text}")
                                valid_segments_count += 1

                            except Exception as e:
                                print(f"Error transcribing segment: {e}")
                                rejected_segments += 1

                        print(
                            f"  -> Extracted {valid_segments_count} segments "
                            f"(rejected={rejected_segments}, similarity_reject={rejected_by_similarity})."
                        )

                        # Explicit Cleanup
                        del full_audio
                        del diarization
                        if resampler_f5:
                            del resampler_f5
                        torch.cuda.empty_cache()

                    except Exception as e:
                        print(f"Failed to process chunk {f}: {e}")
                        torch.cuda.empty_cache()
            finally:
                audit_fh.close()

        # 3. Save Output
        self.generate_arrow(metadata_lines)

    def generate_arrow(self, metadata_lines):
        # Same as before
        data_dicts = []
        durations = []

        print(f"Building dataset from {len(metadata_lines)} segments...")
        for line in tqdm(metadata_lines):
            parts = line.split("|")
            if len(parts) < 2:
                continue
            wav_path = parts[0]
            text = parts[1]
            try:
                info = sf.info(wav_path)
                data_dicts.append({"audio_path": wav_path, "text": text, "duration": info.duration})
                durations.append(info.duration)
            except Exception:
                pass

        if not data_dicts:
            print("Error: No valid data found!")
            return

        ds = Dataset.from_list(data_dicts)
        ds.save_to_disk(str(self.output_dir / "raw"))

        with open(self.output_dir / "duration.json", "w") as f:
            json.dump({"duration": durations}, f)
        print(f"Saved dataset to {self.output_dir / 'raw'}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--trusted_dir", required=True)
    parser.add_argument("--untrusted_dir")
    parser.add_argument("--output_dir", required=True)
    parser.add_argument("--threshold", type=float, default=0.35)
    parser.add_argument("--device", help="Force specific device (e.g. 'cpu', 'cuda')")
    parser.add_argument("--skip_trusted", action="store_true", help="Skip Phase 1 (Trusted/Reference building)")
    parser.add_argument(
        "--speaker_selection",
        choices=["closest", "threshold_all"],
        default=os.environ.get("SPEAKER_SELECTION", "closest"),
        help="Speaker selection strategy (default: closest speaker only)",
    )
    parser.add_argument(
        "--speaker_margin",
        type=float,
        default=float(os.environ.get("SPEAKER_MARGIN", "0.05")),
        help="Minimum distance margin vs 2nd closest speaker (closest mode)",
    )
    parser.add_argument(
        "--min_speaker_share",
        type=float,
        default=float(os.environ.get("SPEAKER_MIN_SHARE", "0.4")),
        help="Minimum share of diarized speech for selected speaker (closest mode)",
    )
    seg_thr_default = None
    seg_thr_env = os.environ.get("SEGMENT_THRESHOLD")
    if seg_thr_env:
        try:
            seg_thr_default = float(seg_thr_env)
        except ValueError:
            seg_thr_default = None
    parser.add_argument(
        "--segment_verify",
        action="store_true",
        default=env_flag("SEGMENT_VERIFY", False),
        help="Enable segment-level speaker verification",
    )
    parser.add_argument(
        "--segment_threshold",
        type=float,
        default=seg_thr_default,
        help="Distance threshold for segment verification (default: use --threshold)",
    )
    args = parser.parse_args()

    DataPreparer(args.output_dir, device=args.device).run(
        args.trusted_dir,
        args.untrusted_dir,
        args.threshold,
        args.skip_trusted,
        speaker_selection=args.speaker_selection,
        speaker_margin=args.speaker_margin,
        min_speaker_share=args.min_speaker_share,
        segment_verify=args.segment_verify,
        segment_threshold=args.segment_threshold,
    )