| """Self-contained, pickle-free model folders (for Hugging Face Hub or local use). |
| |
| out/ |
| config.json MiniJEV config + calibration temperatures |
| model.safetensors MiniJEV weights (without the text encoder) |
| text_encoder/ pretrained text encoder (fine-tuned weights if it was trained) + tokenizer |
| |
| minijev export-hf --checkpoint runs/x/checkpoints/final.pt --out hf/minijev-general |
| dec = TorchDecider.from_pretrained("markuz89/minijev-general") # or a local folder |
| |
| Loading never unpickles anything: weights are safetensors, the config is JSON. |
| """ |
| from __future__ import annotations |
|
|
| import copy |
| import json |
| import shutil |
| from pathlib import Path |
|
|
| FORMAT = "minijev-v2" |
|
|
|
|
| def save_pretrained(checkpoint: str | Path, out: str | Path, extra: dict | None = None) -> Path: |
| import torch |
| from safetensors.torch import save_file |
|
|
| from minijev.training.supervised import make_model |
|
|
| out = Path(out) |
| out.mkdir(parents=True, exist_ok=True) |
| ck = torch.load(checkpoint, map_location="cpu", weights_only=False) |
| cfg = copy.deepcopy(ck["config"]) |
| model = make_model(cfg) |
| model.load_state_dict(ck["model"]) |
| te = cfg["model"].get("text_encoder") or {} |
| if model.backbone is not None: |
| enc_dir = out / "text_encoder" |
| model.backbone.model.save_pretrained(enc_dir, safe_serialization=True) |
| tok_src = _tokenizer_file(te.get("tokenizer") or te["name"]) |
| shutil.copy(tok_src, enc_dir / "tokenizer.json") |
| cfg["model"]["text_encoder"] = {**te, "name": "text_encoder", "tokenizer": "text_encoder"} |
| sd, seen = {}, set() |
| for k, v in model.state_dict().items(): |
| if k.startswith("backbone.") or v.data_ptr() in seen: |
| continue |
| seen.add(v.data_ptr()) |
| sd[k] = v.detach().cpu().contiguous() |
| save_file(sd, out / "model.safetensors", metadata={"format": FORMAT}) |
| meta = {"format": FORMAT, "config": cfg, "temperature": float(model.log_temperature.exp()), |
| "task_temperatures": ck.get("task_temperatures") or {}, **(extra or {})} |
| (out / "config.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False, default=str)) |
| return out |
|
|
|
|
| def _tokenizer_file(name: str) -> Path: |
| p = Path(name) |
| if p.is_dir(): |
| return p / "tokenizer.json" |
| from huggingface_hub import hf_hub_download |
| return Path(hf_hub_download(name, "tokenizer.json")) |
|
|
|
|
| def resolve(path_or_repo: str | Path) -> Path: |
| p = Path(path_or_repo) |
| if p.is_dir(): |
| return p |
| from huggingface_hub import snapshot_download |
| return Path(snapshot_download(str(path_or_repo))) |
|
|
|
|
| def load_pretrained(path_or_repo: str | Path, device: str = "cpu"): |
| """-> (model, cfg, task_temperatures) with local paths filled in.""" |
| from safetensors.torch import load_file |
|
|
| from minijev.training.supervised import make_model |
|
|
| root = resolve(path_or_repo) |
| meta = json.loads((root / "config.json").read_text()) |
| if meta.get("format") != FORMAT: |
| raise ValueError(f"{root} is not a {FORMAT} folder") |
| cfg = meta["config"] |
| te = cfg["model"].get("text_encoder") or {} |
| if te.get("type") == "pretrained": |
| cfg["model"]["text_encoder"] = {**te, "name": str(root / te["name"]), |
| "tokenizer": str(root / te.get("tokenizer", te["name"]))} |
| model = make_model(cfg) |
| weights = load_file(root / "model.safetensors", device="cpu") |
| res = model.load_state_dict(weights, strict=False) |
| ptr = {v.data_ptr() for k, v in model.state_dict().items() if k in weights} |
| full = model.state_dict() |
| missing = [k for k in res.missing_keys if not k.startswith(("backbone.", "core.scalar_head.")) |
| and full[k].data_ptr() not in ptr] |
| if missing or res.unexpected_keys: |
| raise RuntimeError(f"weights mismatch: missing={missing[:10]} unexpected={res.unexpected_keys[:10]}") |
| return model.to(device).eval(), cfg, meta.get("task_temperatures") or {} |
|
|
