| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
| from decision_lab.model_config import ScratchConfig |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, config: ScratchConfig): |
| super().__init__() |
| self.heads = config.heads |
| self.head_size = config.hidden_size // config.heads |
| self.dropout = config.dropout |
| self.qkv = nn.Linear(config.hidden_size, 3 * config.hidden_size) |
| self.output = nn.Linear(config.hidden_size, config.hidden_size) |
|
|
| def forward(self, hidden, attention_mask): |
| batch, length, width = hidden.shape |
| qkv = self.qkv(hidden).view(batch, length, 3, self.heads, self.head_size) |
| query, key, value = qkv.permute(2, 0, 3, 1, 4).unbind(0) |
| attended = F.scaled_dot_product_attention( |
| query, |
| key, |
| value, |
| attn_mask=attention_mask[:, None, None, :].bool(), |
| dropout_p=self.dropout if self.training else 0.0, |
| ) |
| return self.output(attended.transpose(1, 2).reshape(batch, length, width)) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, config: ScratchConfig): |
| super().__init__() |
| self.attention_norm = nn.LayerNorm(config.hidden_size) |
| self.attention = Attention(config) |
| self.feedforward_norm = nn.LayerNorm(config.hidden_size) |
| self.feedforward = nn.Sequential( |
| nn.Linear(config.hidden_size, config.intermediate_size), |
| nn.GELU(), |
| nn.Dropout(config.dropout), |
| nn.Linear(config.intermediate_size, config.hidden_size), |
| ) |
| self.dropout = nn.Dropout(config.dropout) |
|
|
| def forward(self, hidden, attention_mask): |
| hidden = hidden + self.dropout(self.attention(self.attention_norm(hidden), attention_mask)) |
| return hidden + self.dropout(self.feedforward(self.feedforward_norm(hidden))) |
|
|
|
|
| class ScratchDecisionModel(nn.Module): |
| def __init__(self, config: ScratchConfig): |
| super().__init__() |
| self.config = config |
| self.tokens = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_id) |
| self.positions = nn.Embedding(config.max_length, config.hidden_size) |
| self.blocks = nn.ModuleList(Block(config) for _ in range(config.layers)) |
| self.final_norm = nn.LayerNorm(config.hidden_size) |
| self.head = nn.Linear(config.hidden_size, config.classes) |
| self.apply(self.initialize) |
| with torch.no_grad(): |
| self.tokens.weight[config.pad_id].zero_() |
|
|
| def initialize(self, module): |
| if isinstance(module, (nn.Linear, nn.Embedding)): |
| nn.init.normal_(module.weight, std=self.config.initializer_std) |
| if isinstance(module, nn.Linear) and module.bias is not None: |
| nn.init.zeros_(module.bias) |
|
|
| def forward(self, input_ids, attention_mask): |
| length = input_ids.shape[1] |
| if length > self.config.max_length: |
| raise ValueError("Input exceeds scratch model's position limit") |
| if not torch.all(attention_mask[:, 0]): |
| raise ValueError("The first decision token cannot be padding") |
| positions = torch.arange(length, device=input_ids.device) |
| hidden = self.tokens(input_ids) + self.positions(positions) |
| for block in self.blocks: |
| hidden = block(hidden, attention_mask) |
| return self.head(self.final_norm(hidden[:, 0])) |
|
|
| def loss(self, logits, labels): |
| return F.cross_entropy(logits.float(), labels) |
|
|
| def class_logits(self, logits): |
| return logits |
|
|
