Class RecurrentGPT2Block

gpt_utils_systematicity.py:141–185

class RecurrentGPT2Block(nn.Module):
    def __init__(self, config, num_iterations):
        super().__init__()
        self.config = config
        self.num_iterations = num_iterations
        self.blocks = nn.ModuleList([GPT2Block(config) for _ in range(config.n_layer)])

        self.token_embedding = nn.Embedding(config.vocab_size, config.n_embd)
        self.position_embedding = nn.Embedding(config.n_positions, config.n_embd)

        embd_pdrop = getattr(config, "embd_pdrop", 0.1)
        self.dropout = nn.Dropout(embd_pdrop)

        layer_norm_epsilon = getattr(config, "layer_norm_epsilon", 1e-5)
        self.ln_f = nn.LayerNorm(config.n_embd, eps=layer_norm_epsilon)

        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
        self.lm_head.weight = self.token_embedding.weight

    def forward(self, input_ids, attention_mask=None):

        batch_size, seq_len = input_ids.size()
        device = input_ids.device

        position_ids = torch.arange(0, seq_len, dtype=torch.long, device=device)
        position_ids = position_ids.unsqueeze(0).expand(batch_size, seq_len)

        token_embeds = self.token_embedding(input_ids)
        pos_embeds = self.position_embedding(position_ids)
        hidden_states = token_embeds + pos_embeds
        hidden_states = self.dropout(hidden_states)

        if attention_mask is not None:
            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
            attention_mask = attention_mask.to(dtype=hidden_states.dtype)
            attention_mask = (1.0 - attention_mask) * -10000.0

        for _ in range(self.num_iterations):
            for block in self.blocks:
                hidden_states = block(hidden_states, attention_mask=attention_mask)[0]

        hidden_states = self.ln_f(hidden_states)
        logits = self.lm_head(hidden_states)

        return type("Output", (object,), {"logits": logits})