| import math | |
| import torch | |
| import torch.nn as nn | |
| from transformer_lens.HookedTransformer import HookedTransformer | |
| from transformer_lens.HookedTransformerConfig import HookedTransformerConfig | |
| from transformer_lens.train import HookedTransformerTrainConfig, train | |
| from transformers import PretrainedConfig, PreTrainedModel | |
| def generate_config( | |
| n_ctx, | |
| d_model, | |
| d_head, | |
| n_heads, | |
| d_mlp, | |
| n_layers, | |
| attention_dir, | |
| act_fn, | |
| d_vocab, | |
| d_vocab_out, | |
| use_attn_result, | |
| device, | |
| use_hook_tokens, | |
| ): | |
| return HookedTransformerConfig( | |
| n_ctx=n_ctx, | |
| d_model=d_model, | |
| d_head=d_head, | |
| n_heads=n_heads, | |
| d_mlp=d_mlp, | |
| n_layers=n_layers, | |
| attention_dir=attention_dir, | |
| act_fn=act_fn, | |
| d_vocab=d_vocab, | |
| d_vocab_out=d_vocab_out, | |
| use_attn_result=use_attn_result, | |
| device=device, | |
| use_hook_tokens=use_hook_tokens, | |
| ) | |
| def generate_model(config): | |
| return HookedTransformer(config) | |
| def train_model(model, n_epochs, batch_size, lr, dataset): | |
| train_cfg = HookedTransformerTrainConfig( | |
| num_epochs=n_epochs, batch_size=128, lr=0.001, device="cuda:0" | |
| ) | |
| return train(model, train_cfg, dataset) | |
| class ScaledDotProductAttention(nn.Module): | |
| def __init__(self, scale): | |
| super().__init__() | |
| self.scale = scale | |
| def forward(self, q, k, v, mask=None): | |
| attn = torch.matmul(q, k.transpose(-2, -1)) * 1 / self.scale | |
| if mask is not None: | |
| attn = attn.masked_fill(mask == 0, float("-inf")) | |
| attn = torch.softmax(attn, dim=-1) | |
| out = torch.matmul(attn, v) | |
| return out, attn | |
| class MultiHeadAttention(nn.Module): | |
| def __init__(self, n_heads, d_model): | |
| super().__init__() | |
| assert d_model % n_heads == 0, "d_model should be divisible by n_heads" | |
| self.d_model = d_model | |
| self.n_heads = n_heads | |
| self.depth = d_model // n_heads | |
| self.wq = nn.Linear(d_model, d_model) | |
| self.wk = nn.Linear(d_model, d_model) | |
| self.wv = nn.Linear(d_model, d_model) | |
| self.dense = nn.Linear(d_model, d_model) | |
| self.attn = ScaledDotProductAttention(scale=math.sqrt(self.depth)) | |
| def forward(self, q, k, v, mask=None): | |
| batch_size = q.size(0) | |
| q = self.wq(q).view(batch_size, -1, self.n_heads, self.depth).transpose(1, 2) | |
| k = self.wk(k).view(batch_size, -1, self.n_heads, self.depth).transpose(1, 2) | |
| v = self.wv(v).view(batch_size, -1, self.n_heads, self.depth).transpose(1, 2) | |
| attn_out, _ = self.attn(q, k, v, mask=mask) | |
| attn_out = ( | |
| attn_out.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model) | |
| ) | |
| out = self.dense(attn_out) | |
| return out | |
| class TransformerEncoderLayer(nn.Module): | |
| def __init__(self, d_model, n_heads, ff_dim, dropout=0.1): | |
| super().__init__() | |
| self.attn = MultiHeadAttention(n_heads, d_model) | |
| self.ff = nn.Sequential( | |
| nn.Linear(d_model, ff_dim), | |
| nn.ReLU(), | |
| nn.Linear(ff_dim, d_model), | |
| ) | |
| self.ln1 = nn.LayerNorm(d_model) | |
| self.ln2 = nn.LayerNorm(d_model) | |
| self.dropout = nn.Dropout(dropout) | |
| def forward(self, x, mask=None): | |
| attn_out = self.attn(x, x, x, mask=mask) | |
| x = self.ln1(x + self.dropout(attn_out)) | |
| ff_out = self.ff(x) | |
| x = self.ln2(x + self.dropout(ff_out)) | |
| return x | |
| class TransformerClassifierConfig(PretrainedConfig): | |
| model_type = "transformer-checker" | |
| def __init__( | |
| self, | |
| in_dim=512, | |
| d_model=256, | |
| n_heads=8, | |
| ff_dim=2048, | |
| n_layers=6, | |
| n_classes=2, | |
| **kwargs, | |
| ): | |
| self.in_dim = in_dim | |
| self.d_model = d_model | |
| self.n_heads = n_heads | |
| self.ff_dim = ff_dim | |
| self.n_layers = n_layers | |
| self.n_classes = n_classes | |
| super().__init__(**kwargs) | |
| class TransformerClassifier(PreTrainedModel): | |
| config_class = TransformerClassifierConfig | |
| def __init__(self, config: TransformerClassifierConfig): | |
| super().__init__(config) | |
| self.embedding = nn.Linear(config.in_dim, config.d_model) | |
| self.encoders = nn.ModuleList( | |
| [ | |
| TransformerEncoderLayer(config.d_model, config.n_heads, config.ff_dim) | |
| for _ in range(config.n_layers) | |
| ] | |
| ) | |
| self.classifier = nn.Linear(config.d_model, config.n_classes) | |
| def forward(self, x, mask=None): | |
| x = self.embedding(x) | |
| for encoder in self.encoders: | |
| x = encoder(x, mask=mask) | |
| x = self.classifier(x[:, 0]) | |
| return x | |