| """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() |
|
|
