| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| import torch |
| from safetensors.torch import load_file |
|
|
| from modeling_supplychain_jepa import ( |
| ModelConfig, |
| SemanticConceptSlotPredictor, |
| StatePoolPredictor, |
| SupplyChainJEPACore, |
| ) |
|
|
|
|
| def _sha256(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| for block in iter(lambda: handle.read(1024 * 1024), b""): |
| digest.update(block) |
| return digest.hexdigest() |
|
|
|
|
| def verify_release(directory: str | Path) -> None: |
| root = Path(directory) |
| for line in (root / "CHECKSUMS.sha256").read_text().splitlines(): |
| expected, name = line.split(" ", 1) |
| actual = _sha256(root / name) |
| if actual != expected: |
| raise ValueError(f"Checksum mismatch for {name}: {actual} != {expected}") |
|
|
|
|
| def load_supplychain_jepa(directory: str | Path, device: str = "cpu", verify: bool = True): |
| root = Path(directory) |
| if verify: |
| verify_release(root) |
| config = json.loads((root / "config.json").read_text()) |
| model_config = ModelConfig(**config["model"]) |
| model = SupplyChainJEPACore(model_config) |
| model.load_state_dict(load_file(root / config["weights"]["core"], device="cpu"), strict=True) |
|
|
| raw_aux = load_file(root / config["weights"]["state_auxiliary"], device="cpu") |
| pool_cfg = config["state_pool_predictor"] |
| state_pool_predictor = None |
| if pool_cfg["enabled"]: |
| state_pool_predictor = StatePoolPredictor( |
| model_config.d_model, |
| query_vocab_size=max(pool_cfg["query_vocabulary"].values()) + 1, |
| mask_vocab_size=max(pool_cfg["mask_vocabulary"].values()) + 1, |
| hidden_multiplier=pool_cfg["hidden_multiplier"], |
| dropout=pool_cfg["dropout"], |
| ) |
| state_pool_predictor.load_state_dict({k.split(".", 1)[1]: v for k, v in raw_aux.items() if k.startswith("state_pool_predictor.")}, strict=True) |
|
|
| slot_cfg = config["semantic_concept_slot_predictor"] |
| semantic_concept_slot_predictor = None |
| if slot_cfg["enabled"]: |
| slot_state = {k.split(".", 1)[1]: v for k, v in raw_aux.items() if k.startswith("semantic_concept_slot_predictor.")} |
| semantic_concept_slot_predictor = SemanticConceptSlotPredictor( |
| model_config.d_model, |
| concept_count=slot_state["concept_queries"].shape[1], |
| nhead=model_config.nhead, |
| decoder_layers=slot_cfg["decoder_layers"], |
| dropout=slot_cfg["dropout"], |
| ) |
| semantic_concept_slot_predictor.load_state_dict(slot_state, strict=True) |
|
|
| grounding_state = {k.split(".", 1)[1]: v for k, v in raw_aux.items() if k.startswith("grounding_head.")} |
| grounding_head = torch.nn.Linear(model_config.d_model, grounding_state["weight"].shape[0]) |
| grounding_head.load_state_dict(grounding_state, strict=True) |
|
|
| modules = [model, state_pool_predictor, semantic_concept_slot_predictor, grounding_head] |
| for module in modules: |
| if module is not None: |
| module.to(device).eval() |
| for parameter in module.parameters(): |
| parameter.requires_grad_(False) |
| return { |
| "model": model, |
| "state_pool_predictor": state_pool_predictor, |
| "semantic_concept_slot_predictor": semantic_concept_slot_predictor, |
| "grounding_head": grounding_head, |
| "config": config, |
| "tensorizer": json.loads((root / "tensorizer.json").read_text()), |
| "schema": json.loads((root / "schema.json").read_text()), |
| "normalization": json.loads((root / "normalization.json").read_text()), |
| } |
|
|
