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