"""Strict, torch-free HF checkpoint import for the native runtime. Supported: standard dense causal-LM checkpoints (llama/qwen3) with full attention and default RoPE, the SurjoLabs Surjo hybrid family (SurjoForCausalLM: prelude/recurrent/coda XSA + GDN-2 linear attention), plus the FWKV Myosotis recurrent family (FWKVLanguageModel: per-layer WKV exponential-recurrence + GELU FFN, factorized tied head). The only ignored tensors are default-derived rotary ``inv_freq`` buffers (llama/qwen3); all trainable tensors must match the native operators exactly. Surjo/FWKV checkpoints carry no rotary buffers. No modeling code is executed. """ from __future__ import annotations import json import math import os import struct from dataclasses import dataclass from pathlib import Path from typing import Any import numpy as np from huggingface_hub import snapshot_download from safetensors import safe_open class ModelImportError(ValueError): """A checkpoint cannot be faithfully imported into the native runtime.""" class UnsupportedModelError(ModelImportError): """The checkpoint requires an unsupported architecture or operator.""" class InvalidCheckpointError(ModelImportError): """Checkpoint metadata or tensor data is incomplete or inconsistent.""" @dataclass class LoadedModel: """Native inputs plus a tokenizer; source is the resolved local snapshot. For Hub imports, revision is the resolved commit, not a mutable branch name. Local directories retain the caller's revision as informational metadata. Tied embedding/head entries refer to the very same NumPy array. Config includes explicit non-null generation_config.json EOS/BOS/PAD overrides, but not its sampling defaults. """ config: dict weights: dict tokenizer: Any source: str revision: str | None _ARCHITECTURES = {"llama": "LlamaForCausalLM", "qwen3": "Qwen3ForCausalLM"} _SUPPORT = ( "Implemented: standard LlamaForCausalLM (llama) and Qwen3ForCausalLM (qwen3). " "Beta targets not implemented: negative, cma, ember2, spark2a, blaze." ) def validate_config(config: dict) -> dict: """Return normalized HF keys, rejecting unsupported native semantics. Dimensions/context must be explicit; HF defaults are used for optional operator settings. Explicit Qwen3 head_dim need not equal hidden_size/heads. No Transformers AutoConfig or remote modeling code is executed. """ if not isinstance(config, dict): raise UnsupportedModelError("config must be a JSON object. " + _SUPPORT) config = dict(config) model_type = config.get("model_type") if not isinstance(model_type, str) or model_type not in _ARCHITECTURES: raise UnsupportedModelError(f"Unsupported model_type {model_type!r}. {_SUPPORT}") architecture = _ARCHITECTURES[model_type] if config.get("architectures", [architecture]) != [architecture]: raise UnsupportedModelError( f"Expected architectures=[{architecture!r}], got {config.get('architectures')!r}. {_SUPPORT}" ) auto_map = config.get("auto_map") or {} if not isinstance(auto_map, dict) or set(auto_map) - {"AutoTokenizer"}: raise UnsupportedModelError( "Custom auto_map config/modeling overrides are unsupported, even with " "trust_remote_code=True; only AutoTokenizer is allowed. " + _SUPPORT ) config["architectures"] = [architecture] # Keep these safety bounds aligned with native Config::validate. for key, maximum in (("vocab_size", 2_000_000), ("hidden_size", 65536), ("intermediate_size", 262144), ("num_hidden_layers", 256), ("num_attention_heads", 512), ("max_position_embeddings", 1048576)): value = config.get(key) if type(value) is not int or not 0 < value <= maximum: raise UnsupportedModelError(f"{key} must be an integer in [1, {maximum}], got {value!r}") for key in ("eos_token_id", "bos_token_id", "pad_token_id"): value = config.get(key) if value is None: continue tokens = value if key == "eos_token_id" and isinstance(value, list) else [value] if not tokens or any(type(token) is not int or not 0 <= token < config["vocab_size"] for token in tokens): expected = "an integer or a nonempty list of integers" if key == "eos_token_id" else "an integer" raise InvalidCheckpointError(f"{key} must be {expected} in [0, vocab_size={config['vocab_size']}); got {value!r}") if isinstance(value, list): config[key] = list(dict.fromkeys(tokens)) heads = config["num_attention_heads"] hidden = config["hidden_size"] kv_heads = config.get("num_key_value_heads", 32 if model_type == "qwen3" else heads) if kv_heads is None: kv_heads = heads if type(kv_heads) is not int or kv_heads <= 0 or heads % kv_heads: raise UnsupportedModelError("num_key_value_heads must be positive and divide num_attention_heads") config["num_key_value_heads"] = kv_heads head_dim = config.get("head_dim", 128 if model_type == "qwen3" else None) if head_dim is None: if hidden % heads: raise UnsupportedModelError("hidden_size must divide evenly by num_attention_heads when head_dim is absent") head_dim = hidden // heads if type(head_dim) is not int or not 0 < head_dim <= 1024 or head_dim % 2 or heads * head_dim > 65536: raise UnsupportedModelError("head_dim must be an even integer in [2, 1024] with heads * head_dim <= 65536") config["head_dim"] = head_dim for key in ("attention_bias", "mlp_bias", "qkv_bias", "use_bias", "bias", "is_encoder_decoder", "add_cross_attention", "use_sliding_window"): if config.get(key, False) is not False: raise UnsupportedModelError(f"{key} must be false; this operator is not implemented") if config.get("hidden_act", "silu") != "silu": raise UnsupportedModelError("hidden_act must be 'silu' (SwiGLU)") config["hidden_act"] = "silu" config["attention_bias"] = False config["mlp_bias"] = False # Old Qwen3 configs retain an unused window size when use_sliding_window=False. window = config.get("sliding_window") if window is not None and not (model_type == "qwen3" and config.get("use_sliding_window") is False): raise UnsupportedModelError("sliding_window attention is not implemented") layer_types = config.get("layer_types") if layer_types is not None and layer_types != ["full_attention"] * config["num_hidden_layers"]: raise UnsupportedModelError("layer_types must contain exactly one 'full_attention' per layer") config["use_sliding_window"] = False config["sliding_window"] = None for key in ("quantization_config", "compression_config"): if config.get(key) not in (None, {}): raise UnsupportedModelError(f"Nonempty {key} is not implemented; provide unquantized safetensors") for key in ("num_local_experts", "num_experts", "num_experts_per_tok"): if key in config and config[key] not in (None, 0): raise UnsupportedModelError(f"{key}: mixture-of-experts operators are not implemented") for key, default in (("partial_rotary_factor", 1.0), ("rotary_pct", 1.0), ("pretraining_tp", 1)): if type(config.get(key, default)) not in (int, float) or config.get(key, default) != default: raise UnsupportedModelError(f"{key} must be {default}; this variant is not implemented") config["pretraining_tp"] = 1 if config.get("rope_type", "default") != "default": raise UnsupportedModelError("rope_type must be 'default'") if config.get("attention_multiplier") is not None: raise UnsupportedModelError("Custom attention_multiplier is not implemented") if config.get("rope_interleaved", False) is not False: raise UnsupportedModelError("rope_interleaved is not implemented") theta_values = [] if "rope_theta" in config: theta_values.append(config["rope_theta"]) for key in ("rope_scaling", "rope_parameters"): rope = config.get(key) if rope is None: continue if not isinstance(rope, dict): raise UnsupportedModelError(f"{key} must be null or a default RoPE object") if any(rope.get(k, "default") != "default" for k in ("type", "rope_type")): raise UnsupportedModelError(f"{key}: only default RoPE is implemented, got {rope!r}") if set(rope) - {"type", "rope_type", "rope_theta"}: raise UnsupportedModelError(f"{key}: scaled, partial, or per-layer RoPE parameters are not implemented") if "rope_theta" in rope: theta_values.append(rope["rope_theta"]) if not theta_values: theta_values = [10000.0] for theta in theta_values: if type(theta) not in (int, float) or not float(np.finfo(np.float32).smallest_subnormal) <= theta <= float(np.finfo(np.float32).max): raise UnsupportedModelError("rope_theta must be finite, positive, and representable in FP32") if any(theta != theta_values[0] for theta in theta_values): raise UnsupportedModelError("Conflicting rope_theta values in legacy and new RoPE configuration") config["rope_theta"] = float(theta_values[0]) config["rope_scaling"] = None # The native interface consumes theta at the top level only. config["rope_parameters"] = {"rope_type": "default"} eps = config.get("rms_norm_eps", 1e-6) if type(eps) not in (int, float) or not float(np.finfo(np.float32).smallest_subnormal) <= eps <= float(np.finfo(np.float32).max): raise UnsupportedModelError("rms_norm_eps must be finite, positive, and representable in FP32") config["rms_norm_eps"] = float(eps) tied = config.get("tie_word_embeddings", False) if type(tied) is not bool: raise UnsupportedModelError("tie_word_embeddings must be a boolean") config["tie_word_embeddings"] = tied return config _SURJO_ARCHITECTURE = "SurjoForCausalLM" _SURJO_SUPPORT = ( "Surjo support in this snapshot: config validation, execution planning, " "and shapes/weights import. Native GDN-2/XSA execution is the next phase." ) def validate_surjo_config(config: dict) -> dict: """Validate a SurjoLabs Surjo hybrid config; return a normalized copy. Topology is prelude XSA layers, recurrent GDN-2 groups rerun for ``recurrent_passes`` passes, then coda XSA layers (Surjo-50m: 1/8/1 layers, 2 groups, 2 passes, 3 GDN per XSA). No weights are read and no modeling code is executed here; native execution is a later phase. """ if not isinstance(config, dict): raise UnsupportedModelError("config must be a JSON object. " + _SURJO_SUPPORT) config = dict(config) if config.get("model_type") != "surjo": raise UnsupportedModelError( f"Unsupported model_type {config.get('model_type')!r}. " + _SURJO_SUPPORT) if config.get("architectures", [_SURJO_ARCHITECTURE]) != [_SURJO_ARCHITECTURE]: raise UnsupportedModelError( f"Expected architectures=[{_SURJO_ARCHITECTURE!r}], " f"got {config.get('architectures')!r}. " + _SURJO_SUPPORT) auto_map = config.get("auto_map") or {} if not isinstance(auto_map, dict): raise UnsupportedModelError("auto_map must be an object. " + _SURJO_SUPPORT) for name, reference in auto_map.items(): refs = [reference] if isinstance(reference, str) else reference if not isinstance(refs, (list, tuple)) or not refs or any( not isinstance(r, str) or "--" in r or not (r.startswith("configuration_surjo.") or r.startswith("modeling_surjo.")) for r in refs): raise UnsupportedModelError( "Surjo auto_map must reference code within the pinned snapshot " "(configuration_surjo/modeling_surjo modules only); it is never " "executed by the native runtime. " + _SURJO_SUPPORT) for key, maximum in (("vocab_size", 2_000_000), ("hidden_size", 65536), ("intermediate_size", 262144), ("num_hidden_layers", 256), ("num_attention_heads", 512), ("max_position_embeddings", 1048576)): value = config.get(key) if type(value) is not int or not 0 < value <= maximum: raise InvalidCheckpointError( f"{key} must be an integer in [1, {maximum}], got {value!r}") topo = {} for key, maximum in (("prelude_layers", 256), ("recurrent_layers", 256), ("coda_layers", 256), ("num_groups", 64), ("recurrent_passes", 64), ("gdn_per_xsa", 64)): value = config.get(key) if type(value) is not int or not 0 <= value <= maximum: raise InvalidCheckpointError( f"{key} must be an integer in [0, {maximum}], got {value!r}") topo[key] = value if topo["num_groups"] == 0 or topo["recurrent_passes"] == 0: raise InvalidCheckpointError("num_groups and recurrent_passes must be positive") group_block = topo["num_groups"] * (topo["gdn_per_xsa"] + 1) if (topo["prelude_layers"] + group_block + topo["coda_layers"] != config["num_hidden_layers"]): raise InvalidCheckpointError( "prelude_layers + num_groups*(gdn_per_xsa+1) + coda_layers must equal " "num_hidden_layers") if topo["recurrent_layers"] != group_block: raise InvalidCheckpointError( "recurrent_layers must equal num_groups*(gdn_per_xsa+1)") heads = config["num_attention_heads"] hidden = config["hidden_size"] kv_heads = config.get("num_key_value_heads", heads) if kv_heads is None: kv_heads = heads if type(kv_heads) is not int or kv_heads <= 0 or heads % kv_heads: raise InvalidCheckpointError( "num_key_value_heads must be positive and divide num_attention_heads") config["num_key_value_heads"] = kv_heads head_dim = config.get("head_dim") if type(head_dim) is not int or not 0 < head_dim <= 1024 or head_dim % 2: raise InvalidCheckpointError("head_dim must be an even integer in [2, 1024]") config["head_dim"] = head_dim for key in ("gdn_num_v_heads", "gdn_key_dim", "gdn_value_dim"): value = config.get(key) if type(value) is not int or not 0 < value <= 4096: raise InvalidCheckpointError(f"{key} must be an integer in [1, 4096], got {value!r}") kernel = config.get("gdn_conv_kernel_size") if type(kernel) is not int or kernel < 1 or kernel > 64: raise InvalidCheckpointError( f"gdn_conv_kernel_size must be an integer in [1, 64], got {kernel!r}") for key in ("gdn_allow_neg_eigval", "xsa_projection"): if type(config.get(key)) is not bool: raise InvalidCheckpointError(f"{key} must be a boolean") if config.get("attention_bias", False) is not False: raise UnsupportedModelError("attention_bias must be false; this operator is not implemented") dropout = config.get("attention_dropout", 0.0) if type(dropout) not in (int, float) or dropout != 0: raise UnsupportedModelError("attention_dropout must be 0.0") if config.get("hidden_act", "silu") != "silu": raise UnsupportedModelError("hidden_act must be 'silu' (SwiGLU)") config["hidden_act"] = "silu" eps = config.get("rms_norm_eps", 1e-6) if (type(eps) not in (int, float) or not float(np.finfo(np.float32).smallest_subnormal) <= eps <= float(np.finfo(np.float32).max)): raise InvalidCheckpointError("rms_norm_eps must be finite, positive, and representable in FP32") config["rms_norm_eps"] = float(eps) theta = config.get("rope_theta", 10000.0) if (type(theta) not in (int, float) or not float(np.finfo(np.float32).smallest_subnormal) <= theta <= float(np.finfo(np.float32).max)): raise InvalidCheckpointError("rope_theta must be finite, positive, and representable in FP32") config["rope_theta"] = float(theta) rope = config.get("rope_parameters", {"rope_type": "default"}) if not isinstance(rope, dict) or rope.get("rope_type", "default") != "default": raise UnsupportedModelError("rope_parameters: only default RoPE is supported") if rope.get("partial_rotary_factor", 1.0) != 1.0: raise UnsupportedModelError("partial_rotary_factor must be 1.0") if config.get("partial_rotary_factor", 1.0) != 1.0: raise UnsupportedModelError("partial_rotary_factor must be 1.0") tied = config.get("tie_word_embeddings", False) if type(tied) is not bool: raise InvalidCheckpointError("tie_word_embeddings must be a boolean") config["tie_word_embeddings"] = tied for key in ("eos_token_id", "bos_token_id", "pad_token_id"): value = config.get(key) if value is None: continue tokens = value if key == "eos_token_id" and isinstance(value, list) else [value] if not tokens or any(type(t) is not int or not 0 <= t < config["vocab_size"] for t in tokens): raise InvalidCheckpointError(f"{key} is outside [0, vocab_size)") for key in ("quantization_config", "compression_config"): if config.get(key) not in (None, {}): raise UnsupportedModelError(f"Nonempty {key} is not implemented; provide unquantized safetensors") expected_types = ["full_attention"] * topo["prelude_layers"] for _ in range(topo["num_groups"]): expected_types += ["linear_attention"] * topo["gdn_per_xsa"] + ["full_attention"] expected_types += ["full_attention"] * topo["coda_layers"] if config.get("layer_types") is not None and config["layer_types"] != expected_types: raise InvalidCheckpointError("layer_types does not match the prelude/group/coda topology") config["layer_types"] = expected_types config["architectures"] = [_SURJO_ARCHITECTURE] return config def surjo_num_xsa_slots(config: dict) -> int: """KV slot count for a validated Surjo config (prelude + groups*passes + coda).""" return (config["prelude_layers"] + config["num_groups"] * config["recurrent_passes"] + config["coda_layers"]) def surjo_exec_plan(config: dict) -> list: """Execution plan for a validated Surjo config: ``(layer_idx, pass, slot)``. GDN layers carry slot -1 (recurrent state is keyed by (layer, pass)); XSA layers carry their KV slot. Passes share physical layers with distinct recurrent state. """ if not isinstance(config, dict) or config.get("model_type") != "surjo": raise InvalidCheckpointError("surjo_exec_plan requires a Surjo config") prelude = config["prelude_layers"] groups = config["num_groups"] per = config["gdn_per_xsa"] passes = config["recurrent_passes"] coda = config["coda_layers"] recurrent_base = prelude + groups * (per + 1) if prelude + groups * (per + 1) + coda != config["num_hidden_layers"]: raise InvalidCheckpointError("topology fields disagree with num_hidden_layers") plan = [] slot = 0 for li in range(prelude): plan.append((li, 0, slot)) slot += 1 for p in range(passes): for g in range(groups): base = prelude + g * (per + 1) for j in range(per): plan.append((base + j, p, -1)) plan.append((base + per, p, slot)) slot += 1 for li in range(recurrent_base, recurrent_base + coda): plan.append((li, 0, slot)) slot += 1 return plan def _is_surjo_xsa_layer(config: dict, layer: int) -> bool: """True for XSA (full_attention) layers in a validated Surjo config.""" types = config.get("layer_types") if types is not None: return types[layer] == "full_attention" prelude = config["prelude_layers"] per = config["gdn_per_xsa"] groups = config["num_groups"] if layer < prelude: return True if layer >= prelude + groups * (per + 1): return True return (layer - prelude) % (per + 1) == per def _expected_shapes_surjo(config: dict) -> dict[str, tuple[int, ...]]: """Exact trainable tensor shapes for a validated Surjo hybrid config. Mirrors ``modeling_surjo.py`` naming verbatim (verified against the SurjoLabs/Surjo-50m ``model.safetensors`` header: 178 F32 tensors, no ``lm_head`` entry when tied, no rotary buffers): - XSA layers ``self_attn.{q,k,v,o}_proj`` + ``q_norm/k_norm``. - GDN layers ``linear_attn.{q,k,v,b,w}_proj``, ``q/k/v_conv.conv`` 3D ``(dim, 1, kernel)``, ``f_proj.0/1``, ``g_proj.0/1(.bias)``, ``A_log``, ``dt_bias``, ``o_norm/o_proj``. - Per-layer ``input/post_attention_layernorm`` + SwiGLU MLP. """ hidden = config["hidden_size"] intermediate = config["intermediate_size"] vocab = config["vocab_size"] heads = config["num_attention_heads"] kv_heads = config["num_key_value_heads"] head_dim = config["head_dim"] gdn_v_heads = config["gdn_num_v_heads"] gdn_k_dim = config["gdn_key_dim"] gdn_v_dim = config["gdn_value_dim"] kernel = config["gdn_conv_kernel_size"] query = heads * head_dim kv = kv_heads * head_dim key_total = heads * gdn_k_dim value_total = gdn_v_heads * gdn_v_dim shapes: dict[str, tuple[int, ...]] = { "model.embed_tokens.weight": (vocab, hidden), "model.norm.weight": (hidden,), "lm_head.weight": (vocab, hidden), } for layer in range(config["num_hidden_layers"]): prefix = f"model.layers.{layer}." shapes[prefix + "input_layernorm.weight"] = (hidden,) shapes[prefix + "post_attention_layernorm.weight"] = (hidden,) shapes[prefix + "mlp.gate_proj.weight"] = (intermediate, hidden) shapes[prefix + "mlp.up_proj.weight"] = (intermediate, hidden) shapes[prefix + "mlp.down_proj.weight"] = (hidden, intermediate) if _is_surjo_xsa_layer(config, layer): shapes[prefix + "self_attn.q_proj.weight"] = (query, hidden) shapes[prefix + "self_attn.k_proj.weight"] = (kv, hidden) shapes[prefix + "self_attn.v_proj.weight"] = (kv, hidden) shapes[prefix + "self_attn.o_proj.weight"] = (hidden, query) shapes[prefix + "self_attn.q_norm.weight"] = (head_dim,) shapes[prefix + "self_attn.k_norm.weight"] = (head_dim,) else: shapes[prefix + "linear_attn.q_proj.weight"] = (key_total, hidden) shapes[prefix + "linear_attn.k_proj.weight"] = (key_total, hidden) shapes[prefix + "linear_attn.v_proj.weight"] = (value_total, hidden) shapes[prefix + "linear_attn.q_conv.conv.weight"] = (key_total, 1, kernel) shapes[prefix + "linear_attn.k_conv.conv.weight"] = (key_total, 1, kernel) shapes[prefix + "linear_attn.v_conv.conv.weight"] = (value_total, 1, kernel) shapes[prefix + "linear_attn.f_proj.0.weight"] = (gdn_v_dim, hidden) shapes[prefix + "linear_attn.f_proj.1.weight"] = (key_total, gdn_v_dim) shapes[prefix + "linear_attn.b_proj.weight"] = (key_total, hidden) shapes[prefix + "linear_attn.w_proj.weight"] = (value_total, hidden) shapes[prefix + "linear_attn.A_log"] = (heads,) shapes[prefix + "linear_attn.dt_bias"] = (key_total,) shapes[prefix + "linear_attn.g_proj.0.weight"] = (gdn_v_dim, hidden) shapes[prefix + "linear_attn.g_proj.1.weight"] = (value_total, gdn_v_dim) shapes[prefix + "linear_attn.g_proj.1.bias"] = (value_total,) shapes[prefix + "linear_attn.o_norm.weight"] = (gdn_v_dim,) shapes[prefix + "linear_attn.o_proj.weight"] = (hidden, value_total) return shapes _FWKV_ARCHITECTURE = "FWKVLanguageModel" _FWKV_SUPPORT = ( "FWKV support in this snapshot: config validation, shapes/weights import, " "native WKV-recurrence execution." ) def validate_fwkv_config(config: dict) -> dict: """Validate an FWKV Myosotis recurrent config; return a normalized copy. Topology is ``n_layers`` identical FWKVBlock layers (Myosotis-1-base: 13 layers, d_model 768, d_emb 192, ffn_mult 4, vocab 50257, wkv_floor 0.1, factorized tied head). No weights are read and no modeling code is executed here. """ if not isinstance(config, dict): raise UnsupportedModelError("config must be a JSON object. " + _FWKV_SUPPORT) config = dict(config) if config.get("model_type") != "fwkv": raise UnsupportedModelError( f"Unsupported model_type {config.get('model_type')!r}. " + _FWKV_SUPPORT) if config.get("architectures", [_FWKV_ARCHITECTURE]) != [_FWKV_ARCHITECTURE]: raise UnsupportedModelError( f"Expected architectures=[{_FWKV_ARCHITECTURE!r}], " f"got {config.get('architectures')!r}. " + _FWKV_SUPPORT) auto_map = config.get("auto_map") or {} if not isinstance(auto_map, dict): raise UnsupportedModelError("auto_map must be an object. " + _FWKV_SUPPORT) for name, reference in auto_map.items(): refs = [reference] if isinstance(reference, str) else reference if not isinstance(refs, (list, tuple)) or not refs or any( not isinstance(r, str) or "--" in r or not (r.startswith("configuration_fwkv.") or r.startswith("modeling_fwkv.")) for r in refs): raise UnsupportedModelError( "FWKV auto_map must reference code within the pinned snapshot " "(configuration_fwkv/modeling_fwkv modules only); it is never " "executed by the native runtime. " + _FWKV_SUPPORT) for key, maximum in (("vocab_size", 2_000_000), ("d_model", 8192), ("d_emb", 2048), ("n_layers", 128), ("ffn_mult", 16)): value = config.get(key) if type(value) is not int or not 0 < value <= maximum: raise InvalidCheckpointError( f"{key} must be an integer in [1, {maximum}], got {value!r}") floor = config.get("wkv_floor", 0.1) if type(floor) not in (int, float) or not 0 <= float(floor) < 1: raise InvalidCheckpointError( f"wkv_floor must be a number in [0, 1), got {floor!r}") config["wkv_floor"] = float(floor) tied = config.get("tie_word_embeddings", True) if tied is not True: raise UnsupportedModelError( "FWKV requires tie_word_embeddings=true (factorized tied head)") config["tie_word_embeddings"] = True for key in ("eos_token_id", "bos_token_id", "pad_token_id"): value = config.get(key) if value is None: continue tokens = value if key == "eos_token_id" and isinstance(value, list) else [value] if not tokens or any(type(t) is not int or not 0 <= t < config["vocab_size"] for t in tokens): raise InvalidCheckpointError(f"{key} is outside [0, vocab_size)") for key in ("quantization_config", "compression_config"): if config.get(key) not in (None, {}): raise UnsupportedModelError(f"Nonempty {key} is not implemented; provide unquantized safetensors") # Normalized aliases for engine/server convenience (native FWKV keeps d_*). config["num_hidden_layers"] = config["n_layers"] config["hidden_size"] = config["d_model"] config["intermediate_size"] = config["ffn_mult"] * config["d_model"] config["max_position_embeddings"] = int(config.get("seq_len", 1024)) config["architectures"] = [_FWKV_ARCHITECTURE] return config def _expected_shapes_fwkv(config: dict) -> dict[str, tuple[int, ...]]: """Exact trainable tensor shapes for a validated FWKV config. Mirrors ``modeling_fwkv.py`` naming verbatim (verified against the FWKV/Myosotis-1-base ``model.safetensors`` header: 148 F32 tensors, factorized tied head ``shared.weight`` + ``shared.proj.weight``, per-layer ``proj_k/v/r/out``, ``w``, ``ffn.0/2``, ``norm_wkv/norm_ffn`` weight+bias, final ``norm`` weight+bias). """ d = config["d_model"] e = config["d_emb"] vocab = config["vocab_size"] ffn = config["ffn_mult"] * d shapes: dict[str, tuple[int, ...]] = { "shared.weight": (vocab, e), "shared.proj.weight": (d, e), "norm.weight": (d,), "norm.bias": (d,), } for layer in range(config["n_layers"]): prefix = f"blocks.{layer}." shapes[prefix + "proj_k.weight"] = (d, d) shapes[prefix + "proj_v.weight"] = (d, d) shapes[prefix + "proj_r.weight"] = (d, d) shapes[prefix + "proj_out.weight"] = (d, d) shapes[prefix + "w"] = (d,) shapes[prefix + "ffn.0.weight"] = (ffn, d) shapes[prefix + "ffn.2.weight"] = (d, ffn) shapes[prefix + "norm_wkv.weight"] = (d,) shapes[prefix + "norm_wkv.bias"] = (d,) shapes[prefix + "norm_ffn.weight"] = (d,) shapes[prefix + "norm_ffn.bias"] = (d,) return shapes def _unique_object(pairs: list[tuple[str, Any]]) -> dict: result = {} for key, value in pairs: if key in result: raise InvalidCheckpointError(f"Duplicate JSON key: {key}") result[key] = value return result def _read_json(path: Path) -> dict: try: with path.open(encoding="utf-8") as stream: value = json.load(stream, object_pairs_hook=_unique_object) except (OSError, ValueError) as exc: raise InvalidCheckpointError(f"Cannot read {path.name}: {exc}") from exc if not isinstance(value, dict): raise InvalidCheckpointError(f"{path.name} must contain a JSON object") return value def _standalone_chat_template(root: Path) -> Any: """Chat template shipped as a file instead of in ``tokenizer_config.json``. ``AutoTokenizer`` only reads the ``chat_template`` key of ``tokenizer_config.json``, but some checkpoints ship the Jinja as a sibling ``chat_template.jinja`` (the snapshot download already fetches those). Without this fallback the shipped template is silently ignored and chat requests fail with "no chat template". A ``str`` is a single template, a ``dict`` a set of named ones -- both are accepted by ``apply_chat_template``. The config key always wins; this only runs when it is absent. """ candidates: list[Path] = [] for pattern in ("chat_template*.jinja", "chat_template*.json"): candidates.extend(sorted(root.glob(pattern))) templates_dir = root / "chat_templates" if templates_dir.is_dir(): candidates.extend(sorted(templates_dir.glob("*.jinja"))) for path in candidates: try: text = path.read_text(encoding="utf-8") except OSError: continue if path.suffix == ".json": try: parsed = json.loads(text) except ValueError: continue if isinstance(parsed, (str, dict)) and parsed: return parsed continue if text.strip(): return text return None def _expected_shapes(config: dict) -> dict[str, tuple[int, ...]]: hidden, intermediate = config["hidden_size"], config["intermediate_size"] head_dim = config["head_dim"] query, kv = config["num_attention_heads"] * head_dim, config["num_key_value_heads"] * head_dim shapes = { "model.embed_tokens.weight": (config["vocab_size"], hidden), "model.norm.weight": (hidden,), "lm_head.weight": (config["vocab_size"], hidden), } for layer in range(config["num_hidden_layers"]): prefix = f"model.layers.{layer}." for name, shape in { "input_layernorm.weight": (hidden,), "post_attention_layernorm.weight": (hidden,), "self_attn.q_proj.weight": (query, hidden), "self_attn.k_proj.weight": (kv, hidden), "self_attn.v_proj.weight": (kv, hidden), "self_attn.o_proj.weight": (hidden, query), "mlp.gate_proj.weight": (intermediate, hidden), "mlp.up_proj.weight": (intermediate, hidden), "mlp.down_proj.weight": (hidden, intermediate), }.items(): shapes[prefix + name] = shape if config["model_type"] == "qwen3": shapes[prefix + "self_attn.q_norm.weight"] = (head_dim,) shapes[prefix + "self_attn.k_norm.weight"] = (head_dim,) return shapes def _checkpoint_path(root: Path, name: str) -> Path: # Restrict index entries to root filenames, including on Windows. HF cache # symlinks into the repository's blobs directory are legitimate, not traversal. if not isinstance(name, str) or not name.endswith(".safetensors") or any(c in name for c in "/\\:"): raise InvalidCheckpointError(f"Unsafe shard filename: {name!r}; expected a root .safetensors filename") path = root / name resolved = path.resolve() # HF cache symlinks snapshots// files into a blob store. Layouts vary # by hub version (legacy models--X/{snapshots,blobs}, hub>=1.x global # hub/blobs sharded by hash prefix), so match structurally: root lives # under a "snapshots" tree and the target under a "blobs" tree. A bare # filename (enforced above) cannot traversal-escape lexically; this only # blesses cache-style symlinks, and a hostile dir named snapshots/blobs # gains nothing (it would only re-read its own files). root_parts = root.resolve().parts cache_blob = "snapshots" in root_parts and "blobs" in resolved.parts if not resolved.is_relative_to(root) and not cache_blob: raise InvalidCheckpointError(f"Shard escapes checkpoint directory: {name!r}") if not path.is_file(): raise InvalidCheckpointError(f"Missing safetensors shard: {name}") return path def _load_weights(root: Path, config: dict) -> dict[str, np.ndarray]: if config.get("model_type") == "surjo": expected = _expected_shapes_surjo(config) # Surjo computes RoPE frequencies in forward(); checkpoints carry # no inv_freq buffers (verified: 178 F32 tensors, tied lm_head omitted). buffers: set[str] = set() elif config.get("model_type") == "fwkv": expected = _expected_shapes_fwkv(config) # FWKV has no attention KV, no RoPE, no rotary buffers (verified: # 148 F32 tensors, factorized tied head). buffers = set() else: expected = _expected_shapes(config) buffers = {"model.rotary_emb.inv_freq"} | { f"model.layers.{i}.self_attn.rotary_emb.inv_freq" for i in range(config["num_hidden_layers"]) } if sum(math.prod(shape) * 4 for shape in expected.values()) > 64 * 1024**3: raise UnsupportedModelError("FP32 source weights exceed the native 64 GiB safety limit") index_path = root / "model.safetensors.index.json" weight_map = None if index_path.is_file(): weight_map = _read_json(index_path).get("weight_map") if not isinstance(weight_map, dict) or not weight_map: raise InvalidCheckpointError("safetensors index must contain a nonempty weight_map") for name in weight_map.values(): _checkpoint_path(root, name) filenames = sorted(set(weight_map.values())) else: filenames = ["model.safetensors"] unreferenced = {p.name for p in root.glob("*.safetensors")} - set(filenames) if unreferenced: raise InvalidCheckpointError(f"Unreferenced/ambiguous safetensors files: {sorted(unreferenced)}") seen = set() weights = {} for filename in filenames: path = _checkpoint_path(root, filename) try: with safe_open(path, framework="numpy") as tensors, path.open("rb") as raw: # safe_open validates the header and offsets first. Read its raw # metadata as well to detect duplicate JSON keys and decode BF16 # without depending on NumPy BF16 or importing a torch runtime. header_size = struct.unpack(" 100_000_000: raise InvalidCheckpointError("Safetensors header exceeds 100 MB") header = json.loads(raw.read(header_size), object_pairs_hook=_unique_object) for name in tensors.keys(): if name in seen: raise InvalidCheckpointError(f"Duplicate tensor across shards: {name}") seen.add(name) if weight_map is not None and weight_map.get(name) != filename: raise InvalidCheckpointError(f"Tensor {name} is absent from the index or mapped to the wrong shard") if name not in expected and name not in buffers: raise InvalidCheckpointError(f"Unexpected tensor {name}; no unsupported trainable weights may be ignored") tensor = tensors.get_slice(name) shape = tuple(tensor.get_shape()) if name in expected: wanted = expected[name] else: # Buffer-only path (dense RoPE inv_freq); FWKV/Surjo # have empty buffer sets so this never triggers there. wanted = (config["head_dim"] // 2,) if shape != wanted: raise InvalidCheckpointError(f"Wrong shape for {name}: expected {wanted}, got {shape}") dtype = tensor.get_dtype() if dtype == "BF16": raw.seek(8 + header_size + header[name]["data_offsets"][0]) bits = np.fromfile(raw, dtype=" LoadedModel: """Import local or Hub safetensors without ever loading a Transformers model. Downloaded checkpoints/tokenizers share one pinned snapshot. Python files are downloaded only with explicit trust_remote_code=True, which authorizes tokenizer code only, never custom modeling operators. Trusted Python code itself is not sandboxed. Cross-repository tokenizer code is rejected because it cannot be pinned to this snapshot. No pickle/.bin fallback is provided. Only explicit, non-null EOS/BOS/PAD settings from generation_config.json override config.json; sampling defaults are deliberately not imported. """ if type(trust_remote_code) is not bool or type(local_files_only) is not bool: raise TypeError("trust_remote_code and local_files_only must be booleans") source = os.fspath(model_id_or_path) candidate = Path(source).expanduser() if candidate.is_dir(): root = candidate.resolve() resolved_revision = revision else: if candidate.exists() or candidate.is_absolute() or source.startswith((".", "~")) or "\\" in source: raise InvalidCheckpointError(f"Not a checkpoint directory: {source}") patterns = [ "*.safetensors", "*.safetensors.index.json", "config.json", "generation_config.json", "tokenizer*.json", "special_tokens_map.json", "added_tokens.json", "vocab.json", "vocab.txt", "merges.txt", "*.model", "*.tiktoken", "chat_template*.json", "chat_template*.jinja", "chat_templates/*.jinja", ] if trust_remote_code: patterns.append("*.py") root = Path(snapshot_download( repo_id=source, revision=revision, cache_dir=cache_dir, local_files_only=local_files_only, allow_patterns=patterns, )).resolve() resolved_revision = root.name config = _read_json(root / "config.json") generation_path = root / "generation_config.json" if generation_path.is_file(): generation_config = _read_json(generation_path) for key in ("eos_token_id", "bos_token_id", "pad_token_id"): if generation_config.get(key) is not None: config[key] = generation_config[key] # Surjo base checkpoints ship use_cache=false; the native session always # caches (XSA KV slots + GDN recurrent/conv state). The flag is preserved # in config for fidelity and does not disable native caching. if config.get("model_type") == "surjo": config = validate_surjo_config(config) elif config.get("model_type") == "fwkv": config = validate_fwkv_config(config) else: config = validate_config(config) weights = _load_weights(root, config) tokenizer_config = root / "tokenizer_config.json" tokenizer_metadata = _read_json(tokenizer_config) if tokenizer_config.is_file() else {} for metadata in (config, tokenizer_metadata): auto_map = metadata.get("auto_map") or {} references = auto_map.get("AutoTokenizer", []) if isinstance(auto_map, dict) else auto_map if isinstance(references, str): references = [references] if not isinstance(references, (list, tuple)) or any( reference is not None and (not isinstance(reference, str) or "--" in reference) for reference in references ): raise UnsupportedModelError("Tokenizer auto_map must reference code within the pinned snapshot, not another repository") from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained( str(root), trust_remote_code=trust_remote_code, local_files_only=True, ) if not getattr(tokenizer, "chat_template", None): fallback = _standalone_chat_template(root) if fallback is not None: tokenizer.chat_template = fallback return LoadedModel(config=config, weights=weights, tokenizer=tokenizer, source=str(root), revision=resolved_revision)