File size: 6,696 Bytes
fab62cc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 | import math
from typing import Optional, Tuple, Dict, Any
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, GenerationMixin, AutoConfig, AutoModel, AutoModelForCausalLM
try:
from .configuration_gpt2 import GPT2CustomConfig
except ImportError:
from configuration_gpt2 import GPT2CustomConfig
class CausalSelfAttention(nn.Module):
def __init__(self, config):
super().__init__()
assert config.n_embd % config.n_head == 0
self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=True)
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=True)
self.attn_dropout = nn.Dropout(config.attn_pdrop)
self.resid_dropout = nn.Dropout(config.resid_pdrop)
self.n_head = config.n_head
self.n_embd = config.n_embd
# Causal mask buffer
self.register_buffer("bias", torch.tril(torch.ones(config.n_positions, config.n_positions))
.view(1, 1, config.n_positions, config.n_positions))
def forward(self, x):
B, T, C = x.size()
q, k, v = self.c_attn(x).split(self.n_embd, dim=2)
k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
att = F.softmax(att, dim=-1)
att = self.attn_dropout(att)
y = att @ v # (B, nh, T, hs)
y = y.transpose(1, 2).contiguous().view(B, T, C)
y = self.resid_dropout(self.c_proj(y))
return y
class MLP(nn.Module):
def __init__(self, config):
super().__init__()
self.c_fc = nn.Linear(config.n_embd, config.n_inner, bias=True)
self.c_proj = nn.Linear(config.n_inner, config.n_embd, bias=True)
self.act = nn.GELU(approximate="tanh")
self.dropout = nn.Dropout(config.resid_pdrop)
def forward(self, x):
return self.dropout(self.c_proj(self.act(self.c_fc(x))))
class GPT2Block(nn.Module):
def __init__(self, config):
super().__init__()
self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
self.attn = CausalSelfAttention(config)
self.ln_2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
self.mlp = MLP(config)
def forward(self, x):
x = x + self.attn(self.ln_1(x))
x = x + self.mlp(self.ln_2(x))
return x
class GPT2CustomModel(PreTrainedModel):
config_class = GPT2CustomConfig
base_model_prefix = "transformer"
def __init__(self, config):
super().__init__(config)
self.wte = nn.Embedding(config.vocab_size, config.n_embd)
self.wpe = nn.Embedding(config.n_positions, config.n_embd)
self.drop = nn.Dropout(config.embd_pdrop)
self.h = nn.ModuleList([GPT2Block(config) for _ in range(config.n_layer)])
self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
self.post_init()
def _init_weights(self, module):
if isinstance(module, nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
def forward(self, input_ids, **kwargs):
device = input_ids.device
B, T = input_ids.size()
pos = torch.arange(0, T, dtype=torch.long, device=device).unsqueeze(0)
x = self.wte(input_ids) + self.wpe(pos)
x = self.drop(x)
for block in self.h:
x = block(x)
x = self.ln_f(x)
from transformers.modeling_outputs import BaseModelOutputWithPast
return BaseModelOutputWithPast(last_hidden_state=x)
class GPT2CustomLMHeadModel(PreTrainedModel, GenerationMixin):
config_class = GPT2CustomConfig
base_model_prefix = "transformer"
_tied_weights_keys = ["lm_head.weight"]
def __init__(self, config):
super().__init__(config)
self.transformer = GPT2CustomModel(config)
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
self.transformer.wte.weight = self.lm_head.weight # Weight tying
self.post_init()
def get_input_embeddings(self):
return self.transformer.wte
def set_input_embeddings(self, new_embeddings):
self.transformer.wte = new_embeddings
def get_output_embeddings(self):
return self.lm_head
def set_output_embeddings(self, new_embeddings):
self.lm_head = new_embeddings
def tie_weights(self, **kwargs):
self.lm_head.weight = self.transformer.wte.weight
def _init_weights(self, module):
if isinstance(module, nn.Linear):
torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
if module.bias is not None:
torch.nn.init.zeros_(module.bias)
elif isinstance(module, nn.Embedding):
torch.nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
def forward(self, input_ids, labels=None, **kwargs):
transformer_outputs = self.transformer(input_ids)
hidden_states = transformer_outputs.last_hidden_state
logits = self.lm_head(hidden_states)
loss = None
if labels is not None:
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1), ignore_index=-100)
from transformers.modeling_outputs import CausalLMOutputWithPast
return CausalLMOutputWithPast(
loss=loss,
logits=logits
)
def prepare_inputs_for_generation(self, input_ids, **kwargs):
return {"input_ids": input_ids}
# Register configuration and model for auto mapping
AutoConfig.register("gpt2_custom", GPT2CustomConfig)
AutoModel.register(GPT2CustomConfig, GPT2CustomModel)
AutoModelForCausalLM.register(GPT2CustomConfig, GPT2CustomLMHeadModel)
|