Spaces:
Build error
Build error
| import torch | |
| import torch.nn as nn | |
| from torch.nn import functional as F | |
| from .block import Block | |
| import sys | |
| import os | |
| # Add the project root to Python path | |
| sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) | |
| from src.config.model_config import GPTConfig | |
| class GPT(nn.Module): | |
| def __init__(self, config): | |
| super().__init__() | |
| 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), | |
| h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]), | |
| ln_f = nn.LayerNorm(config.n_embd), | |
| )) | |
| self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False) | |
| # weight sharing | |
| self.transformer.wte.weight = self.lm_head.weight | |
| # weight initialization | |
| self.apply(self._init_weights) | |
| def _init_weights(self, module): | |
| if isinstance(module, nn.Linear): | |
| std = 0.02 | |
| if hasattr(module, 'NANGPT_SCALE_INIT'): | |
| std *= (2 * self.config.n_layer) ** -0.5 | |
| torch.nn.init.normal_(module.weight, mean=0.0, std=std) | |
| 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) | |
| def forward(self, idx, targets=None): | |
| B, T = idx.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=idx.device) | |
| pos_emb = self.transformer.wpe(pos) | |
| tok_emb = self.transformer.wte(idx) | |
| x = tok_emb + pos_emb | |
| for block in self.transformer.h: | |
| x = block(x) | |
| x = self.transformer.ln_f(x) | |
| logits = self.lm_head(x) | |
| loss = None | |
| if targets is not None: | |
| loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1)) | |
| return logits, loss | |
| def from_pretrained(cls, model_type): | |
| """Loads pretrained GPT-2 model weights from huggingface""" | |
| assert model_type in {'gpt2', 'gpt2-medium', 'gpt2-large', 'gpt2-xl'} | |
| from transformers import GPT2LMHeadModel | |
| print("loading weights from pretrained gpt: %s" % model_type) | |
| # n_layer, n_head and n_embd are determined from model_type | |
| config_args = { | |
| 'gpt2': dict(n_layer=12, n_head=12, n_embd=768), # 124M params | |
| 'gpt2-medium': dict(n_layer=24, n_head=16, n_embd=1024), # 350M params | |
| 'gpt2-large': dict(n_layer=36, n_head=20, n_embd=1280), # 774M params | |
| 'gpt2-xl': dict(n_layer=48, n_head=25, n_embd=1600), # 1558M params | |
| }[model_type] | |
| config_args['vocab_size'] = 50257 # always 50257 for GPT model checkpoints | |
| config_args['block_size'] = 1024 # always 1024 for GPT model checkpoints | |
| # create a from-scratch initialized minGPT model | |
| config = GPTConfig(**config_args) | |
| model = GPT(config) | |
| sd = model.state_dict() | |
| sd_keys = sd.keys() | |
| sd_keys = [k for k in sd_keys if not k.endswith('.attn.bias')] # discard this mask / buffer, not a param | |
| # init a huggingface/transformers model | |
| model_hf = GPT2LMHeadModel.from_pretrained(model_type) | |
| sd_hf = model_hf.state_dict() | |
| # copy while ensuring all of the parameters are aligned and match in names and shapes | |
| sd_keys_hf = sd_hf.keys() | |
| sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.masked_bias')] # ignore these, just a buffer | |
| sd_keys_hf = [k for k in sd_keys_hf if not k.endswith('.attn.bias')] # same, just the mask (buffer) | |
| transposed = ['attn.c_attn.weight', 'attn.c_proj.weight', 'mlp.c_fc.weight', 'mlp.c_proj.weight'] | |
| assert len(sd_keys_hf) == len(sd_keys), f"mismatched keys: {len(sd_keys_hf)} != {len(sd_keys)}" | |
| for k in sd_keys_hf: | |
| if any(k.endswith(w) for w in transposed): | |
| # special treatment for the Conv1D weights we need to transpose | |
| assert sd_hf[k].shape[::-1] == sd[k].shape | |
| with torch.no_grad(): | |
| sd[k].copy_(sd_hf[k].t()) | |
| else: | |
| # vanilla copy over the other parameters | |
| assert sd_hf[k].shape == sd[k].shape | |
| with torch.no_grad(): | |
| sd[k].copy_(sd_hf[k]) | |
| return model |