sample.py

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
SRC = ROOT / "src"
if str(SRC) not in sys.path:
    sys.path.insert(0, str(SRC))

import torch

from nano_transformer import Tokenizer, TransformerLM
from nano_transformer.inference import branch_kv_caches, decode_step, prefill, sample_next_token
from nano_transformer.training import load_checkpoint


def default_device() -> str:
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Sample from a trained Nano-Transformer checkpoint.")
    parser.add_argument("--run-dir", type=Path, required=True)
    parser.add_argument("--checkpoint", type=Path)
    parser.add_argument("--prompt", type=str, default="")
    parser.add_argument("--max-new-tokens", type=int, default=64)
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--strategy", choices=("greedy", "sample", "top_k", "top_p"), default="greedy")
    parser.add_argument("--top-k", type=int)
    parser.add_argument("--top-p", type=float)
    parser.add_argument("--device", type=str, default=default_device())
    parser.add_argument("--num-samples", type=int, default=1)
    parser.add_argument("--eos-token-id", type=int)
    parser.add_argument("--no-kv-cache", action="store_true")
    return parser.parse_args()


def load_run_config(run_dir: Path) -> dict:
    with (run_dir / "run_config.json").open("r", encoding="utf-8") as f:
        return json.load(f)


def infer_bos_token_id(tokenizer: Tokenizer) -> int | None:
    for special_token in tokenizer.special_tokens:
        token_id = tokenizer.bytes_to_token_id.get(special_token.encode("utf-8"))
        if token_id is not None:
            return token_id
    return None


def main() -> None:
    args = parse_args()
    run_config = load_run_config(args.run_dir)
    checkpoint = args.checkpoint or args.run_dir / "latest.pt"

    tokenizer = Tokenizer.from_files(
        vocab_filepath=run_config["vocab_path"],
        merges_filepath=run_config["merges_path"],
        special_tokens=run_config["special_tokens"],
    )
    model = TransformerLM.from_config(run_config["model_config"]).to(args.device)
    load_checkpoint(checkpoint, model, optimizer=None)
    model.eval()

    prompt_tokens = tokenizer.encode(args.prompt)
    bos_token_id = infer_bos_token_id(tokenizer) if len(prompt_tokens) == 0 else None
    if len(prompt_tokens) == 0 and bos_token_id is None and args.max_new_tokens > 0:
        raise ValueError(
            "Empty prompts require either a tokenizer special token to use as a BOS seed or a non-empty --prompt."
        )
    prompt = torch.tensor(prompt_tokens, dtype=torch.long, device=args.device)

    if args.max_new_tokens == 0:
        decoded_prompt = tokenizer.decode(prompt_tokens)
        if args.num_samples == 1:
            print(decoded_prompt)
            return
        for idx in range(1, args.num_samples + 1):
            print(f"=== Sample {idx} ===")
            print(decoded_prompt)
        return

    if args.num_samples == 1:
        from nano_transformer import generate

        generated = generate(
            model,
            prompt,
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature,
            strategy=args.strategy,
            top_k=args.top_k,
            top_p=args.top_p,
            bos_token_id=bos_token_id,
            eos_token_id=args.eos_token_id,
            use_kv_cache=not args.no_kv_cache,
        )
        print(tokenizer.decode(generated.tolist()))
        return

    if args.no_kv_cache:
        from nano_transformer import generate

        prompt_batch = prompt.unsqueeze(0).expand(args.num_samples, -1).clone()
        generated = generate(
            model,
            prompt_batch,
            max_new_tokens=args.max_new_tokens,
            temperature=args.temperature,
            strategy=args.strategy,
            top_k=args.top_k,
            top_p=args.top_p,
            bos_token_id=bos_token_id,
            eos_token_id=args.eos_token_id,
            use_kv_cache=False,
        )
        for idx, sample in enumerate(generated.tolist(), start=1):
            print(f"=== Sample {idx} ===")
            print(tokenizer.decode(sample))
        return

    generated = prompt.unsqueeze(0).expand(args.num_samples, -1).clone()
    prefill_prompt = prompt.unsqueeze(0)
    strip_prefix_tokens = 0
    if generated.size(1) == 0:
        generated = torch.full((args.num_samples, 1), bos_token_id, dtype=torch.long, device=args.device)
        prefill_prompt = generated[:1]
        strip_prefix_tokens = 1

    logits, kv_caches = prefill(
        model,
        prefill_prompt,
        total_seq_len=prefill_prompt.size(1) + args.max_new_tokens,
    )
    logits = logits.expand(args.num_samples, -1, -1).clone()
    kv_caches = branch_kv_caches(kv_caches, args.num_samples)
    finished = torch.zeros(args.num_samples, dtype=torch.bool, device=args.device)

    for _ in range(args.max_new_tokens):
        next_token = sample_next_token(
            logits[:, -1, :],
            temperature=args.temperature,
            strategy=args.strategy,
            top_k=args.top_k,
            top_p=args.top_p,
        )
        if args.eos_token_id is not None:
            eos_tensor = torch.full_like(next_token, args.eos_token_id)
            next_token = torch.where(finished, eos_tensor, next_token)
        generated = torch.cat((generated, next_token.unsqueeze(-1)), dim=-1)
        if args.eos_token_id is not None:
            finished = finished | (next_token == args.eos_token_id)
            if torch.all(finished):
                break
        logits, kv_caches = decode_step(model, next_token.unsqueeze(-1), kv_caches)

    for idx, sample in enumerate(generated.tolist(), start=1):
        print(f"=== Sample {idx} ===")
        print(tokenizer.decode(sample[strip_prefix_tokens:]))


if __name__ == "__main__":
    main()