Spaces:
Runtime error
Runtime error
| """ | |
| Phase K - text generation from an Abyss-1 checkpoint. | |
| Apllies temoerature scaling, top-k filtering, and top-p (nucleus) | |
| filtering in that order, then samples from the resulting distribution. | |
| Usage (from the_final_product/ directory) | |
| python -m ml.generate | |
| python -m ml.generate --prompt "The dragon" --max_tokens 300 | |
| python -m ml.generate --checkpoint ml/checkpoints/ckpt.pt --temperature 1.0 | |
| """ | |
| import argparse | |
| import sys | |
| import torch | |
| import torch.nn.functional as F | |
| from ml.model import AbyssConfig, TheAbyssModel | |
| from ml.tokenizer.tokenizer import AbyssTokenizer | |
| TOKENIZER_DIR = "ml/tokenizer/vocab" | |
| def sample(model, idx, max_new_tokens, temperature, top_k, top_p): | |
| """Autoregressively yield one token id at a time.""" | |
| block_size = model.config.block_size | |
| for _ in range(max_new_tokens): | |
| idx_cond = idx if idx.size(1) <= block_size else idx[:, -block_size:] | |
| logits, _ = model(idx_cond) | |
| logits = logits[:, -1, :] | |
| # 1. Temperature | |
| logits = logits / temperature | |
| #2. top-k: mask everything bellow the k-th highest logit | |
| if top_k > 0: | |
| k = min(top_k, logits.size(-1)) | |
| kth_val = torch.topk(logits, k).values[:, -1:] | |
| logits = logits.masked_fill(logits < kth_val, float("-inf")) | |
| # 3. top-p (nucleus): keep smallest set whose cumulative prob >= top_p | |
| if 0.0 < top_p <1.0: | |
| sorted_logits, sorted_idx = torch.sort(logits, descending=True) | |
| probs_sorted = F.softmax(sorted_logits, dim=-1) | |
| cum_probs = torch.cumsum(probs_sorted, dim=-1) | |
| # Shift right so the token that first pushes cumsum over top_p is kept | |
| remove = (cum_probs - probs_sorted) > top_p | |
| sorted_logits = sorted_logits.masked_fill(remove, float("-inf")) | |
| logits = torch.zeros_like(logits).scatter_(1, sorted_idx, sorted_logits) | |
| probs = F.softmax(logits, dim=-1) | |
| next_tok = torch.multinomial(probs, num_samples=1) | |
| yield next_tok.item() | |
| idx = torch.cat([idx, next_tok], dim=1) | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Generate text from Abyss-1") | |
| parser.add_argument("--checkpoint", default="ml/checkpoints/ckpt.pt") | |
| parser.add_argument("--prompt", default="Once upon a time") | |
| parser.add_argument("--max_tokens", type=int, default=200) | |
| parser.add_argument("--temperature", type=float, default=1.0) | |
| parser.add_argument("--top_k", type=int, default=75) | |
| parser.add_argument("--top_p", type=float, default=0.95) | |
| parser.add_argument("--device", | |
| default="cuda" if torch.cuda.is_available() else "cpu") | |
| args = parser.parse_args() | |
| # Load Checkpoint | |
| print(f"loading {args.checkpoint} ...", flush=True) | |
| ck = torch.load(args.checkpoint, map_location=args.device, weights_only=False) | |
| config = ck["config"] | |
| model = TheAbyssModel(config).to(args.device) | |
| model.load_state_dict(ck["model"]) | |
| model.eval() | |
| print(f" iter {ck['iter']} | best_val {ck['best_val']:.4f} " | |
| f"| params {model.num_params():,}", flush=True) | |
| # Load Tokenizer | |
| tok = AbyssTokenizer.load(TOKENIZER_DIR) | |
| # Encode prompt and Generate | |
| ids = tok.encode(args.prompt) | |
| idx = torch.tensor([ids], dtype=torch.long, device=args.device) | |
| print(f"\n[temp={args.temperature} top_k={args.top_k} top_p={args.top_p}]\n") | |
| sys.stdout.write(args.prompt) | |
| sys.stdout.flush() | |
| for tok_id in sample(model, idx, args.max_tokens, | |
| args.temperature, args.top_k,args.top_p): | |
| sys.stdout.write(tok.decode([tok_id])) | |
| sys.stdout.flush() | |
| print() # Trailing Newline | |
| if __name__ == "__main__": | |
| main() |