Initial release of CompactNameClassifier (INT8 ONNX + PyTorch + Tokenizer)

fefc1b3 verified 15 days ago

4.06 kB

import torch
import torch.nn as nn

class CompactNameClassifier(nn.Module):
    """
    Ultra-compact transformer binary classifier trained from scratch.
    Features:
      - Factorized Embeddings (vocab_size -> emb_dim -> hidden_dim) for ~50% embedding weight reduction
      - Pre-LN TransformerEncoder for fast, stable training dynamics
      - Masked Dual-Pooling: Concat(MeanPool, MaxPool) for strong anomaly/noise sensitivity
      - Single-task binary classification head: REAL_NAME (1) vs NON_NAME (0)
      - Zero BertPooler / token_type_ids bloat for clean ONNX execution
    """
    def __init__(
        self,
        vocab_size=4000,
        emb_dim=64,
        hidden_dim=128,
        num_layers=2,
        num_heads=4,
        dim_feedforward=256,
        max_len=20,
        dropout=0.1
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.emb_dim = emb_dim
        self.hidden_dim = hidden_dim
        self.max_len = max_len

        # Factorized Token Embeddings: cuts parameters from 4000x128=512k to 4000x64+64x128=264k
        self.token_emb = nn.Embedding(vocab_size, emb_dim, padding_idx=0)
        self.emb_proj = nn.Linear(emb_dim, hidden_dim, bias=False)
        self.pos_emb = nn.Embedding(max_len, hidden_dim)

        # Transformer Encoder with Pre-LN (norm_first=True)
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_dim,
            nhead=num_heads,
            dim_feedforward=dim_feedforward,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
            norm_first=True
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
        self.norm = nn.LayerNorm(hidden_dim)
        self.dropout = nn.Dropout(dropout)

        # Dual Pooling: 128 (MeanPool) + 128 (MaxPool) = 256
        self.classifier = nn.Linear(hidden_dim * 2, 2)

        self._init_weights()

    def _init_weights(self):
        """Standard truncated normal initialization for training from scratch."""
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Embedding):
                nn.init.trunc_normal_(m.weight, std=0.02)
                if m.padding_idx is not None:
                    with torch.no_grad():
                        m.weight[m.padding_idx].fill_(0)
            elif isinstance(m, nn.LayerNorm):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, input_ids, attention_mask):
        # input_ids: [batch_size, seq_len] (int32 or int64)
        # attention_mask: [batch_size, seq_len] (int32 or int64, 1 for real tokens, 0 for pad)
        seq_len = input_ids.size(1)
        positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)

        # Embeddings: token + positional
        x = self.emb_proj(self.token_emb(input_ids.long())) + self.pos_emb(positions)

        # PyTorch TransformerEncoder expects src_key_padding_mask=True for PAD tokens
        padding_mask = (attention_mask == 0)
        hidden = self.encoder(x, src_key_padding_mask=padding_mask)
        hidden = self.norm(hidden) # [batch_size, seq_len, hidden_dim]

        # 1. Masked Mean Pooling
        mask_expanded = attention_mask.unsqueeze(-1).float()
        sum_embeddings = torch.sum(hidden * mask_expanded, dim=1)
        sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
        mean_pooled = sum_embeddings / sum_mask

        # 2. Masked Max Pooling (padded positions set to -1e9 so they are never max)
        hidden_masked = hidden.masked_fill(padding_mask.unsqueeze(-1), -1e9)
        max_pooled = torch.max(hidden_masked, dim=1).values

        # Concatenate features
        features = torch.cat([mean_pooled, max_pooled], dim=-1)
        features = self.dropout(features)
        logits = self.classifier(features)
        return logits

	import torch
	import torch.nn as nn

	class CompactNameClassifier(nn.Module):
	"""
	Ultra-compact transformer binary classifier trained from scratch.
	Features:
	- Factorized Embeddings (vocab_size -> emb_dim -> hidden_dim) for ~50% embedding weight reduction
	- Pre-LN TransformerEncoder for fast, stable training dynamics
	- Masked Dual-Pooling: Concat(MeanPool, MaxPool) for strong anomaly/noise sensitivity
	- Single-task binary classification head: REAL_NAME (1) vs NON_NAME (0)
	- Zero BertPooler / token_type_ids bloat for clean ONNX execution
	"""
	def __init__(
	self,
	vocab_size=4000,
	emb_dim=64,
	hidden_dim=128,
	num_layers=2,
	num_heads=4,
	dim_feedforward=256,
	max_len=20,
	dropout=0.1
	):
	super().__init__()
	self.vocab_size = vocab_size
	self.emb_dim = emb_dim
	self.hidden_dim = hidden_dim
	self.max_len = max_len

	# Factorized Token Embeddings: cuts parameters from 4000x128=512k to 4000x64+64x128=264k
	self.token_emb = nn.Embedding(vocab_size, emb_dim, padding_idx=0)
	self.emb_proj = nn.Linear(emb_dim, hidden_dim, bias=False)
	self.pos_emb = nn.Embedding(max_len, hidden_dim)

	# Transformer Encoder with Pre-LN (norm_first=True)
	encoder_layer = nn.TransformerEncoderLayer(
	d_model=hidden_dim,
	nhead=num_heads,
	dim_feedforward=dim_feedforward,
	dropout=dropout,
	activation="gelu",
	batch_first=True,
	norm_first=True
	)
	self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
	self.norm = nn.LayerNorm(hidden_dim)
	self.dropout = nn.Dropout(dropout)

	# Dual Pooling: 128 (MeanPool) + 128 (MaxPool) = 256
	self.classifier = nn.Linear(hidden_dim * 2, 2)

	self._init_weights()

	def _init_weights(self):
	"""Standard truncated normal initialization for training from scratch."""
	for m in self.modules():
	if isinstance(m, nn.Linear):
	nn.init.trunc_normal_(m.weight, std=0.02)
	if m.bias is not None:
	nn.init.zeros_(m.bias)
	elif isinstance(m, nn.Embedding):
	nn.init.trunc_normal_(m.weight, std=0.02)
	if m.padding_idx is not None:
	with torch.no_grad():
	m.weight[m.padding_idx].fill_(0)
	elif isinstance(m, nn.LayerNorm):
	nn.init.ones_(m.weight)
	nn.init.zeros_(m.bias)

	def forward(self, input_ids, attention_mask):
	# input_ids: [batch_size, seq_len] (int32 or int64)
	# attention_mask: [batch_size, seq_len] (int32 or int64, 1 for real tokens, 0 for pad)
	seq_len = input_ids.size(1)
	positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)

	# Embeddings: token + positional
	x = self.emb_proj(self.token_emb(input_ids.long())) + self.pos_emb(positions)

	# PyTorch TransformerEncoder expects src_key_padding_mask=True for PAD tokens
	padding_mask = (attention_mask == 0)
	hidden = self.encoder(x, src_key_padding_mask=padding_mask)
	hidden = self.norm(hidden) # [batch_size, seq_len, hidden_dim]

	# 1. Masked Mean Pooling
	mask_expanded = attention_mask.unsqueeze(-1).float()
	sum_embeddings = torch.sum(hidden * mask_expanded, dim=1)
	sum_mask = torch.clamp(mask_expanded.sum(dim=1), min=1e-9)
	mean_pooled = sum_embeddings / sum_mask

	# 2. Masked Max Pooling (padded positions set to -1e9 so they are never max)
	hidden_masked = hidden.masked_fill(padding_mask.unsqueeze(-1), -1e9)
	max_pooled = torch.max(hidden_masked, dim=1).values

	# Concatenate features
	features = torch.cat([mean_pooled, max_pooled], dim=-1)
	features = self.dropout(features)
	logits = self.classifier(features)
	return logits