| """Generate text from a trained Eather-500M checkpoint.""" |
| import argparse |
| import json |
| from pathlib import Path |
| import torch |
| import torch.nn.functional as F |
| from tokenizers import Tokenizer |
| from model import ModelConfig, Eather500M |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser() |
| p.add_argument("--checkpoint", required=True) |
| p.add_argument("--tokenizer", default="tokenizer.json") |
| p.add_argument("--prompt", default="Hello!") |
| p.add_argument("--chat", action="store_true") |
| p.add_argument("--system", default="You are a helpful assistant.") |
| p.add_argument("--max-new-tokens", type=int, default=128) |
| p.add_argument("--temperature", type=float, default=0.8) |
| p.add_argument("--top-k", type=int, default=50) |
| args = p.parse_args() |
| tok = Tokenizer.from_file(args.tokenizer) |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| ck = torch.load(args.checkpoint, map_location="cpu", weights_only=False) |
| cfg = ModelConfig(**ck["config"]) |
| model = Eather500M(cfg) |
| model.load_state_dict(ck["model"]) |
| del ck |
| model.to(device).eval() |
| prompt = args.prompt |
| if args.chat: |
| prompt = f"<|system|>{args.system}<|end|><|user|>{args.prompt}<|end|><|assistant|>" |
| ids = tok.encode(prompt, add_special_tokens=False).ids |
| if not ids: |
| ids = [tok.token_to_id("<|bos|>")] |
| stop = {tok.token_to_id("<|end|>"), tok.token_to_id("<|eos|>")} |
| result = [] |
| print("\nOutput:\n", end="", flush=True) |
| with torch.inference_mode(): |
| for _ in range(args.max_new_tokens): |
| inp = torch.tensor([ids[-cfg.max_seq_len:]], device=device) |
| with torch.autocast(device_type="cuda", dtype=torch.bfloat16, |
| enabled=(device == "cuda" and torch.cuda.is_bf16_supported())): |
| logits, _ = model(inp) |
| next_logits = logits[0, -1].float() |
| if args.temperature <= 0: |
| next_id = next_logits.argmax().item() |
| else: |
| next_logits /= args.temperature |
| if args.top_k > 0: |
| v, _ = torch.topk(next_logits, min(args.top_k, next_logits.numel())) |
| next_logits[next_logits < v[-1]] = -float("inf") |
| next_id = torch.multinomial(F.softmax(next_logits, dim=-1), 1).item() |
| if next_id in stop: |
| break |
| ids.append(next_id) |
| result.append(next_id) |
| print(tok.decode(result, skip_special_tokens=True)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
