modeling_firstname_gender.py

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig
from transformers.modeling_outputs import CausalLMOutput


_SAGE_ATTN = None
_SAGE_AVAILABLE = False
_SAGE_IMPORT_ERROR = None

try:
    from sageattention import sageattn
    _SAGE_ATTN = sageattn
    _SAGE_AVAILABLE = True
except Exception as e:
    _SAGE_IMPORT_ERROR = repr(e)
    _SAGE_ATTN = None
    _SAGE_AVAILABLE = False


class FirstNameGenderConfig(PretrainedConfig):
    model_type = "firstname_gender"

    def __init__(
        self,
        vocab_size=32768,
        ctx_len=20,
        n_layer=4,
        n_head=4,
        n_embd=256,
        dropout=0.0,
        bos_token_id=None,
        eos_token_id=None,
        pad_token_id=0,
        F_ID=42,
        M_ID=49,
        attention_backend="sage",
        **kwargs,
    ):
        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            **kwargs,
        )

        self.vocab_size = vocab_size
        self.ctx_len = ctx_len
        self.max_position_embeddings = ctx_len
        self.n_layer = n_layer
        self.n_head = n_head
        self.n_embd = n_embd
        self.hidden_size = n_embd
        self.num_hidden_layers = n_layer
        self.num_attention_heads = n_head
        self.dropout = dropout
        self.F_ID = F_ID
        self.M_ID = M_ID
        self.attention_backend = attention_backend


def sageattention_available():
    return bool(_SAGE_AVAILABLE and _SAGE_ATTN is not None)


def sageattention_import_error():
    return _SAGE_IMPORT_ERROR


class CausalSelfAttention(nn.Module):
    def __init__(self, config):
        super().__init__()

        if config.n_embd % config.n_head != 0:
            raise ValueError("n_embd must be divisible by n_head")

        self.config = config
        self.n_head = config.n_head
        self.head_dim = config.n_embd // config.n_head

        self.qkv = nn.Linear(config.n_embd, 3 * config.n_embd, bias=False)
        self.proj = nn.Linear(config.n_embd, config.n_embd, bias=False)
        self.dropout = nn.Dropout(config.dropout)

        mask = torch.tril(torch.ones(config.ctx_len, config.ctx_len))
        self.register_buffer(
            "mask",
            mask.view(1, 1, config.ctx_len, config.ctx_len),
            persistent=False,
        )

    def _forward_sage(self, q, k, v):
        y = _SAGE_ATTN(
            q,
            k,
            v,
            tensor_layout="HND",
            is_causal=True,
        )

        return y

    def _forward_pytorch(self, q, k, v, t):
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        scores = scores.masked_fill(
            self.mask[:, :, :t, :t] == 0,
            torch.finfo(scores.dtype).min,
        )

        att = F.softmax(scores, dim=-1)
        att = self.dropout(att)

        y = att @ v

        return y

    def forward(self, x):
        b, t, c = x.shape

        qkv = self.qkv(x)
        q, k, v = qkv.chunk(3, dim=-1)

        q = q.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()
        k = k.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()
        v = v.view(b, t, self.n_head, self.head_dim).transpose(1, 2).contiguous()

        use_sage = (
            getattr(self.config, "attention_backend", "sage") == "sage"
            and sageattention_available()
            and q.is_cuda
        )

        if use_sage:
            try:
                y = self._forward_sage(q, k, v)
            except Exception:
                y = self._forward_pytorch(q, k, v, t)
        else:
            y = self._forward_pytorch(q, k, v, t)

        y = y.transpose(1, 2).contiguous().view(b, t, c)
        y = self.proj(y)

        return y


class MLP(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
        self.proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
        self.dropout = nn.Dropout(config.dropout)

    def forward(self, x):
        x = self.fc(x)
        x = F.gelu(x)
        x = self.proj(x)
        x = self.dropout(x)

        return x


class Block(nn.Module):
    def __init__(self, config):
        super().__init__()

        self.ln1 = nn.LayerNorm(config.n_embd)
        self.attn = CausalSelfAttention(config)
        self.ln2 = nn.LayerNorm(config.n_embd)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        x = x + self.mlp(self.ln2(x))

        return x


class FirstNameGenderForCausalLM(PreTrainedModel):
    config_class = FirstNameGenderConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    all_tied_weights_keys = {}
    
    def __init__(self, config):
        super().__init__(config)

        self.config = config

        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
        self.pos_emb = nn.Embedding(config.ctx_len, config.n_embd)
        self.drop = nn.Dropout(config.dropout)

        self.blocks = nn.ModuleList(
            [Block(config) for _ in range(config.n_layer)]
        )

        self.ln_f = nn.LayerNorm(config.n_embd)
        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)

        self.head.weight = self.tok_emb.weight

    def get_input_embeddings(self):
        return self.tok_emb

    def set_input_embeddings(self, value):
        self.tok_emb = value
        self.head.weight = self.tok_emb.weight

    def get_output_embeddings(self):
        return self.head

    def set_output_embeddings(self, new_embeddings):
        self.head = new_embeddings

    def get_attention_backend(self):
        if getattr(self.config, "attention_backend", "sage") == "sage" and sageattention_available():
            return "sageattention"

        return "pytorch"

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None,
        **kwargs,
    ):
        if input_ids is None:
            raise ValueError("input_ids is required")

        b, t = input_ids.shape

        if t > self.config.ctx_len:
            input_ids = input_ids[:, -self.config.ctx_len:]
            t = input_ids.shape[1]

            if labels is not None:
                labels = labels[:, -self.config.ctx_len:]

        pos = torch.arange(
            0,
            t,
            dtype=torch.long,
            device=input_ids.device,
        ).unsqueeze(0)

        x = self.tok_emb(input_ids) + self.pos_emb(pos)
        x = self.drop(x)

        for block in self.blocks:
            x = block(x)

        x = self.ln_f(x)
        logits = self.head(x)

        loss = None

        if labels is not None:
            loss = F.cross_entropy(
                logits.reshape(-1, logits.size(-1)),
                labels.reshape(-1),
                ignore_index=-100,
            )

        return CausalLMOutput(
            loss=loss,
            logits=logits,
        )

    @torch.no_grad()
    def predict_gender(self, input_ids):
        out = self.forward(input_ids=input_ids)
        logits = out.logits

        non_pad = input_ids.ne(self.config.pad_token_id)
        lengths = non_pad.sum(dim=1).clamp(min=1)
        last_pos = lengths - 1

        batch_idx = torch.arange(input_ids.size(0), device=input_ids.device)
        last_logits = logits[batch_idx, last_pos, :]

        fm_logits = torch.stack(
            [
                last_logits[:, self.config.F_ID],
                last_logits[:, self.config.M_ID],
            ],
            dim=-1,
        )

        probs = F.softmax(fm_logits.float(), dim=-1)
        pred_idx = probs.argmax(dim=-1)

        return pred_idx, probs