File size: 2,531 Bytes
7da1a03 | 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 | """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()
|