inkling_mlx/config.py

"""Configuration for the Inkling multimodal model (MLX port).

Mirrors ``thinkingmachines/Inkling`` ``config.json`` and the transformers PR #47347
reference (``InklingConfig`` / ``InklingTextConfig`` / ``InklingVisionConfig`` /
``InklingAudioConfig``). We parse the *checkpoint* config layout (top-level
``text_config`` / ``vision_config`` / ``audio_config`` / ``mtp_config``), not the
flattened transformers layout.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any


def _get(d: dict, *names, default=None):
    for n in names:
        if n in d and d[n] is not None:
            return d[n]
    return default


@dataclass
class TextConfig:
    hidden_size: int = 6144
    num_hidden_layers: int = 66
    vocab_size: int = 201024
    unpadded_vocab_size: int | None = 200058

    # global (full) attention
    num_attention_heads: int = 64
    num_key_value_heads: int = 8
    head_dim: int = 128
    # sliding-window attention
    swa_num_attention_heads: int = 64
    swa_num_key_value_heads: int = 16
    swa_head_dim: int = 128
    sliding_window_size: int = 512

    # relative-position logits
    d_rel: int = 16
    rel_extent: int = 1024
    log_scaling_n_floor: int | None = 128000
    log_scaling_alpha: float = 0.1

    rms_norm_eps: float = 1e-6
    use_embed_norm: bool = True

    # short convolution
    sconv_kernel_size: int = 4

    # dense vs MoE MLP
    dense_mlp_idx: int = 2
    dense_intermediate_size: int = 24576  # dense MLP intermediate
    moe_intermediate_size: int = 3072     # per-expert intermediate

    # MoE routing
    n_routed_experts: int = 256
    num_experts_per_tok: int = 6
    n_shared_experts: int = 2
    shared_expert_sink: bool = True
    route_scale: float = 8.0
    use_gate_bias: bool = True
    norm_after_topk: bool = True
    use_global_scale: bool = True

    logits_mup_width_multiplier: float = 24.0
    hidden_act: str = "silu"

    max_position_embeddings: int = 1048576

    # which layer indices use sliding-window ("local") attention
    local_layer_ids: list[int] = field(default_factory=list)

    # MTP head (dropped for inference)
    num_mtp_layers: int | None = None

    @property
    def layer_types(self) -> list[str]:
        local = set(self.local_layer_ids)
        return [
            "hybrid_sliding" if i in local else "hybrid"
            for i in range(self.num_hidden_layers)
        ]

    @property
    def mlp_layer_types(self) -> list[str]:
        return [
            "dense" if i < self.dense_mlp_idx else "sparse"
            for i in range(self.num_hidden_layers)
        ]

    @classmethod
    def from_dict(cls, tc: dict) -> "TextConfig":
        return cls(
            hidden_size=_get(tc, "hidden_size", default=6144),
            num_hidden_layers=_get(tc, "num_hidden_layers", default=66),
            vocab_size=_get(tc, "vocab_size", default=201024),
            unpadded_vocab_size=_get(tc, "unpadded_vocab_size"),
            num_attention_heads=_get(tc, "num_attention_heads", default=64),
            num_key_value_heads=_get(tc, "num_key_value_heads", default=8),
            head_dim=_get(tc, "head_dim", default=128),
            swa_num_attention_heads=_get(tc, "swa_num_attention_heads", default=64),
            swa_num_key_value_heads=_get(tc, "swa_num_key_value_heads", default=16),
            swa_head_dim=_get(tc, "swa_head_dim", default=128),
            sliding_window_size=_get(tc, "sliding_window_size", default=512),
            d_rel=_get(tc, "d_rel", default=16),
            rel_extent=_get(tc, "rel_extent", default=1024),
            log_scaling_n_floor=_get(tc, "log_scaling_n_floor"),
            log_scaling_alpha=_get(tc, "log_scaling_alpha", default=0.1),
            rms_norm_eps=_get(tc, "rms_norm_eps", default=1e-6),
            use_embed_norm=_get(tc, "use_embed_norm", default=True),
            sconv_kernel_size=_get(tc, "sconv_kernel_size", default=4),
            dense_mlp_idx=_get(tc, "dense_mlp_idx", default=2),
            dense_intermediate_size=_get(tc, "dense_intermediate_size", default=24576),
            # checkpoint labels the *MoE* intermediate as `intermediate_size`
            moe_intermediate_size=_get(tc, "intermediate_size", default=3072),
            n_routed_experts=_get(tc, "n_routed_experts", default=256),
            num_experts_per_tok=_get(tc, "num_experts_per_tok", default=6),
            n_shared_experts=_get(tc, "n_shared_experts", default=2),
            shared_expert_sink=_get(tc, "shared_expert_sink", default=True),
            route_scale=_get(tc, "route_scale", default=8.0),
            use_gate_bias=_get(tc, "use_gate_bias", default=True),
            norm_after_topk=_get(tc, "norm_after_topk", default=True),
            use_global_scale=_get(tc, "use_global_scale", default=True),
            logits_mup_width_multiplier=_get(tc, "logits_mup_width_multiplier", default=24.0),
            max_position_embeddings=_get(tc, "model_max_length", "max_position_embeddings", default=1048576),
            local_layer_ids=list(_get(tc, "local_layer_ids", default=[]) or []),
        )


@dataclass
class VisionConfig:
    text_hidden_size: int = 6144
    patch_size: int = 40
    temporal_patch_size: int = 2
    num_channels: int = 3
    n_layers: int = 4
    rms_norm_eps: float = 1e-6
    use_vision_norm: bool = True

    @classmethod
    def from_dict(cls, vc: dict, text_hidden: int) -> "VisionConfig":
        return cls(
            text_hidden_size=text_hidden,
            patch_size=_get(vc, "patch_size", default=40),
            temporal_patch_size=_get(vc, "temporal_patch_size", default=2),
            num_channels=_get(vc, "n_channels", "num_channels", default=3),
            n_layers=_get(vc, "n_layers", "num_hidden_layers", default=4),
            rms_norm_eps=_get(vc, "rms_norm_eps", default=1e-6),
            use_vision_norm=_get(vc, "use_vision_norm", default=True),
        )


@dataclass
class AudioConfig:
    text_hidden_size: int = 6144
    n_mel_bins: int = 80
    mel_vocab_size: int = 16
    rms_norm_eps: float = 1e-6

    @classmethod
    def from_dict(cls, ac: dict, text_hidden: int) -> "AudioConfig":
        return cls(
            text_hidden_size=text_hidden,
            n_mel_bins=_get(ac, "n_mel_bins", default=80),
            mel_vocab_size=_get(ac, "mel_vocab_size", default=16),
            rms_norm_eps=_get(ac, "rms_norm_eps", default=1e-6),
        )


@dataclass
class InklingConfig:
    text: TextConfig
    vision: VisionConfig
    audio: AudioConfig
    image_token_id: int = 200054
    audio_token_id: int = 200053
    image_bos_token_id: int = 200005
    audio_bos_token_id: int = 200020
    eos_token_id: int = 200006
    model_type: str = "inkling_mm_model"

    @classmethod
    def from_dict(cls, cfg: dict) -> "InklingConfig":
        tc = dict(cfg.get("text_config", {}))
        mtp = cfg.get("mtp_config") or {}
        if mtp.get("num_nextn_predict_layers") is not None:
            tc.setdefault("num_mtp_layers", mtp.get("num_nextn_predict_layers"))
        text = TextConfig.from_dict(tc)
        vision = VisionConfig.from_dict(cfg.get("vision_config", {}) or {}, text.hidden_size)
        audio = AudioConfig.from_dict(cfg.get("audio_config", {}) or {}, text.hidden_size)
        return cls(
            text=text,
            vision=vision,
            audio=audio,
            image_token_id=_get(cfg, "image_token_id", default=200054),
            audio_token_id=_get(cfg, "audio_token_id", default=200053),
            image_bos_token_id=_get(cfg, "image_bos_token_id", default=200005),
            audio_bos_token_id=_get(cfg, "audio_bos_token_id", default=200020),
            eos_token_id=_get(cfg, "eos_token_id", default=200006),
            model_type=_get(cfg, "model_type", default="inkling_mm_model"),
        )

    @property
    def raw(self) -> dict[str, Any]:
        return {"model_type": self.model_type}