bmb_normalizer/runtime.py

"""Hash-verified inference runtime for the public V3 artifact."""
from __future__ import annotations

import hashlib
import json
from pathlib import Path
import threading

from .edit_policy_v3 import inference_policy
from .edit_tags_v3 import render_tag_predictions, tokens


class InputError(ValueError):
    """A safe, user-correctable inference request error."""


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


class V3Runtime:
    def __init__(self, root: Path, device: str = "cuda:0"):
        import torch
        from transformers import AutoModelForTokenClassification, AutoTokenizer

        root = root.resolve()
        manifest = json.loads((root / "release-manifest.json").read_text())
        if manifest.get("product_version") != 3 or manifest.get("system_prompt") is not None:
            raise RuntimeError("V3 release identity mismatch")
        for name, expected in manifest["model_files"].items():
            path = root / name
            if path.stat().st_size != expected["bytes"] or _digest(path) != expected["sha256"]:
                raise RuntimeError(f"V3 model artifact mismatch: {name}")
        for name in ("labels.json", "candidate-index.json"):
            expected = manifest["runtime_files"][name]
            path = root / name
            if path.stat().st_size != expected["bytes"] or _digest(path) != expected["sha256"]:
                raise RuntimeError(f"V3 runtime artifact mismatch: {name}")
        policy_expected = manifest["runtime_files"]["edit_policy_v3.py"]["sha256"]
        if _digest(Path(__file__).with_name("edit_policy_v3.py")) != policy_expected:
            raise RuntimeError("V3 deployed policy differs from release")
        if not torch.cuda.is_available() or not device.startswith("cuda:"):
            raise RuntimeError("A CUDA GPU is required for this release")

        self.labels = json.loads((root / "labels.json").read_text())
        raw_index = json.loads((root / "candidate-index.json").read_text())
        self.index = {
            first: [(tuple(key), int(label_id)) for key, label_id in values]
            for first, values in raw_index.items()
        }
        self.threshold = float(manifest["selected_threshold"])
        self.torch = torch
        self.device = device
        self.tokenizer = AutoTokenizer.from_pretrained(root, local_files_only=True)
        self.model = AutoModelForTokenClassification.from_pretrained(
            root, local_files_only=True, torch_dtype=torch.float16,
        ).to(device).eval()
        self.lock = threading.Lock()
        self.gpu = f"{device} {torch.cuda.get_device_name(int(device.rsplit(':', 1)[1]))}"

    def run(self, text: str) -> dict:
        source_tokens = tokens(text)
        if not source_tokens:
            return {"output": text, "route": "exact_copy", "status": "identity",
                    "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}}
        words = [str(token["text"]) for token in source_tokens]
        encoded = self.tokenizer(
            words, is_split_into_words=True, return_tensors="pt", truncation=False,
        )
        if encoded.input_ids.shape[1] > 256:
            raise InputError("Input exceeds the verified 256-token budget; nothing was truncated.")
        encoded = encoded.to(self.device)
        allowed, forced = inference_policy(text, self.index, self.labels)
        selected = [self.labels[0]] * len(words)
        with self.lock, self.torch.inference_mode():
            logits = self.model(**encoded).logits[0]
        seen = set()
        for token_index, word_id in enumerate(encoded.word_ids()):
            if word_id is None or word_id in seen:
                continue
            seen.add(word_id)
            ids = sorted(allowed[word_id])
            probabilities = logits[token_index, ids].softmax(-1)
            confidence, local_index = probabilities.max(-1)
            chosen = forced[word_id] if forced[word_id] is not None else ids[int(local_index)]
            if forced[word_id] is None and chosen != 0 and float(confidence) < self.threshold:
                chosen = 0
            selected[word_id] = chosen if isinstance(chosen, dict) else self.labels[chosen]
        output = render_tag_predictions(text, selected)
        count = int(encoded.input_ids.shape[1])
        return {
            "output": output,
            "route": "source_conditioned_edit_tagger",
            "status": "normalized" if output != text else "identity",
            "product_version": 3,
            "guard_scope": "licensed source edits plus deterministic semantic and immutable-span policy",
            "usage": {"prompt_tokens": count, "completion_tokens": 0, "total_tokens": count},
        }