Spaces:
Sleeping
Sleeping
| """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.""" | |
| 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/<rev>/ 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("<Q", raw.read(8))[0] | |
| if header_size > 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="<u2", count=math.prod(shape)) | |
| if bits.size != math.prod(shape): | |
| raise InvalidCheckpointError(f"Truncated BF16 tensor: {name}") | |
| # BF16 is the upper 16 bits of IEEE float32, not float16. | |
| value = (bits.astype(np.uint32) << 16).view(np.float32).reshape(shape) | |
| elif dtype in {"F16", "F32", "F64"}: | |
| with np.errstate(over="ignore", invalid="ignore"): | |
| value = np.ascontiguousarray(tensors.get_tensor(name), dtype=np.float32) | |
| else: | |
| raise InvalidCheckpointError(f"Unsupported tensor dtype {dtype} for {name}; expected floating-point weights") | |
| if not np.isfinite(value).all(): | |
| raise InvalidCheckpointError(f"Nonfinite or FP32-overflowing values in tensor {name}") | |
| if name in buffers: | |
| inv_freq = 1.0 / config["rope_theta"] ** ( | |
| np.arange(0, config["head_dim"], 2, dtype=np.float32) / config["head_dim"] | |
| ) | |
| if not np.allclose(value, inv_freq, rtol=1e-5, atol=1e-8): | |
| raise InvalidCheckpointError(f"Rotary buffer {name} does not match default RoPE") | |
| else: | |
| weights[name] = value | |
| except InvalidCheckpointError: | |
| raise | |
| except Exception as exc: | |
| raise InvalidCheckpointError(f"Cannot load {filename}: {exc}") from exc | |
| if weight_map is not None and set(weight_map) != seen: | |
| raise InvalidCheckpointError(f"Indexed tensors missing from shards: {sorted(set(weight_map) - seen)}") | |
| if config["tie_word_embeddings"] and "model.embed_tokens.weight" in weights: | |
| embedding = weights["model.embed_tokens.weight"] | |
| head = weights.get("lm_head.weight") | |
| if head is not None and not np.array_equal(head, embedding): | |
| raise InvalidCheckpointError("Tied lm_head.weight differs from model.embed_tokens.weight") | |
| weights["lm_head.weight"] = embedding | |
| missing = set(expected) - set(weights) | |
| if missing: | |
| raise InvalidCheckpointError(f"Missing required tensors: {sorted(missing)}") | |
| return weights | |
| def import_model( | |
| model_id_or_path: str | os.PathLike, | |
| revision: str | None = None, | |
| trust_remote_code: bool = False, | |
| local_files_only: bool = False, | |
| cache_dir: str | os.PathLike | None = None, | |
| ) -> 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) | |