modeling_talkie.py

from __future__ import annotations

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.cache_utils import Cache, DynamicCache
from transformers import GenerationMixin
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast

try:
    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
except ImportError:  # pragma: no cover - compatibility with older Transformers.
    ALL_ATTENTION_FUNCTIONS = None

from .configuration_talkie import TalkieConfig


def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    dropout: float = 0.0,
    scaling: float | None = None,
    is_causal: bool | None = None,
    **kwargs,
) -> tuple[torch.Tensor, None]:
    del kwargs
    is_causal = is_causal if is_causal is not None else getattr(module, "is_causal", True)
    output = F.scaled_dot_product_attention(
        query,
        key,
        value,
        attn_mask=attention_mask,
        dropout_p=dropout,
        scale=scaling,
        is_causal=is_causal and attention_mask is None,
    )
    return output.transpose(1, 2).contiguous(), None


def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    d = x.shape[3] // 2
    x1 = x[..., :d]
    x2 = x[..., d:]
    y1 = x1 * cos + x2 * sin
    y2 = x1 * (-sin) + x2 * cos
    return torch.cat([y1, y2], 3).type_as(x)


class HeadGain(nn.Module):
    def __init__(self, n_head: int):
        super().__init__()
        self.head_g = nn.Parameter(torch.ones([n_head]))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * self.head_g.type_as(x).view(1, 1, -1, 1)


class WeightGain(nn.Module):
    def __init__(self):
        super().__init__()
        self.w_g = nn.Parameter(torch.ones(1))

    def forward(self, w: torch.Tensor) -> torch.Tensor:
        return w * self.w_g.type_as(w)


class ActGain(nn.Module):
    def __init__(self, init_value: float):
        super().__init__()
        self.a_g = nn.Parameter(torch.ones(1) * init_value)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * self.a_g.type_as(x)


class CausalSelfAttention(nn.Module):
    is_causal = True

    def __init__(self, config: TalkieConfig, layer_idx: int):
        super().__init__()
        self.config = config
        self.layer_idx = layer_idx
        self.n_head = config.n_head
        self.head_dim = config.head_dim
        n_state = config.n_embd

        self.attn_query = nn.Linear(n_state, n_state, bias=False)
        self.attn_key = nn.Linear(n_state, n_state, bias=False)
        self.attn_value = nn.Linear(n_state, n_state, bias=False)
        self.attn_resid = nn.Linear(n_state, n_state, bias=False)
        self.head_gain = HeadGain(config.n_head)

    def forward(
        self,
        x: torch.Tensor,
        cos_sin: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor:
        bsz, seq_len, _ = x.size()
        q = self.attn_query(x).view(bsz, seq_len, self.n_head, self.head_dim)
        k = self.attn_key(x).view(bsz, seq_len, self.n_head, self.head_dim)
        v = self.attn_value(x).view(bsz, seq_len, self.n_head, self.head_dim)

        cos, sin = cos_sin
        q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
        q, k = F.rms_norm(q, (q.size(-1),)), F.rms_norm(k, (k.size(-1),))
        q = self.head_gain(q)

        key_states = k.transpose(1, 2)
        value_states = v.transpose(1, 2)
        if kwargs.get("past_key_values") is not None:
            key_states, value_states = kwargs["past_key_values"].update(
                key_states, value_states, self.layer_idx
            )

        if ALL_ATTENTION_FUNCTIONS is None:
            attention_interface = eager_attention_forward
        elif hasattr(ALL_ATTENTION_FUNCTIONS, "get_interface"):
            attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
                self.config._attn_implementation, eager_attention_forward
            )
        else:  # pragma: no cover - compatibility with older Transformers.
            attention_interface = ALL_ATTENTION_FUNCTIONS.get(
                self.config._attn_implementation, eager_attention_forward
            )
        is_causal = attention_mask is None and key_states.shape[-2] == q.shape[1]
        y, _ = attention_interface(
            self,
            q.transpose(1, 2),
            key_states,
            value_states,
            attention_mask,
            is_causal=is_causal,
            **kwargs,
        )
        y = y.contiguous().view_as(x)
        return self.attn_resid(y)


class MLP(nn.Module):
    def __init__(self, config: TalkieConfig):
        super().__init__()
        n_state = config.n_embd
        n_mlp = int(round(((8 / 3) * n_state) / 128) * 128)

        self.mlp_gate = nn.Linear(n_state, n_mlp, bias=False)
        self.mlp_linear = nn.Linear(n_state, n_mlp, bias=False)
        self.mlp_resid = nn.Linear(n_mlp, n_state, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = F.silu(self.mlp_gate(x)) * self.mlp_linear(x)
        return self.mlp_resid(x)


class Block(nn.Module):
    def __init__(self, config: TalkieConfig, layer_idx: int):
        super().__init__()
        self.attn = CausalSelfAttention(config, layer_idx)
        self.attn_gain = ActGain((2 * config.n_layer) ** -0.5)
        self.mlp = MLP(config)
        self.mlp_gain = ActGain((2 * config.n_layer) ** -0.5)
        self.embed_skip = ActGain(0.0)

    def forward(
        self,
        e_x: torch.Tensor,
        x: torch.Tensor,
        cos_sin: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None = None,
        **kwargs,
    ) -> torch.Tensor:
        x = x + self.attn_gain(
            self.attn(F.rms_norm(x, (x.shape[-1],)), cos_sin, attention_mask, **kwargs)
        )
        x = x + self.mlp_gain(self.mlp(F.rms_norm(x, (x.shape[-1],))))
        x = x + self.embed_skip(e_x)
        return x


class TalkiePreTrainedModel(PreTrainedModel):
    config_class = TalkieConfig
    base_model_prefix = ""
    supports_gradient_checkpointing = False
    _supports_sdpa = True
    _supports_attention_backend = True
    _no_split_modules = ["Block"]
    _tied_weights_keys = None

    def _init_weights(self, module: nn.Module) -> None:
        return


class TalkieModel(TalkiePreTrainedModel, GenerationMixin):
    def __init__(self, config: TalkieConfig):
        super().__init__(config)
        self.embed = nn.Embedding(config.vocab_size, config.n_embd)
        self.blocks = nn.ModuleList([Block(config, i) for i in range(config.n_layer)])

        cos, sin = self._precompute_rotary_embeddings(
            config.max_position_embeddings, config.head_dim, config.rope_base
        )
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)
        self._rotary_initialized = cos.device.type != "meta"
        self.post_init()

    def _precompute_rotary_embeddings(
        self, seq_len: int, head_dim: int, base: int
    ) -> tuple[torch.Tensor, torch.Tensor]:
        device = self.embed.weight.device if hasattr(self, "embed") else "cpu"
        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
        inv_freq = 1.0 / (base ** (channel_range / head_dim))
        t = torch.arange(seq_len, dtype=torch.float32, device=device)
        freqs = torch.outer(t, inv_freq)
        cos, sin = freqs.cos(), freqs.sin()
        cos, sin = cos.bfloat16(), sin.bfloat16()
        cos, sin = cos[None, :, None, :], sin[None, :, None, :]
        return cos, sin

    def _ensure_rotary_embeddings(self, seq_len: int) -> None:
        device = self.embed.weight.device
        needs_init = (
            not self._rotary_initialized
            or self.cos.device != device
            or self.sin.device != device
            or self.cos.shape[1] < seq_len
        )
        if needs_init:
            max_seq_len = max(seq_len, self.config.max_position_embeddings)
            cos, sin = self._precompute_rotary_embeddings(
                max_seq_len, self.config.head_dim, self.config.rope_base
            )
            self.cos = cos.to(device=device)
            self.sin = sin.to(device=device)
            self._rotary_initialized = True

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

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

    def _position_ids(
        self,
        input_ids: torch.LongTensor,
        position_ids: torch.LongTensor | None = None,
        cache_position: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
    ) -> torch.LongTensor:
        batch_size, seq_len = input_ids.shape
        if position_ids is not None:
            if position_ids.dim() == 1:
                position_ids = position_ids.unsqueeze(0)
            return position_ids.to(device=input_ids.device, dtype=torch.long)
        if cache_position is not None:
            if cache_position.dim() == 1:
                cache_position = cache_position.unsqueeze(0)
            if cache_position.shape[0] == 1 and batch_size != 1:
                cache_position = cache_position.expand(batch_size, -1)
            return cache_position.to(device=input_ids.device, dtype=torch.long)
        past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
        position_ids = torch.arange(seq_len, device=input_ids.device, dtype=torch.long) + past_seen
        return position_ids.unsqueeze(0)

    def _attention_mask(
        self,
        attention_mask: torch.Tensor | None,
        input_ids: torch.Tensor,
        position_ids: torch.Tensor,
        past_key_values: Cache | None,
        dtype: torch.dtype,
    ) -> torch.Tensor | None:
        if attention_mask is not None and attention_mask.dim() >= 4:
            return attention_mask
        batch_size, query_length = input_ids.shape
        past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
        key_length = past_seen + query_length

        if attention_mask is not None and attention_mask.dim() != 2:
            return attention_mask
        if attention_mask is not None:
            if attention_mask.shape[-1] == query_length and past_seen:
                prefix = torch.ones(
                    attention_mask.shape[0],
                    past_seen,
                    dtype=attention_mask.dtype,
                    device=attention_mask.device,
                )
                attention_mask = torch.cat([prefix, attention_mask], dim=-1)
            key_length = attention_mask.shape[-1]
            has_padding = not bool(torch.all(attention_mask == 1))
        else:
            has_padding = False

        if attention_mask is None and past_seen == 0:
            return None

        key_positions = torch.arange(key_length, device=input_ids.device, dtype=torch.long)
        future_mask = key_positions.view(1, 1, 1, key_length) > position_ids.view(
            batch_size, 1, query_length, 1
        )
        if attention_mask is not None and has_padding:
            padding_mask = attention_mask[:, None, None, :].to(device=input_ids.device) == 0
            mask = future_mask | padding_mask
        else:
            mask = future_mask

        if not bool(mask.any()):
            return None
        min_value = torch.finfo(dtype).min
        causal_mask = torch.zeros(
            batch_size, 1, query_length, key_length, dtype=dtype, device=input_ids.device
        )
        return causal_mask.masked_fill(mask, min_value)

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: Cache | None = None,
        use_cache: bool | None = None,
        return_dict: bool | None = None,
        **kwargs,
    ) -> BaseModelOutputWithPast | tuple[torch.Tensor, ...]:
        cache_position = kwargs.pop("cache_position", None)
        if input_ids is None and inputs_embeds is None:
            raise ValueError("input_ids or inputs_embeds is required")
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("provide only one of input_ids or inputs_embeds")
        if input_ids is None:
            input_ids = torch.empty(
                inputs_embeds.shape[:2],
                dtype=torch.long,
                device=inputs_embeds.device,
            )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        if use_cache and past_key_values is None:
            past_key_values = DynamicCache(config=self.config)

        position_ids = self._position_ids(input_ids, position_ids, cache_position, past_key_values)
        needed_seq_len = int(position_ids.max().item()) + 1
        self._ensure_rotary_embeddings(needed_seq_len)
        if needed_seq_len > self.cos.shape[1]:
            raise ValueError(
                f"Sequence length {needed_seq_len} exceeds max_position_embeddings "
                f"{self.cos.shape[1]}"
            )

        cos = self.cos[0, position_ids, :, :]
        sin = self.sin[0, position_ids, :, :]
        cos_sin = cos, sin
        x = inputs_embeds if inputs_embeds is not None else self.embed(input_ids)
        x = F.rms_norm(x, (x.shape[-1],))
        attention_mask = self._attention_mask(attention_mask, input_ids, position_ids, past_key_values, x.dtype)
        e_x = x
        for block in self.blocks:
            x = block(
                e_x,
                x,
                cos_sin,
                attention_mask=attention_mask,
                past_key_values=past_key_values if use_cache else None,
                **kwargs,
            )
        x = F.rms_norm(x, (x.shape[-1],))
        past_key_values = past_key_values if use_cache else None
        use_return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        if use_return_dict:
            return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values)
        output = (x,)
        return output + ((past_key_values,) if past_key_values is not None else ())


class TalkieForCausalLM(TalkieModel):
    _tied_weights_keys = None

    def __init__(self, config: TalkieConfig):
        super().__init__(config)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.post_init()

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

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

    def forward(
        self,
        input_ids: torch.LongTensor | None = None,
        attention_mask: torch.Tensor | None = None,
        inputs_embeds: torch.FloatTensor | None = None,
        labels: torch.LongTensor | None = None,
        return_dict: bool | None = None,
        past_key_values: Cache | None = None,
        use_cache: bool | None = None,
        position_ids: torch.LongTensor | None = None,
        logits_to_keep: int | torch.Tensor = 0,
        **kwargs,
    ) -> CausalLMOutputWithPast | tuple[torch.Tensor, ...]:
        if input_ids is None and inputs_embeds is None:
            raise ValueError("input_ids or inputs_embeds is required")
        cache_position = kwargs.pop("cache_position", None)
        outputs = super().forward(
            input_ids,
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            cache_position=cache_position,
            return_dict=True,
            **kwargs,
        )
        hidden_states = outputs.last_hidden_state
        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
        logits = self.lm_head(hidden_states[:, slice_indices, :]).float()
        if self.config.logit_scale != 1.0:
            logits = logits * self.config.logit_scale

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        use_return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        if use_return_dict:
            return CausalLMOutputWithPast(
                loss=loss,
                logits=logits,
                past_key_values=outputs.past_key_values,
            )
        output = (logits,)
        if outputs.past_key_values is not None:
            output += (outputs.past_key_values,)
        return ((loss,) + output) if loss is not None else output


__all__ = ["TalkieConfig", "TalkieForCausalLM", "TalkieModel"]