SupplyChain-JEPA

SupplyChain-JEPA is a predictive world model for structured supply-chain data. It maps partially observed entity-time tokens into latent operational states and predicts both hidden/future state representations and action-conditioned future representations.

Code and documentation: GitHub repository · Documentation index · Getting started

This repository contains the selected pretrained model release. Downstream task heads and task-specific adaptations are intentionally kept separate from the shared pretrained weights.

Code, documentation, and weights

Where	What it holds
GitHub: htizhang/SupplyChain-JEPA	The supplychain_jepa Python package (verified loader, tensorizer, model, training, and inference API), adapters for the ten pretraining sources, the three evaluation protocols, tests, and all documentation
This Hugging Face repository	The pretrained weights and the fitted tensor interface (tensorizer, schema, normalization statistics), plus the release manifest, checksums, and a standalone loader

Both repositories are public. Downloading the weights needs no Hugging Face account or login; signing in is optional. Raw datasets are not distributed in either place and remain under their original terms.

Documentation

The guides explain the method for operations research and operations management readers, with small worked examples and no deep-learning background assumed. Read the four conceptual guides in order, or start with the practical guide.

	Guide	What it covers
1	Data to episodes	How ten real and simulated sources become canonical episodes and tokens
2	State learning	How the model learns a latent operating state from records that arrive late, censored, or not at all
3	Dynamics learning	How paired simulator runs teach what a replenishment plan does to the future
4	Using the model	How three applications turn the frozen state into a demand estimate, a review list, and a plan choice
Practical	Getting started	Installation, loading the pinned weights, CPU smoke checks, preparing a new dataset, and reproducing the examples

The documentation index also lists the technical reference: architecture, training, data contract, data adapters, evaluation, and limitations.

Loading

Recommended: the supplychain_jepa package

Install the package from GitHub with Python 3.11 to 3.13 and uv. The Quick start gives the full installation steps.

git clone https://github.com/htizhang/SupplyChain-JEPA.git
cd SupplyChain-JEPA
uv sync --all-extras

from supplychain_jepa import load_pretrained

bundle = load_pretrained(device="cpu")
model = bundle.model

load_pretrained pins this repository's immutable revision 493786fda0855b72f61235068c62cc7b2a11fbec by default, verifies the release checksums, and returns the loaded modules together with the fitted tensor interface. The package also provides encode_state, predict_masked_state, predict_dynamics, and compare_actions over the loaded bundle; Using the model shows how the three applications use them.

Standalone loader

The files in this repository can also be loaded without the package:

from huggingface_hub import snapshot_download

snapshot = snapshot_download(
    repo_id="htizhang/SupplyChain-JEPA",
    revision="493786fda0855b72f61235068c62cc7b2a11fbec",
)

import sys
sys.path.insert(0, snapshot)
from load_model import load_supplychain_jepa

bundle = load_supplychain_jepa(snapshot)
model = bundle["model"]
state_predictor = model.state_predictor
dynamics_predictor = model.dynamics_predictor
state_pool_predictor = bundle["state_pool_predictor"]

Always pin an immutable repository commit when using the model in an experiment. The release reruns linked below use this pinned revision.

Architecture

The release follows one continuous weight chain:

1. a context encoder for visible operational history;
2. a stop-gradient target encoder defining latent prediction targets;
3. a lightweight State predictor for hidden or future operational state;
4. an action-conditioned Dynamics predictor for counterfactual future state;
5. optional semantic-state and grounding heads.

The core contains 73,629,696 parameters. The context encoder width is 384 with six Transformer layers. Inputs use a unified entity-time token schema spanning numerical values, categorical values, observation status, time, entity roles, and local topology indices.

The Architecture and Training documents give the component sizes, objectives, and training configuration.

Files

• model.safetensors: encoder chain plus State and Dynamics predictors.
• state_auxiliary.safetensors: pooled-state, semantic-state, and grounding heads.
• config.json: architecture and component configuration.
• tensorizer.json: fitted feature and vocabulary mapping.
• schema.json: canonical tensor interface.
• normalization.json: numerical normalization statistics.
• release_manifest.json: source checkpoint hashes, component tensor hashes, architecture, and parameter counts.
• modeling_supplychain_jepa.py: model definition.
• load_model.py: strict loader with checksum verification.
• requirements.txt: Python requirements for the standalone loader.
• CHECKSUMS.sha256: release-file integrity manifest.

Evaluation

Three applications test whether one pretrained representation can support distinct, partially observed tasks. Each keeps a conventional baseline fixed and fits a small task readout on limited labeled data; plan selection also adapts a copy of the pretrained Dynamics predictor. All three were rerun against the pinned revision.

Application	Metric	Conventional baseline	SupplyChain-JEPA system	Improvement
Stockout demand recovery	WAPE, lower is better	0.416952	0.399117	4.28% lower
Early SLA-risk ranking under delayed reports	AUPRC, higher is better	0.386170	0.412200	6.74% higher
Replenishment-plan selection	Regret, lower is better	9.274	5.883	36.6% lower

For plan selection, a same-architecture matched-random diagnostic obtains regret 6.119; the pretrained system is 3.85% lower and wins three of five world splits, a weaker attribution result than the comparison with the fixed baseline. Evaluation documents the protocols, including a negative boundary result.

Intended use

The model is intended for research on partially observed supply-chain state estimation, demand recovery, operational risk ranking, action evaluation, and optimization-support systems. Dataset-specific adapters are required to map raw records into the canonical token schema. The GitHub repository includes adapters for the ten pretraining sources, and Getting started explains how to prepare a new dataset.

Limitations

This is a research model rather than a production decision system. It does not enforce feasibility, inventory balance, service constraints, or safety policies by itself. Operational deployment should place explicit constraints and an optimization or control layer around model predictions. See Limitations for the remaining boundaries.

Terms

See the code repository's licensing notice. Datasets and third-party software remain under their own terms; see Datasets and Third-party notices.