checkpoint.py

1# saves and loads everything needed to use a trained model later:2# weights, model shape, and the tokenizer (talk.py cannot guess any of these)34import json5import os6from dataclasses import asdict7import numpy as np8from tokenizer import Tokenizer9from transformer import ModelConfig1011# one checkpoint is three files that share a prefix, e.g. "checkpoints/model"12#   prefix + ".npz"       every weight array, stored under the same keys as the params dict13#   prefix + ".json"      the ModelConfig (vocab_size, emb_dim, ...), needed to rebuild the shapes14#   prefix + ".tok.json"  the tokenizer: merges + itos1516def save_checkpoint(prefix, params, cfg: ModelConfig, tok: Tokenizer):17    folder = os.path.dirname(prefix)18    if folder:19        os.makedirs(folder, exist_ok=True)2021    np.savez(prefix + ".npz", **params) # keys like "b0.wq" are kept as they are2223    with open(prefix + ".json", "w", encoding="utf-8") as f:24        json.dump(asdict(cfg), f)2526    # stoi is not saved, it is just itos turned around, so load rebuilds it27    # json has no tuples, so each merge's "pair" is stored as a list, merge_tokens only28    # uses pair[0] and pair[1], so a list works exactly the same29    with open(prefix + ".tok.json", "w", encoding="utf-8") as f:30        json.dump({"merges": tok.merges, "itos": tok.itos}, f, ensure_ascii=False)3132def load_checkpoint(prefix):33    with np.load(prefix + ".npz") as z:34        params = {k: z[k] for k in z.files} # copy out, the file closes after this block3536    with open(prefix + ".json", encoding="utf-8") as f:37        cfg = ModelConfig(**json.load(f))3839    with open(prefix + ".tok.json", encoding="utf-8") as f:40        d = json.load(f)4142    stoi = {}43    for i in range(len(d["itos"])):44        stoi[d["itos"][i]] = i4546    return params, cfg, Tokenizer(d["merges"], stoi, d["itos"])