| """ |
| Builds a new TinyTransformerLM, warm-starting the token embeddings |
| from the previously trained MLP model (tinylm.pkl) for any word that |
| exists in both vocabularies. Everything else (attention weights, |
| positional embeddings, new vocab rows, FF head) starts randomly |
| initialized, since there is no equivalent to transfer from the MLP. |
| """ |
| import pickle |
| import re |
| import numpy as np |
| from model_transformer import TinyTransformerLM |
| import config |
|
|
| print("Loading old MLP checkpoint for warm-start...") |
| with open("tinylm_v1_mlp_backup.pkl", "rb") as f: |
| old_ckpt = pickle.load(f) |
|
|
| old_model = old_ckpt["model"] |
| old_word_to_idx = old_ckpt["word_to_idx"] |
| old_embeddings = old_model.embeddings |
|
|
| print(f"Old model: vocab={len(old_word_to_idx)}, embed_dim={old_model.embed_dim}") |
|
|
| |
| with open("dataset.txt", "r", encoding="utf-8") as f: |
| text = f.read().lower() |
| words = re.findall(r"[a-z0-9']+|[.!?]", text) |
| words = [w for w in words if w.strip()] |
|
|
| vocab = sorted(set(words)) |
| word_to_idx = {w: i for i, w in enumerate(vocab)} |
| idx_to_word = {i: w for i, w in enumerate(vocab)} |
| PAD = "<pad>" |
| word_to_idx[PAD] = len(word_to_idx) |
| idx_to_word[len(idx_to_word)] = PAD |
| vocab.append(PAD) |
|
|
| print(f"New vocab size: {len(vocab)}") |
|
|
| CONTEXT = config.CONTEXT_SIZE |
| new_model = TinyTransformerLM( |
| vocab_size=len(vocab), |
| embed_dim=config.EMBED_DIM, |
| hidden_dim=config.HIDDEN_DIM, |
| context_size=CONTEXT, |
| ) |
| print(f"New model params: {new_model.num_params():,}") |
|
|
| |
| |
| |
| old_E = old_model.embed_dim |
| new_E = new_model.embed_dim |
| copy_dim = min(old_E, new_E) |
|
|
| transferred = 0 |
| for w, old_idx in old_word_to_idx.items(): |
| if w in word_to_idx: |
| new_idx = word_to_idx[w] |
| new_model.embeddings[new_idx, :copy_dim] = old_embeddings[old_idx, :copy_dim] |
| transferred += 1 |
|
|
| print(f"Transferred embeddings for {transferred}/{len(word_to_idx)} words " |
| f"(copied {copy_dim} of {new_E} embedding dims, rest stays randomly initialized)") |
|
|
| |
| with open("tinylm.pkl", "wb") as f: |
| pickle.dump({ |
| "model": new_model, |
| "word_to_idx": word_to_idx, |
| "idx_to_word": idx_to_word, |
| "context_size": CONTEXT, |
| "pad_token": PAD, |
| "architecture": "transformer", |
| }, f) |
|
|
| print("Warm-started checkpoint saved to tinylm.pkl") |
|
|
