reproducibility/sequence_repair/optimization.py

"""Explicit adapter or full-encoder optimization and checkpoint ownership."""

# ruff: noqa: PLC0415

import hashlib
import json
import math
from pathlib import Path

from .model import load_model, save_adapter
from .task_head import task_head_scope, tensor_digest


def configure_trainable_tail(model, last_layers):
    """Restrict a complete ModernBERT checkpoint to its tail and entire task head."""
    if (
        type(model).__name__ != "ModernBertForSequenceClassification"
        or hasattr(model, "peft_config")
        or hasattr(model, "get_base_model")
    ):
        raise ValueError("Tail training requires a complete ModernBERT classifier")
    layers = model.model.layers
    if type(last_layers) is not int or not 1 <= last_layers <= len(layers):
        raise ValueError("Trainable tail must contain one through all encoder layers")
    start = len(layers) - last_layers
    modules = [*layers[start:], model.head, model.classifier]
    selected = {
        id(parameter) for module in modules for parameter in module.parameters()
    }
    named = dict(model.named_parameters())
    if not selected or not selected <= {id(parameter) for parameter in named.values()}:
        raise ValueError("Trainable modules do not belong to the classifier")
    if any(
        id(parameter) in selected for parameter in model.model.final_norm.parameters()
    ):
        raise ValueError("Encoder final normalization must remain frozen")
    model.requires_grad_(False)
    for module in modules:
        module.requires_grad_(True)
    tensors = {
        name: {"shape": list(parameter.shape), "numel": parameter.numel()}
        for name, parameter in named.items()
        if parameter.requires_grad
    }
    return {
        "architecture": type(model).__name__,
        "layer_indices": list(range(start, len(layers))),
        "trainable_tensors": tensors,
        "trainable_parameters": sum(item["numel"] for item in tensors.values()),
        "encoder_final_norm_trainable": False,
    }


def frozen_parameter_receipt(model):
    """Hash frozen parameters once so an optional restricted update is auditable."""
    return {
        name: {
            "shape": list(parameter.shape),
            "dtype": str(parameter.dtype),
            "sha256": tensor_digest(parameter),
        }
        for name, parameter in model.named_parameters()
        if not parameter.requires_grad
    }


def assert_frozen_parameters_preserved(expected, model):
    if frozen_parameter_receipt(model) != expected:
        raise ValueError("Frozen parameter scope or tensor values changed")


def validate_method(method, adapter, fresh_head):
    if method not in {"lora", "full"}:
        raise ValueError("Training method must be lora or full")
    if method == "lora" and (adapter is None or fresh_head):
        raise ValueError("LoRA continuation requires its initialized adapter")
    if method == "full" and adapter is not None:
        raise ValueError(
            "Full training starts from a complete checkpoint, not an adapter"
        )


def load_trainable_model(base, contract, adapter, method, fresh_head=False):
    validate_method(method, adapter, fresh_head)
    model, tokenizer, labels, ids = load_model(
        base, contract, adapter, trainable=method == "lora"
    )
    if method == "full":
        if type(model).__name__ != "ModernBertForSequenceClassification":
            raise ValueError("Full sequence training requires a ModernBERT classifier")
        model.requires_grad_(True)
        if fresh_head:
            from torch import nn

            head_parameter = next(model.head.parameters())
            model.head = type(model.head)(model.config).to(
                device=head_parameter.device, dtype=head_parameter.dtype
            )
            classifier = model.classifier
            model.classifier = nn.Linear(
                classifier.in_features,
                classifier.out_features,
                bias=classifier.bias is not None,
                device=classifier.weight.device,
                dtype=classifier.weight.dtype,
            )
            # Fresh modules let HF apply its native initialization without
            # mutating the initialization guards on loaded checkpoint tensors.
            model.head.apply(model._init_weights)
            # ModernBERT initializes the final classifier on its owning task
            # module, not when visiting the bare Linear child.
            model._init_weights(model)
    return model, tokenizer, labels, ids


def optimizer_groups(model, learning_rate, head_learning_rate=None):
    rates = [learning_rate]
    if head_learning_rate is not None:
        rates.append(head_learning_rate)
    if any(not math.isfinite(rate) or rate <= 0 for rate in rates):
        raise ValueError("Learning rates must be finite and positive")
    base = model.get_base_model() if hasattr(model, "get_base_model") else model
    head_ids = {
        id(parameter)
        for module in (base.head, base.classifier)
        for parameter in module.parameters()
        if parameter.requires_grad
    }
    encoder, head = [], []
    for parameter in model.parameters():
        if parameter.requires_grad:
            (head if id(parameter) in head_ids else encoder).append(parameter)
    groups = []
    for name, parameters, rate in (
        ("encoder", encoder, learning_rate),
        ("task_head", head, head_learning_rate or learning_rate),
    ):
        if parameters:
            groups.append(
                {"name": name, "params": parameters, "lr": rate, "initial_lr": rate}
            )
    if not groups:
        raise ValueError("No trainable parameters")
    return groups


def initial_artifact_receipt(base, adapter):
    """Hash the actual local initialization, including every weight shard."""
    root = Path(base)
    if not root.is_dir():
        raise ValueError("Training requires a downloaded local base checkpoint")
    names = {"config.json", "tokenizer.json", "tokenizer_config.json"}
    for path in root.iterdir():
        if (
            path.suffix == ".safetensors"
            or path.name
            in {
                "model.safetensors.index.json",
                "pytorch_model.bin",
                "pytorch_model.bin.index.json",
            }
            or path.name.startswith("pytorch_model-")
        ):
            names.add(path.name)
    files = {}
    for name in sorted(names):
        path = root / name
        if path.is_file():
            with path.open("rb") as stream:
                files[name] = hashlib.file_digest(stream, "sha256").hexdigest()
    result = {"base_files": files}
    if adapter is not None:
        for filename in ("adapter_model.safetensors", "adapter_config.json"):
            path = Path(adapter) / filename
            if filename == "adapter_config.json" and not path.is_file():
                result[filename] = None
                continue
            with path.open("rb") as stream:
                result[filename] = hashlib.file_digest(stream, "sha256").hexdigest()
    return result


def verify_full_checkpoint(model, checkpoint):
    """Require every tensor of a complete sequence checkpoint to be preserved."""
    if type(model).__name__ != "ModernBertForSequenceClassification":
        raise ValueError("Full checkpoints must own the actual encoder weights")
    from safetensors.torch import load_file

    saved = load_file(str(Path(checkpoint) / "model.safetensors"), device="cpu")
    state = model.state_dict()
    if saved.keys() != state.keys():
        raise ValueError("Full checkpoint omitted or introduced model tensors")
    import torch

    for name, value in state.items():
        if saved[name].dtype != value.dtype or not torch.equal(
            saved[name], value.detach().cpu()
        ):
            raise ValueError(f"Full checkpoint changed tensor: {name}")


def save_training_checkpoint(model, tokenizer, destination, method, base_id, revision):
    if method == "lora":
        save_adapter(model, tokenizer, destination, base_id, revision)
        return
    if method != "full" or hasattr(model, "peft_config"):
        raise ValueError("Full checkpoints must own the actual encoder weights")
    destination.mkdir(parents=True, exist_ok=True)
    expected = task_head_scope(model)
    model.save_pretrained(destination, safe_serialization=True)
    verify_full_checkpoint(model, destination)
    tokenizer.save_pretrained(destination)
    (destination / "task-head.json").write_text(json.dumps(expected, indent=2) + "\n")
    (destination / "training-origin.json").write_text(
        json.dumps(
            {"method": method, "base_model": base_id, "base_revision": revision},
            indent=2,
        )
        + "\n"
    )