neuron-tai/packages/node/meshnet_node/safetensors_selection.py

"""Layer-aware SafeTensors snapshot file selection."""

from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any

INDEX_FILENAME = "model.safetensors.index.json"

_LAYER_RE = re.compile(
    r"(?:^|\.)"
    r"(?:model\.layers|layers|h|blocks|decoder\.layers|encoder\.layers)"
    r"\.(\d+)(?:\.|$)"
)

METADATA_FILENAMES = {
    INDEX_FILENAME,
    "config.json",
    "generation_config.json",
    "preprocessor_config.json",
    "special_tokens_map.json",
    "tokenizer.json",
    "tokenizer.model",
    "tokenizer_config.json",
    "vocab.json",
    "merges.txt",
    "added_tokens.json",
}

_METADATA_PREFIXES = ("config.", "tokenizer.", "tokenizer_", "vocab.")

_HEAD_MARKERS = (
    "embed",
    "embedding",
    "embed_tokens",
    "wte",
    "wpe",
)

_TAIL_EXACT = {
    "lm_head.weight",
    "lm_head.bias",
    "model.norm.weight",
    "model.norm.bias",
    "transformer.ln_f.weight",
    "transformer.ln_f.bias",
    "decoder.final_layer_norm.weight",
    "decoder.final_layer_norm.bias",
}

_TAIL_MARKERS = (
    ".lm_head.",
    ".norm.",
    ".ln_f.",
    ".final_layer_norm.",
)


def select_safetensors_files_for_layers(
    model_dir: str | Path,
    start_layer: int,
    end_layer: int,
    *,
    total_layers: int | None = None,
) -> list[str]:
    """Return relative snapshot files needed for an inclusive layer range.

    The returned list always includes root-level config/tokenizer metadata and
    the SafeTensors index. Weight shard files are included only when at least one
    tensor in the index belongs to the assigned layer range, or when the tensor
    is needed by the head/tail shard.
    """
    if start_layer < 0:
        raise ValueError("start_layer must be non-negative")
    if end_layer < start_layer:
        raise ValueError("end_layer must be greater than or equal to start_layer")

    root = Path(model_dir)
    index_path = root / INDEX_FILENAME
    try:
        index = json.loads(index_path.read_text(encoding="utf-8"))
    except FileNotFoundError as exc:
        raise FileNotFoundError(f"missing SafeTensors index: {index_path}") from exc

    weight_map = index.get("weight_map")
    if not isinstance(weight_map, dict):
        raise ValueError(f"{INDEX_FILENAME} must contain a weight_map object")

    inferred_total_layers = total_layers if total_layers is not None else _read_total_layers(root)
    selected = _metadata_files(root)
    selected |= select_files_for_layers_from_index(
        weight_map, start_layer, end_layer, total_layers=inferred_total_layers
    )
    return sorted(selected)


def select_files_for_layers_from_index(
    weight_map: dict[str, str],
    start_layer: int,
    end_layer: int,
    *,
    total_layers: int | None = None,
) -> set[str]:
    """Pure variant of the weight-file selection: takes an already-parsed
    ``weight_map`` (no local snapshot directory needed), so callers that only
    have the index fetched over the network — not a full local snapshot — can
    still compute which shard files they need. Combine the result with
    ``METADATA_FILENAMES`` for a complete download pattern set.
    """
    selected: set[str] = set()
    for tensor_name, rel_file in weight_map.items():
        if not isinstance(tensor_name, str) or not isinstance(rel_file, str):
            continue
        if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, total_layers):
            selected.add(_normalise_relative_file(rel_file))
    return selected


def select_tensor_names_for_layers_from_index(
    weight_map: dict[str, str],
    start_layer: int,
    end_layer: int,
    *,
    total_layers: int | None = None,
) -> set[str]:
    """Pure variant that returns checkpoint tensor names instead of file paths."""
    selected: set[str] = set()
    for tensor_name, rel_file in weight_map.items():
        if not isinstance(tensor_name, str) or not isinstance(rel_file, str):
            continue
        if _tensor_belongs_to_range(tensor_name, start_layer, end_layer, total_layers):
            selected.add(tensor_name)
    return selected


def _tensor_belongs_to_range(
    tensor_name: str,
    start_layer: int,
    end_layer: int,
    total_layers: int | None,
) -> bool:
    layer = _layer_index(tensor_name)
    if layer is not None:
        return start_layer <= layer <= end_layer

    if start_layer == 0 and _is_head_tensor(tensor_name):
        return True

    if total_layers is not None and end_layer >= total_layers - 1 and _is_tail_tensor(tensor_name):
        return True

    return False


def _layer_index(tensor_name: str) -> int | None:
    match = _LAYER_RE.search(tensor_name)
    if match is None:
        return None
    return int(match.group(1))


def _is_head_tensor(tensor_name: str) -> bool:
    lowered = tensor_name.lower()
    return any(marker in lowered for marker in _HEAD_MARKERS)


def _is_tail_tensor(tensor_name: str) -> bool:
    lowered = tensor_name.lower()
    return lowered in _TAIL_EXACT or any(marker in lowered for marker in _TAIL_MARKERS)


def _metadata_files(root: Path) -> set[str]:
    files = {INDEX_FILENAME}
    for path in root.iterdir():
        if not path.is_file():
            continue
        name = path.name
        if name in METADATA_FILENAMES or name.startswith(_METADATA_PREFIXES):
            files.add(name)
    return files


def _read_total_layers(root: Path) -> int | None:
    config_path = root / "config.json"
    if not config_path.exists():
        return None
    config = json.loads(config_path.read_text(encoding="utf-8"))
    return layers_from_config_dict(config)


def layers_from_config_dict(config: dict[str, Any]) -> int | None:
    for key in ("num_hidden_layers", "num_layers", "n_layer", "n_layers"):
        value = config.get(key)
        if isinstance(value, int) and value > 0:
            return value

    text_config = config.get("text_config")
    if isinstance(text_config, dict):
        return layers_from_config_dict(text_config)

    return None


def _normalise_relative_file(rel_file: str) -> str:
    path = Path(rel_file)
    if path.is_absolute() or ".." in path.parts:
        raise ValueError(f"unsafe relative file in {INDEX_FILENAME}: {rel_file}")
    return path.as_posix()