speculators.models

@classmethod
def from_training_args(
    cls,
    verifier_config: PretrainedConfig,
    **kwargs,
) -> "Eagle3DraftModel":
    """Create Eagle3 model from training arguments.

    Args:
        verifier_config: Verifier model configuration
        **kwargs: Training arguments with Eagle3-specific params
            - num_layers: Number of decoder layers
            - norm_before_residual: Whether to normalize before residual connection
            - t2d: Target-to-draft vocabulary mapping tensor
            - d2t: Draft-to-target vocabulary mapping tensor
            - ttt_steps: Number of TTT steps
            - verifier_name_or_path: Path to verifier model

    Returns:
        Initialized Eagle3DraftModel
    """
    config = Eagle3SpeculatorConfig(
        transformer_layer_config=verifier_config,
        draft_vocab_size=kwargs["draft_vocab_size"],
        norm_before_residual=kwargs["norm_before_residual"],
        embed_requires_grad=kwargs.get("embed_requires_grad", False),
        speculators_config=SpeculatorsConfig(
            algorithm="eagle3",
            proposal_methods=[
                GreedyTokenProposalConfig(
                    speculative_tokens=kwargs["ttt_steps"],
                )
            ],
            default_proposal_method="greedy",
            verifier=VerifierConfig.from_config(
                verifier_config, name_or_path=kwargs["verifier_name_or_path"]
            ),
        ),
    )

    return cls(config=config, t2d=kwargs.get("t2d"), d2t=kwargs.get("d2t"))

@staticmethod
def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:
    """Get training and validation kwargs for Eagle3.

    Args:
        **kwargs: Training arguments

    Returns:
        Tuple of (train_call_kwargs, val_call_kwargs)
    """
    train_kwargs = {
        "use_off_policy_tokens": kwargs["use_off_policy_tokens"],
        "ttt_steps": kwargs["ttt_steps"],
        "ttt_step_loss_decay": kwargs["ttt_step_loss_decay"],
    }
    val_kwargs = {
        "use_off_policy_tokens": False,
        "ttt_steps": kwargs["ttt_steps"],
        "ttt_step_loss_decay": kwargs["ttt_step_loss_decay"],
    }
    return train_kwargs, val_kwargs

@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] = LlamaConfig
        if "model_type" in value:
            config_class = AutoConfig.for_model(
                model_type=value["model_type"]
            ).__class__
        return config_class(**value)
    return value

def attach_verifier(
    self,
    verifier: str | os.PathLike | PreTrainedModel,
    mode: Literal["full", "train_only"] | None = None,
):
    """
    Attach a verifier model to the EagleSpeculator for speculative decoding.
    Utilizes the verifier's embed_tokens, rotary_emb, and lm_head layers
    for the speculator's forward pass and generation methods.
    Additionally, for `generate`, it uses the verifier's hidden states
    to generate speculative token predictions.

    If mode is "full", the verifier is fully integrated for use with
    both `generate` and `forward` methods.

    If mode is "train_only", only the verifier's layers required for a forward pass
    are attached, allowing for better resource utilization during training.
    `generate` will not be available until a full verifier is attached.

    Example:
        ```python
        # Load and attach a verifier
        verifier = EagleSpeculator(...)

        # For generation
        speculator.attach_verifier(verifier)
        outputs = speculator.generate(input_ids)
        speculator.detach_verifier()

        # For training
        speculator.attach_verifier(verifier, mode="train_only")
        outputs = speculator(input_ids, hidden_states)
        speculator.detach_verifier()
        ```

    :param verifier: The verifier model to attach. This can be a path to a local
        model directory, a Hugging Face model identifier, or an instance of
        PreTrainedModel. If a path or identifier is provided, the model will be
        loaded automatically. If an instance is provided, it will be used directly.
    :param mode: The mode for attaching the verifier. Can be "full" or "train_only".
        If None, defaults to "full". In "train_only" mode, only the layers
        required for a forward pass are attached, and the speculator cannot
        perform generation until a full verifier is attached.
    :return: The PreTrainedModel instance for the verifier that was attached.
    """
    super().attach_verifier(verifier=verifier, mode=mode)

    if self.verifier_attachment_mode == "train_only":
        verifier_model = self.resolve_verifier(verifier)
    elif self.verifier_attachment_mode == "full":
        verifier_model = cast("PreTrainedModel", self.verifier)
    else:
        return

    if hasattr(verifier_model, "model"):
        self.embed_tokens = verifier_model.model.embed_tokens  # type: ignore[assignment,union-attr]
        self.rotary_emb = verifier_model.model.rotary_emb  # type: ignore[assignment,union-attr]
    else:
        # Bare model structure
        self.embed_tokens = verifier_model.embed_tokens  # type: ignore[assignment,attr-defined]
        self.rotary_emb = verifier_model.rotary_emb  # type: ignore[assignment,attr-defined]

    # lm_head is always at the top level of the verifier
    self.lm_head = verifier_model.lm_head  # type: ignore[assignment,attr-defined]

def detach_verifier(self):
    """
    Removes the reference to the attached verifier model and frees up the
    associated memory. After calling this method, the speculator will not
    be able to perform forward passes or generation until a new verifier
    is attached.
    """
    super().detach_verifier()

    del self.embed_tokens
    self.embed_tokens = None
    del self.rotary_emb
    self.rotary_emb = None
    del self.lm_head
    self.lm_head = None

def forward(
    self,
    input_ids: torch.LongTensor,
    hidden_states: torch.FloatTensor,
    attention_mask: torch.Tensor | None = None,
    position_ids: torch.LongTensor | None = None,
    past_key_values: tuple[tuple[torch.FloatTensor]] | None = None,
    use_cache: bool | None = None,
    output_attentions: bool | None = None,
    output_hidden_states: bool | None = None,  # noqa: ARG002
    return_dict: bool | None = None,
) -> torch.FloatTensor | CausalLMOutputWithPast:
    """
    Execute the forward pass for speculative token generation.

    Processes input tokens and verifier hidden states through the EAGLE architecture
    to generate candidate tokens for speculative decoding. The method combines input
    embeddings with verifier hidden states via a fusion layer, processes them
    through a transformer decoder layer, and produces logits for next token
    prediction.

    :param input_ids: Token IDs for the current input sequence. Shape: (batch_size,
        sequence_length). These represent the tokens that will be converted to
        embeddings and combined with verifier hidden states.
    :param hidden_states: Hidden state representations from the verifier model
        corresponding to the input sequence. Shape: (batch_size, sequence_length,
        hidden_size). These capture the verifier's understanding of the context.
    :param attention_mask: Optional attention mask to avoid attending to padding
        tokens. Shape: (batch_size, sequence_length) for 2D or (batch_size, 1,
        sequence_length, sequence_length) for 4D causal mask.
    :param position_ids: Optional position indices for tokens in the sequence.
        Shape: (batch_size, sequence_length). If None, auto-generated based on
        sequence length and past key values.
    :param past_key_values: Optional cached key-value states from previous forward
        passes for efficient generation. Tuple of layer key-value pairs.
    :param use_cache: Whether to return key-value states for caching in subsequent
        forward passes. Useful for autoregressive generation efficiency.
    :param output_attentions: Whether to return attention weights from the
        transformer layer. Used for analysis and visualization.
    :param output_hidden_states: Whether to return hidden states from the
        transformer layer. Currently not implemented in this model.
    :param return_dict: Whether to return structured CausalLMOutputWithPast instead
        of raw logits. If None, uses config.use_return_dict default.
    :return: Either raw logits tensor (batch_size, sequence_length, vocab_size) if
        return_dict=False, or CausalLMOutputWithPast containing logits, past key
        values, and optional attention weights.
    :raises ValueError: If verifier components (embed_tokens, rotary_emb, lm_head)
        are not attached. Call attach_verifier() before using forward().
    """
    if self.embed_tokens is None or self.rotary_emb is None or self.lm_head is None:
        raise ValueError(
            "Verifier model layers not initialized. "
            "Call `attach_verifier` to set up the model before using forward."
        )

    return_dict = (
        return_dict if return_dict is not None else self.config.use_return_dict
    )

    inputs_embeds = self.embed_tokens(input_ids)
    if self.embedding_layernorm is not None:
        inputs_embeds = self.embedding_layernorm(inputs_embeds)

    hidden_states = self.fusion_fc(
        torch.cat([inputs_embeds, hidden_states], dim=-1)
    )
    hidden_states, attention_mask, position_ids = self._prepare_decoder_inputs(
        hidden_states, attention_mask, position_ids, past_key_values
    )

    cos, sin = self.rotary_emb(hidden_states, position_ids)
    layer_outputs = self.transformer(
        hidden_states,
        attention_mask=attention_mask,
        position_ids=position_ids,
        past_key_value=past_key_values[0] if past_key_values else None,
        output_attentions=output_attentions,
        use_cache=use_cache,
        position_embeddings=(cos, sin),
    )
    hidden_states = layer_outputs[0]

    if self.pre_lm_head_layernorm is not None:
        hidden_states = self.pre_lm_head_layernorm(hidden_states)

    logits = self.lm_head(hidden_states)

    if not return_dict:
        return logits

    return CausalLMOutputWithPast(
        logits=logits,
        past_key_values=layer_outputs[1] if use_cache else None,
        hidden_states=None,
        attentions=None,
    )

@classmethod
def from_training_args(
    cls,
    verifier_config: PretrainedConfig,
    **kwargs,
) -> "EagleSpeculator":
    """Create EAGLE model from training arguments.

    Args:
        verifier_config: Verifier model configuration
        **kwargs: Training arguments with EAGLE-specific params
            - layernorms: Whether to include layer normalization layers
            - fusion_bias: Whether to add bias to fusion layer
            - transformer_layer_architecture: Name of transformer decoder layer
                class
            - verifier_name_or_path: Path to verifier model

    Returns:
        Initialized EagleSpeculator
    """
    config = EagleSpeculatorConfig(
        transformer_layer_config=verifier_config,
        layernorms=kwargs.get("layernorms", False),
        fusion_bias=kwargs.get("fusion_bias", False),
        transformer_layer_architecture=kwargs.get(
            "transformer_layer_architecture", "auto"
        ),
        speculators_config=SpeculatorsConfig(
            algorithm="eagle",
            proposal_methods=[GreedyTokenProposalConfig()],
            default_proposal_method="greedy",
            verifier=VerifierConfig.from_config(
                verifier_config, name_or_path=kwargs["verifier_name_or_path"]
            ),
        ),
    )

    return cls(config=config)

@staticmethod
def get_trainer_kwargs(**kwargs) -> tuple[dict, dict]:  # noqa: ARG004
    """Get training and validation kwargs for EAGLE.

    EAGLE doesn't require any special forward pass arguments during training,
    so this returns empty dictionaries.

    Args:
        **kwargs: Training arguments (unused)

    Returns:
        Tuple of (train_call_kwargs, val_call_kwargs), both empty dicts
    """
    return {}, {}

@model_validator(mode="after")
def check_add_architectures(self) -> Self:
    """
    Automatically adds the transformer layer architecture to the
    architectures list if it's not already present.

    :return: The validated configuration instance with updated architectures
    """
    if (
        self.transformer_layer_architecture != "auto"
        and self.transformer_layer_architecture not in self.architectures
    ):
        self.architectures.append(self.transformer_layer_architecture)

    return self

@field_serializer("transformer_layer_config")
def serialize_transformer_layer_config(self, value: PretrainedConfig) -> dict:
    """
    Serialize the transformer_layer_config to a dictionary for JSON storage.

    Converts the PretrainedConfig object to its dictionary representation
    using to_diff_dict() to only include non-default values.

    :param value: The PretrainedConfig instance to serialize
    :return: Dictionary representation of the transformer layer configuration
    """
    return value.to_diff_dict()

@field_validator("transformer_layer_config", mode="before")
@classmethod
def validate_transformer_layer_config(cls, value: Any) -> PretrainedConfig:
    """
    Validate and convert transformer_layer_config to a PretrainedConfig instance.

    Accepts either a dictionary that can be converted to a PretrainedConfig
    or an existing PretrainedConfig instance.

    :param value: The value to validate (dict or PretrainedConfig)
    :return: A validated PretrainedConfig instance
    :raises ValueError: If the value cannot be converted to a PretrainedConfig
    """
    if isinstance(value, dict):
        return AutoConfig.for_model(**value)
    if isinstance(value, PretrainedConfig):
        return value

    raise ValueError(
        "transformer_layer_config must be a PretrainedConfig instance or a "
        "dictionary that can be converted to a PretrainedConfig."
    )