File size: 6,323 Bytes
c776e2a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | """Try Walker -- load the Walker 1B model and generate text.
A minimal, self-contained entry point for users. It reuses the model definition
inside train_walker_1b_continue.py, so put this file next to that trainer (or
pass --trainer). The architecture is inferred from the saved checkpoint, so any
Walker checkpoint works (ongoing training or best).
Examples:
python try_walker.py
python try_walker.py --prompt "The migration of the Arctic tern"
python try_walker.py --ckpt checkpoints/ckpt_ongoing.pt --greedy
python try_walker.py --device cpu --max-new 32
"""
import argparse
import importlib.util
import os
import sys
from pathlib import Path
import torch
def _first_existing(paths):
for p in paths:
if not p:
continue
cand = Path(p).expanduser()
if cand.exists():
return str(cand.resolve())
return None
def locate_trainer(explicit):
here = Path(__file__).resolve().parent
cwd = Path.cwd()
found = _first_existing([
explicit,
os.environ.get("TB_TRAINER"),
here / "train_walker_1b_continue.py",
cwd / "train_walker_1b_continue.py",
here / "walker_space" / "train_walker_1b_continue.py",
here / "CPT" / "train_walker_1b_continue.py",
here / "Experiment" / "train_walker_1b_continue.py",
])
if not found:
sys.exit(
"Could not find train_walker_1b_continue.py.\n"
"Put this script next to it, or pass --trainer PATH."
)
return found
def locate_ckpt(explicit):
here = Path(__file__).resolve().parent
cwd = Path.cwd()
found = _first_existing([
explicit,
os.environ.get("TB_CKPT_PATH"),
os.environ.get("TB_CKPT"),
"checkpoints/ckpt_ongoing.pt",
here / "checkpoints" / "ckpt_ongoing.pt",
cwd / "checkpoints" / "ckpt_ongoing.pt",
])
if not found:
sys.exit(
"Could not find a checkpoint.\n"
"Pass --ckpt PATH, for example --ckpt checkpoints/ckpt_ongoing.pt"
)
return found
def load_trainer(path):
spec = importlib.util.spec_from_file_location("walker_trainer", path)
mod = importlib.util.module_from_spec(spec)
spec.loader.exec_module(mod)
return mod
def pick_device(requested):
if requested:
return requested
return "cuda" if torch.cuda.is_available() else "cpu"
def build(mod, ckpt, device, tokenizer_name):
print(f"Loading Walker from {ckpt} on {device} ...", flush=True)
model, tok, cfg = mod.build_model_for_inference(
ckpt, device=device, tokenizer_name=tokenizer_name
)
mod.DEVICE = device
model.eval()
print(f"Loaded: D={cfg.D} L={cfg.L} heads={cfg.n_head} vocab={cfg.vocab} "
f"max_t={cfg.max_t}", flush=True)
return model, tok, cfg
@torch.inference_mode()
def generate_text(mod, model, tok, prompt, device, max_new, temperature, top_p, rep_penalty, seed):
torch.manual_seed(int(seed))
ids = tok(prompt, return_tensors="pt", add_special_tokens=False)["input_ids"]
max_t = int(getattr(model, "max_t", 4096))
keep = max_t - int(max_new) - 1
if ids.shape[1] > keep:
ids = ids[:, -keep:]
ids = ids.to(device)
out = model.generate(
ids,
max_new=int(max_new),
temperature=float(temperature),
top_p=float(top_p),
rep_penalty=float(rep_penalty),
)
return tok.decode(out[0].tolist(), skip_special_tokens=True)
def parse_args():
ap = argparse.ArgumentParser(description="Try the Walker 1B model.")
ap.add_argument("--prompt", default=None, help="Text prompt. Omit for interactive mode.")
ap.add_argument("--ckpt", default=None, help="Checkpoint path (default: checkpoints/ckpt_ongoing.pt).")
ap.add_argument("--trainer", default=None, help="Path to train_walker_1b_continue.py.")
ap.add_argument("--device", default=None, help="cuda or cpu (default: auto).")
ap.add_argument("--tokenizer", default=None, help="Tokenizer name (default: from checkpoint, else gpt2).")
ap.add_argument("--max-new", type=int, default=64, help="Number of new tokens to generate.")
ap.add_argument("--temperature", type=float, default=0.8)
ap.add_argument("--top-p", type=float, default=0.9)
ap.add_argument("--rep-penalty", type=float, default=1.1)
ap.add_argument("--seed", type=int, default=0)
ap.add_argument("--greedy", action="store_true", help="Deterministic decoding (temperature 0).")
return ap.parse_args()
def main():
args = parse_args()
temperature = 0.0 if args.greedy else float(args.temperature)
trainer = locate_trainer(args.trainer)
ckpt = locate_ckpt(args.ckpt)
device = pick_device(args.device)
if device == "cpu":
print("Note: running the 1B model on CPU will be very slow.", flush=True)
mod = load_trainer(trainer)
model, tok, cfg = build(mod, ckpt, device, args.tokenizer)
if args.prompt is not None:
text = generate_text(mod, model, tok, args.prompt, device,
args.max_new, temperature, args.top_p, args.rep_penalty, args.seed)
print("\n" + text + "\n")
return
print("Ready. Type a prompt, or :quit to exit.")
print("Commands: :temp X :top-p X :max-new N")
while True:
try:
prompt = input("prompt> ")
except (EOFError, KeyboardInterrupt):
print()
break
if not prompt.strip():
continue
if prompt.strip() in (":quit", ":q", ":exit"):
break
if prompt.startswith(":temp "):
temperature = float(prompt.split(None, 1)[1])
print(f"temperature = {temperature}")
continue
if prompt.startswith(":top-p "):
args.top_p = float(prompt.split(None, 1)[1])
print(f"top_p = {args.top_p}")
continue
if prompt.startswith(":max-new "):
args.max_new = int(prompt.split(None, 1)[1])
print(f"max_new = {args.max_new}")
continue
text = generate_text(mod, model, tok, prompt, device,
args.max_new, temperature, args.top_p, args.rep_penalty, args.seed)
print("\n" + text + "\n", flush=True)
if __name__ == "__main__":
main()
|