Upload folder using huggingface_hub

config.json

@@ -0,0 +1,30 @@
+{
+  "F_ID": 42,
+  "M_ID": 49,
+  "architectures": [
+    "FirstNameGenderForCausalLM"
+  ],
+  "attention_backend": "pytorch",
+  "auto_map": {
+    "AutoConfig": "modeling_firstname_gender.FirstNameGenderConfig",
+    "AutoModel": "modeling_firstname_gender.FirstNameGenderForCausalLM",
+    "AutoModelForCausalLM": "modeling_firstname_gender.FirstNameGenderForCausalLM"
+  },
+  "bos_token_id": 2,
+  "ctx_len": 20,
+  "dropout": 0.0,
+  "dtype": "float32",
+  "eos_token_id": 3,
+  "hidden_size": 384,
+  "max_position_embeddings": 20,
+  "model_type": "firstname_gender",
+  "n_embd": 384,
+  "n_head": 4,
+  "n_layer": 4,
+  "num_attention_heads": 4,
+  "num_hidden_layers": 4,
+  "pad_token_id": 0,
+  "transformers_version": "5.10.2",
+  "use_cache": false,
+  "vocab_size": 32768
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:706d118f19b2808209911f99dbe6ac495ead204370cf8358fd3210bde43117c1
+size 129036560

modeling_firstname_gender.py

@@ -0,0 +1,297 @@
+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
+
+    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

tokenizer_config.json

@@ -0,0 +1,20 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "[BOS]",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "[EOS]",
+  "is_local": false,
+  "local_files_only": false,
+  "lowercase": false,
+  "mask_token": "[MASK]",
+  "max_length": null,
+  "model_max_length": 20,
+  "pad_to_multiple_of": null,
+  "pad_token": "[PAD]",
+  "pad_token_type_id": 0,
+  "padding_side": "right",
+  "tokenizer_class": "TokenizersBackend",
+  "truncation_side": "right",
+  "unk_token": "[UNK]",
+  "vocab_size": 32768
+}

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:05727dced653438df160bb8641a235f1f6a29b67eb049c96044874005d6eb0e8
+size 5201