AR / ar /eval /generate.py

Upload AR-4B pretraining codebase, configs, and cluster scripts

2ae6691 verified 18 days ago

2.32 kB

"""Generate text from an exported checkpoint (quick quality check)."""
import argparse

import torch

from ar.model.config import ARConfig
from ar.model.model import ARModel
from ar.model.mask import build_doc_block_mask
from ar.model.rope import build_rope_cache
from ar.data.tokenizer import Tokenizer

SEQ = 256  # fixed length, so FlexAttention compiles once instead of for every new length


def load_model(path, device):
    cfg = ARConfig()
    with torch.device(device):
        model = ARModel(cfg)                     # real init also builds correct RoPE tables
    model.load_state_dict(torch.load(path, map_location=device))
    model = model.to(torch.bfloat16).eval()    
    return model


@torch.no_grad()
def generate(model, tok, prompt, device, max_new=150, temperature=0.8, top_k=50):
    ids = tok.encode(prompt)
    n = len(ids)
    # fixed-size buffer; slots after the current position are placeholders
    # that the causal mask hides from earlier tokens
    buf = torch.full((1, SEQ), tok.eot_id, dtype=torch.long, device=device)
    buf[0, :n] = torch.tensor(ids, device=device)
    seg = torch.zeros(1, SEQ, dtype=torch.int32, device=device)   # one document
    mask = build_doc_block_mask(seg, SEQ, device=device)

    pos = n
    for pos in range(n, min(SEQ, n + max_new)):
        logits = model(buf, mask)[0, pos - 1].float() / temperature
        # keep only the top_k most likely tokens, then sample one of them
        kth = torch.topk(logits, top_k).values[-1]
        logits[logits < kth] = -float("inf")
        nxt = torch.multinomial(torch.softmax(logits, -1), 1)
        buf[0, pos] = nxt
        if nxt.item() == tok.eot_id:              # model decided the text is over
            break
    return tok.decode(buf[0, : pos + 1].tolist())


if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True)
    args = p.parse_args()

    device = "cuda"
    tok = Tokenizer()
    model = load_model(args.model, device)
    torch.manual_seed(0)

    prompts = [
        "The capital of France is",
        "Once upon a time, there was a little",
        "def fibonacci(n):",
        "Question: What is 12 + 7?\nAnswer:",
    ]
    for pr in prompts:
        print("=" * 60, flush=True)
        print(generate(model, tok, pr, device), flush=True)

	"""Generate text from an exported checkpoint (quick quality check)."""
	import argparse

	import torch

	from ar.model.config import ARConfig
	from ar.model.model import ARModel
	from ar.model.mask import build_doc_block_mask
	from ar.model.rope import build_rope_cache
	from ar.data.tokenizer import Tokenizer

	SEQ = 256 # fixed length, so FlexAttention compiles once instead of for every new length


	def load_model(path, device):
	cfg = ARConfig()
	with torch.device(device):
	model = ARModel(cfg) # real init also builds correct RoPE tables
	model.load_state_dict(torch.load(path, map_location=device))
	model = model.to(torch.bfloat16).eval()
	return model


	@torch.no_grad()
	def generate(model, tok, prompt, device, max_new=150, temperature=0.8, top_k=50):
	ids = tok.encode(prompt)
	n = len(ids)
	# fixed-size buffer; slots after the current position are placeholders
	# that the causal mask hides from earlier tokens
	buf = torch.full((1, SEQ), tok.eot_id, dtype=torch.long, device=device)
	buf[0, :n] = torch.tensor(ids, device=device)
	seg = torch.zeros(1, SEQ, dtype=torch.int32, device=device) # one document
	mask = build_doc_block_mask(seg, SEQ, device=device)

	pos = n
	for pos in range(n, min(SEQ, n + max_new)):
	logits = model(buf, mask)[0, pos - 1].float() / temperature
	# keep only the top_k most likely tokens, then sample one of them
	kth = torch.topk(logits, top_k).values[-1]
	logits[logits < kth] = -float("inf")
	nxt = torch.multinomial(torch.softmax(logits, -1), 1)
	buf[0, pos] = nxt
	if nxt.item() == tok.eot_id: # model decided the text is over
	break
	return tok.decode(buf[0, : pos + 1].tolist())


	if __name__ == "__main__":
	p = argparse.ArgumentParser()
	p.add_argument("--model", required=True)
	args = p.parse_args()

	device = "cuda"
	tok = Tokenizer()
	model = load_model(args.model, device)
	torch.manual_seed(0)

	prompts = [
	"The capital of France is",
	"Once upon a time, there was a little",
	"def fibonacci(n):",
	"Question: What is 12 + 7?\nAnswer:",
	]
	for pr in prompts:
	print("=" * 60, flush=True)
	print(generate(model, tok, pr, device), flush=True)