File size: 1,677 Bytes
921bf58 | 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 | """Load the inference-only FP32 Purrence package without training state."""
import argparse
import json
from pathlib import Path
from dataclasses import fields
import torch
from safetensors.torch import load_file
from purrence_runtime.config import ModelConfig
from purrence_runtime.model import RecurrentLM
from purrence_runtime.tokenizer import Sub3MTokenizer
from purrence_runtime.scorer import Scorer
def load_model(directory=None, device="cpu"):
root = Path(directory) if directory else Path(__file__).resolve().parent
values = json.loads((root / "config.json").read_text())
config = ModelConfig(**{f.name: values[f.name] for f in fields(ModelConfig)})
model = RecurrentLM(config)
state = load_file(str(root / "model.safetensors"), device="cpu")
model.load_state_dict(state, strict=True)
if any(not torch.isfinite(p).all() for p in model.parameters()):
raise ValueError("Nonfinite model weights")
if model.parameter_count() != 2987712:
raise ValueError("Inference parameter count differs from the qualified model")
model = model.to(device).float().eval()
tokenizer = Sub3MTokenizer(root / "tokenizer")
return model, tokenizer
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("--directory", type=Path)
p.add_argument("--device", default="cpu")
p.add_argument("--prompt", default="Once upon a time")
p.add_argument("--max-new-tokens", type=int, default=32)
args = p.parse_args()
torch.set_num_threads(2)
model, tokenizer = load_model(args.directory, args.device)
print(Scorer(model, tokenizer).generate(args.prompt, max_new_tokens=args.max_new_tokens))
|