| """Generate text from a trained checkpoint to verify English quality.""" |
| import torch |
| import tiktoken |
| from config import Config |
| from model import build_model |
|
|
| def generate(model, cfg, prompt_ids, max_tokens=200, temperature=0.8, top_k=40): |
| """Greedy or sampled generation from the base model (no Medusa, no cache).""" |
| model.eval() |
| enc = tiktoken.get_encoding("gpt2") |
| ids = torch.tensor([prompt_ids], device="cuda", dtype=torch.long) |
| |
| with torch.no_grad(): |
| for _ in range(max_tokens): |
| |
| x = ids[:, -1024:] |
| h = model(x) |
| logits = model.lm_head(h[:, -1]) |
| |
| if temperature == 0: |
| next_id = logits.argmax(dim=-1, keepdim=True) |
| else: |
| logits = logits / temperature |
| if top_k > 0: |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) |
| logits[logits < v[:, [-1]]] = float('-inf') |
| probs = torch.softmax(logits, dim=-1) |
| next_id = torch.multinomial(probs, 1) |
| |
| ids = torch.cat([ids, next_id], dim=1) |
| |
| return ids[0].tolist() |
|
|
| if __name__ == "__main__": |
| import sys |
| ckpt_path = sys.argv[1] if len(sys.argv) > 1 else "checkpoints/base/best.pt" |
| |
| ckpt = torch.load(ckpt_path, map_location="cuda") |
| cfg = Config.v5_500m() |
| model = build_model(cfg, "cuda") |
| model.load_state_dict(ckpt["model"]) |
| model.eval() |
| |
| print(f"Loaded: {ckpt_path} (step {ckpt.get('step', '?')}, loss {ckpt.get('loss', '?'):.4f})") |
| |
| enc = tiktoken.get_encoding("gpt2") |
| |
| prompts = [ |
| "The quick brown fox", |
| "In a galaxy far far away", |
| "The capital of France is", |
| "Once upon a time", |
| ] |
| |
| for prompt in prompts: |
| prompt_ids = enc.encode_ordinary(prompt) |
| out_ids = generate(model, cfg, prompt_ids, max_tokens=100, temperature=0.7, top_k=40) |
| text = enc.decode(out_ids) |
| print(f"\n--- Prompt: {prompt!r} ---") |
| print(text) |
|
|
