"""Strict loader preserving the complete trained UniTok state, including semantic heads."""
import json
from pathlib import Path
from types import SimpleNamespace
import torch
from safetensors.torch import load_file
from native.models.unitok import UniTok


def load_model(path=None, dtype=torch.float32, device='cpu'):
    root = Path(path or Path(__file__).parent)
    cfg = json.loads((root / 'config.json').read_text())
    cfg.update(device='cpu', grad_ckpt=False)
    # The native hybrid embedder runs a CPU shape probe in its constructor.
    model = UniTok(SimpleNamespace(**cfg))
    model.load_state_dict(load_file(root / 'model.safetensors'), strict=True, assign=True)
    return model.to(device=device, dtype=dtype).eval()