1import sys2import numpy as np3import config4from checkpoint import load_checkpoint5from generate import generate67defmain(prefix="checkpoints/model"):8 params, cfg, tok =load_checkpoint(prefix)9 rng = np.random.default_rng(config.seed)10print(f"loaded {prefix} (V={cfg.vocab_size})")11print("Type the start of a text and press Enter. Empty line quits.")1213whileTrue:14try:15 prompt =input("> ")16exceptEOFError:17break1819if prompt.strip()=="":20break2122 ids = tok.encode(prompt)2324 out =generate(params, cfg, ids, config.max_new_tokens, config.seq_len,25 config.temperature, config.top_k, config.top_p, rng)2627print(tok.decode(out[len(ids):]))# only the new tokens, the prompt is already on screen2829if__name__=="__main__":30main(*sys.argv[1:])