     1→"""GPU-accelerated audio processing - semantically equivalent to CPU path.
     2→
     3→This module provides a drop-in replacement for segment processing that:
     4→1. Uses EXACT same chunking logic (make_chunk_windows)
     5→2. Uses EXACT same metrics computation (compute_audio_metrics with pyloudnorm)
     6→3. Only accelerates I/O-bound operations:
     7→   - Parallel FLAC decoding
     8→   - GPU-batched resampling
     9→   - Parallel WAV encoding
    10→
    11→The GPU path produces IDENTICAL outputs to the CPU path.
    12→"""
    13→
    14→from __future__ import annotations
    15→
    16→import io
    17→import math
    18→import random
    19→import threading
    20→from concurrent.futures import ThreadPoolExecutor
    21→from dataclasses import dataclass
    22→from typing import Any, Optional
    23→
    24→import numpy as np
    25→import soundfile as sf
    26→
    27→# Import exact same functions from prepare_shards for equivalence
    28→from prepare_shards import (
    29→    ChunkConfig,
    30→    SilenceConfig,
    31→    SegmentOutput,
    32→    make_chunk_windows,
    33→    compute_audio_metrics,
    34→    compute_quality_score,
    35→    duration_bucket,
    36→    classify_native_bucket,
    37→)
    38→
    39→# GPU acceleration (optional)
    40→try:
    41→    import torch
    42→    import torchaudio
    43→    HAS_TORCH = True
    44→except ImportError:
    45→    HAS_TORCH = False
    46→    torch = None
    47→    torchaudio = None
    48→
    49→
    50→@dataclass
    51→class SegmentInput:
    52→    """Input segment for processing."""
    53→    video_id: str
    54→    language: str
    55→    segment_bytes: bytes
    56→    segment_meta: dict[str, Any]
    57→    speaker: str
    58→    index: int
    59→    is_val: bool
    60→    rng_seed: int
    61→
    62→
    63→@dataclass
    64→class OutputSegment:
    65→    """Final output ready for shard writing."""
    66→    sample_rate: int
    67→    audio_bytes: bytes
    68→    metadata: dict[str, Any]
    69→    is_val: bool
    70→
    71→
    72→class GPUResampler:
    73→    """GPU-accelerated resampler using torchaudio."""
    74→
    75→    def __init__(self, device: str = "cuda"):
    76→        self.use_gpu = HAS_TORCH and torch.cuda.is_available()
    77→        self.device = torch.device(device if self.use_gpu else "cpu")
    78→        self._resamplers: dict[tuple[int, int], Any] = {}
    79→        self._lock = threading.Lock()
    80→
    81→        if self.use_gpu:
    82→            print(f"GPUResampler: Using {torch.cuda.get_device_name()}")
    83→
    84→    def resample(self, audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray:
    85→        """Resample audio using GPU if available."""
    86→        if src_sr == dst_sr:
    87→            return audio
    88→
    89→        if not self.use_gpu:
    90→            # CPU fallback using scipy
    91→            from scipy import signal
    92→            gcd = math.gcd(src_sr, dst_sr)
    93→            up = dst_sr // gcd
    94→            down = src_sr // gcd
    95→            return signal.resample_poly(audio, up, down).astype(np.float32, copy=False)
    96→
    97→        # GPU path
    98→        resampler = self._get_resampler(src_sr, dst_sr)
    99→        audio_t = torch.from_numpy(audio).unsqueeze(0).to(self.device)
   100→        with torch.no_grad():
   101→            resampled_t = resampler(audio_t)
   102→        return resampled_t.squeeze(0).cpu().numpy()
   103→
   104→    def resample_batch(
   105→        self, audios: list[np.ndarray], src_sr: int, dst_sr: int
   106→    ) -> list[np.ndarray]:
   107→        """Resample multiple audio arrays to the same target rate."""
   108→        if src_sr == dst_sr:
   109→            return audios
   110→
   111→        if not self.use_gpu or not audios:
   112→            return [self.resample(a, src_sr, dst_sr) for a in audios]
   113→
   114→        # Batch on GPU
   115→        resampler = self._get_resampler(src_sr, dst_sr)
   116→
   117→        # Pad to max length for batching
   118→        lengths = [a.shape[0] for a in audios]
   119→        max_len = max(lengths)
   120→        batch = np.zeros((len(audios), max_len), dtype=np.float32)
   121→        for i, a in enumerate(audios):
   122→            batch[i, :a.shape[0]] = a
   123→
   124→        batch_t = torch.from_numpy(batch).to(self.device)
   125→        with torch.no_grad():
   126→            resampled_t = resampler(batch_t)
   127→        resampled_batch = resampled_t.cpu().numpy()
   128→
   129→        # Trim to actual lengths
   130→        ratio = dst_sr / src_sr
   131→        results = []
   132→        for i, orig_len in enumerate(lengths):
   133→            new_len = int(orig_len * ratio)
   134→            results.append(resampled_batch[i, :new_len])
   135→
   136→        return results
   137→
   138→    def _get_resampler(self, src_sr: int, dst_sr: int):
   139→        """Get or create a resampler for the given rate pair."""
   140→        key = (src_sr, dst_sr)
   141→        if key not in self._resamplers:
   142→            with self._lock:
   143→                if key not in self._resamplers:
   144→                    self._resamplers[key] = torchaudio.transforms.Resample(
   145→                        src_sr, dst_sr, dtype=torch.float32
   146→                    ).to(self.device)
   147→        return self._resamplers[key]
   148→
   149→
   150→class EquivalentBatchProcessor:
   151→    """Batch processor that produces IDENTICAL outputs to CPU path.
   152→
   153→    This processor:
   154→    1. Decodes segments in parallel (CPU ThreadPool)
   155→    2. Applies EXACT same chunking logic as CPU path
   156→    3. Computes metrics using EXACT same functions (including pyloudnorm LUFS)
   157→    4. Resamples using GPU (same algorithm, just faster)
   158→    5. Encodes in parallel (CPU ThreadPool)
   159→
   160→    The only difference from CPU is parallelism/hardware, not algorithms.
   161→    """
   162→
   163→    def __init__(
   164→        self,
   165→        decode_workers: int = 8,
   166→        encode_workers: int = 8,
   167→        output_format: str = "wav",
   168→        chunk_cfg: Optional[ChunkConfig] = None,
   169→        silence_cfg: Optional[SilenceConfig] = None,
   170→    ):
   171→        self.decode_workers = decode_workers
   172→        self.encode_workers = encode_workers
   173→        self.output_format = output_format
   174→        self.chunk_cfg = chunk_cfg or ChunkConfig(
   175→            enabled=True, min_sec=6.0, max_sec=10.0, threshold_sec=12.0, all_buckets=False
   176→        )
   177→        self.silence_cfg = silence_cfg or SilenceConfig(
   178→            enabled=True, frame_sec=0.05, hop_sec=0.01,
   179→            search_sec=0.5, percentile=20.0, db_drop=20.0
   180→        )
   181→
   182→        self.resampler = GPUResampler()
   183→
   184→        self.stats = {
   185→            "decoded": 0,
   186→            "chunks_emitted": 0,
   187→            "segments_chunked": 0,
   188→            "dropped_too_short": 0,
   189→            "dropped_too_long": 0,
   190→            "dropped_clipped": 0,
   191→        }
   192→
   193→    def process_segments(
   194→        self, segments: list[SegmentInput]
   195→    ) -> list[OutputSegment]:
   196→        """Process segments and return outputs - equivalent to CPU path.
   197→
   198→        This method mirrors the exact logic of process_segment() in prepare_shards.py.
   199→        """
   200→        if not segments:
   201→            return []
   202→
   203→        # Step 1: Parallel decode
   204→        decoded = self._decode_parallel(segments)
   205→
   206→        # Step 2: Process each segment (chunking + metrics + resample + encode)
   207→        all_outputs: list[OutputSegment] = []
   208→
   209→        for seg_input, audio, sr in decoded:
   210→            outputs = self._process_single_segment(seg_input, audio, sr)
   211→            all_outputs.extend(outputs)
   212→
   213→        return all_outputs
   214→
   215→    def _decode_parallel(
   216→        self, segments: list[SegmentInput]
   217→    ) -> list[tuple[SegmentInput, np.ndarray, int]]:
   218→        """Decode segments in parallel."""
   219→        def decode_one(seg: SegmentInput) -> Optional[tuple[SegmentInput, np.ndarray, int]]:
   220→            try:
   221→                audio, sr = sf.read(io.BytesIO(seg.segment_bytes), dtype="float32")
   222→                if audio.ndim > 1:
   223→                    audio = np.mean(audio, axis=1)
   224→                return (seg, audio.astype(np.float32), sr)
   225→            except Exception:
   226→                return None
   227→
   228→        with ThreadPoolExecutor(max_workers=self.decode_workers) as pool:
   229→            results = list(pool.map(decode_one, segments))
   230→
   231→        decoded = [r for r in results if r is not None]
   232→        self.stats["decoded"] += len(decoded)
   233→        return decoded
   234→
   235→    def _process_single_segment(
   236→        self,
   237→        seg: SegmentInput,
   238→        audio: np.ndarray,
   239→        sr: int,
   240→    ) -> list[OutputSegment]:
   241→        """Process a single segment - mirrors process_segment() exactly."""
   242→        outputs: list[OutputSegment] = []
   243→
   244→        # Create RNG with same seed as CPU path
   245→        rng = random.Random(seg.rng_seed)
   246→
   247→        # Get chunk windows - EXACT same logic as CPU
   248→        windows = make_chunk_windows(audio, sr, self.chunk_cfg, self.silence_cfg, rng)
   249→
   250→        music_stats = seg.segment_meta.get("music_stats") or {}
   251→        has_music = bool(music_stats.get("has_music", False))
   252→
   253→        sample_id_base = f"{seg.video_id}__{seg.speaker}_{seg.index:04d}"
   254→        use_chunk_suffix = len(windows) > 1
   255→        if use_chunk_suffix:
   256→            self.stats["segments_chunked"] += 1
   257→
   258→        segment_start = seg.segment_meta.get("start")
   259→        segment_end = seg.segment_meta.get("end")
   260→        original_start = seg.segment_meta.get("original_start")
   261→        original_end = seg.segment_meta.get("original_end")
   262→
   263→        # Process 16k outputs for all chunk windows
   264→        for chunk_idx, (offset_sec, chunk_len_sec) in enumerate(windows):
   265→            start_idx = int(round(offset_sec * sr))
   266→            end_idx = int(round((offset_sec + chunk_len_sec) * sr))
   267→            audio_slice = audio[start_idx:end_idx]
   268→            if audio_slice.size == 0:
   269→                continue
   270→
   271→            slice_duration = float(audio_slice.shape[0] / sr)
   272→
   273→            # Compute metrics - EXACT same function as CPU
   274→            slice_metrics = compute_audio_metrics(audio_slice, sr)
   275→
   276→            chunk_start = float(segment_start) + offset_sec if segment_start is not None else None
   277→            chunk_end = chunk_start + slice_duration if chunk_start is not None else None
   278→            chunk_original_start = (
   279→                float(original_start) + offset_sec if original_start is not None else None
   280→            )
   281→            chunk_original_end = (
   282→                chunk_original_start + slice_duration if chunk_original_start is not None else None
   283→            )
   284→
   285→            sample_id = sample_id_base
   286→            chunk_index = None
   287→            if use_chunk_suffix:
   288→                sample_id = f"{sample_id_base}_c{chunk_idx:02d}"
   289→                chunk_index = chunk_idx
   290→
   291→            # Emit 16k output
   292→            chunk_outputs = self._emit_outputs(
   293→                audio_slice=audio_slice,
   294→                sr=sr,
   295→                metrics=slice_metrics,
   296→                sample_id=sample_id,
   297→                video_id=seg.video_id,
   298→                language=seg.language,
   299→                segment_index=seg.index,
   300→                chunk_index=chunk_index,
   301→                start=chunk_start,
   302→                end=chunk_end,
   303→                original_start=chunk_original_start,
   304→                original_end=chunk_original_end,
   305→                target_srs=[16000],
   306→                has_music=has_music,
   307→                music_stats=music_stats,
   308→                speaker=seg.speaker,
   309→                is_val=seg.is_val,
   310→            )
   311→            outputs.extend(chunk_outputs)
   312→
   313→            if chunk_outputs:
   314→                self.stats["chunks_emitted"] += 1
   315→
   316→        # High-rate buckets: use full segment unless chunk_all_buckets
   317→        full_metrics = compute_audio_metrics(audio, sr)
   318→        effective_nyquist = full_metrics["effective_nyquist_hz"]
   319→        targets_high: list[int] = []
   320→        if effective_nyquist >= 8000:
   321→            targets_high.append(24000)
   322→        if effective_nyquist >= 12000:
   323→            targets_high.append(32000)
   324→        if effective_nyquist >= 16000:
   325→            targets_high.append(44100)
   326→
   327→        if targets_high:
   328→            if self.chunk_cfg.all_buckets and use_chunk_suffix:
   329→                # Emit high-rate for each chunk
   330→                for chunk_idx, (offset_sec, chunk_len_sec) in enumerate(windows):
   331→                    start_idx = int(round(offset_sec * sr))
   332→                    end_idx = int(round((offset_sec + chunk_len_sec) * sr))
   333→                    audio_slice = audio[start_idx:end_idx]
   334→                    if audio_slice.size == 0:
   335→                        continue
   336→
   337→                    slice_duration = float(audio_slice.shape[0] / sr)
   338→                    slice_metrics = compute_audio_metrics(audio_slice, sr)
   339→                    slice_effective_nyquist = slice_metrics["effective_nyquist_hz"]
   340→
   341→                    slice_targets: list[int] = []
   342→                    if slice_effective_nyquist >= 8000:
   343→                        slice_targets.append(24000)
   344→                    if slice_effective_nyquist >= 12000:
   345→                        slice_targets.append(32000)
   346→                    if slice_effective_nyquist >= 16000:
   347→                        slice_targets.append(44100)
   348→                    if not slice_targets:
   349→                        continue
   350→
   351→                    chunk_start = (
   352→                        float(segment_start) + offset_sec if segment_start is not None else None
   353→                    )
   354→                    chunk_end = chunk_start + slice_duration if chunk_start is not None else None
   355→                    chunk_original_start = (
   356→                        float(original_start) + offset_sec if original_start is not None else None
   357→                    )
   358→                    chunk_original_end = (
   359→                        chunk_original_start + slice_duration
   360→                        if chunk_original_start is not None
   361→                        else None
   362→                    )
   363→
   364→                    sample_id = f"{sample_id_base}_c{chunk_idx:02d}"
   365→                    chunk_outputs = self._emit_outputs(
   366→                        audio_slice=audio_slice,
   367→                        sr=sr,
   368→                        metrics=slice_metrics,
   369→                        sample_id=sample_id,
   370→                        video_id=seg.video_id,
   371→                        language=seg.language,
   372→                        segment_index=seg.index,
   373→                        chunk_index=chunk_idx,
   374→                        start=chunk_start,
   375→                        end=chunk_end,
   376→                        original_start=chunk_original_start,
   377→                        original_end=chunk_original_end,
   378→                        target_srs=slice_targets,
   379→                        has_music=has_music,
   380→                        music_stats=music_stats,
   381→                        speaker=seg.speaker,
   382→                        is_val=seg.is_val,
   383→                    )
   384→                    outputs.extend(chunk_outputs)
   385→            else:
   386→                # Emit high-rate for full segment
   387→                high_outputs = self._emit_outputs(
   388→                    audio_slice=audio,
   389→                    sr=sr,
   390→                    metrics=full_metrics,
   391→                    sample_id=sample_id_base,
   392→                    video_id=seg.video_id,
   393→                    language=seg.language,
   394→                    segment_index=seg.index,
   395→                    chunk_index=None,
   396→                    start=segment_start,
   397→                    end=segment_end,
   398→                    original_start=original_start,
   399→                    original_end=original_end,
   400→                    target_srs=targets_high,
   401→                    has_music=has_music,
   402→                    music_stats=music_stats,
   403→                    speaker=seg.speaker,
   404→                    is_val=seg.is_val,
   405→                )
   406→                outputs.extend(high_outputs)
   407→
   408→        return outputs
   409→
   410→    def _emit_outputs(
   411→        self,
   412→        audio_slice: np.ndarray,
   413→        sr: int,
   414→        metrics: dict[str, float],
   415→        sample_id: str,
   416→        video_id: str,
   417→        language: str,
   418→        segment_index: int,
   419→        chunk_index: Optional[int],
   420→        start: Optional[float],
   421→        end: Optional[float],
   422→        original_start: Optional[float],
   423→        original_end: Optional[float],
   424→        target_srs: list[int],
   425→        has_music: bool,
   426→        music_stats: dict,
   427→        speaker: str,
   428→        is_val: bool,
   429→    ) -> list[OutputSegment]:
   430→        """Emit outputs for target sample rates - mirrors emit_outputs() exactly."""
   431→        outputs: list[OutputSegment] = []
   432→
   433→        slice_duration = float(audio_slice.shape[0] / sr)
   434→        bucket = duration_bucket(slice_duration)
   435→        if bucket is None:
   436→            if slice_duration < 3.0:
   437→                self.stats["dropped_too_short"] += 1
   438→            else:
   439→                self.stats["dropped_too_long"] += 1
   440→            return outputs
   441→
   442→        if metrics["clipping_fraction"] > 0.1:
   443→            self.stats["dropped_clipped"] += 1
   444→            return outputs
   445→
   446→        duration_bucket_name, base_weight = bucket
   447→        sampling_weight = base_weight * (0.5 if has_music else 1.0)
   448→        quality_score = compute_quality_score(metrics, has_music)
   449→        effective_nyquist = metrics["effective_nyquist_hz"]
   450→
   451→        base_metadata = {
   452→            "id": sample_id,
   453→            "video_id": video_id,
   454→            "language": language,
   455→            "segment_index": segment_index,
   456→            "chunk_index": chunk_index,
   457→            "start": start,
   458→            "end": end,
   459→            "original_start": original_start,
   460→            "original_end": original_end,
   461→            "duration_sec": slice_duration,
   462→            "native_sr": sr,
   463→            "effective_nyquist_hz": effective_nyquist,
   464→            "spectral_rolloff_hz": metrics["spectral_rolloff_hz"],
   465→            "high_freq_energy_ratio": metrics["high_freq_energy_ratio"],
   466→            "is_truly_wideband": bool(effective_nyquist >= 12000),
   467→            "native_sr_bucket": classify_native_bucket(effective_nyquist),
   468→            "loudness_lufs": metrics["loudness_lufs"],
   469→            "loudness_rms_db": metrics["loudness_rms_db"],
   470→            "clipping_fraction": metrics["clipping_fraction"],
   471→            "dc_offset": metrics["dc_offset"],  # Include dc_offset!
   472→            "snr_proxy": metrics["snr_proxy"],
   473→            "quality_score": quality_score,
   474→            "music_prob": float(music_stats.get("music_mean", 0.0)),
   475→            "music_stats": music_stats,
   476→            "speaker": speaker,
   477→            "vad_speech_regions": [[0.0, slice_duration]],
   478→            "duration_bucket": duration_bucket_name,
   479→            "sampling_weight": sampling_weight,
   480→        }
   481→
   482→        # Resample and encode for each target rate
   483→        # Use GPU resampling for speed
   484→        for target_sr in target_srs:
   485→            resampled = self.resampler.resample(audio_slice, sr, target_sr)
   486→            resampled = np.clip(resampled, -1.0, 1.0)
   487→            audio_bytes = self._encode_audio(resampled, target_sr)
   488→
   489→            meta = dict(base_metadata)
   490→            meta["sample_rate"] = target_sr
   491→            meta["resample_target_sr"] = target_sr
   492→
   493→            outputs.append(OutputSegment(
   494→                sample_rate=target_sr,
   495→                audio_bytes=audio_bytes,
   496→                metadata=meta,
   497→                is_val=is_val,
   498→            ))
   499→
   500→        return outputs
   501→
   502→    def _encode_audio(self, audio: np.ndarray, sr: int) -> bytes:
   503→        """Encode audio to WAV or FLAC."""
   504→        buf = io.BytesIO()
   505→        if self.output_format == "wav":
   506→            sf.write(buf, audio, sr, format="WAV", subtype="PCM_16")
   507→        else:
   508→            sf.write(buf, audio, sr, format="FLAC", subtype="PCM_16")
   509→        return buf.getvalue()
   510→
   511→
   512→def benchmark():
   513→    """Benchmark the equivalent batch processor."""
   514→    import time
   515→
   516→    print("=" * 60)
   517→    print("Equivalent Batch Processor Benchmark")
   518→    print("=" * 60)
   519→
   520→    # Generate test audio
   521→    sr = 48000
   522→    duration = 15.0  # Long enough to trigger chunking
   523→    num_samples = int(sr * duration)
   524→    t = np.linspace(0, duration, num_samples, dtype=np.float32)
   525→
   526→    # Create test segments
   527→    test_segments = []
   528→    for i in range(20):
   529→        # Mix of frequencies for realistic spectral content
   530→        audio = (
   531→            np.sin(2 * np.pi * 440 * t) * 0.3 +
   532→            np.sin(2 * np.pi * 1000 * t) * 0.2 +
   533→            np.sin(2 * np.pi * 5000 * t) * 0.1 +
   534→            np.random.randn(num_samples).astype(np.float32) * 0.02
   535→        ).astype(np.float32)
   536→
   537→        buf = io.BytesIO()
   538→        sf.write(buf, audio, sr, format="FLAC")
   539→        flac_bytes = buf.getvalue()
   540→
   541→        test_segments.append(SegmentInput(
   542→            video_id=f"test_{i}",
   543→            language="english",
   544→            segment_bytes=flac_bytes,
   545→            segment_meta={"start": 0.0, "end": duration, "original_start": 0.0, "original_end": duration},
   546→            speaker="SPEAKER_00",
   547→            index=i,
   548→            is_val=i % 10 == 0,
   549→            rng_seed=42 + i,
   550→        ))
   551→
   552→    print(f"Created {len(test_segments)} test segments ({duration}s each)")
   553→    print(f"Chunking enabled: segments > 12s will be split into 6-10s windows")
   554→
   555→    # Test processor with default chunking config
   556→    chunk_cfg = ChunkConfig(
   557→        enabled=True,
   558→        min_sec=6.0,
   559→        max_sec=10.0,
   560→        threshold_sec=12.0,
   561→        all_buckets=False,
   562→    )
   563→    silence_cfg = SilenceConfig(
   564→        enabled=True,
   565→        frame_sec=0.05,
   566→        hop_sec=0.01,
   567→        search_sec=0.5,
   568→        percentile=20.0,
   569→        db_drop=20.0,
   570→    )
   571→
   572→    processor = EquivalentBatchProcessor(
   573→        decode_workers=8,
   574→        encode_workers=8,
   575→        output_format="wav",
   576→        chunk_cfg=chunk_cfg,
   577→        silence_cfg=silence_cfg,
   578→    )
   579→    print(f"GPU resampling: {processor.resampler.use_gpu}")
   580→
   581→    start = time.perf_counter()
   582→    outputs = processor.process_segments(test_segments)
   583→    elapsed = time.perf_counter() - start
   584→
   585→    print(f"\nResults:")
   586→    print(f"  Time: {elapsed:.2f}s")
   587→    print(f"  Input segments: {len(test_segments)}")
   588→    print(f"  Output samples: {len(outputs)}")
   589→    print(f"  Segments chunked: {processor.stats['segments_chunked']}")
   590→    print(f"  Chunks emitted: {processor.stats['chunks_emitted']}")
   591→    print(f"  Throughput: {len(test_segments)/elapsed:.1f} input segs/sec")
   592→    print(f"  Stats: {processor.stats}")
   593→
   594→    # Count outputs by sample rate
   595→    sr_counts = {}
   596→    for out in outputs:
   597→        sr_counts[out.sample_rate] = sr_counts.get(out.sample_rate, 0) + 1
   598→    print(f"\nOutputs by sample rate:")
   599→    for sr, count in sorted(sr_counts.items()):
   600→        print(f"  {sr}: {count}")
   601→
   602→    # Show sample output metadata
   603→    if outputs:
   604→        out = outputs[0]
   605→        print(f"\nSample output metadata:")
   606→        for key in ["id", "duration_sec", "loudness_lufs", "dc_offset", "snr_proxy",
   607→                    "effective_nyquist_hz", "native_sr_bucket", "chunk_index"]:
   608→            if key in out.metadata:
   609→                print(f"  {key}: {out.metadata[key]}")
   610→
   611→
   612→if __name__ == "__main__":
   613→    benchmark()
   614→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

"""GPU-accelerated audio processing - semantically equivalent to CPU path.

This module provides a drop-in replacement for segment processing that:
1. Uses EXACT same chunking logic (make_chunk_windows)
2. Uses EXACT same metrics computation (compute_audio_metrics with pyloudnorm)
3. Only accelerates I/O-bound operations:
   - Parallel FLAC decoding
   - GPU-batched resampling
   - Parallel WAV encoding

The GPU path produces IDENTICAL outputs to the CPU path.
"""

from __future__ import annotations

import io
import math
import random
import threading
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import Any, Optional

import numpy as np
import soundfile as sf

# Import exact same functions from prepare_shards for equivalence
from prepare_shards import (
    ChunkConfig,
    SilenceConfig,
    SegmentOutput,
    make_chunk_windows,
    compute_audio_metrics,
    compute_quality_score,
    duration_bucket,
    classify_native_bucket,
)

# GPU acceleration (optional)
try:
    import torch
    import torchaudio
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False
    torch = None
    torchaudio = None

@dataclass
class SegmentInput:
    """Input segment for processing."""
    video_id: str
    language: str
    segment_bytes: bytes
    segment_meta: dict[str, Any]
    speaker: str
    index: int
    is_val: bool
    rng_seed: int

@dataclass
class OutputSegment:
    """Final output ready for shard writing."""
    sample_rate: int
    audio_bytes: bytes
    metadata: dict[str, Any]
    is_val: bool

class GPUResampler:
    """GPU-accelerated resampler using torchaudio."""

    def __init__(self, device: str = "cuda"):
        self.use_gpu = HAS_TORCH and torch.cuda.is_available()
        self.device = torch.device(device if self.use_gpu else "cpu")
        self._resamplers: dict[tuple[int, int], Any] = {}
        self._lock = threading.Lock()

        if self.use_gpu:
            print(f"GPUResampler: Using {torch.cuda.get_device_name()}")

    def resample(self, audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray:
        """Resample audio using GPU if available."""
        if src_sr == dst_sr:
            return audio

        if not self.use_gpu:
            # CPU fallback using scipy
            from scipy import signal
            gcd = math.gcd(src_sr, dst_sr)
            up = dst_sr // gcd
            down = src_sr // gcd
            return signal.resample_poly(audio, up, down).astype(np.float32, copy=False)

        # GPU path
        resampler = self._get_resampler(src_sr, dst_sr)
        audio_t = torch.from_numpy(audio).unsqueeze(0).to(self.device)
        with torch.no_grad():
            resampled_t = resampler(audio_t)
        return resampled_t.squeeze(0).cpu().numpy()

    def resample_batch(
        self, audios: list[np.ndarray], src_sr: int, dst_sr: int
    ) -> list[np.ndarray]:
        """Resample multiple audio arrays to the same target rate."""
        if src_sr == dst_sr:
            return audios

        if not self.use_gpu or not audios:
            return [self.resample(a, src_sr, dst_sr) for a in audios]

        # Batch on GPU
        resampler = self._get_resampler(src_sr, dst_sr)

        # Pad to max length for batching
        lengths = [a.shape[0] for a in audios]
        max_len = max(lengths)
        batch = np.zeros((len(audios), max_len), dtype=np.float32)
        for i, a in enumerate(audios):
            batch[i, :a.shape[0]] = a

        batch_t = torch.from_numpy(batch).to(self.device)
        with torch.no_grad():
            resampled_t = resampler(batch_t)
        resampled_batch = resampled_t.cpu().numpy()

        # Trim to actual lengths
        ratio = dst_sr / src_sr
        results = []
        for i, orig_len in enumerate(lengths):
            new_len = int(orig_len * ratio)
            results.append(resampled_batch[i, :new_len])

        return results

    def _get_resampler(self, src_sr: int, dst_sr: int):
        """Get or create a resampler for the given rate pair."""
        key = (src_sr, dst_sr)
        if key not in self._resamplers:
            with self._lock:
                if key not in self._resamplers:
                    self._resamplers[key] = torchaudio.transforms.Resample(
                        src_sr, dst_sr, dtype=torch.float32
                    ).to(self.device)
        return self._resamplers[key]

class EquivalentBatchProcessor:
    """Batch processor that produces IDENTICAL outputs to CPU path.

    This processor:
    1. Decodes segments in parallel (CPU ThreadPool)
    2. Applies EXACT same chunking logic as CPU path
    3. Computes metrics using EXACT same functions (including pyloudnorm LUFS)
    4. Resamples using GPU (same algorithm, just faster)
    5. Encodes in parallel (CPU ThreadPool)

    The only difference from CPU is parallelism/hardware, not algorithms.
    """

    def __init__(
        self,
        decode_workers: int = 8,
        encode_workers: int = 8,
        output_format: str = "wav",
        chunk_cfg: Optional[ChunkConfig] = None,
        silence_cfg: Optional[SilenceConfig] = None,
    ):
        self.decode_workers = decode_workers
        self.encode_workers = encode_workers
        self.output_format = output_format
        self.chunk_cfg = chunk_cfg or ChunkConfig(
            enabled=True, min_sec=6.0, max_sec=10.0, threshold_sec=12.0, all_buckets=False
        )
        self.silence_cfg = silence_cfg or SilenceConfig(
            enabled=True, frame_sec=0.05, hop_sec=0.01,
            search_sec=0.5, percentile=20.0, db_drop=20.0
        )

        self.resampler = GPUResampler()

        self.stats = {
            "decoded": 0,
            "chunks_emitted": 0,
            "segments_chunked": 0,
            "dropped_too_short": 0,
            "dropped_too_long": 0,
            "dropped_clipped": 0,
        }

    def process_segments(
        self, segments: list[SegmentInput]
    ) -> list[OutputSegment]:
        """Process segments and return outputs - equivalent to CPU path.

        This method mirrors the exact logic of process_segment() in prepare_shards.py.
        """
        if not segments:
            return []

        # Step 1: Parallel decode
        decoded = self._decode_parallel(segments)

        # Step 2: Process each segment (chunking + metrics + resample + encode)
        all_outputs: list[OutputSegment] = []

        for seg_input, audio, sr in decoded:
            outputs = self._process_single_segment(seg_input, audio, sr)
            all_outputs.extend(outputs)

        return all_outputs

    def _decode_parallel(
        self, segments: list[SegmentInput]
    ) -> list[tuple[SegmentInput, np.ndarray, int]]:
        """Decode segments in parallel."""
        def decode_one(seg: SegmentInput) -> Optional[tuple[SegmentInput, np.ndarray, int]]:
            try:
                audio, sr = sf.read(io.BytesIO(seg.segment_bytes), dtype="float32")
                if audio.ndim > 1:
                    audio = np.mean(audio, axis=1)
                return (seg, audio.astype(np.float32), sr)
            except Exception:
                return None

        with ThreadPoolExecutor(max_workers=self.decode_workers) as pool:
            results = list(pool.map(decode_one, segments))

        decoded = [r for r in results if r is not None]
        self.stats["decoded"] += len(decoded)
        return decoded

    def _process_single_segment(
        self,
        seg: SegmentInput,
        audio: np.ndarray,
        sr: int,
    ) -> list[OutputSegment]:
        """Process a single segment - mirrors process_segment() exactly."""
        outputs: list[OutputSegment] = []

        # Create RNG with same seed as CPU path
        rng = random.Random(seg.rng_seed)

        # Get chunk windows - EXACT same logic as CPU
        windows = make_chunk_windows(audio, sr, self.chunk_cfg, self.silence_cfg, rng)

        music_stats = seg.segment_meta.get("music_stats") or {}
        has_music = bool(music_stats.get("has_music", False))

        sample_id_base = f"{seg.video_id}__{seg.speaker}_{seg.index:04d}"
        use_chunk_suffix = len(windows) > 1
        if use_chunk_suffix:
            self.stats["segments_chunked"] += 1

        segment_start = seg.segment_meta.get("start")
        segment_end = seg.segment_meta.get("end")
        original_start = seg.segment_meta.get("original_start")
        original_end = seg.segment_meta.get("original_end")

        # Process 16k outputs for all chunk windows
        for chunk_idx, (offset_sec, chunk_len_sec) in enumerate(windows):
            start_idx = int(round(offset_sec * sr))
            end_idx = int(round((offset_sec + chunk_len_sec) * sr))
            audio_slice = audio[start_idx:end_idx]
            if audio_slice.size == 0:
                continue

            slice_duration = float(audio_slice.shape[0] / sr)

            # Compute metrics - EXACT same function as CPU
            slice_metrics = compute_audio_metrics(audio_slice, sr)

            chunk_start = float(segment_start) + offset_sec if segment_start is not None else None
            chunk_end = chunk_start + slice_duration if chunk_start is not None else None
            chunk_original_start = (
                float(original_start) + offset_sec if original_start is not None else None
            )
            chunk_original_end = (
                chunk_original_start + slice_duration if chunk_original_start is not None else None
            )

            sample_id = sample_id_base
            chunk_index = None
            if use_chunk_suffix:
                sample_id = f"{sample_id_base}_c{chunk_idx:02d}"
                chunk_index = chunk_idx

            # Emit 16k output
            chunk_outputs = self._emit_outputs(
                audio_slice=audio_slice,
                sr=sr,
                metrics=slice_metrics,
                sample_id=sample_id,
                video_id=seg.video_id,
                language=seg.language,
                segment_index=seg.index,
                chunk_index=chunk_index,
                start=chunk_start,
                end=chunk_end,
                original_start=chunk_original_start,
                original_end=chunk_original_end,
                target_srs=[16000],
                has_music=has_music,
                music_stats=music_stats,
                speaker=seg.speaker,
                is_val=seg.is_val,
            )
            outputs.extend(chunk_outputs)

            if chunk_outputs:
                self.stats["chunks_emitted"] += 1

        # High-rate buckets: use full segment unless chunk_all_buckets
        full_metrics = compute_audio_metrics(audio, sr)
        effective_nyquist = full_metrics["effective_nyquist_hz"]
        targets_high: list[int] = []
        if effective_nyquist >= 8000:
            targets_high.append(24000)
        if effective_nyquist >= 12000:
            targets_high.append(32000)
        if effective_nyquist >= 16000:
            targets_high.append(44100)

        if targets_high:
            if self.chunk_cfg.all_buckets and use_chunk_suffix:
                # Emit high-rate for each chunk
                for chunk_idx, (offset_sec, chunk_len_sec) in enumerate(windows):
                    start_idx = int(round(offset_sec * sr))
                    end_idx = int(round((offset_sec + chunk_len_sec) * sr))
                    audio_slice = audio[start_idx:end_idx]
                    if audio_slice.size == 0:
                        continue

                    slice_duration = float(audio_slice.shape[0] / sr)
                    slice_metrics = compute_audio_metrics(audio_slice, sr)
                    slice_effective_nyquist = slice_metrics["effective_nyquist_hz"]

                    slice_targets: list[int] = []
                    if slice_effective_nyquist >= 8000:
                        slice_targets.append(24000)
                    if slice_effective_nyquist >= 12000:
                        slice_targets.append(32000)
                    if slice_effective_nyquist >= 16000:
                        slice_targets.append(44100)
                    if not slice_targets:
                        continue

                    chunk_start = (
                        float(segment_start) + offset_sec if segment_start is not None else None
                    )
                    chunk_end = chunk_start + slice_duration if chunk_start is not None else None
                    chunk_original_start = (
                        float(original_start) + offset_sec if original_start is not None else None
                    )
                    chunk_original_end = (
                        chunk_original_start + slice_duration
                        if chunk_original_start is not None
                        else None
                    )

                    sample_id = f"{sample_id_base}_c{chunk_idx:02d}"
                    chunk_outputs = self._emit_outputs(
                        audio_slice=audio_slice,
                        sr=sr,
                        metrics=slice_metrics,
                        sample_id=sample_id,
                        video_id=seg.video_id,
                        language=seg.language,
                        segment_index=seg.index,
                        chunk_index=chunk_idx,
                        start=chunk_start,
                        end=chunk_end,
                        original_start=chunk_original_start,
                        original_end=chunk_original_end,
                        target_srs=slice_targets,
                        has_music=has_music,
                        music_stats=music_stats,
                        speaker=seg.speaker,
                        is_val=seg.is_val,
                    )
                    outputs.extend(chunk_outputs)
            else:
                # Emit high-rate for full segment
                high_outputs = self._emit_outputs(
                    audio_slice=audio,
                    sr=sr,
                    metrics=full_metrics,
                    sample_id=sample_id_base,
                    video_id=seg.video_id,
                    language=seg.language,
                    segment_index=seg.index,
                    chunk_index=None,
                    start=segment_start,
                    end=segment_end,
                    original_start=original_start,
                    original_end=original_end,
                    target_srs=targets_high,
                    has_music=has_music,
                    music_stats=music_stats,
                    speaker=seg.speaker,
                    is_val=seg.is_val,
                )
                outputs.extend(high_outputs)

        return outputs

    def _emit_outputs(
        self,
        audio_slice: np.ndarray,
        sr: int,
        metrics: dict[str, float],
        sample_id: str,
        video_id: str,
        language: str,
        segment_index: int,
        chunk_index: Optional[int],
        start: Optional[float],
        end: Optional[float],
        original_start: Optional[float],
        original_end: Optional[float],
        target_srs: list[int],
        has_music: bool,
        music_stats: dict,
        speaker: str,
        is_val: bool,
    ) -> list[OutputSegment]:
        """Emit outputs for target sample rates - mirrors emit_outputs() exactly."""
        outputs: list[OutputSegment] = []

        slice_duration = float(audio_slice.shape[0] / sr)
        bucket = duration_bucket(slice_duration)
        if bucket is None:
            if slice_duration < 3.0:
                self.stats["dropped_too_short"] += 1
            else:
                self.stats["dropped_too_long"] += 1
            return outputs

        if metrics["clipping_fraction"] > 0.1:
            self.stats["dropped_clipped"] += 1
            return outputs

        duration_bucket_name, base_weight = bucket
        sampling_weight = base_weight * (0.5 if has_music else 1.0)
        quality_score = compute_quality_score(metrics, has_music)
        effective_nyquist = metrics["effective_nyquist_hz"]

        base_metadata = {
            "id": sample_id,
            "video_id": video_id,
            "language": language,
            "segment_index": segment_index,
            "chunk_index": chunk_index,
            "start": start,
            "end": end,
            "original_start": original_start,
            "original_end": original_end,
            "duration_sec": slice_duration,
            "native_sr": sr,
            "effective_nyquist_hz": effective_nyquist,
            "spectral_rolloff_hz": metrics["spectral_rolloff_hz"],
            "high_freq_energy_ratio": metrics["high_freq_energy_ratio"],
            "is_truly_wideband": bool(effective_nyquist >= 12000),
            "native_sr_bucket": classify_native_bucket(effective_nyquist),
            "loudness_lufs": metrics["loudness_lufs"],
            "loudness_rms_db": metrics["loudness_rms_db"],
            "clipping_fraction": metrics["clipping_fraction"],
            "dc_offset": metrics["dc_offset"],  # Include dc_offset!
            "snr_proxy": metrics["snr_proxy"],
            "quality_score": quality_score,
            "music_prob": float(music_stats.get("music_mean", 0.0)),
            "music_stats": music_stats,
            "speaker": speaker,
            "vad_speech_regions": [[0.0, slice_duration]],
            "duration_bucket": duration_bucket_name,
            "sampling_weight": sampling_weight,
        }

        # Resample and encode for each target rate
        # Use GPU resampling for speed
        for target_sr in target_srs:
            resampled = self.resampler.resample(audio_slice, sr, target_sr)
            resampled = np.clip(resampled, -1.0, 1.0)
            audio_bytes = self._encode_audio(resampled, target_sr)

            meta = dict(base_metadata)
            meta["sample_rate"] = target_sr
            meta["resample_target_sr"] = target_sr

            outputs.append(OutputSegment(
                sample_rate=target_sr,
                audio_bytes=audio_bytes,
                metadata=meta,
                is_val=is_val,
            ))

        return outputs

    def _encode_audio(self, audio: np.ndarray, sr: int) -> bytes:
        """Encode audio to WAV or FLAC."""
        buf = io.BytesIO()
        if self.output_format == "wav":
            sf.write(buf, audio, sr, format="WAV", subtype="PCM_16")
        else:
            sf.write(buf, audio, sr, format="FLAC", subtype="PCM_16")
        return buf.getvalue()

def benchmark():
    """Benchmark the equivalent batch processor."""
    import time

    print("=" * 60)
    print("Equivalent Batch Processor Benchmark")
    print("=" * 60)

    # Generate test audio
    sr = 48000
    duration = 15.0  # Long enough to trigger chunking
    num_samples = int(sr * duration)
    t = np.linspace(0, duration, num_samples, dtype=np.float32)

    # Create test segments
    test_segments = []
    for i in range(20):
        # Mix of frequencies for realistic spectral content
        audio = (
            np.sin(2 * np.pi * 440 * t) * 0.3 +
            np.sin(2 * np.pi * 1000 * t) * 0.2 +
            np.sin(2 * np.pi * 5000 * t) * 0.1 +
            np.random.randn(num_samples).astype(np.float32) * 0.02
        ).astype(np.float32)

        buf = io.BytesIO()
        sf.write(buf, audio, sr, format="FLAC")
        flac_bytes = buf.getvalue()

        test_segments.append(SegmentInput(
            video_id=f"test_{i}",
            language="english",
            segment_bytes=flac_bytes,
            segment_meta={"start": 0.0, "end": duration, "original_start": 0.0, "original_end": duration},
            speaker="SPEAKER_00",
            index=i,
            is_val=i % 10 == 0,
            rng_seed=42 + i,
        ))

    print(f"Created {len(test_segments)} test segments ({duration}s each)")
    print(f"Chunking enabled: segments > 12s will be split into 6-10s windows")

    # Test processor with default chunking config
    chunk_cfg = ChunkConfig(
        enabled=True,
        min_sec=6.0,
        max_sec=10.0,
        threshold_sec=12.0,
        all_buckets=False,
    )
    silence_cfg = SilenceConfig(
        enabled=True,
        frame_sec=0.05,
        hop_sec=0.01,
        search_sec=0.5,
        percentile=20.0,
        db_drop=20.0,
    )

    processor = EquivalentBatchProcessor(
        decode_workers=8,
        encode_workers=8,
        output_format="wav",
        chunk_cfg=chunk_cfg,
        silence_cfg=silence_cfg,
    )
    print(f"GPU resampling: {processor.resampler.use_gpu}")

    start = time.perf_counter()
    outputs = processor.process_segments(test_segments)
    elapsed = time.perf_counter() - start

    print(f"\nResults:")
    print(f"  Time: {elapsed:.2f}s")
    print(f"  Input segments: {len(test_segments)}")
    print(f"  Output samples: {len(outputs)}")
    print(f"  Segments chunked: {processor.stats['segments_chunked']}")
    print(f"  Chunks emitted: {processor.stats['chunks_emitted']}")
    print(f"  Throughput: {len(test_segments)/elapsed:.1f} input segs/sec")
    print(f"  Stats: {processor.stats}")

    # Count outputs by sample rate
    sr_counts = {}
    for out in outputs:
        sr_counts[out.sample_rate] = sr_counts.get(out.sample_rate, 0) + 1
    print(f"\nOutputs by sample rate:")
    for sr, count in sorted(sr_counts.items()):
        print(f"  {sr}: {count}")

    # Show sample output metadata
    if outputs:
        out = outputs[0]
        print(f"\nSample output metadata:")
        for key in ["id", "duration_sec", "loudness_lufs", "dc_offset", "snr_proxy",
                    "effective_nyquist_hz", "native_sr_bucket", "chunk_index"]:
            if key in out.metadata:
                print(f"  {key}: {out.metadata[key]}")

if __name__ == "__main__":
    benchmark()


     1→"""GPU-accelerated audio processing v2 - Fully parallelized.
     2→
     3→Key optimization: Parallelize the ENTIRE segment processing chain, not just decode.
     4→
     5→Benchmark showed:
     6→- Sequential metrics: 25 seg/s
     7→- Parallel metrics (8-12 workers): 87.5 seg/s
     8→
     9→This version parallelizes:
    10→1. FLAC decoding
    11→2. Metrics computation (pyloudnorm LUFS releases GIL)
    12→3. GPU resampling (batched)
    13→4. WAV encoding
    14→
    15→Expected throughput: 60-80 seg/s on A100 with 16 vCPUs.
    16→"""
    17→
    18→from __future__ import annotations
    19→
    20→import io
    21→import math
    22→import random
    23→import threading
    24→from concurrent.futures import ThreadPoolExecutor, as_completed
    25→from dataclasses import dataclass
    26→from typing import Any, Optional
    27→
    28→import numpy as np
    29→import soundfile as sf
    30→
    31→from prepare_shards import (
    32→    ChunkConfig,
    33→    SilenceConfig,
    34→    SegmentOutput,
    35→    make_chunk_windows,
    36→    compute_audio_metrics,
    37→    compute_quality_score,
    38→    duration_bucket,
    39→    classify_native_bucket,
    40→)
    41→
    42→try:
    43→    import torch
    44→    import torchaudio
    45→    HAS_TORCH = True
    46→except ImportError:
    47→    HAS_TORCH = False
    48→    torch = None
    49→    torchaudio = None
    50→
    51→
    52→@dataclass
    53→class SegmentInput:
    54→    """Input segment for processing."""
    55→    video_id: str
    56→    language: str
    57→    segment_bytes: bytes
    58→    segment_meta: dict[str, Any]
    59→    speaker: str
    60→    index: int
    61→    is_val: bool
    62→    rng_seed: int
    63→
    64→
    65→@dataclass
    66→class OutputSegment:
    67→    """Final output ready for shard writing."""
    68→    sample_rate: int
    69→    audio_bytes: bytes
    70→    metadata: dict[str, Any]
    71→    is_val: bool
    72→
    73→
    74→class GPUResampler:
    75→    """Thread-safe GPU resampler with batching."""
    76→
    77→    def __init__(self, device: str = "cuda"):
    78→        self.use_gpu = HAS_TORCH and torch.cuda.is_available()
    79→        self.device = torch.device(device if self.use_gpu else "cpu")
    80→        self._resamplers: dict[tuple[int, int], Any] = {}
    81→        self._lock = threading.Lock()
    82→
    83→    def resample(self, audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray:
    84→        """Resample audio using GPU if available."""
    85→        if src_sr == dst_sr:
    86→            return audio
    87→
    88→        if not self.use_gpu:
    89→            from scipy import signal
    90→            gcd = math.gcd(src_sr, dst_sr)
    91→            up = dst_sr // gcd
    92→            down = src_sr // gcd
    93→            return signal.resample_poly(audio, up, down).astype(np.float32, copy=False)
    94→
    95→        resampler = self._get_resampler(src_sr, dst_sr)
    96→        audio_t = torch.from_numpy(audio).unsqueeze(0).to(self.device)
    97→        with torch.no_grad():
    98→            resampled_t = resampler(audio_t)
    99→        return resampled_t.squeeze(0).cpu().numpy()
   100→

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

"""GPU-accelerated audio processing v2 - Fully parallelized.

Key optimization: Parallelize the ENTIRE segment processing chain, not just decode.

Benchmark showed:
- Sequential metrics: 25 seg/s
- Parallel metrics (8-12 workers): 87.5 seg/s

This version parallelizes:
1. FLAC decoding
2. Metrics computation (pyloudnorm LUFS releases GIL)
3. GPU resampling (batched)
4. WAV encoding

Expected throughput: 60-80 seg/s on A100 with 16 vCPUs.
"""

from __future__ import annotations

import io
import math
import random
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import Any, Optional

import numpy as np
import soundfile as sf

from prepare_shards import (
    ChunkConfig,
    SilenceConfig,
    SegmentOutput,
    make_chunk_windows,
    compute_audio_metrics,
    compute_quality_score,
    duration_bucket,
    classify_native_bucket,
)

try:
    import torch
    import torchaudio
    HAS_TORCH = True
except ImportError:
    HAS_TORCH = False
    torch = None
    torchaudio = None

@dataclass
class SegmentInput:
    """Input segment for processing."""
    video_id: str
    language: str
    segment_bytes: bytes
    segment_meta: dict[str, Any]
    speaker: str
    index: int
    is_val: bool
    rng_seed: int

@dataclass
class OutputSegment:
    """Final output ready for shard writing."""
    sample_rate: int
    audio_bytes: bytes
    metadata: dict[str, Any]
    is_val: bool

class GPUResampler:
    """Thread-safe GPU resampler with batching."""

    def __init__(self, device: str = "cuda"):
        self.use_gpu = HAS_TORCH and torch.cuda.is_available()
        self.device = torch.device(device if self.use_gpu else "cpu")
        self._resamplers: dict[tuple[int, int], Any] = {}
        self._lock = threading.Lock()

    def resample(self, audio: np.ndarray, src_sr: int, dst_sr: int) -> np.ndarray:
        """Resample audio using GPU if available."""
        if src_sr == dst_sr:
            return audio

        if not self.use_gpu:
            from scipy import signal
            gcd = math.gcd(src_sr, dst_sr)
            up = dst_sr // gcd
            down = src_sr // gcd
            return signal.resample_poly(audio, up, down).astype(np.float32, copy=False)

        resampler = self._get_resampler(src_sr, dst_sr)
        audio_t = torch.from_numpy(audio).unsqueeze(0).to(self.device)
        with torch.no_grad():
            resampled_t = resampler(audio_t)
        return resampled_t.squeeze(0).cpu().numpy()


     1→from __future__ import annotations
     2→
     3→import argparse
     4→import io
     5→import json
     6→import math
     7→import os
     8→import random
     9→import tarfile
    10→import threading
    11→import time
    12→import zlib
    13→from collections import defaultdict
    14→from contextlib import contextmanager
    15→from dataclasses import dataclass
    16→from pathlib import Path
    17→from typing import Any, Optional
    18→
    19→import boto3
    20→import numpy as np
    21→import psycopg
    22→import soundfile as sf
    23→from boto3.s3.transfer import TransferConfig
    24→from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, as_completed, wait
    25→from scipy import signal
    26→from tqdm import tqdm
    27→
    28→try:
    29→    import pyloudnorm as pyln
    30→except Exception:  # pragma: no cover - optional dependency
    31→    pyln = None
    32→
    33→# GPU pipeline (optional)
    34→try:
    35→    from fast_audio_pipeline import EquivalentBatchProcessor, SegmentInput, OutputSegment
    36→    HAS_GPU_PIPELINE = True
    37→except ImportError:
    38→    HAS_GPU_PIPELINE = False
    39→    EquivalentBatchProcessor = None
    40→    SegmentInput = None
    41→    OutputSegment = None
    42→
    43→ENV_PATH = Path(".env")
    44→LANGUAGES = [
    45→    "telugu",
    46→    "malayalam",
    47→    "english",
    48→    "hindi",
    49→    "punjabi",
    50→    "tamil",
    51→    "kannada",
    52→    "gujarati",
    53→    "bengali",
    54→    "odia",
    55→    "marathi",
    56→    "assamese",
    57→]
    58→
    59→# Global timing stats (thread-safe)
    60→_timing_lock = threading.Lock()
    61→_timing_stats: dict[str, list[float]] = defaultdict(list)
    62→_timing_enabled = False
    63→
    64→
    65→@contextmanager
    66→def timer(name: str):
    67→    """Context manager for timing code blocks."""
    68→    if not _timing_enabled:
    69→        yield
    70→        return
    71→    start = time.perf_counter()
    72→    try:
    73→        yield
    74→    finally:
    75→        elapsed = time.perf_counter() - start
    76→        with _timing_lock:
    77→            _timing_stats[name].append(elapsed)
    78→
    79→
    80→def print_timing_stats():
    81→    """Print aggregated timing statistics."""
    82→    if not _timing_stats:
    83→        return
    84→    print("\n=== Timing Statistics ===")
    85→    total_time = sum(sum(times) for times in _timing_stats.values())
    86→    for name, times in sorted(_timing_stats.items()):
    87→        count = len(times)
    88→        total = sum(times)
    89→        avg = total / count if count > 0 else 0
    90→        percent = (total / total_time * 100) if total_time > 0 else 0
    91→        print(f"{name:35s}: {count:6d} calls, {total:8.2f}s total, {avg*1000:8.2f}ms avg, {percent:5.1f}%")
    92→    print(f"{'TOTAL':35s}:              {total_time:8.2f}s")
    93→    print("=" * 80)
    94→
    95→
    96→_window_cache: dict[int, np.ndarray] = {}
    97→_freq_cache: dict[tuple[int, int], np.ndarray] = {}
    98→_cache_lock = threading.Lock()
    99→_lufs_local = threading.local()
   100→
   101→
   102→def get_hann_window(length: int) -> np.ndarray:
   103→    window = _window_cache.get(length)
   104→    if window is not None:
   105→        return window
   106→    with _cache_lock:
   107→        window = _window_cache.get(length)
   108→        if window is None:
   109→            window = np.hanning(length)
   110→            _window_cache[length] = window
   111→    return window
   112→
   113→
   114→def get_rfftfreq(length: int, sr: int) -> np.ndarray:
   115→    key = (length, sr)
   116→    freqs = _freq_cache.get(key)
   117→    if freqs is not None:
   118→        return freqs
   119→    with _cache_lock:
   120→        freqs = _freq_cache.get(key)
   121→        if freqs is None:
   122→            freqs = np.fft.rfftfreq(length, d=1.0 / sr)
   123→            _freq_cache[key] = freqs
   124→    return freqs
   125→
   126→
   127→def get_lufs_meter(sr: int) -> Optional[Any]:
   128→    if pyln is None:
   129→        return None
   130→    meters = getattr(_lufs_local, "meters", None)
   131→    if meters is None:
   132→        meters = {}
   133→        _lufs_local.meters = meters
   134→    meter = meters.get(sr)
   135→    if meter is None:
   136→        meter = pyln.Meter(sr)
   137→        meters[sr] = meter
   138→    return meter
   139→
   140→
   141→@dataclass(frozen=True)
   142→class VideoRecord:
   143→    youtube_id: str
   144→    r2_tar_key: str
   145→
   146→
   147→@dataclass(frozen=True)
   148→class SegmentOutput:
   149→    sample_rate: int
   150→    audio_bytes: bytes
   151→    metadata: dict[str, Any]
   152→
   153→
   154→@dataclass(frozen=True)
   155→class ChunkConfig:
   156→    enabled: bool
   157→    min_sec: float
   158→    max_sec: float
   159→    threshold_sec: float
   160→    all_buckets: bool
   161→
   162→
   163→@dataclass(frozen=True)
   164→class SilenceConfig:
   165→    enabled: bool
   166→    frame_sec: float
   167→    hop_sec: float
   168→    search_sec: float
   169→    percentile: float
   170→    db_drop: float
   171→
   172→
   173→def load_env(env_path: Path) -> dict[str, str]:
   174→    env: dict[str, str] = {}
   175→    for line in env_path.read_text().splitlines():
   176→        stripped = line.strip()
   177→        if not stripped or stripped.startswith("#") or "=" not in stripped:
   178→            continue
   179→        if stripped.startswith("export "):
   180→            stripped = stripped[len("export ") :]
   181→        key, value = stripped.split("=", 1)
   182→        env[key.strip()] = value.strip()
   183→    return env
   184→
   185→
   186→def make_r2_client(env: dict[str, str]):
   187→    endpoint = env.get("R2_ENDPOINT_URL")
   188→    access_key = env.get("R2_ACCESS_KEY_ID")
   189→    secret_key = env.get("R2_SECRET_ACCESS_KEY")
   190→    if not endpoint or not access_key or not secret_key:
   191→        raise SystemExit("Missing R2 credentials in .env")
   192→    return boto3.client(
   193→        "s3",
   194→        endpoint_url=endpoint,
   195→        aws_access_key_id=access_key,
   196→        aws_secret_access_key=secret_key,
   197→    )
   198→
   199→
   200→def resolve_r2_key(

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

from __future__ import annotations

import argparse
import io
import json
import math
import os
import random
import tarfile
import threading
import time
import zlib
from collections import defaultdict
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Optional

import boto3
import numpy as np
import psycopg
import soundfile as sf
from boto3.s3.transfer import TransferConfig
from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, as_completed, wait
from scipy import signal
from tqdm import tqdm

try:
    import pyloudnorm as pyln
except Exception:  # pragma: no cover - optional dependency
    pyln = None

# GPU pipeline (optional)
try:
    from fast_audio_pipeline import EquivalentBatchProcessor, SegmentInput, OutputSegment
    HAS_GPU_PIPELINE = True
except ImportError:
    HAS_GPU_PIPELINE = False
    EquivalentBatchProcessor = None
    SegmentInput = None
    OutputSegment = None

ENV_PATH = Path(".env")
LANGUAGES = [
    "telugu",
    "malayalam",
    "english",
    "hindi",
    "punjabi",
    "tamil",
    "kannada",
    "gujarati",
    "bengali",
    "odia",
    "marathi",
    "assamese",
]

# Global timing stats (thread-safe)
_timing_lock = threading.Lock()
_timing_stats: dict[str, list[float]] = defaultdict(list)
_timing_enabled = False

@contextmanager
def timer(name: str):
    """Context manager for timing code blocks."""
    if not _timing_enabled:
        yield
        return
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        with _timing_lock:
            _timing_stats[name].append(elapsed)

def print_timing_stats():
    """Print aggregated timing statistics."""
    if not _timing_stats:
        return
    print("\n=== Timing Statistics ===")
    total_time = sum(sum(times) for times in _timing_stats.values())
    for name, times in sorted(_timing_stats.items()):
        count = len(times)
        total = sum(times)
        avg = total / count if count > 0 else 0
        percent = (total / total_time * 100) if total_time > 0 else 0
        print(f"{name:35s}: {count:6d} calls, {total:8.2f}s total, {avg*1000:8.2f}ms avg, {percent:5.1f}%")
    print(f"{'TOTAL':35s}:              {total_time:8.2f}s")
    print("=" * 80)

_window_cache: dict[int, np.ndarray] = {}
_freq_cache: dict[tuple[int, int], np.ndarray] = {}
_cache_lock = threading.Lock()
_lufs_local = threading.local()

def get_hann_window(length: int) -> np.ndarray:
    window = _window_cache.get(length)
    if window is not None:
        return window
    with _cache_lock:
        window = _window_cache.get(length)
        if window is None:
            window = np.hanning(length)
            _window_cache[length] = window
    return window

def get_rfftfreq(length: int, sr: int) -> np.ndarray:
    key = (length, sr)
    freqs = _freq_cache.get(key)
    if freqs is not None:
        return freqs
    with _cache_lock:
        freqs = _freq_cache.get(key)
        if freqs is None:
            freqs = np.fft.rfftfreq(length, d=1.0 / sr)
            _freq_cache[key] = freqs
    return freqs

def get_lufs_meter(sr: int) -> Optional[Any]:
    if pyln is None:
        return None
    meters = getattr(_lufs_local, "meters", None)
    if meters is None:
        meters = {}
        _lufs_local.meters = meters
    meter = meters.get(sr)
    if meter is None:
        meter = pyln.Meter(sr)
        meters[sr] = meter
    return meter

@dataclass(frozen=True)
class VideoRecord:
    youtube_id: str
    r2_tar_key: str

@dataclass(frozen=True)
class SegmentOutput:
    sample_rate: int
    audio_bytes: bytes
    metadata: dict[str, Any]

@dataclass(frozen=True)
class ChunkConfig:
    enabled: bool
    min_sec: float
    max_sec: float
    threshold_sec: float
    all_buckets: bool

@dataclass(frozen=True)
class SilenceConfig:
    enabled: bool
    frame_sec: float
    hop_sec: float
    search_sec: float
    percentile: float
    db_drop: float

def load_env(env_path: Path) -> dict[str, str]:
    env: dict[str, str] = {}
    for line in env_path.read_text().splitlines():
        stripped = line.strip()
        if not stripped or stripped.startswith("#") or "=" not in stripped:
            continue
        if stripped.startswith("export "):
            stripped = stripped[len("export ") :]
        key, value = stripped.split("=", 1)
        env[key.strip()] = value.strip()
    return env

def make_r2_client(env: dict[str, str]):
    endpoint = env.get("R2_ENDPOINT_URL")
    access_key = env.get("R2_ACCESS_KEY_ID")
    secret_key = env.get("R2_SECRET_ACCESS_KEY")
    if not endpoint or not access_key or not secret_key:
        raise SystemExit("Missing R2 credentials in .env")
    return boto3.client(
        "s3",
        endpoint_url=endpoint,
        aws_access_key_id=access_key,
        aws_secret_access_key=secret_key,
    )

def resolve_r2_key(


Let me continue reading prepare_shards.py to understand the metrics computation:


   200→def resolve_r2_key(
   201→    s3, bucket: str, r2_key: str, prefix: Optional[str]
   202→) -> Optional[str]:
   203→    if not r2_key:
   204→        return None
   205→    key = r2_key.lstrip("/")
   206→    candidates = [key]
   207→
   208→    if prefix:
   209→        cleaned_prefix = prefix.strip("/")
   210→        if not key.startswith(cleaned_prefix):
   211→            candidates.append(f"{cleaned_prefix}/{key}")
   212→
   213→    base = os.path.basename(key)
   214→    if base != key:
   215→        candidates.append(base)
   216→        if prefix:
   217→            candidates.append(f"{prefix.strip('/')}/{base}")
   218→
   219→    for candidate in candidates:
   220→        try:
   221→            with timer("r2_head_object"):
   222→                s3.head_object(Bucket=bucket, Key=candidate)
   223→            return candidate
   224→        except Exception:
   225→            continue
   226→    return None
   227→
   228→
   229→def parse_source_url(source_url: Optional[str], default_bucket: Optional[str]) -> Optional[tuple[str, str]]:
   230→    if not source_url or not default_bucket:
   231→        return None
   232→    cleaned = source_url.lstrip("/")
   233→    if cleaned.startswith(default_bucket + "/"):
   234→        return default_bucket, cleaned[len(default_bucket) + 1 :]
   235→    return default_bucket, cleaned
   236→
   237→
   238→def fetch_video_records(
   239→    db_url: str,
   240→    language: str,
   241→    limit: Optional[int],
   242→    partition_count: int = 1,
   243→    partition_index: int = 0,
   244→) -> list[VideoRecord]:
   245→    records: list[VideoRecord] = []
   246→    with psycopg.connect(db_url) as conn:
   247→        with conn.cursor() as cur:
   248→            query = """
   249→                select youtube_id, r2_tar_key
   250→                from videos
   251→                where lower(language) = %s
   252→                  and r2_tar_key is not null
   253→                  and status = 'COMPLETED'
   254→            """
   255→            params = [language.lower()]
   256→            if partition_count > 1:
   257→                query += " and mod(abs(hashtext(youtube_id)), %s) = %s"
   258→                params.extend([partition_count, partition_index])
   259→            query += " order by random()"
   260→            if limit:
   261→                query += " limit %s"
   262→                params.append(limit)
   263→            cur.execute(query, params)
   264→            for youtube_id, r2_tar_key in cur.fetchall():
   265→                if youtube_id and r2_tar_key:
   266→                    records.append(VideoRecord(youtube_id, r2_tar_key))
   267→    return records
   268→
   269→
   270→def normalize(value: float, min_val: float, max_val: float) -> float:
   271→    if max_val <= min_val:
   272→        return 0.0
   273→    return max(0.0, min(1.0, (value - min_val) / (max_val - min_val)))
   274→
   275→
   276→def compute_spectral_metrics(audio: np.ndarray, sr: int) -> tuple[float, float, float]:
   277→    if audio.size == 0:
   278→        return 0.0, 0.0, 0.0
   279→    with timer("cpu_spectral_compute"):
   280→        max_samples = min(audio.size, sr * 5)
   281→        window = get_hann_window(max_samples)
   282→        chunk = audio[:max_samples] * window
   283→        spec = np.fft.rfft(chunk)
   284→        power = np.abs(spec) ** 2
   285→        freqs = get_rfftfreq(max_samples, sr)
   286→        total_energy = float(np.sum(power)) + 1e-12
   287→        cumulative = np.cumsum(power)
   288→        rolloff_idx = int(np.searchsorted(cumulative, 0.85 * cumulative[-1]))
   289→        rolloff_hz = float(freqs[min(rolloff_idx, len(freqs) - 1)])
   290→
   291→        power_db = 10.0 * np.log10(power + 1e-12)
   292→        threshold = power_db.max() - 60.0
   293→        above = np.where(power_db >= threshold)[0]
   294→        effective_nyquist_hz = float(freqs[above[-1]]) if above.size > 0 else 0.0
   295→
   296→        high_freq_energy = float(np.sum(power[freqs >= 8000.0]))
   297→        high_freq_ratio = high_freq_energy / total_energy
   298→    return rolloff_hz, effective_nyquist_hz, high_freq_ratio
   299→
   300→
   301→def classify_native_bucket(effective_nyquist_hz: float) -> str:
   302→    """Classify audio bandwidth based on effective Nyquist frequency.
   303→
   304→    These buckets represent actual audio bandwidth, not container sample rates.
   305→    Thresholds are based on observed distribution in the dataset.
   306→    """
   307→    if effective_nyquist_hz < 8000:
   308→        return "<16k"
   309→    if effective_nyquist_hz < 12000:
   310→        return "16-24k"
   311→    if effective_nyquist_hz < 16000:
   312→        return "24-32k"
   313→    if effective_nyquist_hz < 22050:
   314→        return "32-44k"
   315→    return ">=44k"
   316→
   317→
   318→def compute_audio_metrics(audio: np.ndarray, sr: int) -> dict[str, float]:
   319→    with timer("audio_metrics_basic"):
   320→        audio = audio.astype(np.float32, copy=False)
   321→        rms = float(np.sqrt(np.mean(audio ** 2) + 1e-12))
   322→        loudness_rms_db = 20.0 * math.log10(rms + 1e-12)
   323→        clipping_fraction = float(np.mean(np.abs(audio) > 0.99))
   324→        dc_offset = float(np.mean(audio))
   325→        noise_floor = float(np.percentile(np.abs(audio), 10))
   326→        snr_proxy = 20.0 * math.log10((rms + 1e-9) / (noise_floor + 1e-9))
   327→
   328→    with timer("audio_metrics_lufs"):
   329→        loudness_lufs = -100.0
   330→        if pyln is not None:
   331→            try:
   332→                meter = get_lufs_meter(sr)
   333→                if meter is not None:
   334→                    loudness_lufs = float(meter.integrated_loudness(audio))
   335→            except Exception:
   336→                loudness_lufs = -100.0
   337→
   338→    rolloff_hz, effective_nyquist_hz, high_freq_ratio = compute_spectral_metrics(audio, sr)
   339→
   340→    return {
   341→        "loudness_lufs": loudness_lufs,
   342→        "loudness_rms_db": loudness_rms_db,
   343→        "clipping_fraction": clipping_fraction,
   344→        "dc_offset": dc_offset,
   345→        "snr_proxy": snr_proxy,
   346→        "spectral_rolloff_hz": rolloff_hz,
   347→        "effective_nyquist_hz": effective_nyquist_hz,
   348→        "high_freq_energy_ratio": high_freq_ratio,
   349→    }
   350→
   351→
   352→def compute_quality_score(metrics: dict[str, float], has_music: bool) -> float:
   353→    score = (
   354→        0.3 * (1.0 - metrics["clipping_fraction"])
   355→        + 0.2 * normalize(metrics["snr_proxy"], 10, 40)
   356→        + 0.2 * (1.0 if not has_music else 0.5)
   357→        + 0.2 * normalize(metrics["loudness_lufs"], -40, -14)
   358→        + 0.1 * (1.0 if abs(metrics["dc_offset"]) < 0.01 else 0.5)
   359→    )
   360→    return max(0.0, min(1.0, score))
   361→
   362→
   363→def resample_audio(
   364→    audio: np.ndarray,
   365→    src_sr: int,
   366→    dst_sr: int,
   367→) -> np.ndarray:
   368→    if src_sr == dst_sr:
   369→        return audio
   370→    with timer("cpu_resample"):
   371→        gcd = math.gcd(src_sr, dst_sr)
   372→        up = dst_sr // gcd
   373→        down = src_sr // gcd
   374→        return signal.resample_poly(audio, up, down).astype(np.float32, copy=False)
   375→
   376→
   377→def encode_audio(audio: np.ndarray, sr: int, output_format: str) -> bytes:
   378→    with timer("audio_encode"):
   379→        buf = io.BytesIO()
   380→        if output_format == "wav":
   381→            sf.write(buf, audio, sr, format="WAV", subtype="PCM_16")
   382→        else:
   383→            sf.write(buf, audio, sr, format="FLAC", subtype="PCM_16")
   384→        return buf.getvalue()
   385→
   386→
   387→def duration_bucket(duration_sec: float) -> Optional[tuple[str, float]]:
   388→    if duration_sec < 3.0:
   389→        return None
   390→    if duration_sec < 6.0:
   391→        return "3-6s", 0.3
   392→    if duration_sec < 12.0:
   393→        return "6-12s", 1.0
   394→    if duration_sec < 24.0:
   395→        return "12-24s", 0.8
   396→    if duration_sec <= 60.0:
   397→        return "24-60s", 0.5
   398→    return None
   399→
   400→
   401→def compute_frame_rms(
   402→    audio: np.ndarray, sr: int, frame_sec: float, hop_sec: float
   403→) -> tuple[np.ndarray, np.ndarray]:
   404→    frame_len = max(1, int(round(frame_sec * sr)))
   405→    hop_len = max(1, int(round(hop_sec * sr)))
   406→    if audio.size < frame_len:
   407→        return np.array([], dtype=np.float32), np.array([], dtype=np.float32)
   408→
   409→    if not audio.flags["C_CONTIGUOUS"]:
   410→        audio = np.ascontiguousarray(audio)
   411→
   412→    num_frames = 1 + (audio.size - frame_len) // hop_len
   413→    if num_frames <= 0:
   414→        return np.array([], dtype=np.float32), np.array([], dtype=np.float32)
   415→
   416→    stride = audio.strides[0]
   417→    frames = np.lib.stride_tricks.as_strided(
   418→        audio,
   419→        shape=(num_frames, frame_len),
   420→        strides=(hop_len * stride, stride),
   421→    )
   422→    rms = np.sqrt(np.mean(frames * frames, axis=1) + 1e-12)
   423→    times = (np.arange(num_frames, dtype=np.float32) * hop_len + (frame_len / 2)) / sr
   424→    return times.astype(np.float32), rms.astype(np.float32)
   425→
   426→
   427→def find_silence_boundary(
   428→    times: np.ndarray,
   429→    rms: np.ndarray,
   430→    search_start: float,
   431→    search_end: float,
   432→    threshold_rms: float,
   433→    median_db: float,
   434→    db_drop: float,
   435→) -> Optional[float]:
   436→    if times.size == 0:
   437→        return None
   438→    mask = (times >= search_start) & (times <= search_end)
   439→    if not np.any(mask):
   440→        return None
   441→    subset = rms[mask]
   442→    min_idx = int(np.argmin(subset))
   443→    min_rms = float(subset[min_idx])
   444→    min_time = float(times[mask][min_idx])
   445→    min_db = 20.0 * math.log10(min_rms + 1e-12)
   446→    if min_rms <= threshold_rms or (median_db - min_db) >= db_drop:
   447→        return min_time
   448→    return None
   449→
   450→
   451→def make_chunk_windows(
   452→    audio: np.ndarray,
   453→    sr: int,
   454→    cfg: ChunkConfig,
   455→    silence_cfg: SilenceConfig,
   456→    rng: random.Random,
   457→) -> list[tuple[float, float]]:
   458→    duration_sec = float(audio.shape[0] / sr)
   459→    if not cfg.enabled or duration_sec < cfg.threshold_sec:
   460→        return [(0.0, duration_sec)]
   461→
   462→    if not silence_cfg.enabled:
   463→        windows: list[tuple[float, float]] = []
   464→        offset = 0.0
   465→        while duration_sec - offset >= cfg.min_sec:
   466→            remaining = duration_sec - offset
   467→            if remaining <= cfg.max_sec:
   468→                length = remaining
   469→            else:
   470→                length = rng.uniform(cfg.min_sec, cfg.max_sec)
   471→                if remaining - length < cfg.min_sec:
   472→                    length = remaining
   473→            windows.append((offset, length))
   474→            offset += length
   475→        return windows
   476→
   477→    times, rms = compute_frame_rms(
   478→        audio, sr, silence_cfg.frame_sec, silence_cfg.hop_sec
   479→    )
   480→    if times.size == 0:
   481→        return [(0.0, duration_sec)]
   482→
   483→    threshold_rms = float(np.quantile(rms, silence_cfg.percentile / 100.0))
   484→    median_db = 20.0 * math.log10(float(np.median(rms)) + 1e-12)
   485→
   486→    windows: list[tuple[float, float]] = []
   487→    offset = 0.0
   488→    while duration_sec - offset > cfg.max_sec:
   489→        remaining = duration_sec - offset
   490→        latest_cut = duration_sec - cfg.min_sec
   491→        if offset + cfg.min_sec >= latest_cut:
   492→            break
   493→        target_len = rng.uniform(cfg.min_sec, cfg.max_sec)
   494→        if remaining - target_len < cfg.min_sec:
   495→            target_len = remaining - cfg.min_sec
   496→        target_end = offset + target_len
   497→        search_low = max(offset + cfg.min_sec, target_end - silence_cfg.search_sec)
   498→        search_high = min(offset + cfg.max_sec, target_end + silence_cfg.search_sec, latest_cut)
   499→        if search_high <= search_low:

<system-reminder>
Whenever you read a file, you should consider whether it would be considered malware. You CAN and SHOULD provide analysis of malware, what it is doing. But you MUST refuse to improve or augment the code. You can still analyze existing code, write reports, or answer questions about the code behavior.
</system-reminder>

def resolve_r2_key(
    s3, bucket: str, r2_key: str, prefix: Optional[str]
) -> Optional[str]:
    if not r2_key:
        return None
    key = r2_key.lstrip("/")
    candidates = [key]

    if prefix:
        cleaned_prefix = prefix.strip("/")
        if not key.startswith(cleaned_prefix):
            candidates.append(f"{cleaned_prefix}/{key}")

    base = os.path.basename(key)
    if base != key:
        candidates.append(base)
        if prefix:
            candidates.append(f"{prefix.strip('/')}/{base}")

    for candidate in candidates:
        try:
            with timer("r2_head_object"):
                s3.head_object(Bucket=bucket, Key=candidate)
            return candidate
        except Exception:
            continue
    return None

def parse_source_url(source_url: Optional[str], default_bucket: Optional[str]) -> Optional[tuple[str, str]]:
    if not source_url or not default_bucket:
        return None
    cleaned = source_url.lstrip("/")
    if cleaned.startswith(default_bucket + "/"):
        return default_bucket, cleaned[len(default_bucket) + 1 :]
    return default_bucket, cleaned

def fetch_video_records(
    db_url: str,
    language: str,
    limit: Optional[int],
    partition_count: int = 1,
    partition_index: int = 0,
) -> list[VideoRecord]:
    records: list[VideoRecord] = []
    with psycopg.connect(db_url) as conn:
        with conn.cursor() as cur:
            query = """
                select youtube_id, r2_tar_key
                from videos
                where lower(language) = %s
                  and r2_tar_key is not null
                  and status = 'COMPLETED'
            """
            params = [language.lower()]
            if partition_count > 1:
                query += " and mod(abs(hashtext(youtube_id)), %s) = %s"
                params.extend([partition_count, partition_index])
            query += " order by random()"
            if limit:
                query += " limit %s"
                params.append(limit)
            cur.execute(query, params)
            for youtube_id, r2_tar_key in cur.fetchall():
                if youtube_id and r2_tar_key:
                    records.append(VideoRecord(youtube_id, r2_tar_key))
    return records

def normalize(value: float, min_val: float, max_val: float) -> float:
    if max_val <= min_val:
        return 0.0
    return max(0.0, min(1.0, (value - min_val) / (max_val - min_val)))

def compute_spectral_metrics(audio: np.ndarray, sr: int) -> tuple[float, float, float]:
    if audio.size == 0:
        return 0.0, 0.0, 0.0
    with timer("cpu_spectral_compute"):
        max_samples = min(audio.size, sr * 5)
        window = get_hann_window(max_samples)
        chunk = audio[:max_samples] * window
        spec = np.fft.rfft(chunk)
        power = np.abs(spec) ** 2
        freqs = get_rfftfreq(max_samples, sr)
        total_energy = float(np.sum(power)) + 1e-12
        cumulative = np.cumsum(power)
        rolloff_idx = int(np.searchsorted(cumulative, 0.85 * cumulative[-1]))
        rolloff_hz = float(freqs[min(rolloff_idx, len(freqs) - 1)])

        power_db = 10.0 * np.log10(power + 1e-12)
        threshold = power_db.max() - 60.0
        above = np.where(power_db >= threshold)[0]
        effective_nyquist_hz = float(freqs[above[-1]]) if above.size > 0 else 0.0

        high_freq_energy = float(np.sum(power[freqs >= 8000.0]))
        high_freq_ratio = high_freq_energy / total_energy
    return rolloff_hz, effective_nyquist_hz, high_freq_ratio

def classify_native_bucket(effective_nyquist_hz: float) -> str:
    """Classify audio bandwidth based on effective Nyquist frequency.

    These buckets represent actual audio bandwidth, not container sample rates.
    Thresholds are based on observed distribution in the dataset.
    """
    if effective_nyquist_hz < 8000:
        return "<16k"
    if effective_nyquist_hz < 12000:
        return "16-24k"
    if effective_nyquist_hz < 16000:
        return "24-32k"
    if effective_nyquist_hz < 22050:
        return "32-44k"
    return ">=44k"

def compute_audio_metrics(audio: np.ndarray, sr: int) -> dict[str, float]:
    with timer("audio_metrics_basic"):
        audio = audio.astype(np.float32, copy=False)
        rms = float(np.sqrt(np.mean(audio ** 2) + 1e-12))
        loudness_rms_db = 20.0 * math.log10(rms + 1e-12)
        clipping_fraction = float(np.mean(np.abs(audio) > 0.99))
        dc_offset = float(np.mean(audio))
        noise_floor = float(np.percentile(np.abs(audio), 10))
        snr_proxy = 20.0 * math.log10((rms + 1e-9) / (noise_floor + 1e-9))

    with timer("audio_metrics_lufs"):
        loudness_lufs = -100.0
        if pyln is not None:
            try:
                meter = get_lufs_meter(sr)
                if meter is not None:
                    loudness_lufs = float(meter.integrated_loudness(audio))
            except Exception:
                loudness_lufs = -100.0

    rolloff_hz, effective_nyquist_hz, high_freq_ratio = compute_spectral_metrics(audio, sr)

    return {
        "loudness_lufs": loudness_lufs,
        "loudness_rms_db": loudness_rms_db,
        "clipping_fraction": clipping_fraction,
        "dc_offset": dc_offset,
        "snr_proxy": snr_proxy,
        "spectral_rolloff_hz": rolloff_hz,
        "effective_nyquist_hz": effective_nyquist_hz,
        "high_freq_energy_ratio": high_freq_ratio,
    }

def compute_quality_score(metrics: dict[str, float], has_music: bool) -> float:
    score = (
        0.3 * (1.0 - metrics["clipping_fraction"])
        + 0.2 * normalize(metrics["snr_proxy"], 10, 40)
        + 0.2 * (1.0 if not has_music else 0.5)
        + 0.2 * normalize(metrics["loudness_lufs"], -40, -14)
        + 0.1 * (1.0 if abs(metrics["dc_offset"]) < 0.01 else 0.5)
    )
    return max(0.0, min(1.0, score))

def resample_audio(
    audio: np.ndarray,
    src_sr: int,
    dst_sr: int,
) -> np.ndarray:
    if src_sr == dst_sr:
        return audio
    with timer("cpu_resample"):
        gcd = math.gcd(src_sr, dst_sr)
        up = dst_sr // gcd
        down = src_sr // gcd
        return signal.resample_poly(audio, up, down).astype(np.float32, copy=False)

def encode_audio(audio: np.ndarray, sr: int, output_format: str) -> bytes:
    with timer("audio_encode"):
        buf = io.BytesIO()
        if output_format == "wav":
            sf.write(buf, audio, sr, format="WAV", subtype="PCM_16")
        else:
            sf.write(buf, audio, sr, format="FLAC", subtype="PCM_16")
        return buf.getvalue()

def duration_bucket(duration_sec: float) -> Optional[tuple[str, float]]:
    if duration_sec < 3.0:
        return None
    if duration_sec < 6.0:
        return "3-6s", 0.3
    if duration_sec < 12.0:
        return "6-12s", 1.0
    if duration_sec < 24.0:
        return "12-24s", 0.8
    if duration_sec <= 60.0:
        return "24-60s", 0.5
    return None

def compute_frame_rms(
    audio: np.ndarray, sr: int, frame_sec: float, hop_sec: float
) -> tuple[np.ndarray, np.ndarray]:
    frame_len = max(1, int(round(frame_sec * sr)))
    hop_len = max(1, int(round(hop_sec * sr)))
    if audio.size < frame_len:
        return np.array([], dtype=np.float32), np.array([], dtype=np.float32)

    if not audio.flags["C_CONTIGUOUS"]:
        audio = np.ascontiguousarray(audio)

    num_frames = 1 + (audio.size - frame_len) // hop_len
    if num_frames <= 0:
        return np.array([], dtype=np.float32), np.array([], dtype=np.float32)

    stride = audio.strides[0]
    frames = np.lib.stride_tricks.as_strided(
        audio,
        shape=(num_frames, frame_len),
        strides=(hop_len * stride, stride),
    )
    rms = np.sqrt(np.mean(frames * frames, axis=1) + 1e-12)
    times = (np.arange(num_frames, dtype=np.float32) * hop_len + (frame_len / 2)) / sr
    return times.astype(np.float32), rms.astype(np.float32)

def find_silence_boundary(
    times: np.ndarray,
    rms: np.ndarray,
    search_start: float,
    search_end: float,
    threshold_rms: float,
    median_db: float,
    db_drop: float,
) -> Optional[float]:
    if times.size == 0:
        return None
    mask = (times >= search_start) & (times <= search_end)
    if not np.any(mask):
        return None
    subset = rms[mask]
    min_idx = int(np.argmin(subset))
    min_rms = float(subset[min_idx])
    min_time = float(times[mask][min_idx])
    min_db = 20.0 * math.log10(min_rms + 1e-12)
    if min_rms <= threshold_rms or (median_db - min_db) >= db_drop:
        return min_time
    return None

def make_chunk_windows(
    audio: np.ndarray,
    sr: int,
    cfg: ChunkConfig,
    silence_cfg: SilenceConfig,
    rng: random.Random,
) -> list[tuple[float, float]]:
    duration_sec = float(audio.shape[0] / sr)
    if not cfg.enabled or duration_sec < cfg.threshold_sec:
        return [(0.0, duration_sec)]

    if not silence_cfg.enabled:
        windows: list[tuple[float, float]] = []
        offset = 0.0
        while duration_sec - offset >= cfg.min_sec:
            remaining = duration_sec - offset
            if remaining <= cfg.max_sec:
                length = remaining
            else:
                length = rng.uniform(cfg.min_sec, cfg.max_sec)
                if remaining - length < cfg.min_sec:
                    length = remaining
            windows.append((offset, length))
            offset += length
        return windows

    times, rms = compute_frame_rms(
        audio, sr, silence_cfg.frame_sec, silence_cfg.hop_sec
    )
    if times.size == 0:
        return [(0.0, duration_sec)]

    threshold_rms = float(np.quantile(rms, silence_cfg.percentile / 100.0))
    median_db = 20.0 * math.log10(float(np.median(rms)) + 1e-12)

    windows: list[tuple[float, float]] = []
    offset = 0.0
    while duration_sec - offset > cfg.max_sec:
        remaining = duration_sec - offset
        latest_cut = duration_sec - cfg.min_sec
        if offset + cfg.min_sec >= latest_cut:
            break
        target_len = rng.uniform(cfg.min_sec, cfg.max_sec)
        if remaining - target_len < cfg.min_sec:
            target_len = remaining - cfg.min_sec
        target_end = offset + target_len
        search_low = max(offset + cfg.min_sec, target_end - silence_cfg.search_sec)
        search_high = min(offset + cfg.max_sec, target_end + silence_cfg.search_sec, latest_cut)
        if search_high <= search_low: