modeling_addition.py

from __future__ import annotations

from typing import Any

import torch
import torch.nn.functional as F
from torch import nn
from transformers import GenerationMixin, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast

from .addition_transformer import AdditionTransformer
from .configuration_addition import AdditionConfig


class AdditionForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = AdditionConfig
    base_model_prefix = "model"
    main_input_name = "input_ids"

    def __init__(self, config: AdditionConfig) -> None:
        super().__init__(config)
        self.model = AdditionTransformer(config, vocab_size=config.vocab_size)
        self.post_init()

    def _init_weights(self, module: nn.Module) -> None:
        # AdditionTransformer owns initialization so research and exported models remain identical.
        return None

    def get_input_embeddings(self) -> nn.Embedding:
        return self.model.token_embedding

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.model.token_embedding = value

    def get_output_embeddings(self) -> nn.Linear:
        return self.model.unembedding

    def set_output_embeddings(self, value: nn.Linear) -> None:
        self.model.unembedding = value

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        labels: torch.Tensor | None = None,
        past_key_values: Any | None = None,
        use_cache: bool | None = None,
        output_attentions: bool | None = None,
        output_hidden_states: bool | None = None,
        return_dict: bool | None = None,
        **kwargs: Any,
    ) -> CausalLMOutputWithPast | tuple[torch.Tensor, ...]:
        if kwargs:
            names = ", ".join(sorted(kwargs))
            raise TypeError(f"Unsupported model inputs: {names}")
        if past_key_values is not None or use_cache:
            raise ValueError("AdditionForCausalLM does not implement a key-value cache.")
        if attention_mask is not None:
            if attention_mask.shape != input_ids.shape:
                raise ValueError("attention_mask must have the same shape as input_ids.")
            if not bool(torch.all(attention_mask != 0)):
                raise ValueError("Padding is unsupported; attention_mask must contain only ones.")

        return_dict = self.config.return_dict if return_dict is None else return_dict
        output_attentions = bool(output_attentions)
        output_hidden_states = bool(output_hidden_states)
        core_output = self.model(
            input_ids,
            return_activations=output_attentions or output_hidden_states,
        )
        loss = None
        if labels is not None:
            if labels.shape != input_ids.shape:
                raise ValueError("labels must have the same shape as input_ids.")
            loss = F.cross_entropy(
                core_output.logits[:, :-1, :].contiguous().view(-1, self.config.vocab_size),
                labels[:, 1:].contiguous().view(-1),
                ignore_index=-100,
            )

        hidden_states = None
        if output_hidden_states:
            hidden_states = (
                core_output.block_outputs[0].residual_pre_attention,
                *(block.residual_after_mlp for block in core_output.block_outputs),
            )
        attentions = (
            tuple(block.attention_pattern for block in core_output.block_outputs)
            if output_attentions
            else None
        )

        if not return_dict:
            values = (core_output.logits, None, hidden_states, attentions)
            return ((loss,) + values) if loss is not None else values
        return CausalLMOutputWithPast(
            loss=loss,
            logits=core_output.logits,
            past_key_values=None,
            hidden_states=hidden_states,
            attentions=attentions,
        )

    def prepare_inputs_for_generation(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        **kwargs: Any,
    ) -> dict[str, torch.Tensor | None | bool]:
        return {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "use_cache": False,
        }

    def analysis_tensors(self, detach: bool = True) -> dict[str, torch.Tensor]:
        return self.model.analysis_tensors(detach=detach)

    def symmetrized_mlp_tensor(self, detach: bool = True, layer: int = 0) -> torch.Tensor:
        return self.model.symmetrized_mlp_tensor(detach=detach, layer=layer)


AdditionForCausalLM.register_for_auto_class("AutoModelForCausalLM")