#!/usr/bin/env python3
"""
HF-compatible single-language Hawk / RG-LRU model, for lm-eval.

Self-contained `trust_remote_code` modeling file. The building blocks are the
SAME code used at training time (De et al., 2024, Griffin/Hawk; arXiv:2402.19427),
so the per-language export state_dict maps 1:1 onto this module's parameters
(top-level attribute names wte / layers / norm_f / lm_head match the export keys
exactly -- no renaming, no transpose). Exposes the standard
forward(input_ids, labels=None) -> CausalLMOutputWithPast that lm-eval expects.

Register via config.json:
  "model_type": "hawk_rglru",
  "architectures": ["HawkForCausalLM"],
  "auto_map": {
    "AutoConfig": "modeling_hawk.HawkConfig",
    "AutoModelForCausalLM": "modeling_hawk.HawkForCausalLM",
    "AutoModelForSequenceClassification": "modeling_hawk.HawkForSequenceClassification"
  }
"""

from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import (
    CausalLMOutputWithPast, SequenceClassifierOutput,
)


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return self.weight * (x * norm)


class RGLRU(nn.Module):
    """Real-Gated Linear Recurrent Unit (De et al., 2024)."""

    def __init__(self, width: int, c: float = 8.0):
        super().__init__()
        self.width = width
        self.c = c
        self.input_gate = nn.Linear(width, width)
        self.recur_gate = nn.Linear(width, width)
        lam = torch.empty(width).uniform_(2.197, 6.907)
        self.log_lambda = nn.Parameter(lam)

    def forward(self, x):                       # x: (B, T, W)
        B, T, W = x.shape
        r = torch.sigmoid(self.recur_gate(x))
        i = torch.sigmoid(self.input_gate(x))
        log_a = -F.softplus(-self.log_lambda)
        log_a_t = self.c * r * log_a
        a_t = torch.exp(log_a_t)
        mult = torch.sqrt(torch.clamp(-torch.expm1(2.0 * log_a_t), min=1e-8))
        gated_x = mult * (i * x)
        h = torch.zeros(B, W, device=x.device, dtype=x.dtype)
        outs = []
        for t in range(T):
            h = a_t[:, t] * h + gated_x[:, t]
            outs.append(h)
        return torch.stack(outs, dim=1)


class RecurrentBlock(nn.Module):
    def __init__(self, d_model: int, d_rnn: int, conv_kernel: int = 4, rglru_c: float = 8.0):
        super().__init__()
        self.conv_kernel = conv_kernel
        self.in_gate = nn.Linear(d_model, d_rnn)
        self.in_recur = nn.Linear(d_model, d_rnn)
        self.conv = nn.Conv1d(d_rnn, d_rnn, conv_kernel, groups=d_rnn,
                              padding=conv_kernel - 1)
        self.rglru = RGLRU(d_rnn, rglru_c)
        self.out = nn.Linear(d_rnn, d_model)

    def forward(self, x):
        gate = F.gelu(self.in_gate(x))
        rec = self.in_recur(x).transpose(1, 2)
        rec = self.conv(rec)[..., : x.size(1)]
        rec = self.rglru(rec.transpose(1, 2))
        return self.out(gate * rec)


class MLPBlock(nn.Module):
    def __init__(self, d_model: int, expansion: int = 3):
        super().__init__()
        hidden = expansion * d_model
        self.gate = nn.Linear(d_model, hidden)
        self.up = nn.Linear(d_model, hidden)
        self.down = nn.Linear(hidden, d_model)

    def forward(self, x):
        return self.down(F.gelu(self.gate(x)) * self.up(x))


class HawkLayer(nn.Module):
    def __init__(self, d_model, d_rnn, conv_kernel, mlp_expansion, eps, rglru_c=8.0):
        super().__init__()
        self.norm1 = RMSNorm(d_model, eps)
        self.recur = RecurrentBlock(d_model, d_rnn, conv_kernel, rglru_c)
        self.norm2 = RMSNorm(d_model, eps)
        self.mlp = MLPBlock(d_model, mlp_expansion)

    def forward(self, x):
        x = x + self.recur(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x


class HawkConfig(PretrainedConfig):
    model_type = "hawk_rglru"

    def __init__(self, vocab_size: int = 16384, n_layer: int = 12, n_embd: int = 768,
                 rnn_width: Optional[int] = None, conv_kernel: int = 4,
                 mlp_expansion: int = 3, rmsnorm_eps: float = 1e-6, rglru_c: float = 8.0,
                 max_position_embeddings: int = 1024, tie_word_embeddings: bool = True,
                 bos_token_id: int = 2, eos_token_id: int = 3, pad_token_id: int = 1,
                 **kwargs):
        self.vocab_size = vocab_size
        self.n_layer = n_layer
        self.n_embd = n_embd
        self.rnn_width = rnn_width
        self.conv_kernel = conv_kernel
        self.mlp_expansion = mlp_expansion
        self.rmsnorm_eps = rmsnorm_eps
        self.rglru_c = rglru_c
        self.max_position_embeddings = max_position_embeddings
        self.auto_map = {
            "AutoConfig": "modeling_hawk.HawkConfig",
            "AutoModelForCausalLM": "modeling_hawk.HawkForCausalLM",
            "AutoModelForSequenceClassification": "modeling_hawk.HawkForSequenceClassification",
        }
        super().__init__(tie_word_embeddings=tie_word_embeddings,
                         bos_token_id=bos_token_id, eos_token_id=eos_token_id,
                         pad_token_id=pad_token_id, **kwargs)

    @property
    def d_rnn(self):
        return self.rnn_width if self.rnn_width is not None else self.n_embd


class HawkForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = HawkConfig
    _tied_weights_keys = {"lm_head.weight"}

    def __init__(self, config: HawkConfig):
        super().__init__(config)
        d_rnn = config.d_rnn
        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
        self.layers = nn.ModuleList([
            HawkLayer(config.n_embd, d_rnn, config.conv_kernel,
                      config.mlp_expansion, config.rmsnorm_eps, config.rglru_c)
            for _ in range(config.n_layer)])
        self.norm_f = RMSNorm(config.n_embd, config.rmsnorm_eps)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.post_init()
        self.lm_head.weight = self.wte.weight   # hard-tie: survives tie flag + 4.x/5.x

    def get_input_embeddings(self):
        return self.wte

    def set_input_embeddings(self, new):
        self.wte = new

    def get_output_embeddings(self):
        return self.lm_head

    def forward(self, input_ids: torch.LongTensor,
                attention_mask: Optional[torch.Tensor] = None,
                labels: Optional[torch.LongTensor] = None,
                **kwargs) -> CausalLMOutputWithPast:
        x = self.wte(input_ids)
        for layer in self.layers:
            x = layer(x)
        x = self.norm_f(x)
        logits = self.lm_head(x)
        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)
        return CausalLMOutputWithPast(loss=loss, logits=logits)

    # ----- generation support -------------------------------------------------
    # This backbone is stateless across calls (no KV cache): each step recomputes
    # the full prefix. generate() is therefore correct but O(T) per new token.
    # We override prepare_inputs_for_generation to always pass the full sequence
    # and never request a cache, so HF's default cache plumbing stays out of the
    # way regardless of the installed transformers version.
    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
        return {"input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False}

    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, **kwargs):
        # No cache to carry forward; just keep attention_mask growing if present.
        if model_kwargs.get("attention_mask") is not None:
            am = model_kwargs["attention_mask"]
            model_kwargs["attention_mask"] = torch.cat(
                [am, am.new_ones((am.shape[0], 1))], dim=-1)
        model_kwargs["past_key_values"] = None
        return model_kwargs


class HawkForSequenceClassification(PreTrainedModel):
    """
    Sequence-classification head on top of the SAME Hawk backbone used for
    causal LM. The backbone attribute names (wte / layers / norm_f) are
    IDENTICAL to HawkForCausalLM, so a CausalLM export state_dict maps 1:1 onto
    the backbone with no renaming. Only `score` is newly initialised, which is
    the expected behaviour when starting a fine-tuning run.

    The pooled representation is read from the hidden state at the last
    non-padding position (right padding, as produced by the BabyLM finetune
    tokenizer), matching the GPT-2 / Mamba sequence-classification convention.
    """

    config_class = HawkConfig

    def __init__(self, config: HawkConfig):
        super().__init__(config)
        self.num_labels = config.num_labels
        d_rnn = config.d_rnn
        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
        self.layers = nn.ModuleList([
            HawkLayer(config.n_embd, d_rnn, config.conv_kernel,
                      config.mlp_expansion, config.rmsnorm_eps, config.rglru_c)
            for _ in range(config.n_layer)])
        self.norm_f = RMSNorm(config.n_embd, config.rmsnorm_eps)
        self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.wte

    def set_input_embeddings(self, new):
        self.wte = new

    def forward(self, input_ids: torch.LongTensor,
                attention_mask: Optional[torch.Tensor] = None,
                labels: Optional[torch.LongTensor] = None,
                **kwargs) -> SequenceClassifierOutput:
        x = self.wte(input_ids)
        for layer in self.layers:
            x = layer(x)
        x = self.norm_f(x)
        logits = self.score(x)                       # (B, T, num_labels)

        B, T = input_ids.shape[:2]
        # Index of the last real token per sequence (assumes right padding).
        if attention_mask is not None:
            last_idx = attention_mask.long().sum(-1) - 1
        elif self.config.pad_token_id is not None:
            last_idx = (input_ids != self.config.pad_token_id).int().sum(-1) - 1
        else:
            last_idx = torch.full((B,), T - 1, device=input_ids.device)
        last_idx = last_idx.clamp(min=0)
        pooled_logits = logits[torch.arange(B, device=input_ids.device), last_idx]

        loss = None
        if labels is not None:
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and labels.dtype in (torch.long, torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = nn.MSELoss()
                loss = (loss_fct(pooled_logits.squeeze(), labels.squeeze())
                        if self.num_labels == 1
                        else loss_fct(pooled_logits, labels))
            elif self.config.problem_type == "single_label_classification":
                loss_fct = nn.CrossEntropyLoss()
                loss = loss_fct(pooled_logits.view(-1, self.num_labels),
                                labels.view(-1))
            else:  # multi_label_classification
                loss_fct = nn.BCEWithLogitsLoss()
                loss = loss_fct(pooled_logits, labels.float())

        return SequenceClassifierOutput(loss=loss, logits=pooled_logits)#!/usr/bin/env python3
"""
HF-compatible single-language Hawk / RG-LRU model, for lm-eval.

Self-contained `trust_remote_code` modeling file. The building blocks are the
SAME code used at training time (De et al., 2024, Griffin/Hawk; arXiv:2402.19427),
so the per-language export state_dict maps 1:1 onto this module's parameters
(top-level attribute names wte / layers / norm_f / lm_head match the export keys
exactly -- no renaming, no transpose). Exposes the standard
forward(input_ids, labels=None) -> CausalLMOutputWithPast that lm-eval expects.

Register via config.json:
  "model_type": "hawk_rglru",
  "architectures": ["HawkForCausalLM"],
  "auto_map": {
    "AutoConfig": "modeling_hawk.HawkConfig",
    "AutoModelForCausalLM": "modeling_hawk.HawkForCausalLM",
    "AutoModelForSequenceClassification": "modeling_hawk.HawkForSequenceClassification"
  }
"""

from typing import Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import (
    CausalLMOutputWithPast, SequenceClassifierOutput,
)


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        norm = torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return self.weight * (x * norm)


class RGLRU(nn.Module):
    """Real-Gated Linear Recurrent Unit (De et al., 2024)."""

    def __init__(self, width: int, c: float = 8.0):
        super().__init__()
        self.width = width
        self.c = c
        self.input_gate = nn.Linear(width, width)
        self.recur_gate = nn.Linear(width, width)
        lam = torch.empty(width).uniform_(2.197, 6.907)
        self.log_lambda = nn.Parameter(lam)

    def forward(self, x):                       # x: (B, T, W)
        B, T, W = x.shape
        r = torch.sigmoid(self.recur_gate(x))
        i = torch.sigmoid(self.input_gate(x))
        log_a = -F.softplus(-self.log_lambda)
        log_a_t = self.c * r * log_a
        a_t = torch.exp(log_a_t)
        mult = torch.sqrt(torch.clamp(-torch.expm1(2.0 * log_a_t), min=1e-8))
        gated_x = mult * (i * x)
        h = torch.zeros(B, W, device=x.device, dtype=x.dtype)
        outs = []
        for t in range(T):
            h = a_t[:, t] * h + gated_x[:, t]
            outs.append(h)
        return torch.stack(outs, dim=1)


class RecurrentBlock(nn.Module):
    def __init__(self, d_model: int, d_rnn: int, conv_kernel: int = 4, rglru_c: float = 8.0):
        super().__init__()
        self.conv_kernel = conv_kernel
        self.in_gate = nn.Linear(d_model, d_rnn)
        self.in_recur = nn.Linear(d_model, d_rnn)
        self.conv = nn.Conv1d(d_rnn, d_rnn, conv_kernel, groups=d_rnn,
                              padding=conv_kernel - 1)
        self.rglru = RGLRU(d_rnn, rglru_c)
        self.out = nn.Linear(d_rnn, d_model)

    def forward(self, x):
        gate = F.gelu(self.in_gate(x))
        rec = self.in_recur(x).transpose(1, 2)
        rec = self.conv(rec)[..., : x.size(1)]
        rec = self.rglru(rec.transpose(1, 2))
        return self.out(gate * rec)


class MLPBlock(nn.Module):
    def __init__(self, d_model: int, expansion: int = 3):
        super().__init__()
        hidden = expansion * d_model
        self.gate = nn.Linear(d_model, hidden)
        self.up = nn.Linear(d_model, hidden)
        self.down = nn.Linear(hidden, d_model)

    def forward(self, x):
        return self.down(F.gelu(self.gate(x)) * self.up(x))


class HawkLayer(nn.Module):
    def __init__(self, d_model, d_rnn, conv_kernel, mlp_expansion, eps, rglru_c=8.0):
        super().__init__()
        self.norm1 = RMSNorm(d_model, eps)
        self.recur = RecurrentBlock(d_model, d_rnn, conv_kernel, rglru_c)
        self.norm2 = RMSNorm(d_model, eps)
        self.mlp = MLPBlock(d_model, mlp_expansion)

    def forward(self, x):
        x = x + self.recur(self.norm1(x))
        x = x + self.mlp(self.norm2(x))
        return x


class HawkConfig(PretrainedConfig):
    model_type = "hawk_rglru"

    def __init__(self, vocab_size: int = 16384, n_layer: int = 12, n_embd: int = 768,
                 rnn_width: Optional[int] = None, conv_kernel: int = 4,
                 mlp_expansion: int = 3, rmsnorm_eps: float = 1e-6, rglru_c: float = 8.0,
                 max_position_embeddings: int = 1024, tie_word_embeddings: bool = True,
                 bos_token_id: int = 2, eos_token_id: int = 3, pad_token_id: int = 1,
                 **kwargs):
        self.vocab_size = vocab_size
        self.n_layer = n_layer
        self.n_embd = n_embd
        self.rnn_width = rnn_width
        self.conv_kernel = conv_kernel
        self.mlp_expansion = mlp_expansion
        self.rmsnorm_eps = rmsnorm_eps
        self.rglru_c = rglru_c
        self.max_position_embeddings = max_position_embeddings
        self.auto_map = {
            "AutoConfig": "modeling_hawk.HawkConfig",
            "AutoModelForCausalLM": "modeling_hawk.HawkForCausalLM",
            "AutoModelForSequenceClassification": "modeling_hawk.HawkForSequenceClassification",
        }
        super().__init__(tie_word_embeddings=tie_word_embeddings,
                         bos_token_id=bos_token_id, eos_token_id=eos_token_id,
                         pad_token_id=pad_token_id, **kwargs)

    @property
    def d_rnn(self):
        return self.rnn_width if self.rnn_width is not None else self.n_embd


class HawkForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = HawkConfig
    _tied_weights_keys = {"lm_head.weight": "wte.weight"}

    def __init__(self, config: HawkConfig):
        super().__init__(config)
        d_rnn = config.d_rnn
        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
        self.layers = nn.ModuleList([
            HawkLayer(config.n_embd, d_rnn, config.conv_kernel,
                      config.mlp_expansion, config.rmsnorm_eps, config.rglru_c)
            for _ in range(config.n_layer)])
        self.norm_f = RMSNorm(config.n_embd, config.rmsnorm_eps)
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.wte

    def set_input_embeddings(self, new):
        self.wte = new

    def get_output_embeddings(self):
        return self.lm_head

    def forward(self, input_ids: torch.LongTensor,
                attention_mask: Optional[torch.Tensor] = None,
                labels: Optional[torch.LongTensor] = None,
                **kwargs) -> CausalLMOutputWithPast:
        x = self.wte(input_ids)
        for layer in self.layers:
            x = layer(x)
        x = self.norm_f(x)
        logits = self.lm_head(x)
        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)
        return CausalLMOutputWithPast(loss=loss, logits=logits)

    # ----- generation support -------------------------------------------------
    # This backbone is stateless across calls (no KV cache): each step recomputes
    # the full prefix. generate() is therefore correct but O(T) per new token.
    # We override prepare_inputs_for_generation to always pass the full sequence
    # and never request a cache, so HF's default cache plumbing stays out of the
    # way regardless of the installed transformers version.
    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
        return {"input_ids": input_ids, "attention_mask": attention_mask, "use_cache": False}

    def _update_model_kwargs_for_generation(self, outputs, model_kwargs, **kwargs):
        # No cache to carry forward; just keep attention_mask growing if present.
        if model_kwargs.get("attention_mask") is not None:
            am = model_kwargs["attention_mask"]
            model_kwargs["attention_mask"] = torch.cat(
                [am, am.new_ones((am.shape[0], 1))], dim=-1)
        model_kwargs["past_key_values"] = None
        return model_kwargs


class HawkForSequenceClassification(PreTrainedModel):
    """
    Sequence-classification head on top of the SAME Hawk backbone used for
    causal LM. The backbone attribute names (wte / layers / norm_f) are
    IDENTICAL to HawkForCausalLM, so a CausalLM export state_dict maps 1:1 onto
    the backbone with no renaming. Only `score` is newly initialised, which is
    the expected behaviour when starting a fine-tuning run.

    The pooled representation is read from the hidden state at the last
    non-padding position (right padding, as produced by the BabyLM finetune
    tokenizer), matching the GPT-2 / Mamba sequence-classification convention.
    """

    config_class = HawkConfig

    def __init__(self, config: HawkConfig):
        super().__init__(config)
        self.num_labels = config.num_labels
        d_rnn = config.d_rnn
        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
        self.layers = nn.ModuleList([
            HawkLayer(config.n_embd, d_rnn, config.conv_kernel,
                      config.mlp_expansion, config.rmsnorm_eps, config.rglru_c)
            for _ in range(config.n_layer)])
        self.norm_f = RMSNorm(config.n_embd, config.rmsnorm_eps)
        self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.wte

    def set_input_embeddings(self, new):
        self.wte = new

    def forward(self, input_ids: torch.LongTensor,
                attention_mask: Optional[torch.Tensor] = None,
                labels: Optional[torch.LongTensor] = None,
                **kwargs) -> SequenceClassifierOutput:
        x = self.wte(input_ids)
        for layer in self.layers:
            x = layer(x)
        x = self.norm_f(x)
        logits = self.score(x)                       # (B, T, num_labels)

        B, T = input_ids.shape[:2]
        # Index of the last real token per sequence (assumes right padding).
        if attention_mask is not None:
            last_idx = attention_mask.long().sum(-1) - 1
        elif self.config.pad_token_id is not None:
            last_idx = (input_ids != self.config.pad_token_id).int().sum(-1) - 1
        else:
            last_idx = torch.full((B,), T - 1, device=input_ids.device)
        last_idx = last_idx.clamp(min=0)
        pooled_logits = logits[torch.arange(B, device=input_ids.device), last_idx]

        loss = None
        if labels is not None:
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and labels.dtype in (torch.long, torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = nn.MSELoss()
                loss = (loss_fct(pooled_logits.squeeze(), labels.squeeze())
                        if self.num_labels == 1
                        else loss_fct(pooled_logits, labels))
            elif self.config.problem_type == "single_label_classification":
                loss_fct = nn.CrossEntropyLoss()
                loss = loss_fct(pooled_logits.view(-1, self.num_labels),
                                labels.view(-1))
            else:  # multi_label_classification
                loss_fct = nn.BCEWithLogitsLoss()
                loss = loss_fct(pooled_logits, labels.float())

        return SequenceClassifierOutput(loss=loss, logits=pooled_logits)