"""NanoGPT model implementation for HuggingFace""" import torch import torch.nn as nn import torch.nn.functional as F import math from transformers import PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions from transformers.utils import logging # Import configuration try: from .configuration_nanogpt import NanoGPTConfig except ImportError: from configuration_nanogpt import NanoGPTConfig logger = logging.get_logger(__name__) class ExactNanoGPTAttention(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.dropout = config.dropout self.flash = hasattr(torch.nn.functional, '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): 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 self.flash: y = torch.nn.functional.scaled_dot_product_attention( q, k, v, attn_mask=None, dropout_p=self.dropout if self.training else 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 class ExactNanoGPTMLP(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): x = self.c_fc(x) x = self.gelu(x) x = self.c_proj(x) x = self.dropout(x) return x class ExactNanoGPTBlock(nn.Module): def __init__(self, config): super().__init__() self.ln_1 = nn.LayerNorm(config.n_embd, bias=config.bias) self.attn = ExactNanoGPTAttention(config) self.ln_2 = nn.LayerNorm(config.n_embd, bias=config.bias) self.mlp = ExactNanoGPTMLP(config) def forward(self, x): x = x + self.attn(self.ln_1(x)) x = x + self.mlp(self.ln_2(x)) return x class NanoGPTPreTrainedModel(PreTrainedModel): """Base class for NanoGPT models""" config_class = NanoGPTConfig base_model_prefix = "transformer" supports_gradient_checkpointing = False _no_split_modules = ["ExactNanoGPTBlock"] def _init_weights(self, module): if isinstance(module, nn.Linear): torch.nn.init.normal_(module.weight, mean=0.0, std=0.02) 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=0.02) elif isinstance(module, nn.LayerNorm): torch.nn.init.zeros_(module.bias) torch.nn.init.ones_(module.weight) class NanoGPTModel(NanoGPTPreTrainedModel): """The main NanoGPT model""" def __init__(self, config): super().__init__(config) self.config = 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([ExactNanoGPTBlock(config) for _ in range(config.n_layer)]), ln_f = nn.LayerNorm(config.n_embd, bias=config.bias), )) self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) # Initialize weights 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 forward(self, input_ids, attention_mask=None, **kwargs): device = input_ids.device b, t = input_ids.size() assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}" pos = torch.arange(0, t, dtype=torch.long, device=device) tok_emb = self.transformer.wte(input_ids) pos_emb = self.transformer.wpe(pos) x = self.transformer.drop(tok_emb + pos_emb) for block in self.transformer.h: x = block(x) x = self.transformer.ln_f(x) logits = self.lm_head(x) return CausalLMOutputWithCrossAttentions(logits=logits) def generate(self, input_ids, max_length=None, max_new_tokens=None, temperature=1.0, top_k=None, do_sample=True, top_p=None, pad_token_id=None, eos_token_id=None, **kwargs): if max_new_tokens is None: max_new_tokens = max_length - input_ids.shape[1] if max_length else 50 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:] with torch.no_grad(): outputs = self(idx_cond) logits = outputs.logits logits = 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') if top_p is not None: sorted_logits, sorted_indices = torch.sort(logits, descending=True) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) sorted_indices_to_remove = cumulative_probs > top_p sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] = 0 indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove) logits[indices_to_remove] = -float('Inf') probs = F.softmax(logits, dim=-1) if do_sample: idx_next = torch.multinomial(probs, num_samples=1) else: _, idx_next = torch.topk(probs, k=1, dim=-1) input_ids = torch.cat((input_ids, idx_next), dim=1) return input_ids # For backward compatibility NanoGPTForCausalLM = NanoGPTModel