bleep/cli.py

"""Command-line interface.

    bleep find audio.wav --recall high
    bleep redact audio.wav                       # writes a review plan, not audio
    bleep redact audio.wav --apply --out red.wav # needs a decided plan
    bleep union audio.wav --asr-spans spans.json
    bleep info

**No command in this module can print what was said.**  ``find`` emits
category + time; ``redact`` emits a plan of categories + times; ``union`` takes
someone else's spans and returns spans.  :mod:`bleep.asr` and
:mod:`bleep.textpii` are not imported here at all, so no code path reachable
from the CLI can produce a transcript.
``tests/test_no_content_path.py`` asserts both the import property and that the
rendered output contains no transcript-like text.

Claim demonstrated by this module: no-content, and the reviewability
requirement (``redact`` defaults to plan-only).
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

from .audio import load_audio
from .redact import MODES, RedactionPlan, Redactor, verify_redaction
from .spans import ANY, CATEGORIES, Span
from .tagger import load_model, measure_rtf, predict_scores
from .threshold import (
    PRESETS,
    OperatingPoint,
    RecallFirstThresholding,
    get_operating_point,
    load_operating_point,
)
from .union import UnionCombiner

DEFAULT_CKPT = Path("checkpoints/bleep-base.pt")

#: Shown on every command's output. Imported from the package so the CLI, the
#: Gradio app and the library cannot drift into saying different things about
#: what the tool does and does not promise.
from . import DISCLAIMER  # noqa: E402


# ---------------------------------------------------------------------------


def _resolve_model(path: str | None):
    p = Path(path) if path else DEFAULT_CKPT
    if not p.exists():
        raise SystemExit(
            f"no checkpoint at {p}.\n"
            f"Train one:   python scripts/train.py --data data --preset bleep-base\n"
            f"Or point at one: bleep ... --model /path/to/bleep-base.pt"
        )
    return load_model(p)


def _operating_points_file(model_path: str | None) -> Path:
    """Where a calibrated operating-point set lives: next to the checkpoint."""
    base = Path(model_path) if model_path else DEFAULT_CKPT
    return base.parent / "operating_points.json"


def _operating_point(args, frame_rate: float):
    """Resolve the decision rule, preferring a calibrated threshold.

    Resolution order: an explicit ``--operating-point-file``, then the
    ``operating_points.json`` written next to the checkpoint by
    ``run_benchmarks.py --calibrate-only``, then failure with the command to
    run.  A preset's *threshold* is not portable across checkpoints, so
    inventing one here would silently produce a decision rule nobody measured.
    """
    if getattr(args, "operating_point_file", None):
        op = load_operating_point(args.operating_point_file)
    else:
        op = get_operating_point(args.recall)
        calibrated = _operating_points_file(getattr(args, "model", None))
        if calibrated.exists():
            try:
                blob = json.loads(calibrated.read_text(encoding="utf-8"))
            except json.JSONDecodeError as e:
                raise SystemExit(f"{calibrated} is not valid JSON: {e}")
            if args.recall in blob:
                op = OperatingPoint.from_dict(blob[args.recall])
    if op.threshold is None:
        raise SystemExit(
            f"operating point {op.name!r} has no calibrated threshold for this "
            f"checkpoint.\n\nCalibrate it on the held-out `calib` split:\n"
            f"  python scripts/run_benchmarks.py --data data "
            f"--checkpoint {args.model or DEFAULT_CKPT} --calibrate-only\n\n"
            f"That writes {_operating_points_file(getattr(args, 'model', None))}, "
            f"which this command picks up automatically."
        )
    if getattr(args, "pre_margin", None) is not None:
        op.pre_margin_s = args.pre_margin
    if getattr(args, "post_margin", None) is not None:
        op.post_margin_s = args.post_margin
    return RecallFirstThresholding(op, frame_rate)


def _detect(args) -> tuple[list[Span], dict]:
    model, feat, extra = _resolve_model(args.model)
    audio = load_audio(args.audio)
    rf = _operating_point(args, model.cfg.frame_rate)
    scores = predict_scores(model, audio.samples, feat, device=args.device,
                            chunk_seconds=args.chunk_seconds)
    typed = rf.decode_typed(scores, duration=audio.duration)
    anyv = rf.decode_any(scores, duration=audio.duration)
    meta = {
        "audio": str(args.audio),
        "duration_s": round(audio.duration, 3),
        "sample_rate": audio.sample_rate,
        "model": extra.get("preset", model.cfg.name),
        "param_count": model.param_count(),
        "operating_point": rf.op.as_dict(),
        "frame_rate_hz": model.cfg.frame_rate,
    }
    # Prefer typed spans where available; fall back to the ANY channel so a
    # region is never dropped merely because no category channel claimed it.
    spans = list(typed)
    for s in anyv:
        if not any(t.overlap(s) > 0.5 * s.duration for t in typed):
            spans.append(s)
    return sorted(spans, key=lambda s: (s.start, s.label)), meta


# ---------------------------------------------------------------------------
# commands
# ---------------------------------------------------------------------------


def cmd_find(args) -> int:
    """Report sensitive regions: category + time only.

    Claim: no-content -- the output is category and time, nothing else.
    """
    spans, meta = _detect(args)
    if args.json:
        print(json.dumps({"meta": meta, "spans": [s.as_dict() for s in spans],
                          "disclaimer": DISCLAIMER}, indent=2))
        return 0
    print(f"{meta['audio']}  ({meta['duration_s']:.2f}s, "
          f"model {meta['model']}, operating point "
          f"{meta['operating_point']['name']} @ thr "
          f"{meta['operating_point']['threshold']})")
    if not spans:
        print("\n  no regions flagged")
    else:
        print(f"\n  {'start':>8}  {'end':>8}  {'dur':>6}  {'score':>5}  category")
        print("  " + "-" * 56)
        for s in spans:
            print(f"  {s.start:8.2f}  {s.end:8.2f}  {s.duration:6.2f}  "
                  f"{s.score:5.2f}  {s.label}")
        total = sum(s.duration for s in spans)
        print("  " + "-" * 56)
        print(f"  {len(spans)} region(s), {total:.2f}s "
              f"({total/meta['duration_s']:.1%} of the recording)")
    print(f"\n{DISCLAIMER}")
    return 0


def cmd_redact(args) -> int:
    """Plan a redaction (default) or apply a decided plan.

    Claim: the reviewability requirement -- plan by default, audio only from a
    decided plan.
    """
    audio = load_audio(args.audio)
    red = Redactor(args.mode)

    if args.apply:
        if not args.plan:
            raise SystemExit(
                "--apply needs --plan PATH (the reviewed plan).\n"
                "Produce one first:  bleep redact audio.wav --plan-out plan.json"
            )
        plan = RedactionPlan.load(args.plan)
        if args.decisions_csv:
            plan.apply_csv_decisions(Path(args.decisions_csv).read_text(encoding="utf-8"))
        try:
            out_audio, report = red.apply(audio, plan,
                                          require_review=not args.no_review)
        except ValueError as e:
            raise SystemExit(f"refused: {e}")
        out_path = Path(args.out or (Path(args.audio).with_suffix("")).name + ".redacted.wav")
        Redactor.write(out_path, out_audio, source_path=args.audio)
        v = verify_redaction(audio, out_audio, plan)
        print(json.dumps({"written": str(out_path), "report": report,
                          "verification": v}, indent=2))
        if not v["ok"]:
            print("WARNING: redaction verification failed", file=sys.stderr)
            return 1
        return 0

    spans, meta = _detect(args)
    plan = red.plan(audio, spans, source_path=str(args.audio),
                    operating_point=meta["operating_point"],
                    auto_approve=args.auto_approve)
    plan_path = Path(args.plan_out or
                     Path(args.audio).with_suffix(".bleep-plan.json").name)
    plan.save(plan_path)
    csv_path = plan_path.with_suffix(".csv")
    csv_path.write_text(plan.to_csv(), encoding="utf-8")
    print(json.dumps(plan.summary(), indent=2))
    print(f"\nplan:  {plan_path}\nreview: {csv_path}")
    print("\nNo audio was modified. Review the CSV (set each `decision` to "
          "approved/rejected), then:")
    print(f"  bleep redact {args.audio} --apply --plan {plan_path} "
          f"--decisions-csv {csv_path} --out redacted.wav")
    print(f"\n{DISCLAIMER}")
    return 0


def cmd_union(args) -> int:
    """Union Bleep's detections with an existing detector's spans.

    Claim: complementarity and containment; also no-content, since text on the
    incoming spans is dropped.
    """
    spans, meta = _detect(args)
    payload = json.loads(Path(args.asr_spans).read_text(encoding="utf-8"))
    raw = payload.get("spans", payload) if isinstance(payload, dict) else payload
    asr = []
    for d in raw:
        # Accept foreign span formats, and drop any text field the other tool
        # attached: Bleep's artefacts stay content-free even when its input
        # was not.
        asr.append(Span(float(d["start"]), float(d["end"]),
                        d.get("label") if d.get("label") in CATEGORIES else ANY,
                        float(d.get("score", 1.0)), "asr"))
    uc = UnionCombiner(args.gap_tolerance)
    bd = uc.combine(asr, spans, duration=meta["duration_s"])
    ok, missing = uc.verify_containment(asr, bd)
    out = {
        "meta": meta,
        "union": [s.as_dict() for s in bd.union],
        "breakdown": bd.as_dict(),
        "containment_holds": ok,
        "containment_violations": [s.as_dict() for s in missing],
        "disclaimer": DISCLAIMER,
    }
    print(json.dumps(out, indent=2))
    return 0 if ok else 1


def cmd_info(args) -> int:
    """Model size, receptive field, measured speed, and the size-band check.

    Claim: on-device speed -- size, receptive field, measured real-time
    factor.
    """
    from .tagger import build, in_target_band

    if args.model or DEFAULT_CKPT.exists():
        model, feat, extra = _resolve_model(args.model)
        info = {
            "checkpoint": str(args.model or DEFAULT_CKPT),
            "preset": extra.get("preset", model.cfg.name),
            "trained_step": extra.get("step"),
            "validation": extra.get("val"),
        }
    else:
        model = build(args.preset)
        feat = None
        info = {"checkpoint": None, "preset": args.preset,
                "note": "no checkpoint found; reporting the untrained architecture"}
    from .features import DEFAULT_FEATURES
    info.update(
        param_count=model.param_count(),
        param_count_millions=round(model.param_count() / 1e6, 2),
        in_target_size_band_0p05_0p15B=in_target_band(model),
        receptive_field_seconds=round(model.cfg.receptive_field_seconds, 2),
        output_frame_rate_hz=model.cfg.frame_rate,
        gmac_per_audio_second=round(
            model.macs_per_output_frame() * model.cfg.frame_rate / 1e9, 3),
    )
    if args.speed:
        info["speed"] = measure_rtf(model, feat or DEFAULT_FEATURES,
                                    seconds=args.speed_seconds,
                                    device=args.device, repeats=3)
    info["operating_points"] = {k: v.as_dict() for k, v in PRESETS.items()}
    info["categories"] = list(CATEGORIES)
    info["disclaimer"] = DISCLAIMER
    print(json.dumps(info, indent=2))
    return 0


# ---------------------------------------------------------------------------


def build_parser() -> argparse.ArgumentParser:
    """Build the argument parser.

    Claim: no-content -- there is no subcommand that could emit a transcript.
    """
    ap = argparse.ArgumentParser(
        prog="bleep",
        description="Transcript-free detection of sensitive speech regions "
                    "(category + time only).",
        epilog=DISCLAIMER,
    )
    sub = ap.add_subparsers(dest="cmd", required=True)

    def common(p):
        """Add the arguments every audio-taking subcommand shares.

        Claim: recall at fixed false-positive rate -- the operating point is selectable on every command.
        """
        p.add_argument("audio", type=Path)
        p.add_argument("--model", default=None, help="checkpoint path")
        p.add_argument("--recall", default="high", choices=sorted(PRESETS),
                       help="operating point (default: high = recall-first)")
        p.add_argument("--operating-point-file", default=None,
                       help="calibrated operating point JSON")
        p.add_argument("--pre-margin", type=float, default=None,
                       help="override leading margin, seconds")
        p.add_argument("--post-margin", type=float, default=None)
        p.add_argument("--device", default="cpu")
        p.add_argument("--chunk-seconds", type=float, default=30.0)

    p = sub.add_parser("find", help="report sensitive regions (category + time)")
    common(p)
    p.add_argument("--json", action="store_true")
    p.set_defaults(func=cmd_find)

    p = sub.add_parser("redact", help="plan a redaction; --apply to write audio")
    common(p)
    p.add_argument("--mode", default="silence", choices=MODES)
    p.add_argument("--plan-out", default=None)
    p.add_argument("--auto-approve", action="store_true",
                   help="mark every region approved in the plan (skips review)")
    p.add_argument("--apply", action="store_true", help="write redacted audio")
    p.add_argument("--plan", default=None, help="reviewed plan to apply")
    p.add_argument("--decisions-csv", default=None)
    p.add_argument("--no-review", action="store_true",
                   help="apply pending regions without review (recorded in the plan)")
    p.add_argument("--out", default=None)
    p.set_defaults(func=cmd_redact)

    p = sub.add_parser("union", help="union Bleep with an existing detector's spans")
    common(p)
    p.add_argument("--asr-spans", required=True,
                   help="JSON list of {start,end[,label,score]} from your detector")
    p.add_argument("--gap-tolerance", type=float, default=0.08)
    p.set_defaults(func=cmd_union)

    p = sub.add_parser("info", help="model size, receptive field, speed")
    p.add_argument("--model", default=None)
    p.add_argument("--preset", default="bleep-base")
    p.add_argument("--speed", action="store_true", help="measure real-time factor")
    p.add_argument("--speed-seconds", type=float, default=60.0)
    p.add_argument("--device", default="cpu")
    p.set_defaults(func=cmd_info)
    return ap


def main(argv: list[str] | None = None) -> int:
    """CLI entry point.

    Claim: no-content.
    """
    args = build_parser().parse_args(argv)
    return args.func(args)


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