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