Planner-Cache / src /pcm /planner /compatibility.py
Dermitio's picture
Upload folder using huggingface_hub
ca3d977 verified
Raw
History Blame Contribute Delete
12.5 kB
"""Explicit model compatibility boundaries for canonical Planner Cache state.
TTL exposes canonical state inside a model forward path. LTL is intentionally
weaker and controls lexical output at a tokenizer or runtime boundary.
"""
from __future__ import annotations
from dataclasses import asdict, dataclass
from enum import Enum
import hashlib
import json
from pathlib import Path
import warnings
import torch
from safetensors import safe_open
from safetensors._safetensors_rust import SafetensorError
from safetensors.torch import load_file, save_file
from pcm.planner.canonical import CANONICAL_P_PROTOCOL
from pcm.planner.split_translator import (
SPLIT_TRANSLATE_FORMAT,
SplitPTranslatePackage,
SplitTranslateConfig,
tensor_checksum,
)
TTL_FORMAT = "planner-cache-ttl-v1"
LTL_FORMAT = "planner-cache-ltl-v1"
TTL_EXTENSION = ".ttl"
LTL_EXTENSION = ".ltl"
FORBIDDEN_ADAPTER_FIELDS = (
"base_model",
"conversation",
"p_cache",
"canonical_values",
"optimizer",
)
class CompatibilityKind(str, Enum):
NATIVE = "native"
TTL = "ttl"
LTL = "ltl"
@dataclass(frozen=True)
class CompatibilityResolution:
kind: CompatibilityKind
support_level: str
artifact: Path | None
class TensorTranslationLayer(SplitPTranslatePackage):
"""Semantic/internal compatibility module backed by the Pythia TTL.
The learned module is unchanged from the proven split translator. The new
container identifies its stronger semantic contract and rejects LTL files.
"""
adapter_class = CompatibilityKind.TTL.value
support_level = "semantic/internal"
extension = TTL_EXTENSION
def save(self, path: str | Path) -> None:
path = Path(path)
if path.suffix != TTL_EXTENSION:
raise ValueError(f"TTL artifacts must use the {TTL_EXTENSION} extension")
state = {name: value.detach().cpu() for name, value in self.state_dict().items()}
if any(field in name.casefold() for name in state for field in FORBIDDEN_ADAPTER_FIELDS):
raise ValueError("TTL contains forbidden model or conversation state")
config = asdict(self.config)
config.pop("format", None)
manifest = {
"format": TTL_FORMAT,
"adapter_class": self.adapter_class,
"support_level": self.support_level,
"canonical_protocol": self.config.canonical_protocol,
"config": config,
"weights_sha256": tensor_checksum(state),
}
save_file(state, str(path), metadata={
"manifest": json.dumps(manifest, sort_keys=True, separators=(",", ":")),
})
@classmethod
def load(
cls,
path: str | Path,
*,
device: str | torch.device = "cpu",
dtype: torch.dtype = torch.float32,
allow_legacy: bool = True,
) -> "TensorTranslationLayer":
path = Path(path)
try:
with safe_open(str(path), framework="pt", device="cpu") as handle:
metadata = handle.metadata()
except SafetensorError as error:
raise ValueError("artifact is not a Tensor Translation Layer") from error
manifest = json.loads(metadata["manifest"]) if "manifest" in metadata else metadata
file_format = manifest.get("format")
if file_format == SPLIT_TRANSLATE_FORMAT:
if not allow_legacy:
raise ValueError("legacy .translate artifact is not an explicit TTL")
warnings.warn(
".translate is deprecated. This semantic adapter is classified as TTL.",
DeprecationWarning,
stacklevel=2,
)
elif file_format != TTL_FORMAT:
raise ValueError("artifact is not a Tensor Translation Layer")
if file_format == TTL_FORMAT:
if manifest.get("adapter_class") != CompatibilityKind.TTL.value:
raise ValueError("TTL adapter class metadata does not match")
if manifest.get("support_level") != "semantic/internal":
raise ValueError("TTL support level metadata does not match")
raw_config = manifest["config"]
raw = json.loads(raw_config) if isinstance(raw_config, str) else dict(raw_config)
raw["attachment_layers"] = tuple(raw["attachment_layers"])
# The neural architecture remains the proven split translator. The
# container format, not the in-memory architecture config, is migrated.
raw["format"] = SPLIT_TRANSLATE_FORMAT
result = cls(SplitTranslateConfig(**raw)).to(device=device, dtype=dtype)
state = load_file(str(path), device=str(device))
if any(field in name.casefold() for name in state for field in FORBIDDEN_ADAPTER_FIELDS):
raise ValueError("TTL contains forbidden model or conversation state")
if tensor_checksum(state) != manifest.get("weights_sha256"):
raise ValueError("TTL weights checksum does not match")
result.load_state_dict({name: value.to(dtype=dtype) for name, value in state.items()})
return result
@dataclass(frozen=True)
class LexicalTranslationConfig:
model_id: str
model_architecture: str
model_sha256: str
runtime: str
runtime_version: str
tokenizer_bundle_sha256: str
canonical_protocol: str = CANONICAL_P_PROTOCOL
format: str = LTL_FORMAT
adapter_class: str = CompatibilityKind.LTL.value
support_level: str = "lexical/output"
control: str = "direct_adaptive_logit_bias"
logit_margin: float = 0.01
parameter_count: int = 0
def __post_init__(self) -> None:
if self.format != LTL_FORMAT or self.adapter_class != CompatibilityKind.LTL.value:
raise ValueError("invalid LTL format or adapter class")
if self.support_level != "lexical/output":
raise ValueError("invalid LTL support level")
if self.canonical_protocol != CANONICAL_P_PROTOCOL:
raise ValueError("unsupported canonical P protocol")
if self.parameter_count != 0:
raise ValueError("the direct adaptive logit-bias LTL has no learned parameters")
if self.logit_margin < 0:
raise ValueError("logit margin must be non-negative")
def _canonical_json(value: object) -> bytes:
return json.dumps(
value, ensure_ascii=False, sort_keys=True, separators=(",", ":"),
).encode("utf-8")
def tokenizer_bundle_checksum(path: str | Path) -> str:
"""Hash the exact tokenizer and metadata bundle used by the Gemma LTL."""
root = Path(path)
names = (
"chat_template.jinja",
"config.json",
"generation_config.json",
"processor_config.json",
"tokenizer_config.json",
"tokenizer.json",
)
digest = hashlib.sha256()
for name in names:
item = root / name
if not item.is_file():
raise FileNotFoundError(f"tokenizer bundle file is missing: {item}")
digest.update(name.encode("utf-8"))
digest.update(b"\0")
digest.update(item.read_bytes())
return digest.hexdigest()
class LexicalTranslationLayer:
"""Metadata-only lexical/output compatibility for llama.cpp runtimes."""
adapter_class = CompatibilityKind.LTL.value
support_level = "lexical/output"
extension = LTL_EXTENSION
def __init__(self, config: LexicalTranslationConfig) -> None:
self.config = config
def target(self, canonical_value: str, *, route_accepted: bool) -> str | None:
"""Return an output target only after the universal router accepts it."""
if not route_accepted:
return None
value = str(canonical_value)
return value if value else None
def token_targets(
self,
canonical_value: str,
tokenizer,
*,
route_accepted: bool,
) -> tuple[int, ...]:
target = self.target(canonical_value, route_accepted=route_accepted)
if target is None:
return ()
encoded = tokenizer(target, add_special_tokens=False).input_ids
return tuple(int(token_id) for token_id in encoded)
def adaptive_bias(self, logits: torch.Tensor, target_token_id: int) -> float:
"""Return the minimum non-negative bias that wins by the configured margin."""
flat = logits.detach().float().flatten()
if target_token_id < 0 or target_token_id >= flat.numel():
raise IndexError("target token is outside the model vocabulary")
masked = flat.clone()
masked[target_token_id] = -torch.inf
required = masked.max() - flat[target_token_id] + self.config.logit_margin
return max(0.0, float(required))
def validate_compatibility(
self,
*,
model_id: str,
model_architecture: str,
model_sha256: str,
runtime: str,
runtime_version: str | None = None,
tokenizer_bundle_sha256: str | None = None,
canonical_protocol: str = CANONICAL_P_PROTOCOL,
) -> None:
mismatches = []
if model_id != self.config.model_id:
mismatches.append("model identifier")
if model_architecture != self.config.model_architecture:
mismatches.append("model architecture")
if model_sha256 != self.config.model_sha256:
mismatches.append("model checksum")
if runtime != self.config.runtime:
mismatches.append("runtime")
if runtime_version is not None and runtime_version != self.config.runtime_version:
mismatches.append("runtime version")
if (
tokenizer_bundle_sha256 is not None
and tokenizer_bundle_sha256 != self.config.tokenizer_bundle_sha256
):
mismatches.append("tokenizer bundle checksum")
if canonical_protocol != self.config.canonical_protocol:
mismatches.append("canonical protocol")
if mismatches:
raise ValueError("incompatible LTL: " + ", ".join(mismatches))
def save(self, path: str | Path) -> None:
path = Path(path)
if path.suffix != LTL_EXTENSION:
raise ValueError(f"LTL artifacts must use the {LTL_EXTENSION} extension")
payload = asdict(self.config)
payload_bytes = _canonical_json(payload)
envelope = {
"format": LTL_FORMAT,
"payload": payload,
"payload_sha256": hashlib.sha256(payload_bytes).hexdigest(),
}
path.write_bytes(_canonical_json(envelope) + b"\n")
@classmethod
def load(cls, path: str | Path) -> "LexicalTranslationLayer":
envelope = json.loads(Path(path).read_text(encoding="utf-8"))
if envelope.get("format") != LTL_FORMAT:
raise ValueError("artifact is not a Lexical Translation Layer")
payload = envelope.get("payload")
if not isinstance(payload, dict):
raise ValueError("LTL payload is missing")
actual = hashlib.sha256(_canonical_json(payload)).hexdigest()
if actual != envelope.get("payload_sha256"):
raise ValueError("LTL checksum does not match")
return cls(LexicalTranslationConfig(**payload))
def classify_compatibility_artifact(path: str | Path) -> CompatibilityKind:
"""Classify modern adapters and the one supported legacy semantic format."""
path = Path(path)
if path.suffix == LTL_EXTENSION:
return CompatibilityKind.LTL
try:
with safe_open(str(path), framework="pt", device="cpu") as handle:
metadata = handle.metadata()
manifest = json.loads(metadata["manifest"]) if "manifest" in metadata else metadata
file_format = manifest.get("format")
except Exception as error:
raise ValueError(f"unrecognized compatibility artifact: {path}") from error
if file_format in (TTL_FORMAT, SPLIT_TRANSLATE_FORMAT):
return CompatibilityKind.TTL
raise ValueError(
"legacy artifact is research-only and has no active TTL or LTL classification"
)
def resolve_compatibility(path: str | Path | None) -> CompatibilityResolution:
if path is None:
return CompatibilityResolution(CompatibilityKind.NATIVE, "native", None)
artifact = Path(path)
kind = classify_compatibility_artifact(artifact)
level = "semantic/internal" if kind is CompatibilityKind.TTL else "lexical/output"
return CompatibilityResolution(kind, level, artifact)