File size: 18,278 Bytes
35d483e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
#!/usr/bin/env python3
"""Measure endpoint-model sensitivity to appended trailing silence.

The report is deliberately aggregate-only.  It never serializes record IDs,
source rows, audio, transcripts, or per-example probabilities.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import os
import statistics
import sys
from collections.abc import Mapping, Sequence
from pathlib import Path
from typing import Any

REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
SOURCE_ROOT = REPOSITORY_ROOT / "src"
if str(SOURCE_ROOT) not in sys.path:
    sys.path.insert(0, str(SOURCE_ROOT))

TRAILING_SILENCE_MS = (0, 200, 400, 800)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--checkpoint", required=True)
    parser.add_argument("--source", required=True, help="JSON/JSONL split manifest")
    parser.add_argument(
        "--source-root",
        required=True,
        help="base directory for manifest source_file entries",
    )
    parser.add_argument("--split", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument(
        "--max-examples",
        type=int,
        help="deterministic manifest-order prefix after split filtering",
    )
    parser.add_argument(
        "--threshold",
        type=float,
        help="default: checkpoint's validation-calibrated threshold, otherwise 0.5",
    )
    parser.add_argument("--device", default="auto")
    return parser.parse_args()


def _resolve(value: str) -> Path:
    path = Path(value).expanduser()
    return path.resolve() if path.is_absolute() else (REPOSITORY_ROOT / path).resolve()


def _portable_path(path: Path) -> str:
    try:
        return path.resolve().relative_to(REPOSITORY_ROOT).as_posix()
    except ValueError:
        return path.name


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(block)
    return digest.hexdigest()


def _field(record: Any, *names: str, default: Any = None) -> Any:
    for name in names:
        if isinstance(record, Mapping) and name in record:
            return record[name]
        if hasattr(record, name):
            return getattr(record, name)
    return default


def _identity_collate(records: list[Any]) -> list[Any]:
    """Keep raw records intact so every silence condition shares one decode."""

    return records


def _device(torch: Any, requested: str) -> Any:
    if requested != "auto":
        return torch.device(requested)
    if torch.cuda.is_available():
        return torch.device("cuda")
    mps = getattr(torch.backends, "mps", None)
    if mps is not None and mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


def _percentile(values: Sequence[float], quantile: float) -> float:
    if not values:
        raise ValueError("percentiles require at least one value")
    if not 0.0 <= quantile <= 1.0:
        raise ValueError("quantile must be in [0, 1]")
    ordered = sorted(float(value) for value in values)
    position = (len(ordered) - 1) * quantile
    lower = math.floor(position)
    upper = math.ceil(position)
    if lower == upper:
        return ordered[lower]
    fraction = position - lower
    return ordered[lower] * (1.0 - fraction) + ordered[upper] * fraction


def _distribution(values: Sequence[float]) -> dict[str, float | int]:
    if not values:
        raise ValueError("distribution requires at least one value")
    finite = [float(value) for value in values]
    if any(not math.isfinite(value) for value in finite):
        raise ValueError("distribution values must be finite")
    return {
        "count": len(finite),
        "mean": statistics.fmean(finite),
        "standard_deviation": statistics.pstdev(finite),
        "minimum": min(finite),
        "p05": _percentile(finite, 0.05),
        "p50": _percentile(finite, 0.50),
        "p95": _percentile(finite, 0.95),
        "maximum": max(finite),
    }


def _probability_shift(
    labels: Sequence[int], baseline: Sequence[float], condition: Sequence[float]
) -> dict[str, Any]:
    if not labels or len(labels) != len(baseline) or len(labels) != len(condition):
        raise ValueError("paired shift inputs must have the same non-zero length")
    deltas = [
        float(candidate) - float(reference)
        for reference, candidate in zip(baseline, condition, strict=True)
    ]
    absolute = [abs(value) for value in deltas]

    def label_distribution(label: int) -> dict[str, float | int] | None:
        selected = [delta for target, delta in zip(labels, deltas, strict=True) if target == label]
        return _distribution(selected) if selected else None

    return {
        "definition": "p_end(condition) - p_end(0ms) on the same examples",
        **_distribution(deltas),
        "mean_absolute": statistics.fmean(absolute),
        "root_mean_square": math.sqrt(statistics.fmean(value * value for value in deltas)),
        "increased_count": sum(value > 0.0 for value in deltas),
        "decreased_count": sum(value < 0.0 for value in deltas),
        "unchanged_count": sum(value == 0.0 for value in deltas),
        "by_label": {
            "HOLD": label_distribution(0),
            "END": label_distribution(1),
        },
    }


def _decision_flips(
    labels: Sequence[int],
    baseline: Sequence[float],
    condition: Sequence[float],
    threshold: float,
) -> dict[str, int | float]:
    if not labels or len(labels) != len(baseline) or len(labels) != len(condition):
        raise ValueError("paired flip inputs must have the same non-zero length")
    baseline_predictions = [int(float(value) >= threshold) for value in baseline]
    condition_predictions = [int(float(value) >= threshold) for value in condition]
    pairs = list(zip(labels, baseline_predictions, condition_predictions, strict=True))
    hold_to_end = sum(before == 0 and after == 1 for _, before, after in pairs)
    end_to_hold = sum(before == 1 and after == 0 for _, before, after in pairs)
    flips = hold_to_end + end_to_hold
    return {
        "count": flips,
        "rate": flips / len(labels),
        "unchanged_count": len(labels) - flips,
        "HOLD_to_END": hold_to_end,
        "END_to_HOLD": end_to_hold,
        "false_interruptions_introduced": sum(
            target == 0 and before == 0 and after == 1 for target, before, after in pairs
        ),
        "false_interruptions_resolved": sum(
            target == 0 and before == 1 and after == 0 for target, before, after in pairs
        ),
        "missed_ends_introduced": sum(
            target == 1 and before == 1 and after == 0 for target, before, after in pairs
        ),
        "missed_ends_resolved": sum(
            target == 1 and before == 0 and after == 1 for target, before, after in pairs
        ),
    }


def _append_silence_and_pad(
    waveforms: Sequence[Any],
    silence_ms: int,
    sample_rate: int,
    max_seconds: float,
    pad_side: str,
    torch: Any,
) -> tuple[Any, Any]:
    """Append valid zeros, suffix-crop to the model window, then batch-pad."""

    if silence_ms < 0:
        raise ValueError("silence_ms cannot be negative")
    if sample_rate <= 0 or max_seconds <= 0.0:
        raise ValueError("sample_rate and max_seconds must be positive")
    if pad_side not in {"left", "right"}:
        raise ValueError("pad_side must be 'left' or 'right'")
    if not waveforms:
        raise ValueError("waveforms cannot be empty")
    target_samples = int(round(max_seconds * sample_rate))
    if target_samples < 1:
        raise ValueError("model window rounds to zero samples")
    silence_samples = int(round(silence_ms * sample_rate / 1000.0))
    prepared: list[Any] = []
    for waveform in waveforms:
        signal = torch.as_tensor(waveform, dtype=torch.float32).flatten()
        if signal.numel() == 0:
            raise ValueError("decoded audio cannot be empty")
        if silence_samples:
            signal = torch.cat((signal, signal.new_zeros(silence_samples)))
        prepared.append(signal[-target_samples:])
    lengths = torch.tensor([signal.numel() for signal in prepared], dtype=torch.long)
    padded = torch.zeros((len(prepared), target_samples), dtype=torch.float32)
    for index, signal in enumerate(prepared):
        if pad_side == "left":
            padded[index, -signal.numel() :] = signal
        else:
            padded[index, : signal.numel()] = signal
    return padded, lengths


def _validate_args(args: argparse.Namespace) -> tuple[Path, Path, Path, Path]:
    if args.batch_size < 1:
        raise SystemExit("--batch-size must be positive")
    if args.max_examples is not None and args.max_examples < 1:
        raise SystemExit("--max-examples must be positive")
    if args.threshold is not None and (
        not math.isfinite(args.threshold) or not 0.0 <= args.threshold <= 1.0
    ):
        raise SystemExit("--threshold must be finite and in [0, 1]")
    checkpoint = _resolve(args.checkpoint)
    source = _resolve(args.source)
    source_root = _resolve(args.source_root)
    output = _resolve(args.output)
    if not checkpoint.is_file():
        raise SystemExit(f"checkpoint does not exist: {checkpoint}")
    if not source.is_file() or source.suffix.lower() not in {".json", ".jsonl"}:
        raise SystemExit("--source must be an existing JSON/JSONL manifest")
    if not source_root.is_dir():
        raise SystemExit(f"--source-root is not a directory: {source_root}")
    return checkpoint, source, source_root, output


def main() -> int:
    args = parse_args()
    checkpoint_path, source_path, source_root, output_path = _validate_args(args)
    try:
        import torch
        from torch.utils.data import DataLoader
    except ImportError as exc:
        raise SystemExit("silence sensitivity analysis requires PyTorch") from exc

    from turn_detection.models import LogMelConfig, LogMelFrontend, load_model_checkpoint
    from turn_detection.training.datasets import build_record_dataloader, decode_audio
    from turn_detection.training.metrics import binary_classification_metrics

    device = _device(torch, args.device)
    try:
        model, checkpoint = load_model_checkpoint(checkpoint_path, map_location=device)
    except (OSError, RuntimeError, ValueError) as exc:
        raise SystemExit(f"cannot load checkpoint: {exc}") from exc
    model.to(device).eval()
    metadata = checkpoint.get("metadata", {})
    if not isinstance(metadata, Mapping):
        metadata = {}
    feature_config = LogMelConfig.from_mapping(metadata.get("feature_config", {}))
    frontend = LogMelFrontend(feature_config).cpu().eval()
    max_seconds = float(metadata.get("max_seconds", 8.0))
    if not math.isfinite(max_seconds) or max_seconds <= 0.0:
        raise SystemExit("checkpoint metadata has an invalid max_seconds")
    raw_threshold = (
        args.threshold if args.threshold is not None else checkpoint.get("threshold", 0.5)
    )
    try:
        threshold = float(raw_threshold)
    except (TypeError, ValueError) as exc:
        raise SystemExit("checkpoint threshold is not numeric") from exc
    if not math.isfinite(threshold) or not 0.0 <= threshold <= 1.0:
        raise SystemExit("effective threshold must be finite and in [0, 1]")

    # Reuse the canonical source routing (including lazy Parquet resolution and
    # deterministic max_examples), but keep raw records so audio is decoded once
    # and the exact same examples are evaluated under every paired condition.
    routed_loader = build_record_dataloader(
        source_path,
        split=args.split,
        frontend=frontend,
        batch_size=args.batch_size,
        max_seconds=max_seconds,
        shuffle=False,
        num_workers=0,
        token=os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN"),
        max_examples=args.max_examples,
        source_root=source_root,
    )
    raw_loader = DataLoader(
        routed_loader.dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=0,
        collate_fn=_identity_collate,
    )
    labels: list[int] = []
    probabilities: dict[int, list[float]] = {silence_ms: [] for silence_ms in TRAILING_SILENCE_MS}
    with torch.inference_mode():
        for raw_records in raw_loader:
            batch_labels: list[int] = []
            waveforms: list[Any] = []
            for record in raw_records:
                raw_label = _field(record, "endpoint", "endpoint_bool", "label", default=None)
                try:
                    numeric_label = float(raw_label)
                except (TypeError, ValueError) as exc:
                    raise SystemExit(
                        "manifest contains a missing or non-numeric endpoint label"
                    ) from exc
                if not math.isfinite(numeric_label) or numeric_label not in (0.0, 1.0):
                    raise SystemExit("manifest endpoint labels must be binary (0 or 1)")
                batch_labels.append(int(numeric_label))
                try:
                    waveform = decode_audio(
                        record,
                        target_sample_rate=feature_config.sample_rate,
                        max_seconds=max_seconds,
                    )
                except (OSError, RuntimeError, TypeError, ValueError) as exc:
                    raise SystemExit(
                        "manifest must resolve to waveform-backed audio for silence analysis: "
                        f"{exc}"
                    ) from exc
                waveforms.append(waveform)
            labels.extend(batch_labels)
            for silence_ms in TRAILING_SILENCE_MS:
                padded, lengths = _append_silence_and_pad(
                    waveforms,
                    silence_ms=silence_ms,
                    sample_rate=feature_config.sample_rate,
                    max_seconds=max_seconds,
                    pad_side=feature_config.pad_side,
                    torch=torch,
                )
                features, attention_mask = frontend(padded, lengths)
                output = model(features.to(device), attention_mask.to(device))
                scores = torch.sigmoid(output.endpoint_logits).detach().cpu().tolist()
                probabilities[silence_ms].extend(float(score) for score in scores)

    if not labels:
        raise SystemExit("manifest split produced no examples")
    if args.max_examples is not None and len(labels) > args.max_examples:
        raise SystemExit("internal error: deterministic max-examples bound was exceeded")
    baseline = probabilities[0]
    conditions: list[dict[str, Any]] = []
    for silence_ms in TRAILING_SILENCE_MS:
        scores = probabilities[silence_ms]
        if len(scores) != len(labels):
            raise SystemExit("internal error: silence-condition predictions are misaligned")
        condition: dict[str, Any] = {
            "trailing_silence_ms": silence_ms,
            "probability_summary": _distribution(scores),
            "classification_metrics": binary_classification_metrics(labels, scores, threshold),
        }
        if silence_ms == 0:
            condition["relative_to_0ms"] = None
        else:
            condition["relative_to_0ms"] = {
                "probability_shift": _probability_shift(labels, baseline, scores),
                "threshold_decision_flips": _decision_flips(labels, baseline, scores, threshold),
            }
        conditions.append(condition)

    run_metadata = metadata.get("run_metadata", {})
    if not isinstance(run_metadata, Mapping):
        run_metadata = {}
    report = {
        "schema_version": 1,
        "analysis": "trailing_silence_sensitivity",
        "checkpoint": _portable_path(checkpoint_path),
        "checkpoint_sha256": _sha256(checkpoint_path),
        "model_parameter_count": sum(parameter.numel() for parameter in model.parameters()),
        "training_status": run_metadata.get("status"),
        "data_scope": metadata.get("data_scope"),
        "data_revision": metadata.get("data_revision"),
        "source_manifest": _portable_path(source_path),
        "source_manifest_sha256": _sha256(source_path),
        "split": args.split,
        "threshold": threshold,
        "sample_rate": feature_config.sample_rate,
        "model_window_seconds": max_seconds,
        "example_count": len(labels),
        "positive_count": sum(labels),
        "negative_count": len(labels) - sum(labels),
        "selection": {
            "shuffle": False,
            "num_workers": 0,
            "policy": "manifest-order prefix after split filtering",
            "max_examples": args.max_examples,
        },
        "perturbation": {
            "conditions_ms": list(TRAILING_SILENCE_MS),
            "sample_value": 0.0,
            "policy": (
                "Decode/resample the original suffix once, append valid zero-valued samples, "
                "then suffix-crop to the checkpoint window before the canonical frontend."
            ),
        },
        "conditions": conditions,
        "privacy": {
            "aggregate_only": True,
            "contains_per_example_rows": False,
            "contains_raw_identifiers": False,
            "contains_audio_or_transcripts": False,
            "note": (
                "The report contains aggregate paired statistics and file-level SHA-256 "
                "provenance only."
            ),
        },
    }
    output_path.parent.mkdir(parents=True, exist_ok=True)
    temporary = output_path.with_name(output_path.name + ".tmp")
    temporary.write_text(
        json.dumps(report, indent=2, sort_keys=True, allow_nan=False) + "\n",
        encoding="utf-8",
    )
    temporary.replace(output_path)
    print(
        json.dumps(
            {
                "example_count": len(labels),
                "output": str(output_path),
                "conditions_ms": list(TRAILING_SILENCE_MS),
            },
            indent=2,
        )
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())