| | import json |
| | import torch |
| | import torch.nn as nn |
| | from transformers import GPT2Tokenizer |
| |
|
| |
|
| | def load_tinylm(model_dir, device="cpu"): |
| | |
| | with open(f"{model_dir}/config.json") as f: |
| | config = json.load(f) |
| |
|
| | VOCAB_SIZE = config["vocab_size"] |
| | EMBED_RANK = config["embed_rank"] |
| | D_MODEL = config["d_model"] |
| | N_HEADS = config["n_heads"] |
| | FFN_DIM = config["ffn_dim"] |
| | N_LAYERS = config["n_layers"] |
| | MAX_SEQ_LEN = config["max_seq_len"] |
| | DROPOUT = config["dropout"] |
| |
|
| | class FactoredEmbedding(nn.Module): |
| | def __init__(self, vocab_size, rank, d_model): |
| | super().__init__() |
| | self.in_proj = nn.Embedding(vocab_size, rank) |
| | self.out_proj = nn.Linear(rank, d_model, bias=False) |
| |
|
| | def forward(self, x): |
| | return self.out_proj(self.in_proj(x)) |
| |
|
| | class TransformerBlock(nn.Module): |
| | def __init__(self): |
| | super().__init__() |
| | self.ln1 = nn.LayerNorm(D_MODEL) |
| | self.attn = nn.MultiheadAttention(D_MODEL, N_HEADS, dropout=DROPOUT, batch_first=True) |
| | self.ln2 = nn.LayerNorm(D_MODEL) |
| | self.ffn = nn.Sequential( |
| | nn.Linear(D_MODEL, FFN_DIM), |
| | nn.GELU(), |
| | nn.Linear(FFN_DIM, D_MODEL), |
| | nn.Dropout(DROPOUT), |
| | ) |
| |
|
| | def forward(self, x, attn_mask=None, key_padding_mask=None): |
| | x_norm = self.ln1(x) |
| | attn_out, _ = self.attn(x_norm, x_norm, x_norm, |
| | attn_mask=attn_mask, |
| | key_padding_mask=key_padding_mask, |
| | is_causal=True) |
| | x = x + attn_out |
| | x = x + self.ffn(self.ln2(x)) |
| | return x |
| |
|
| | class TinyLM(nn.Module): |
| | def __init__(self): |
| | super().__init__() |
| | self.tok_emb = FactoredEmbedding(VOCAB_SIZE, EMBED_RANK, D_MODEL) |
| | self.pos_emb = nn.Embedding(MAX_SEQ_LEN, D_MODEL) |
| | self.drop = nn.Dropout(DROPOUT) |
| | self.blocks = nn.ModuleList([TransformerBlock() for _ in range(N_LAYERS)]) |
| | self.ln_final = nn.LayerNorm(D_MODEL) |
| | self.head_down = nn.Linear(D_MODEL, EMBED_RANK, bias=False) |
| | self.head_vocab = nn.Linear(EMBED_RANK, VOCAB_SIZE, bias=False) |
| | self.head_vocab.weight = nn.Parameter(self.tok_emb.in_proj.weight) |
| |
|
| | def forward(self, idx): |
| | B, T = idx.shape |
| | if T > MAX_SEQ_LEN: |
| | idx = idx[:, :MAX_SEQ_LEN] |
| | T = idx.shape[1] |
| | positions = torch.arange(T, device=idx.device).unsqueeze(0) |
| | x = self.drop(self.tok_emb(idx) + self.pos_emb(positions)) |
| | mask = nn.Transformer.generate_square_subsequent_mask(T, device=idx.device) |
| | for block in self.blocks: |
| | x = block(x, attn_mask=mask) |
| | x = self.ln_final(x) |
| | x = self.head_down(x) |
| | return self.head_vocab(x) |
| |
|
| | |
| | model = TinyLM().to(device) |
| | state_dict = torch.load(f"{model_dir}/pytorch_model.bin", map_location=device) |
| | model.load_state_dict(state_dict) |
| | model.eval() |
| |
|
| | |
| | tokenizer = GPT2Tokenizer.from_pretrained(model_dir) |
| | tokenizer.pad_token = tokenizer.eos_token |
| |
|
| | return model, tokenizer, config |
| |
|
| |
|
| | def generate(model, tokenizer, prompt, max_new_tokens=100, temperature=0.1, top_k=25, device="cpu"): |
| | MAX_SEQ_LEN = model.pos_emb.num_embeddings |
| | model.eval() |
| | ids = tokenizer(prompt, return_tensors="pt").input_ids.to(device) |
| |
|
| | with torch.no_grad(): |
| | for _ in range(max_new_tokens): |
| | idx_cond = ids[:, -MAX_SEQ_LEN:] |
| | logits = model(idx_cond) |
| | logits = logits[:, -1, :] / temperature |
| | if top_k is not None: |
| | values, _ = torch.topk(logits, top_k) |
| | logits[logits < values[:, -1:]] = -float("inf") |
| | probs = torch.softmax(logits, dim=-1) |
| | next_id = torch.multinomial(probs, num_samples=1) |
| | if next_id.item() == tokenizer.eos_token_id: |
| | break |
| | ids = torch.cat([ids, next_id], dim=1) |
| |
|
| | return tokenizer.decode(ids[0], skip_special_tokens=True) |
| |
|
| |
|
| | if __name__ == "__main__": |
| | model, tokenizer, config = load_tinylm("./tinylm") |
| | print("Model loaded!") |
| | print("Use 'module.generate(model, tokenizer, \"Once upon a time\")' to generate.") |
| |
|
