Fu01978/TinyLM

1import json2import torch3import torch.nn as nn4from transformers import GPT2Tokenizer5 6 7def load_tinylm(model_dir, device="cpu"):8    # Load config9    with open(f"{model_dir}/config.json") as f:10        config = json.load(f)11 12    VOCAB_SIZE  = config["vocab_size"]13    EMBED_RANK  = config["embed_rank"]14    D_MODEL     = config["d_model"]15    N_HEADS     = config["n_heads"]16    FFN_DIM     = config["ffn_dim"]17    N_LAYERS    = config["n_layers"]18    MAX_SEQ_LEN = config["max_seq_len"]19    DROPOUT     = config["dropout"]20 21    class FactoredEmbedding(nn.Module):22        def __init__(self, vocab_size, rank, d_model):23            super().__init__()24            self.in_proj  = nn.Embedding(vocab_size, rank)25            self.out_proj = nn.Linear(rank, d_model, bias=False)26 27        def forward(self, x):28            return self.out_proj(self.in_proj(x))29 30    class TransformerBlock(nn.Module):31        def __init__(self):32            super().__init__()33            self.ln1  = nn.LayerNorm(D_MODEL)34            self.attn = nn.MultiheadAttention(D_MODEL, N_HEADS, dropout=DROPOUT, batch_first=True)35            self.ln2  = nn.LayerNorm(D_MODEL)36            self.ffn  = nn.Sequential(37                nn.Linear(D_MODEL, FFN_DIM),38                nn.GELU(),39                nn.Linear(FFN_DIM, D_MODEL),40                nn.Dropout(DROPOUT),41            )42 43        def forward(self, x, attn_mask=None, key_padding_mask=None):44            x_norm = self.ln1(x)45            attn_out, _ = self.attn(x_norm, x_norm, x_norm,46                                    attn_mask=attn_mask,47                                    key_padding_mask=key_padding_mask,48                                    is_causal=True)49            x = x + attn_out50            x = x + self.ffn(self.ln2(x))51            return x52 53    class TinyLM(nn.Module):54        def __init__(self):55            super().__init__()56            self.tok_emb  = FactoredEmbedding(VOCAB_SIZE, EMBED_RANK, D_MODEL)57            self.pos_emb  = nn.Embedding(MAX_SEQ_LEN, D_MODEL)58            self.drop     = nn.Dropout(DROPOUT)59            self.blocks   = nn.ModuleList([TransformerBlock() for _ in range(N_LAYERS)])60            self.ln_final = nn.LayerNorm(D_MODEL)61            self.head_down  = nn.Linear(D_MODEL, EMBED_RANK, bias=False)62            self.head_vocab = nn.Linear(EMBED_RANK, VOCAB_SIZE, bias=False)63            self.head_vocab.weight = nn.Parameter(self.tok_emb.in_proj.weight)64 65        def forward(self, idx):66            B, T = idx.shape67            if T > MAX_SEQ_LEN:68                idx = idx[:, :MAX_SEQ_LEN]69            T = idx.shape[1]70            positions = torch.arange(T, device=idx.device).unsqueeze(0)71            x = self.drop(self.tok_emb(idx) + self.pos_emb(positions))72            mask = nn.Transformer.generate_square_subsequent_mask(T, device=idx.device)73            for block in self.blocks:74                x = block(x, attn_mask=mask)75            x = self.ln_final(x)76            x = self.head_down(x)77            return self.head_vocab(x)78 79    # Build and load weights80    model = TinyLM().to(device)81    state_dict = torch.load(f"{model_dir}/pytorch_model.bin", map_location=device)82    model.load_state_dict(state_dict)83    model.eval()84 85    # Load tokenizer86    tokenizer = GPT2Tokenizer.from_pretrained(model_dir)87    tokenizer.pad_token = tokenizer.eos_token88 89    return model, tokenizer, config90 91 92def generate(model, tokenizer, prompt, max_new_tokens=100, temperature=0.1, top_k=25, device="cpu"):93    MAX_SEQ_LEN = model.pos_emb.num_embeddings94    model.eval()95    ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device)96 97    with torch.no_grad():98        for _ in range(max_new_tokens):99            idx_cond = ids[:, -MAX_SEQ_LEN:]100            logits   = model(idx_cond)101            logits   = logits[:, -1, :] / temperature102            if top_k is not None:103                values, _ = torch.topk(logits, top_k)104                logits[logits < values[:, -1:]] = -float("inf")105            probs   = torch.softmax(logits, dim=-1)106            next_id = torch.multinomial(probs, num_samples=1)107            if next_id.item() == tokenizer.eos_token_id:108                break109            ids = torch.cat([ids, next_id], dim=1)110 111    return tokenizer.decode(ids[0], skip_special_tokens=True)112 113 114if __name__ == "__main__":115    model, tokenizer, config = load_tinylm("./tinylm")116    print("Model loaded!")117    print("Use 'module.generate(model, tokenizer, \"Once upon a time\")' to generate.")118