Blame for src/dlb/rdlm_diagnostics.py

"""Read-only diagnostics for suspicious RDLM quality rows."""

from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
from statistics import mean, median
import sys
from typing import Any, Iterable, Mapping, Sequence

from dlb.io import sha256_file


SCHEMA = "dlb-rdlm-diagnostics-v1"
DEFAULT_STEPS = (1, 2, 4, 8, 16, 32, 1024)
COLLAPSE_WARNINGS = frozenset(
    {
        "low_ppl_with_low_entropy",
        "high_self_bleu",
        "high_duplicate_texts",
        "short_decoded_texts",
        "high_padding_fraction",
    }
)


def _read_json(path: Path, warnings: list[str], label: str) -> dict[str, Any] | None:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError):
        warnings.append(f"invalid_{label}")
        return None
    if not isinstance(value, dict):
        warnings.append(f"invalid_{label}")
        return None
    return value


def _read_samples(path: Path, warnings: list[str]) -> list[dict[str, Any]]:
    samples: list[dict[str, Any]] = []
    try:
        with path.open("r", encoding="utf-8") as source:
            for expected_id, line in enumerate(source):
                try:
                    record = json.loads(line)
                except json.JSONDecodeError:
                    warnings.append("invalid_sample_jsonl")
                    return samples
                if not isinstance(record, dict):
                    warnings.append("invalid_sample_record")
                    return samples
                if record.get("sample_id") != expected_id:
                    warnings.append("sample_id_mismatch")
                samples.append(record)
    except OSError:
        warnings.append("invalid_samples")
    return samples


def _metric(metrics: Mapping[str, Any], name: str) -> float | None:
    aliases = {
        "generative_perplexity": (
            ("generative_perplexity", "perplexity"),
            ("gen_ppl", "perplexity"),
            ("ppl", None),
        ),
        "unigram_entropy": (
            ("unigram_entropy", "mean_entropy"),
            ("entropy", "mean_entropy"),
            ("entropy", None),
        ),
        "self_bleu": (
            ("self_bleu", "score"),
            ("self_bleu_score", None),
        ),
    }[name]
    for key, field in aliases:
        if key not in metrics:
            continue
        value = metrics[key]
        if field is not None:
            if not isinstance(value, Mapping) or field not in value:
                continue
            value = value[field]
        if type(value) in {int, float} and not isinstance(value, bool):
            result = float(value)
            return result if math.isfinite(result) else None
    return None


def _command_override(command: object, key: str) -> str | None:
    if not isinstance(command, list):
        return None
    prefix = key + "="
    values = [
        argument[len(prefix) :]
        for argument in command
        if isinstance(argument, str) and argument.startswith(prefix)
    ]
    return values[0] if len(values) == 1 else None


def _sample_stats(samples: list[dict[str, Any]]) -> dict[str, Any]:
    texts = [item.get("text") for item in samples if isinstance(item.get("text"), str)]
    token_rows = [
        item.get("token_ids")
        for item in samples
        if isinstance(item.get("token_ids"), list)
    ]
    text_lengths = [len(text) for text in texts]
    token_lengths = [len(tokens) for tokens in token_rows]
    flat_tokens = [
        token
        for tokens in token_rows
        for token in tokens
        if type(token) is int and token >= 0
    ]
    total_tokens = len(flat_tokens)
    return {
        "sample_count": len(samples),
        "unique_text_count": len(set(texts)),
        "unique_text_ratio": len(set(texts)) / len(texts) if texts else None,
        "mean_text_chars": mean(text_lengths) if text_lengths else None,
        "median_text_chars": median(text_lengths) if text_lengths else None,
        "mean_token_length": mean(token_lengths) if token_lengths else None,
        "min_token_length": min(token_lengths) if token_lengths else None,
        "max_token_length": max(token_lengths) if token_lengths else None,
        "pad_token_fraction": (
            sum(1 for token in flat_tokens if token == 0) / total_tokens
            if total_tokens
            else None
        ),
        "cls_text_fraction": (
            sum(1 for text in texts if "[CLS]" in text) / len(texts) if texts else None
        ),
        "sep_text_fraction": (
            sum(1 for text in texts if "[SEP]" in text) / len(texts) if texts else None
        ),
    }


def _add_quality_warnings(
    row: dict[str, Any],
    warnings: list[str],
    *,
    expected_sample_count: int,
    low_ppl_threshold: float,
    low_entropy_threshold: float,
    high_self_bleu_threshold: float,
    duplicate_unique_ratio_threshold: float,
    short_text_chars_threshold: float,
    high_padding_fraction_threshold: float,
) -> None:
    if row.get("sample_count") not in {None, expected_sample_count}:
        warnings.append("sample_count_mismatch")
    if row.get("metrics_sample_count") not in {None, expected_sample_count}:
        warnings.append("metrics_sample_count_mismatch")
    if (
        row.get("generative_perplexity") is not None
        and row.get("unigram_entropy") is not None
        and row["generative_perplexity"] < low_ppl_threshold
        and row["unigram_entropy"] < low_entropy_threshold
    ):
        warnings.append("low_ppl_with_low_entropy")
    if (
        row.get("self_bleu") is not None
        and row["self_bleu"] > high_self_bleu_threshold
    ):
        warnings.append("high_self_bleu")
    if (
        row.get("unique_text_ratio") is not None
        and row["unique_text_ratio"] < duplicate_unique_ratio_threshold
    ):
        warnings.append("high_duplicate_texts")
    if (
        row.get("median_text_chars") is not None
        and row["median_text_chars"] < short_text_chars_threshold
    ):
        warnings.append("short_decoded_texts")
    if (
        row.get("pad_token_fraction") is not None
        and row["pad_token_fraction"] > high_padding_fraction_threshold
    ):
        warnings.append("high_padding_fraction")
    if row.get("token_ids_source") == "retokenized":
        warnings.append("retokenized_from_text")


def _row_for_step(
    root: Path,
    *,
    dataset: str,
    model: str,
    step: int,
    expected_sample_count: int,
    low_ppl_threshold: float,
    low_entropy_threshold: float,
    high_self_bleu_threshold: float,
    duplicate_unique_ratio_threshold: float,
    short_text_chars_threshold: float,
    high_padding_fraction_threshold: float,
) -> dict[str, Any]:
    sample_dir = root / "results" / "samples" / dataset / model / f"steps_{step}"
    samples_path = sample_dir / "samples.jsonl"
    metrics_path = root / "results" / "metrics" / dataset / model / f"steps_{step}" / "metrics.json"
    conversion_path = sample_dir / "conversion_metadata.json"
    metadata_path = sample_dir / "run_metadata.json"
    warnings: list[str] = []
    row: dict[str, Any] = {
        "dataset": dataset,
        "model": model,
        "steps": step,
        "sample_dir": str(sample_dir),
        "metrics_path": str(metrics_path),
        "sample_count": None,
        "metrics_sample_count": None,
        "generative_perplexity": None,
        "unigram_entropy": None,
        "self_bleu": None,
        "token_ids_source": None,
        "token_ids_reason": None,
        "command_sampling_steps": None,
        "command_batch_per_gpu": None,
    }

    if samples_path.is_file() and not samples_path.is_symlink():
        samples = _read_samples(samples_path, warnings)
        row.update(_sample_stats(samples))
    else:
        warnings.append("missing_samples")

    if metrics_path.is_file() and not metrics_path.is_symlink():
        metrics_document = _read_json(metrics_path, warnings, "metrics")
        if metrics_document is not None:
            metric_values = metrics_document.get("metrics")
            metric_values = metric_values if isinstance(metric_values, Mapping) else {}
            row["metrics_sample_count"] = metrics_document.get("sample_count")
            row["metrics_partial"] = metrics_document.get("partial")
            row["generative_perplexity"] = _metric(
                metric_values, "generative_perplexity"
            )
            row["unigram_entropy"] = _metric(metric_values, "unigram_entropy")
            row["self_bleu"] = _metric(metric_values, "self_bleu")
            recorded_sha = metrics_document.get("samples_sha256")
            if (
                isinstance(recorded_sha, str)
                and samples_path.is_file()
                and sha256_file(samples_path) != recorded_sha
            ):
                warnings.append("metrics_sample_sha_mismatch")
    else:
        warnings.append("missing_metrics")

    if conversion_path.is_file() and not conversion_path.is_symlink():
        conversion = _read_json(conversion_path, warnings, "conversion_metadata")
        if conversion is not None:
            row["token_ids_source"] = conversion.get("token_ids_source")
            transformation = conversion.get("token_ids_transformation")
            if isinstance(transformation, Mapping):
                row["token_ids_reason"] = transformation.get("reason")
    elif samples_path.is_file():
        warnings.append("missing_conversion_metadata")

    if metadata_path.is_file() and not metadata_path.is_symlink():
        metadata = _read_json(metadata_path, warnings, "run_metadata")
        if metadata is not None:
            row["run_status"] = metadata.get("status")
            command = metadata.get("command")
            row["command_sampling_steps"] = _command_override(
                command, "sampling.steps"
            )
            row["command_batch_per_gpu"] = _command_override(
                command, "sampling.batch_per_gpu"
            )
            if row["command_sampling_steps"] not in {None, str(step)}:
                warnings.append("command_step_mismatch")
    elif samples_path.is_file():
        warnings.append("missing_run_metadata")

    _add_quality_warnings(
        row,
        warnings,
        expected_sample_count=expected_sample_count,
        low_ppl_threshold=low_ppl_threshold,
        low_entropy_threshold=low_entropy_threshold,
        high_self_bleu_threshold=high_self_bleu_threshold,
        duplicate_unique_ratio_threshold=duplicate_unique_ratio_threshold,
        short_text_chars_threshold=short_text_chars_threshold,
        high_padding_fraction_threshold=high_padding_fraction_threshold,
    )
    row["warnings"] = list(dict.fromkeys(warnings))
    return row


def _ppl_nonmonotonic(rows: Sequence[Mapping[str, Any]]) -> bool:
    values = [
        row.get("generative_perplexity")
        for row in rows
        if type(row.get("generative_perplexity")) in {int, float}
    ]
    if len(values) < 3:
        return False
    nondecreasing = all(a <= b for a, b in zip(values, values[1:]))
    nonincreasing = all(a >= b for a, b in zip(values, values[1:]))
    return not nondecreasing and not nonincreasing


def _summary(rows: list[dict[str, Any]]) -> dict[str, Any]:
    suspect_steps = [
        row["steps"]
        for row in rows
        if COLLAPSE_WARNINGS.intersection(row["warnings"])
        or (
            "retokenized_from_text" in row["warnings"]
            and row.get("generative_perplexity") is not None
            and row["generative_perplexity"] < 100.0
        )
    ]
    missing_steps = [
        row["steps"]
        for row in rows
        if "missing_samples" in row["warnings"] or "missing_metrics" in row["warnings"]
    ]
    retokenized_steps = [
        row["steps"] for row in rows if row.get("token_ids_source") == "retokenized"
    ]
    nonmonotonic = _ppl_nonmonotonic(rows)
    if suspect_steps:
        verdict = "low_ppl_points_show_collapse_signals"
    elif missing_steps:
        verdict = "missing_server_artifacts"
    elif nonmonotonic:
        verdict = "ppl_curve_is_nonmonotonic"
    else:
        verdict = "no_obvious_artifact"
    return {
        "verdict": verdict,
        "suspect_steps": suspect_steps,
        "missing_steps": missing_steps,
        "retokenized_steps": retokenized_steps,
        "ppl_curve_nonmonotonic": nonmonotonic,
    }


def diagnose_rdlm(
    root: Path | str,
    *,
    dataset: str = "lm1b",
    model: str = "rdlm",
    steps: Iterable[int] = DEFAULT_STEPS,
    expected_sample_count: int = 1024,
    low_ppl_threshold: float = 100.0,
    low_entropy_threshold: float = 3.8,
    high_self_bleu_threshold: float = 0.8,
    duplicate_unique_ratio_threshold: float = 0.9,
    short_text_chars_threshold: float = 80.0,
    high_padding_fraction_threshold: float = 0.25,
) -> dict[str, Any]:
    """Inspect already-produced RDLM artifacts and summarize suspicious signals."""

    root = Path(root).resolve()
    rows = [
        _row_for_step(
            root,
            dataset=dataset,
            model=model,
            step=step,
            expected_sample_count=expected_sample_count,
            low_ppl_threshold=low_ppl_threshold,
            low_entropy_threshold=low_entropy_threshold,
            high_self_bleu_threshold=high_self_bleu_threshold,
            duplicate_unique_ratio_threshold=duplicate_unique_ratio_threshold,
            short_text_chars_threshold=short_text_chars_threshold,
            high_padding_fraction_threshold=high_padding_fraction_threshold,
        )
        for step in steps
    ]
    return {
        "schema": SCHEMA,
        "root": str(root),
        "dataset": dataset,
        "model": model,
        "expected_sample_count": expected_sample_count,
        "summary": _summary(rows),
        "rows": rows,
    }


def _parse_steps(value: str) -> tuple[int, ...]:
    try:
        steps = tuple(int(part) for part in value.split(",") if part)
    except ValueError as error:
        raise argparse.ArgumentTypeError("steps must be comma-separated integers") from error
    if not steps or any(step <= 0 for step in steps):
        raise argparse.ArgumentTypeError("steps must contain positive integers")
    return steps


def _format(value: object) -> str:
    if value is None:
        return "-"
    if isinstance(value, float):
        return f"{value:.4g}"
    return str(value)


def print_report(report: Mapping[str, Any]) -> None:
    rows = report["rows"]
    print(f"schema\t{report['schema']}")
    print(f"verdict\t{report['summary']['verdict']}")
    print(
        "steps\tsamples\tppl\tentropy\tself_bleu\tunique\tmedian_chars\t"
        "pad_frac\ttoken_source\twarnings"
    )
    for row in rows:
        print(
            "\t".join(
                [
                    str(row["steps"]),
                    _format(row.get("sample_count")),
                    _format(row.get("generative_perplexity")),
                    _format(row.get("unigram_entropy")),
                    _format(row.get("self_bleu")),
                    _format(row.get("unique_text_ratio")),
                    _format(row.get("median_text_chars")),
                    _format(row.get("pad_token_fraction")),
                    _format(row.get("token_ids_source")),
                    ",".join(row["warnings"]),
                ]
            )
        )


def main(argv: Sequence[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=Path.cwd())
    parser.add_argument("--dataset", default="lm1b")
    parser.add_argument("--model", default="rdlm")
    parser.add_argument("--steps", type=_parse_steps, default=DEFAULT_STEPS)
    parser.add_argument("--expected-sample-count", type=int, default=1024)
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args(argv)
    report = diagnose_rdlm(
        args.root,
        dataset=args.dataset,
        model=args.model,
        steps=args.steps,
        expected_sample_count=args.expected_sample_count,
    )
    if args.json:
        json.dump(report, sys.stdout, indent=2, sort_keys=True)
        print()
    else:
        print_report(report)
    return 0


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