| """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)) |
|
|
