File size: 6,665 Bytes
f5f443b 08dbee5 f5f443b 08dbee5 f5f443b 08dbee5 f5f443b 4f331d8 f5f443b 4f331d8 f5f443b 4f331d8 f5f443b 08dbee5 f5f443b | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | import argparse
import json
import math
import time
from pathlib import Path
import torch
from tokenizers import Tokenizer
from .model import LanguageModel, ModelConfig
from .artifacts import digest_file
def load_model(checkpoint_path, tokenizer_path, device="cpu"):
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
if digest_file(tokenizer_path) != checkpoint["tokenizer_sha256"]:
raise ValueError("Tokenizer does not match checkpoint")
tokenizer = Tokenizer.from_file(str(tokenizer_path))
model = LanguageModel(ModelConfig(**checkpoint["config"]))
if tokenizer.get_vocab_size() != model.config.vocab_size:
raise ValueError("Tokenizer vocabulary does not match model")
model.load_state_dict(checkpoint["model"])
step = checkpoint["step"]
del checkpoint
return model.to(device).eval(), tokenizer, step
def load_export(directory, device="cpu"):
from safetensors.torch import load_file
directory = Path(directory)
metadata = json.loads((directory / "config.json").read_text(encoding="utf-8"))
if metadata.get("format") != "tiny-english-safetensors-v1":
raise ValueError("Unsupported exported model format")
if metadata.get("dtype") != "bfloat16":
raise ValueError("This export requires BF16 weights")
if device != "cuda":
raise ValueError("This BF16 export requires a CUDA GPU; use --device cuda")
tokenizer_path = directory / "tokenizer.json"
weights_path = directory / "model.safetensors"
if digest_file(tokenizer_path) != metadata["tokenizer_sha256"]:
raise ValueError("Export tokenizer checksum does not match")
if digest_file(weights_path) != metadata["weights_sha256"]:
raise ValueError("Export weights checksum does not match")
tokenizer = Tokenizer.from_file(str(tokenizer_path))
model = LanguageModel(ModelConfig(**metadata["model_config"]))
if tokenizer.get_vocab_size() != model.config.vocab_size:
raise ValueError("Export tokenizer vocabulary does not match model")
model.load_state_dict(load_file(str(weights_path), device="cpu"), strict=True)
return model.to(device=device, dtype=torch.bfloat16).eval(), tokenizer, metadata["checkpoint_step"]
def complete(model, tokenizer, prompt, max_new_tokens=150, temperature=0.8, top_k=40, seed=42, deadline=None, strict_context=False, stop_at_sentence_end=False, min_new_tokens=0):
if not isinstance(prompt, str) or not prompt.strip():
raise ValueError("Prompt must not be empty")
if not math.isfinite(temperature) or temperature <= 0 or top_k < 1 or max_new_tokens < 1 or min_new_tokens < 0:
raise ValueError("Generation numeric options must be positive and finite")
tokens = tokenizer.encode(prompt).ids
prompt_tokens = len(tokens)
if not tokens:
raise ValueError("Prompt must contain tokens")
if strict_context and prompt_tokens > model.config.context_length:
raise ValueError(f"Prompt has {prompt_tokens} tokens; the maximum is {model.config.context_length}. Shorten the prompt.")
device = model.embedding.weight.device
generator = torch.Generator(device=device).manual_seed(seed)
finish_reason = "length"
with torch.inference_mode():
for _ in range(max_new_tokens):
if deadline is not None and time.monotonic() >= deadline:
raise TimeoutError("Generation exceeded its time limit. Try fewer tokens or a shorter prompt.")
inputs = torch.tensor([tokens[-model.config.context_length:]], device=device)
logits = model(inputs)[0, -1].float() / temperature
values, indices = logits.topk(min(top_k, logits.numel()))
next_token = indices[torch.multinomial(values.softmax(-1), 1, generator=generator)].item()
if next_token == tokenizer.token_to_id("<eos>"):
finish_reason = "eos"
break
tokens.append(next_token)
if stop_at_sentence_end and len(tokens) - prompt_tokens >= min_new_tokens:
ending = tokenizer.decode([next_token]).rstrip()
completion = tokenizer.decode(tokens[prompt_tokens:])
if ending.endswith((".", "!", "?")) and completion.count('"') % 2 == 0:
finish_reason = "sentence"
break
return {
"text": tokenizer.decode(tokens),
"completion": tokenizer.decode(tokens[prompt_tokens:]),
"prompt_tokens": prompt_tokens, "generated_tokens": len(tokens) - prompt_tokens,
"context_length": model.config.context_length, "finish_reason": finish_reason,
"sliding_context": len(tokens) > model.config.context_length,
}
def main():
parser = argparse.ArgumentParser(description="Continue a simple English prompt")
source = parser.add_mutually_exclusive_group(required=True)
source.add_argument("--checkpoint", help="Local training checkpoint")
source.add_argument("--model-dir", type=Path, help="Downloaded/exported safetensors model directory")
parser.add_argument("--tokenizer", help="Required with --checkpoint; inferred for --model-dir")
parser.add_argument("--prompt", default="Once upon a time, a little rabbit found a blue hat.")
parser.add_argument("--max-new-tokens", type=int, default=150)
parser.add_argument("--temperature", type=float, default=0.8)
parser.add_argument("--top-k", type=int, default=40)
parser.add_argument("--device", choices=["cpu", "cuda"], default="cuda")
parser.add_argument("--threads", type=int, default=8)
parser.add_argument("--seed", type=int, default=42)
args = parser.parse_args()
if args.checkpoint and not args.tokenizer:
parser.error("--tokenizer is required with --checkpoint")
if args.model_dir and args.tokenizer:
parser.error("--model-dir uses its own tokenizer.json; omit --tokenizer")
if not math.isfinite(args.temperature) or args.temperature <= 0 or args.top_k < 1 or args.max_new_tokens < 1 or args.threads < 1:
parser.error("Generation numeric options must be positive")
torch.set_num_threads(args.threads)
model, tokenizer, _ = load_export(args.model_dir, args.device) if args.model_dir else load_model(args.checkpoint, args.tokenizer, args.device)
try:
result = complete(model, tokenizer, args.prompt, args.max_new_tokens, args.temperature, args.top_k, args.seed)
except ValueError as error:
parser.error(str(error))
print(result["text"])
if __name__ == "__main__":
main()
|