configuration_multiscreen.py

"""Configuration for the Transformers-compatible Multiscreen model."""

from __future__ import annotations

import math
from typing import Any

try:  # Transformers has historically used both spellings in examples.
    from transformers import PreTrainedConfig
except ImportError:  # pragma: no cover - compatibility fallback for old releases.
    from transformers import PretrainedConfig as PreTrainedConfig  # type: ignore


class MultiscreenConfig(PreTrainedConfig):
    """Configuration for Multiscreen causal language models.

    This mirrors the architecture knobs from ``dieOD/multiscreen-pytorch`` while
    exposing the conventional Transformers names where possible.

    Important aliases
    -----------------
    ``hidden_size`` <-> original ``hidden_dim``
    ``num_hidden_layers`` <-> original ``num_layers``
    ``num_attention_heads`` <-> original ``num_heads``
    ``max_position_embeddings`` <-> original ``max_seq_len``

    Reproducibility controls
    ------------------------
    ``mipe_compute_dtype`` and ``softmask_compute_dtype`` can be ``"fp32"``
    for the numerically safer Transformers port behavior, or ``"reference"``
    to use the incoming tensor dtype like the standalone PyTorch reference.
    ``strict_position_ids`` rejects batch-specific or non-contiguous
    ``position_ids`` because the reference cache API is based on a scalar
    ``start_pos``. ``zero_pad_hidden_states`` can additionally zero padded
    query states after each residual layer; it defaults to ``False`` to keep
    original residual behavior.
    """

    model_type = "multiscreen"
    keys_to_ignore_at_inference = ["past_key_values"]
    _alias_to_primary = {
        "hidden_dim": "hidden_size",
        "num_layers": "num_hidden_layers",
        "num_heads": "num_attention_heads",
        "max_seq_len": "max_position_embeddings",
    }

    def __init__(
        self,
        vocab_size: int = 50_257,
        hidden_size: int | None = None,
        hidden_dim: int | None = None,
        num_hidden_layers: int | None = None,
        num_layers: int | None = None,
        num_attention_heads: int | None = None,
        num_heads: int | None = None,
        key_dim: int = 16,
        value_dim: int = 64,
        max_position_embeddings: int | None = None,
        max_seq_len: int | None = None,
        mipe_threshold: float = 256.0,
        gradient_checkpointing: bool = False,
        use_cache: bool = True,
        labels_are_shifted: bool = False,
        mipe_compute_dtype: str = "fp32",
        softmask_compute_dtype: str = "fp32",
        strict_position_ids: bool = True,
        zero_pad_hidden_states: bool = False,
        initializer_range: float = 0.1,
        bos_token_id: int | None = None,
        eos_token_id: int | None = None,
        pad_token_id: int | None = None,
        tie_word_embeddings: bool = True,
        **kwargs: Any,
    ) -> None:
        # Saved Transformers configs may contain these superclass fields in
        # kwargs. Multiscreen is always a decoder-only model, and passing them
        # through while also setting them below would duplicate keyword args.
        is_decoder = kwargs.pop("is_decoder", True)
        is_encoder_decoder = kwargs.pop("is_encoder_decoder", False)
        kwargs.pop("model_type", None)
        if is_decoder is not True:
            raise ValueError("MultiscreenConfig requires is_decoder=True")
        if is_encoder_decoder is not False:
            raise ValueError("MultiscreenConfig requires is_encoder_decoder=False")

        hidden_size = self._resolve_alias(
            primary=hidden_size,
            alias=hidden_dim,
            default=256,
            primary_name="hidden_size",
            alias_name="hidden_dim",
        )
        num_hidden_layers = self._resolve_alias(
            primary=num_hidden_layers,
            alias=num_layers,
            default=8,
            primary_name="num_hidden_layers",
            alias_name="num_layers",
        )
        num_attention_heads = self._resolve_alias(
            primary=num_attention_heads,
            alias=num_heads,
            default=8,
            primary_name="num_attention_heads",
            alias_name="num_heads",
        )
        max_position_embeddings = self._resolve_alias(
            primary=max_position_embeddings,
            alias=max_seq_len,
            default=256,
            primary_name="max_position_embeddings",
            alias_name="max_seq_len",
        )

        if not bool(tie_word_embeddings):
            raise ValueError(
                "Multiscreen uses normalized tied input/output embeddings; "
                "tie_word_embeddings must be True."
            )

        self.vocab_size = int(vocab_size)
        self.hidden_size = int(hidden_size)
        self.hidden_dim = int(hidden_size)  # original repo alias
        self.num_hidden_layers = int(num_hidden_layers)
        self.num_layers = int(num_hidden_layers)  # original repo alias
        self.num_attention_heads = int(num_attention_heads)
        self.num_heads = int(num_attention_heads)  # original repo alias
        self.key_dim = int(key_dim)
        self.value_dim = int(value_dim)
        self.max_position_embeddings = int(max_position_embeddings)
        self.max_seq_len = int(max_position_embeddings)  # original repo alias
        self.mipe_threshold = float(mipe_threshold)
        self.gradient_checkpointing = bool(gradient_checkpointing)
        self.use_cache = bool(use_cache)
        self.labels_are_shifted = bool(labels_are_shifted)
        self.mipe_compute_dtype = str(mipe_compute_dtype)
        self.softmask_compute_dtype = str(softmask_compute_dtype)
        self.strict_position_ids = bool(strict_position_ids)
        self.zero_pad_hidden_states = bool(zero_pad_hidden_states)
        self.initializer_range = float(initializer_range)

        self._validate()

        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            tie_word_embeddings=tie_word_embeddings,
            is_decoder=True,
            is_encoder_decoder=False,
            use_cache=use_cache,
            **kwargs,
        )

        # Useful when pushing a repo with these Python files to the Hub and
        # loading with trust_remote_code=True.
        if not getattr(self, "auto_map", None):
            self.auto_map = {
                "AutoConfig": "configuration_multiscreen.MultiscreenConfig",
                "AutoModel": "modeling_multiscreen.MultiscreenModel",
                "AutoModelForCausalLM": "modeling_multiscreen.MultiscreenForCausalLM",
            }
        if not getattr(self, "architectures", None):
            self.architectures = ["MultiscreenForCausalLM"]

    @staticmethod
    def _resolve_alias(
        *,
        primary: int | None,
        alias: int | None,
        default: int,
        primary_name: str,
        alias_name: str,
    ) -> int:
        if primary is None and alias is None:
            return default
        if primary is None:
            return int(alias)  # type: ignore[arg-type]
        if alias is None:
            return int(primary)
        if int(primary) != int(alias):
            raise ValueError(
                f"Conflicting values for {primary_name}={primary} and "
                f"{alias_name}={alias}. Use only one or make them equal."
            )
        return int(primary)

    @classmethod
    def from_psi(
        cls,
        psi: int,
        vocab_size: int = 50_257,
        max_seq_len: int = 256,
        **overrides: Any,
    ) -> "MultiscreenConfig":
        """Build a paper-style config from the supraparameter Psi.

        The scaling rule used in the reference repo is ``N_L = N_H = Psi`` and
        ``d_E = Psi²``.
        """

        return cls(
            vocab_size=vocab_size,
            hidden_size=psi * psi,
            num_hidden_layers=psi,
            num_attention_heads=psi,
            max_position_embeddings=max_seq_len,
            **overrides,
        )

    def _validate(self) -> None:
        if self.vocab_size <= 0:
            raise ValueError("vocab_size must be positive")
        if self.hidden_size <= 0:
            raise ValueError("hidden_size/hidden_dim must be positive")
        if self.num_hidden_layers <= 0:
            raise ValueError("num_hidden_layers/num_layers must be positive")
        if self.num_attention_heads <= 0:
            raise ValueError("num_attention_heads/num_heads must be positive")
        if self.key_dim < 2:
            raise ValueError("key_dim must be at least 2 because MiPE rotates the first two coordinates")
        if self.value_dim <= 0:
            raise ValueError("value_dim must be positive")
        if self.max_position_embeddings <= 0:
            raise ValueError("max_position_embeddings/max_seq_len must be positive")
        if self.mipe_threshold <= 0:
            raise ValueError("mipe_threshold must be positive")
        allowed_compute_dtypes = {"fp32", "reference"}
        if self.mipe_compute_dtype not in allowed_compute_dtypes:
            raise ValueError(
                "mipe_compute_dtype must be either 'fp32' or 'reference', "
                f"got {self.mipe_compute_dtype!r}"
            )
        if self.softmask_compute_dtype not in allowed_compute_dtypes:
            raise ValueError(
                "softmask_compute_dtype must be either 'fp32' or 'reference', "
                f"got {self.softmask_compute_dtype!r}"
            )
        if self.initializer_range <= 0:
            raise ValueError("initializer_range must be positive")

    @classmethod
    def _normalize_alias_updates(cls, updates: dict[str, Any]) -> dict[str, Any]:
        """Map original-repository aliases to canonical Transformers field names."""

        normalized = dict(updates)
        for alias, primary in cls._alias_to_primary.items():
            if alias not in normalized:
                continue
            alias_value = normalized.pop(alias)
            if primary in normalized and int(normalized[primary]) != int(alias_value):
                raise ValueError(
                    f"Conflicting update values for {primary}={normalized[primary]} "
                    f"and {alias}={alias_value}. Use only one or make them equal."
                )
            normalized[primary] = alias_value
        return normalized

    def clone(self, **updates: Any) -> "MultiscreenConfig":
        """Return a config copy with updated fields.

        ``PreTrainedConfig.to_dict()`` contains both Transformers field names and
        original-repository aliases such as ``hidden_size``/``hidden_dim``.
        Reusing that dictionary directly can create alias conflicts when callers
        update only one spelling, so this method rebuilds from canonical fields
        and canonicalizes alias-style updates.
        """

        normalized_updates = self._normalize_alias_updates(updates)

        # ``is_decoder`` and ``is_encoder_decoder`` are forced by this class and
        # are passed explicitly to ``PreTrainedConfig`` in ``__init__``.
        # Silently accepting the matching values makes clone robust to dicts
        # produced by Transformers; conflicting values should fail loudly.
        if "is_decoder" in normalized_updates:
            if normalized_updates.pop("is_decoder") is not True:
                raise ValueError("MultiscreenConfig requires is_decoder=True")
        if "is_encoder_decoder" in normalized_updates:
            if normalized_updates.pop("is_encoder_decoder") is not False:
                raise ValueError("MultiscreenConfig requires is_encoder_decoder=False")

        data: dict[str, Any] = {
            "vocab_size": self.vocab_size,
            "hidden_size": self.hidden_size,
            "num_hidden_layers": self.num_hidden_layers,
            "num_attention_heads": self.num_attention_heads,
            "key_dim": self.key_dim,
            "value_dim": self.value_dim,
            "max_position_embeddings": self.max_position_embeddings,
            "mipe_threshold": self.mipe_threshold,
            "gradient_checkpointing": self.gradient_checkpointing,
            "use_cache": self.use_cache,
            "labels_are_shifted": self.labels_are_shifted,
            "mipe_compute_dtype": self.mipe_compute_dtype,
            "softmask_compute_dtype": self.softmask_compute_dtype,
            "strict_position_ids": self.strict_position_ids,
            "zero_pad_hidden_states": self.zero_pad_hidden_states,
            "initializer_range": self.initializer_range,
            "bos_token_id": self.bos_token_id,
            "eos_token_id": self.eos_token_id,
            "pad_token_id": self.pad_token_id,
            "tie_word_embeddings": True,
        }

        # Preserve useful PreTrainedConfig extras without reintroducing explicit
        # constructor arguments, aliases, or forced superclass kwargs.
        skip_keys = set(data) | set(self._alias_to_primary) | {
            "model_type",
            "is_decoder",
            "is_encoder_decoder",
            "transformers_version",
        }
        for key, value in self.to_dict().items():
            if key not in skip_keys:
                data[key] = value

        data.update(normalized_updates)
        return self.__class__(**data)

    @property
    def num_params_estimate(self) -> int:
        """Approximate parameter count, following the reference implementation.

        The estimate assumes tied input/output embeddings and ignores the small
        learned scalar parameters.
        """

        embed = self.vocab_size * self.hidden_size
        per_tile = self.hidden_size * (2 * self.key_dim + 3 * self.value_dim)
        total_tiles = self.num_hidden_layers * self.num_attention_heads
        return int(embed + total_tiles * per_tile)

    @property
    def sqrt_hidden_size(self) -> float:
        return math.sqrt(self.hidden_size)