File size: 2,324 Bytes
2ae6691 | 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 | """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) |