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