File size: 4,067 Bytes
45507bb | 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 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 | """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: # tied weights (shared byte embedding) once
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] # tied to a loaded tensor
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 {}
|