|
|
import torch |
|
|
import torch.nn as nn |
|
|
import torch.nn.functional as F |
|
|
import math |
|
|
from transformers import PreTrainedModel, PretrainedConfig |
|
|
from transformers.modeling_outputs import CausalLMOutputWithPast |
|
|
|
|
|
class GPTConfig(PretrainedConfig): |
|
|
model_type = "babylang" |
|
|
|
|
|
def __init__( |
|
|
self, |
|
|
vocab_size=50257, |
|
|
block_size=128, |
|
|
n_layer=6, |
|
|
n_head=6, |
|
|
n_embd=384, |
|
|
dropout=0.0, |
|
|
bias=True, |
|
|
**kwargs |
|
|
): |
|
|
super().__init__(**kwargs) |
|
|
self.vocab_size = vocab_size |
|
|
self.block_size = block_size |
|
|
self.n_layer = n_layer |
|
|
self.n_head = n_head |
|
|
self.n_embd = n_embd |
|
|
self.dropout = dropout |
|
|
self.bias = bias |
|
|
|
|
|
class LayerNorm(nn.Module): |
|
|
def __init__(self, ndim, bias): |
|
|
super().__init__() |
|
|
self.weight = nn.Parameter(torch.ones(ndim)) |
|
|
self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None |
|
|
|
|
|
def forward(self, x): |
|
|
return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5) |
|
|
|
|
|
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=config.bias) |
|
|
self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias) |
|
|
self.attn_dropout = nn.Dropout(config.dropout) |
|
|
self.resid_dropout = nn.Dropout(config.dropout) |
|
|
self.n_head = config.n_head |
|
|
self.n_embd = config.n_embd |
|
|
self.flash = hasattr(F, 'scaled_dot_product_attention') |
|
|
if not self.flash: |
|
|
self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size)).view(1, 1, config.block_size, config.block_size)) |
|
|
|
|
|
def forward(self, x, layer_past=None): |
|
|
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) |
|
|
q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
|
|
v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) |
|
|
|
|
|
if layer_past is not None: |
|
|
past_key, past_value = layer_past |
|
|
k = torch.cat((past_key, k), dim=-2) |
|
|
v = torch.cat((past_value, v), dim=-2) |
|
|
|
|
|
present = (k, v) |
|
|
|
|
|
if self.flash: |
|
|
y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.attn_dropout.p if self.training else 0.0, is_causal=True) |
|
|
else: |
|
|
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 |
|
|
|
|
|
y = y.transpose(1, 2).contiguous().view(B, T, C) |
|
|
y = self.resid_dropout(self.c_proj(y)) |
|
|
return y, present |
|
|
|
|
|
class MLP(nn.Module): |
|
|
def __init__(self, config): |
|
|
super().__init__() |
|
|
self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias) |
|
|
self.gelu = nn.GELU() |
|
|
self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias) |
|
|
self.dropout = nn.Dropout(config.dropout) |
|
|
|
|
|
def forward(self, x): |
|
|
return self.dropout(self.c_proj(self.gelu(self.c_fc(x)))) |
|
|
|
|
|
class Block(nn.Module): |
|
|
def __init__(self, config): |
|
|
super().__init__() |
|
|
self.ln1 = LayerNorm(config.n_embd, config.bias) |
|
|
self.attn = CausalSelfAttention(config) |
|
|
self.ln2 = LayerNorm(config.n_embd, config.bias) |
|
|
self.mlp = MLP(config) |
|
|
|
|
|
def forward(self, x, layer_past=None): |
|
|
x = x + self.attn(self.ln1(x), layer_past=layer_past) |
|
|
x = x + self.mlp(self.ln2(x)) |
|
|
return x |
|
|
|
|
|
class GPT(PreTrainedModel): |
|
|
config_class = GPTConfig |
|
|
base_model_prefix = "babylang" |
|
|
|
|
|
def __init__(self, config): |
|
|
super().__init__(config) |
|
|
self.transformer = nn.ModuleDict(dict( |
|
|
wte=nn.Embedding(config.vocab_size, config.n_embd), |
|
|
wpe=nn.Embedding(config.block_size, config.n_embd), |
|
|
drop=nn.Dropout(config.dropout), |
|
|
h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]), |
|
|
ln_f=LayerNorm(config.n_embd, config.bias), |
|
|
)) |
|
|
self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) |
|
|
self.transformer.wte.weight = self.lm_head.weight |
|
|
|
|
|
self.apply(self._init_weights) |
|
|
for pn, p in self.named_parameters(): |
|
|
if pn.endswith('c_proj.weight'): |
|
|
nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layer)) |
|
|
|
|
|
def _init_weights(self, module): |
|
|
if isinstance(module, nn.Linear): |
|
|
nn.init.normal_(module.weight, mean=0.0, std=0.02) |
|
|
if module.bias is not None: |
|
|
nn.init.zeros_(module.bias) |
|
|
elif isinstance(module, nn.Embedding): |
|
|
nn.init.normal_(module.weight, mean=0.0, std=0.02) |
|
|
|
|
|
def forward(self, input_ids, past_key_values=None, attention_mask=None, labels=None): |
|
|
device = input_ids.device |
|
|
b, t = input_ids.size() |
|
|
assert t <= self.config.block_size |
|
|
pos = torch.arange(0, t, dtype=torch.long, device=device) |
|
|
|
|
|
if past_key_values is not None: |
|
|
pos = pos[-1].unsqueeze(0) |
|
|
|
|
|
tok_emb = self.transformer.wte(input_ids) |
|
|
pos_emb = self.transformer.wpe(pos) |
|
|
x = self.transformer.drop(tok_emb + pos_emb) |
|
|
|
|
|
new_past_key_values = [] |
|
|
for i, block in enumerate(self.transformer.h): |
|
|
x, past = block(x, layer_past=past_key_values[i] if past_key_values is not None else None) |
|
|
new_past_key_values.append(past) |
|
|
|
|
|
x = self.transformer.ln_f(x) |
|
|
logits = self.lm_head(x) |
|
|
|
|
|
loss = None |
|
|
if labels is not None: |
|
|
loss = F.cross_entropy(logits.view(-1, logits.size(-1)), labels.view(-1), ignore_index=-1) |
|
|
|
|
|
return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=new_past_key_values) |
|
|
|
|
|
def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs): |
|
|
if past_key_values: |
|
|
input_ids = input_ids[:, -1].unsqueeze(-1) |
|
|
return {"input_ids": input_ids, "past_key_values": past_key_values} |
|
|
|
|
|
@torch.no_grad() |
|
|
def generate(self, input_ids, max_new_tokens, temperature=1.0, top_k=None): |
|
|
for _ in range(max_new_tokens): |
|
|
idx_cond = input_ids if input_ids.size(1) <= self.config.block_size else input_ids[:, -self.config.block_size:] |
|
|
out = self(idx_cond) |
|
|
logits = out['logits'][:, -1, :] / temperature |
|
|
if top_k is not None: |
|
|
v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
|
|
logits[logits < v[:, [-1]]] = -float('Inf') |
|
|
probs = F.softmax(logits, dim=-1) |
|
|
idx_next = torch.multinomial(probs, num_samples=1) |
|
|
input_ids = torch.cat((input_ids, idx_next), dim=1) |
|
|
return input_ids |
|
|
|
