source/src/vimeml/deployment/validation.py

"""Numerical fixtures, frozen candidate evaluations and Mac-only timing probes."""
import json
import time
from pathlib import Path

import numpy as np
import torch

from vimeml.benchmarks.ajimee import DEVELOPMENT_FORMAT, FORMAT
from vimeml.benchmarks.evaluate_ajimee import load_export, rerank, summarize
from vimeml.deployment.bundle import (fresh_directory, read_json, tree_inventory, write_json)
from vimeml.training.data import file_sha
from vimeml.training.evaluate_ime import score_candidates

SCORING_CASES = [
    {"context": "", "candidates": ["東京", "とうきょう", "東京都"]},
    {"context": "明日", "candidates": ["は", "は晴れ", "の会議"]},
    {"context": "コーヒーを", "candidates": ["飲む", "買う", "飲みたい"]},
    {"context": "ファイル", "candidates": ["名", "名前", "を開く"]},
]


def numeric_error(expected, actual, atol, rtol):
    if expected.shape != actual.shape or not np.isfinite(actual).all():
        raise ValueError("Nonfinite or mismatched comparison output.")
    delta = np.abs(expected.astype(np.float64) - actual.astype(np.float64))
    return {"max_abs": float(delta.max()), "mean_abs": float(delta.mean()),
            "rmse": float(np.sqrt(np.mean(delta ** 2))),
            "outside_tolerance": int((delta > atol + rtol * np.abs(expected)).sum()),
            "values": int(delta.size)}


@torch.inference_mode()
def reference(lm, output):
    if output.exists():
        raise ValueError("Reference output exists; choose a new directory.")
    special = lm.special
    short = lm.prefix_ids("明日の会議までに、")
    rng = np.random.default_rng(716)
    long = rng.integers(4, lm.model.config.vocab_size, size=lm.model.config.context_length).tolist()
    long[0] = special["bos"]
    changed = long.copy()
    changed[16:] = rng.integers(4, lm.model.config.vocab_size, size=len(long) - 16).tolist()
    # Explicit byte pieces exercise byte fallback even if the text is in the vocabulary.
    byte_ids = [lm.processor.piece_to_id(f"<0x{value:02X}>") for value in (0xF0, 0x9F, 0xA6, 0x84)]
    if special["unk"] in byte_ids:
        raise ValueError("Expected byte-fallback tokenizer.")
    fixtures = [
        {"id": "bos", "input_ids": [special["bos"]], "valid_length": 1},
        {"id": "short", "input_ids": short, "valid_length": len(short)},
        {"id": "right-pad", "input_ids": short + [special["pad"]] * (lm.model.config.context_length - len(short)), "valid_length": len(short)},
        {"id": "byte-eos", "input_ids": [special["bos"], *byte_ids, special["eos"]], "valid_length": 6},
        {"id": "long", "input_ids": long, "valid_length": len(long)},
        {"id": "causal-mutation", "input_ids": changed, "valid_length": len(changed)},
    ]
    arrays = {}
    for fixture in fixtures:
        logits = lm.model(torch.tensor([fixture["input_ids"]])).float().cpu().numpy()
        arrays[fixture["id"]] = logits
        fixture["last_top1"] = int(logits[0, fixture["valid_length"] - 1].argmax())
    torch.testing.assert_close(torch.from_numpy(arrays["short"]),
        torch.from_numpy(arrays["right-pad"][:, :len(short)]), atol=3e-5, rtol=3e-4)
    torch.testing.assert_close(torch.from_numpy(arrays["long"][:, :16]),
        torch.from_numpy(arrays["causal-mutation"][:, :16]), atol=3e-5, rtol=3e-4)
    scored, device_candidates = [], []
    for case in SCORING_CASES:
        result = score_candidates(lm, case["context"], case["candidates"])
        scored.append({**case, "scored": result})
        for candidate in result["candidates"]:
            seq = [special["bos"], *lm.processor.encode(case["context"] + candidate["text"], out_type=int)]
            device_candidates.append({"context": case["context"], "text": candidate["text"],
                "input_ids": seq[:-1], "targets": seq[1:],
                "score_start": result["common_prefix_tokens_including_bos"] - 1,
                "log_probability_sum": candidate["log_probability_sum"]})
    fresh_directory(output)
    np.savez_compressed(output / "logits.npz", **arrays)
    write_json(output / "fixtures.json", {"logits": fixtures, "scoring": scored})
    write_json(output / "device-fixtures.json", {"format": "vimeml_device_fixtures_v1",
        "bundle_manifest_sha256": lm.metadata["bundle_manifest_sha256"],
        "logits": fixtures, "candidates": device_candidates})
    write_json(output / "manifest.json", {"format": "vimeml_reference_v1", "status": "complete",
        "model": lm.metadata, "files": tree_inventory(output),
        "note": "Windows CPU FP32 reference. Explicit token IDs test tensor contract; device-side SentencePiece roundtrip/joint tokenization still requires separate integration tests."})


@torch.inference_mode()
def validate(lm, reference_dir, output, atol, rtol, score_atol):
    if output.exists():
        raise ValueError("Alignment output exists.")
    manifest = read_json(reference_dir / "manifest.json")
    if manifest.get("format") != "vimeml_reference_v1" or manifest.get("status") != "complete":
        raise ValueError("Invalid reference manifest.")
    if manifest["model"]["bundle_manifest_sha256"] != lm.metadata["bundle_manifest_sha256"]:
        raise ValueError("Reference belongs to another bundle.")
    for name, entry in manifest["files"].items():
        if file_sha(reference_dir / name) != entry["sha256"]:
            raise ValueError(f"Reference changed: {name}")
    fixtures = read_json(reference_dir / "fixtures.json")
    errors, predictions = [], {}
    with np.load(reference_dir / "logits.npz", allow_pickle=False) as expected:
        for fixture in fixtures["logits"]:
            ids = torch.tensor([fixture["input_ids"]], dtype=torch.long)
            actual = lm.model(ids).float().cpu().numpy()
            predictions[fixture["id"]] = actual
            valid = fixture["valid_length"]
            error = numeric_error(expected[fixture["id"]][:, :valid], actual[:, :valid], atol, rtol)
            error.update(id=fixture["id"], last_top1_equal=int(actual[0, valid - 1].argmax()) == fixture["last_top1"])
            errors.append(error)
    # Explicitly verify causality and right-PAD invariance on the converted model too.
    length = predictions["short"].shape[1]
    invariants = {
        "right_pad": numeric_error(predictions["short"], predictions["right-pad"][:, :length], atol, rtol),
        "causal": numeric_error(predictions["long"][:, :16], predictions["causal-mutation"][:, :16], atol, rtol),
    }
    score_errors = []
    for case in fixtures["scoring"]:
        scored = score_candidates(lm, case["context"], case["candidates"])
        before = case["scored"]["candidates"]
        after = scored["candidates"]
        if scored["common_prefix_tokens_including_bos"] != case["scored"]["common_prefix_tokens_including_bos"]:
            raise ValueError("Joint tokenization/common-prefix contract changed.")
        max_delta = 0.0
        for left, right in zip(before, after):
            if left["token_ids"] != right["token_ids"]:
                raise ValueError("Candidate target tokens changed.")
            max_delta = max(max_delta, abs(left["log_probability_sum"] - right["log_probability_sum"]))
        order = lambda values: [item["text"] for item in sorted(values, key=lambda item: item["log_probability_sum"], reverse=True)]
        score_errors.append({"context": case["context"], "max_sum_abs": max_delta,
                             "order_equal": order(before) == order(after), "scored": scored})
    passed = (all(item["outside_tolerance"] == 0 for item in [*errors, *invariants.values()]) and
              all(item["max_sum_abs"] <= score_atol and item["order_equal"] for item in score_errors))
    fresh_directory(output)
    report = {"format": "vimeml_alignment_v1", "passed": passed, "model": lm.metadata,
        "coreml_manifest_sha256": lm.metadata.get("coreml_manifest_sha256"),
        "reference_manifest_sha256": file_sha(reference_dir / "manifest.json"),
        "tolerances": {"logits_atol": atol, "logits_rtol": rtol, "score_sum_atol": score_atol},
        "logits": errors, "invariants": invariants, "scores": score_errors,
        "note": "Numerical fixture gate, not a quality certification. Review dev/AJIMEE and phrases separately. Quantization may fail FP16 tolerances; preserve failures and inspect before changing thresholds."}
    write_json(output / "alignment.json", report)
    return report


def compare_rows(before, after):
    if len(before) != len(after):
        raise ValueError("Baseline case count changed.")
    changed, max_delta, total_delta, score_count = [], 0.0, 0.0, 0
    for left, right in zip(before, after):
        for key in ("id", "query", "left_context", "answers", "candidates"):
            if left[key] != right[key]:
                raise ValueError(f"Baseline candidate/input mismatch: {key}")
        delta = None
        a = left.get("lm_scores", {}).get("contextual")
        b = right.get("lm_scores", {}).get("contextual")
        if a and b:
            if a["common_prefix_tokens_including_bos"] != b["common_prefix_tokens_including_bos"]:
                raise ValueError("Baseline tokenizer boundary mismatch.")
            deltas = []
            for x, y in zip(a["candidates"], b["candidates"]):
                if x["text"] != y["text"] or x["token_ids"] != y["token_ids"]:
                    raise ValueError("Baseline token targets mismatch.")
                deltas.append(abs(x["log_probability_sum"] - y["log_probability_sum"]))
            delta = max(deltas)
            max_delta = max(max_delta, delta)
            total_delta += sum(deltas)
            score_count += len(deltas)
        old, new = left["orders"]["lm_context_sum"], right["orders"]["lm_context_sum"]
        if old != new or left["fallbacks"] != right["fallbacks"]:
            changed.append({"id": right["id"], "old_order": old, "new_order": new,
                "top1_changed": old[:1] != new[:1], "answers": right["answers"],
                "max_score_sum_abs": delta, "old_fallbacks": left["fallbacks"], "new_fallbacks": right["fallbacks"]})
    return {"order_changed_cases": len(changed), "top1_changed_cases": sum(row["top1_changed"] for row in changed),
            "max_score_sum_abs": max_delta, "mean_score_sum_abs": total_delta / score_count if score_count else None,
            "changes": changed}


def evaluate(lm, benchmark, role, output, baseline=None):
    if output.exists():
        raise ValueError("Evaluation output exists; preserve the baseline.")
    manifest, provenance, rows = load_export(benchmark)
    expected_format = DEVELOPMENT_FORMAT if role == "dev" else FORMAT
    if manifest["format"] != expected_format:
        raise ValueError("Dataset role does not match the frozen benchmark format.")
    started = time.perf_counter()
    rerank(lm, rows)  # Reuse exact frozen scoring, eligibility, stable ties and whole-case fallback.
    comparison = None
    if baseline:
        baseline = Path(baseline)
        metadata_path = baseline / "manifest.json" if (baseline / "manifest.json").exists() else baseline / "metrics.json"
        prior = read_json(metadata_path)
        model = prior["model"]
        for key in ("checkpoint_sha256", "tokenizer_sha256"):
            if model[key] != lm.metadata[key]:
                raise ValueError(f"Baseline identity mismatch: {key}")
        if prior["benchmark_manifest"] != manifest or prior["export_provenance"] != provenance:
            raise ValueError("Baseline uses a different frozen dataset/candidate export.")
        digest = prior.get("files_sha256", {}).get("scores.jsonl")
        if not digest or file_sha(baseline / "scores.jsonl") != digest:
            raise ValueError("Baseline score hash mismatch.")
        prior_rows = [json.loads(line) for line in (baseline / "scores.jsonl").read_text(encoding="utf-8").splitlines()]
        comparison = compare_rows(prior_rows, rows)
    fresh_directory(output)
    (output / "scores.jsonl").write_text("".join(json.dumps(row, ensure_ascii=False, allow_nan=False) + "\n" for row in rows), encoding="utf-8")
    report = {"format": "vimeml_deployment_metrics_v1", "status": "complete", "role": role,
        "model": lm.metadata, "benchmark_manifest": manifest, "export_provenance": provenance,
        "metrics": summarize(rows), "comparison": comparison,
        "files_sha256": {"scores.jsonl": file_sha(output / "scores.jsonl")},
        "elapsed_seconds": time.perf_counter() - started,
        "policy": "LM-only contextual logP sum primary. AzooKey retrieves; scores ignored except original-order ties/fallback. Mean and no-context secondary. No lambda selection on AJIMEE.",
        "note": "All cases remain in denominators. Fixed AJIMEE is an observed comparison, not a fresh blind test. Latency here includes diagnostics/context-free scores; use timing command for workload measurements."}
    write_json(output / "metrics.json", report)
    return report


@torch.inference_mode()
def timing(lm, benchmark, prompts, output, repeats, warmup):
    from vimeml.benchmarks.evaluate_ajimee import eligibility
    from vimeml.tools.phrase_demo import PhraseDemo
    if output.exists():
        raise ValueError("Timing output exists.")
    _, provenance, rows = load_export(benchmark)
    cases = read_json(prompts)
    demo = PhraseDemo(lm, cases)
    times, workloads = [], []
    for row in rows:
        texts = row["orders"]["azookey"]
        reason = eligibility(lm, row["left_context"], texts)
        if reason:
            workloads.append({"id": row["id"], "fallback": reason})
            continue
        def operation():
            return score_candidates(lm, row["left_context"], texts)
        for _ in range(warmup):
            operation()
        elapsed = []
        for _ in range(repeats):
            start = time.perf_counter()
            operation()
            elapsed.append((time.perf_counter() - start) * 1000)
        times.extend(elapsed)
        workloads.append({"id": row["id"], "candidates": len(texts), "elapsed_ms": elapsed})
    phrase_times, phrases = [], []
    for case in cases:
        settings = {"prompt": case["prompt"], "max_tokens": 8, "mode": "beam", "count": 5, "seed": 42}
        for _ in range(warmup):
            demo.suggest(settings)
        elapsed = []
        for _ in range(repeats):
            result = demo.suggest(settings)
            elapsed.append(result["elapsed_ms"])
        phrase_times.extend(elapsed)
        phrases.append({"id": case["id"], "elapsed_ms": elapsed, "returned_count": result["returned_count"]})
    percentiles = lambda values: {"count": len(values), **({"p50_ms": float(np.percentile(values, 50)),
        "p95_ms": float(np.percentile(values, 95)), "max_ms": max(values)} if values else {})}
    fresh_directory(output)
    write_json(output / "timing.json", {"model": lm.metadata, "warmup_per_case": warmup, "repeats": repeats,
        "prompts_sha256": file_sha(prompts), "export_provenance": provenance,
        "load_seconds": getattr(lm.model, "load_seconds", None),
        "reranking": percentiles(times), "phrases": percentiles(phrase_times),
        "reranking_cases": workloads, "phrase_cases": phrases,
        "predict_calls_total": getattr(lm.model, "calls", None),
        "predict_seconds_total": getattr(lm.model, "predict_seconds", None),
        "note": "Mac host timings, full vocab transfer and app-side Python search included. Load may compile/use cache; not true cold iPhone loading. No iPhone/keyboard memory claim."})