ELECTRA › Generator Training

class ElectraGenerator(nn.Module):
    """Small MLM model that generates replacement tokens."""

    def __init__(
        self,
        vocab_size,
        hidden_size,
        num_layers,
        num_heads,
        intermediate_size,
        max_position=512,
        dropout=0.1,
    ):
        super().__init__()

        # Embeddings
        self.token_embedding = nn.Embedding(vocab_size, hidden_size)
        self.position_embedding = nn.Embedding(max_position, hidden_size)
        self.embedding_norm = nn.LayerNorm(hidden_size)
        self.embedding_dropout = nn.Dropout(dropout)

        # Transformer layers
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=hidden_size,
            nhead=num_heads,
            dim_feedforward=intermediate_size,
            dropout=dropout,
            activation="gelu",
            batch_first=True,
        )
        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers)

        # MLM head
        self.mlm_dense = nn.Linear(hidden_size, hidden_size)
        self.mlm_norm = nn.LayerNorm(hidden_size)
        self.mlm_output = nn.Linear(hidden_size, vocab_size)

        self.vocab_size = vocab_size

    def forward(self, input_ids, attention_mask=None):
        batch_size, seq_len = input_ids.shape

        # Embeddings
        positions = torch.arange(seq_len, device=input_ids.device).unsqueeze(0)
        x = self.token_embedding(input_ids) + self.position_embedding(positions)
        x = self.embedding_norm(x)
        x = self.embedding_dropout(x)

        # Create attention mask for transformer
        if attention_mask is not None:
            # Convert to format expected by PyTorch transformer
            src_key_padding_mask = attention_mask == 0
        else:
            src_key_padding_mask = None

        # Encode
        x = self.encoder(x, src_key_padding_mask=src_key_padding_mask)

        # MLM predictions
        x = self.mlm_dense(x)
        x = F.gelu(x)
        x = self.mlm_norm(x)
        logits = self.mlm_output(x)

        return logits

    def sample_replacements(self, input_ids, masked_positions, temperature=1.0):
        """Sample replacement tokens for masked positions."""
        logits = self(input_ids)

        # Get logits only at masked positions
        batch_size = input_ids.shape[0]
        sampled_tokens = input_ids.clone()

        for b in range(batch_size):
            mask_indices = masked_positions[b].nonzero(as_tuple=True)[0]
            if len(mask_indices) > 0:
                masked_logits = logits[b, mask_indices] / temperature
                probs = F.softmax(masked_logits, dim=-1)
                sampled = torch.multinomial(probs, num_samples=1).squeeze(-1)
                sampled_tokens[b, mask_indices] = sampled

        return sampled_tokens

def create_masked_input(input_ids, mask_token_id, mask_prob=0.15):
    """Create masked input and track masked positions."""
    masked_ids = input_ids.clone()
    masked_positions = torch.zeros_like(input_ids, dtype=torch.bool)

    # Sample positions to mask
    rand = torch.rand_like(input_ids, dtype=torch.float)
    masked_positions = rand < mask_prob

    # Replace with mask token
    masked_ids[masked_positions] = mask_token_id

    return masked_ids, masked_positions

# Create a small generator for demonstration
torch.manual_seed(42)
generator = ElectraGenerator(
    vocab_size=30522,
    hidden_size=256,
    num_layers=4,
    num_heads=4,
    intermediate_size=1024,
)

# Sample input
input_ids = torch.randint(100, 1000, (2, 32))  # Batch of 2, length 32
mask_token_id = 103

masked_ids, masked_positions = create_masked_input(
    input_ids, mask_token_id, mask_prob=0.15
)