File size: 2,368 Bytes
689a869 | 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 | """Load one model of this repository on the CPU and run it.
Usage, from the downloaded repository:
python example.py # the Low model, float32
python example.py --tier max --dtype bfloat16
Needs only torch and safetensors (see requirements.txt). Nothing is written to disk.
"""
import argparse
import os
import sys
sys.dont_write_bytecode = True # leave the downloaded folder unchanged
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import torch # noqa: E402
from modeling_dense import DenseDecoder # noqa: E402
TIERS = ["low", "medium", "high", "extra", "max"]
# GPT-NeoX-20B token ids of "The capital of France is". Any copy of that tokenizer gives them, for example
# tokenizers.Tokenizer.from_pretrained("EleutherAI/gpt-neox-20b").encode("The capital of France is").ids
PROMPT_IDS = [510, 5347, 273, 6181, 310]
def main() -> None:
parser = argparse.ArgumentParser(description="Run one model of this repository on the CPU.")
parser.add_argument("--tier", default="low", choices=TIERS, help="model folder to load (default: low)")
parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"],
help="compute precision (the released weights are float32)")
parser.add_argument("--new-tokens", type=int, default=8, help="greedy tokens to generate")
args = parser.parse_args()
model = DenseDecoder.from_pretrained(os.path.join(HERE, args.tier), device="cpu",
dtype=getattr(torch, args.dtype))
cfg = model.config
n_params = sum(p.numel() for p in model.parameters())
print(f"{args.tier}: width {cfg.hidden_size}, {cfg.num_attention_heads} heads, FFN width "
f"{cfg.intermediate_size}, {n_params:,} parameters, {args.dtype}")
ids = torch.tensor([PROMPT_IDS])
with torch.no_grad():
logits, loss = model(ids[:, :-1], labels=ids[:, 1:])
out = model.generate(ids, max_new_tokens=args.new_tokens)
assert logits.shape == (1, len(PROMPT_IDS) - 1, cfg.vocab_size)
assert torch.isfinite(logits).all() and torch.isfinite(loss)
print(f"prompt ids: {PROMPT_IDS}")
print(f"mean next-token loss on the prompt: {loss.item():.4f}")
print(f"greedy continuation ids: {out[0, len(PROMPT_IDS):].tolist()}")
if __name__ == "__main__":
main()
|