config.py

from typing import Any, Literal

from pydantic import Field, field_serializer, field_validator
from transformers import AutoConfig, PretrainedConfig
from transformers.models.qwen3.modeling_qwen3 import (
    Qwen3Config,
)

from speculators import SpeculatorModelConfig

__all__ = [
    "DFlashSpeculatorConfig",
]


@SpeculatorModelConfig.register("dflash")
class DFlashSpeculatorConfig(SpeculatorModelConfig):
    """
    Configuration for DFlash speculator with vocabulary mapping.

    DFlash features vocabulary mapping between draft (64K) and target (128K)
    vocabularies, enabling cross-tokenizer speculation.

    :param transformer_layer_config: Configuration for the transformer decoder layer
    :param draft_vocab_size: Size of draft model vocabulary for speculation
    """

    speculators_model_type: Literal["dflash"] = "dflash"
    architectures: list[str] = Field(
        default_factory=lambda: ["DFlashSpeculator"],
        description="Model architectures that can load these weights",
    )

    transformer_layer_config: PretrainedConfig = Field(
        default_factory=Qwen3Config,
        description="Configuration for the transformer decoder layer",
    )

    draft_vocab_size: int = Field(
        default=32000,
        description="Size of draft model vocabulary for speculation",
    )

    block_size: int = Field(
        default=8,
        description=(
            "Default size of the draft block predicted with a forward pass of the model"
        ),
    )

    max_anchors: int = Field(
        default=256,
        description=(
            "Maximum number of anchor positions to sample during training "
            "(controls memory usage and training efficiency)"
        ),
    )

    target_hidden_size: int | None = Field(
        default=None,
        description="Hidden size of the target model (if different from draft model)",
    )

    aux_hidden_state_layer_ids: list[int] | None = Field(
        default=None,
        description="Layer IDs of the DFlash auxiliary hidden state layers",
    )

    decoder_layer_type: Literal["qwen3", "laguna_xs"] = Field(
        default="qwen3",
        description="Decoder layer implementation used by the DFlash drafter.",
    )

    mask_token_id: int | None = Field(
        default=None,
        description="Token ID used for masking",
    )

    sliding_window_non_causal: bool = Field(
        default=False,
        description="Use non-causal synthetic block attention for sliding-window layers.",
    )

    sliding_window_base: Literal["fixed_anchor", "moving_query"] = Field(
        default="moving_query",
        description=(
            "Base-token sliding-window lower-bound policy. 'moving_query' matches "
            "FlashAttention-style SWA during inference; 'fixed_anchor' preserves "
            "the legacy DFlash training mask."
        ),
    )

    loss_type: Literal["distill", "dflash", "lk", "tv"] = Field(
        default="distill",
        description="DFlash objective. 'lk' uses hard-label LK loss.",
    )

    ce_weight: float | None = Field(
        default=None,
        description="Additive weight for hard-label DFlash CE.",
    )

    tv_weight: float | None = Field(
        default=None,
        description="Additive weight for full-distribution TV loss.",
    )

    kl_weight: float | None = Field(
        default=None,
        description="Additive weight for full-distribution KL distillation.",
    )

    lk_lambda: float = Field(
        default=0.5,
        description="Blend coefficient for hard-label LK loss.",
    )

    tv_temperature: float = Field(
        default=1.0,
        description="Teacher softmax temperature for TV/KL terms.",
    )

    cumacc_weight: bool = Field(
        default=False,
        description="Weight hard-label DFlash CE by draft cumulative acceptance.",
    )

    veri_cum_acc: bool = Field(
        default=False,
        description="Weight DFlash loss by verifier cumulative acceptance.",
    )

    veri_acc_temperature: float = Field(
        default=1.0,
        description="Temperature for verifier cumulative acceptance weighting.",
    )

    static_decay_weight: bool = Field(
        default=True,
        description="Apply DFlash position decay to hard-label CE.",
    )

    kl_distill_weight: float = Field(
        default=0.0,
        description="Back-compatible alias for kl_weight when kl_weight is unset.",
    )

    compile_decoder_layers: bool = Field(
        default=True,
        description=(
            "If True, torch.compile each decoder layer forward during training. "
            "The DFlash loss remains eager."
        ),
    )

    @field_serializer("transformer_layer_config")
    def serialize_transformer_config(self, value: PretrainedConfig) -> dict:
        """Serialize transformer config to dict."""
        return value.to_diff_dict()

    @field_validator("transformer_layer_config", mode="before")
    @classmethod
    def validate_transformer_config(cls, value: Any) -> PretrainedConfig:
        """Validate and convert transformer config."""
        if isinstance(value, dict):
            config_class: type[PretrainedConfig] = Qwen3Config
            if "model_type" in value:
                config_class = AutoConfig.for_model(
                    model_type=value["model_type"]
                ).__class__
            return config_class(**value)
        return value

    @property
    def target_vocab_size(self) -> int:
        """Get target vocabulary size from transformer config."""
        return self.transformer_layer_config.vocab_size

    def resolve_loss_weights(self) -> tuple[float, float, float]:
        if self.loss_type == "tv":
            ce_default, tv_default = 0.0, 1.0
        else:
            ce_default, tv_default = 1.0, 0.0

        ce = ce_default if self.ce_weight is None else self.ce_weight
        tv = tv_default if self.tv_weight is None else self.tv_weight
        kl = self.kl_distill_weight if self.kl_weight is None else self.kl_weight
        return float(ce), float(tv), float(kl)