src/kd_contracts.py

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any

PROVENANCE_SCHEMA_VERSION = 4

_SHARD_SCHEMA = {
    "support": "teacher_topk_plus_other_bucket",
    "layout": "chunked_sample_lists",
    "logprobs_dtype": "float16",
    "ids_dtype": "int32",
    "other_logprob_dtype": "float16",
}

_SPECIAL_TOKEN_ID_FIELDS = (
    "bos_token_id",
    "eos_token_id",
    "pad_token_id",
    "unk_token_id",
    "cls_token_id",
    "sep_token_id",
    "mask_token_id",
    "additional_special_tokens_ids",
)

def normalize_config_revision(value: str | None) -> str | None:
    if value is None:
        return None
    stripped = value.strip()
    return stripped or None

def canonical_revision(value: str | None) -> str:
    return normalize_config_revision(value) or "unversioned"

def sha256_file(path: str | Path, chunk_size: int = 1 << 20) -> str:
    digest = hashlib.sha256()
    with open(path, "rb") as handle:
        while True:
            chunk = handle.read(chunk_size)
            if not chunk:
                break
            digest.update(chunk)
    return digest.hexdigest()

def _special_token_ids(tokenizer) -> dict[str, Any]:
    snapshot: dict[str, Any] = {}
    for field in _SPECIAL_TOKEN_ID_FIELDS:
        value = getattr(tokenizer, field, None)
        if isinstance(value, tuple):
            value = list(value)
        snapshot[field] = value
    return snapshot

def build_tokenizer_contract(tokenizer) -> dict[str, Any]:
    canonical = {
        "tokenizer_class": tokenizer.__class__.__name__,
        "full_vocab_size": len(tokenizer),
        "special_token_ids": _special_token_ids(tokenizer),
        "vocab": dict(sorted(tokenizer.get_vocab().items())),
    }
    encoded = json.dumps(
        canonical,
        sort_keys=True,
        separators=(",", ":"),
        ensure_ascii=True,
    ).encode("utf-8")
    return {
        "tokenizer_class": canonical["tokenizer_class"],
        "full_vocab_size": canonical["full_vocab_size"],
        "special_token_ids": canonical["special_token_ids"],
        "fingerprint": hashlib.sha256(encoded).hexdigest(),
    }

def build_shard_schema() -> dict[str, str]:
    return dict(_SHARD_SCHEMA)

def collect_model_vocab_sizes(model) -> dict[str, int]:
    sizes: dict[str, int] = {}

    config_size = getattr(getattr(model, "config", None), "vocab_size", None)
    if isinstance(config_size, int):
        sizes["config"] = config_size

    input_embeddings = model.get_input_embeddings()
    if input_embeddings is not None and getattr(input_embeddings, "weight", None) is not None:
        sizes["input_embeddings"] = int(input_embeddings.weight.shape[0])

    output_embeddings = model.get_output_embeddings()
    if output_embeddings is not None and getattr(output_embeddings, "weight", None) is not None:
        sizes["output_embeddings"] = int(output_embeddings.weight.shape[0])

    return sizes