Upload folder using huggingface_hub

checkpoint.py

@@ -13,14 +13,8 @@ from .configuration_cue import CueConfig
 JOINT_CONFIG_NAME = "unified_config.json"
 JOINT_WEIGHTS_NAME = "model.pt"
 REFINED_DECODER_STATE = "decoder_state.pt"
-# Newest first; `stage2_config.json` predates the pipeline rename.
-REFINED_CONFIG_NAMES = (
-    "refinement_config.json",
-    "decoder_refinement_config.json",
-    "stage2_config.json",
-)
-# Newest first; `stage1_checkpoint` predates the rename.
-JOINT_CHECKPOINT_KEYS = ("joint_training_checkpoint", "stage1_checkpoint")
+REFINED_CONFIG_NAME = "refinement_config.json"
+JOINT_CHECKPOINT_KEY = "joint_training_checkpoint"
 
 LAYOUT_NATIVE = "native"
 LAYOUT_JOINT = "joint"
@@ -55,16 +49,13 @@ def detect_layout(directory: str | Path) -> str:
         return LAYOUT_NATIVE
     raise ValueError(
         f"{directory} is not a cue-hf, joint training, or refinement checkpoint "
-        f"(looked for config.json, {JOINT_CONFIG_NAME}, {'/'.join(REFINED_CONFIG_NAMES)})"
+        f"(looked for config.json, {JOINT_CONFIG_NAME}, {REFINED_CONFIG_NAME})"
     )
 
 
 def _refined_config_path(directory: Path) -> Path | None:
-    for name in REFINED_CONFIG_NAMES:
-        path = directory / name
-        if path.is_file():
-            return path
-    return None
+    path = directory / REFINED_CONFIG_NAME
+    return path if path.is_file() else None
 
 
 def config_from_joint_dir(directory: str | Path) -> CueConfig:
@@ -91,12 +82,10 @@ def resolve_refined(directory: str | Path) -> tuple[Path, Path]:
     if config_path is None:
         raise ValueError(f"{directory} has no refinement config")
     payload = json.loads(config_path.read_text(encoding="utf-8"))
-    joint = next(
-        (payload[key] for key in JOINT_CHECKPOINT_KEYS if payload.get(key)), None
-    )
+    joint = payload.get(JOINT_CHECKPOINT_KEY)
     if not joint:
         raise ValueError(
-            f"{config_path} has none of {JOINT_CHECKPOINT_KEYS}; cannot find the "
+            f"{config_path} has no {JOINT_CHECKPOINT_KEY}; cannot find the "
             "joint training checkpoint holding the encoder"
         )
     decoder_state = directory / REFINED_DECODER_STATE

configuration_cue.py

@@ -12,6 +12,7 @@ from .encoder_model import (
     DEFAULT_PASSAGE_PREFIX,
     DEFAULT_QUERY_PREFIX,
 )
+from .session_preprocess import normalize_session_preprocess
 
 # Defaults mirror the final CUE joint-training config (ModernBERT turn encoder).
 DEFAULT_ENCODER_CONFIG: dict[str, Any] = {
@@ -59,7 +60,7 @@ class CueConfig(PretrainedConfig):
         style_command_slots: int = 4,
         slot_embedding_dim: int = 64,
         target_max_tokens: int = 2048,
-        session_preprocess: str = "full",
+        session_preprocess: str | None = "full",
         decode_temperature: float = 0.0,
         decode_max_new_tokens: int = 2048,
         slot_dedup_jaccard: float = 0.8,
@@ -81,7 +82,9 @@ class CueConfig(PretrainedConfig):
         self.style_command_slots = int(style_command_slots)
         self.slot_embedding_dim = int(slot_embedding_dim)
         self.target_max_tokens = int(target_max_tokens)
-        self.session_preprocess = str(session_preprocess)
+        # A config exported without an explicit mode carries null, and str(None) would reach
+        # the encoder as the literal "None".
+        self.session_preprocess = normalize_session_preprocess(session_preprocess)
         self.decode_temperature = float(decode_temperature)
         self.decode_max_new_tokens = int(decode_max_new_tokens)
         self.slot_dedup_jaccard = float(slot_dedup_jaccard)

modeling_cue.py

@@ -23,7 +23,7 @@ from .decoder_model import CueDecoder
 from .encoder_model import CueEncoder
 from .example_pool import META_NAME, POOL_DIRNAME, ExamplePool, inject_examples
 from .log import log, warn
-from .sampler_sample import CueSampler, load_sampler
+from .sampler_sample import SAMPLER_NAME, CueSampler, load_sampler
 from .schema import (
     COMMAND_BLOCK_SENTINEL,
     COMMAND_SLOT_SENTINEL,
@@ -156,8 +156,14 @@ class CueModel(PreTrainedModel):
     def attach_sampler(self, checkpoint: str | Path, **kwargs: Any) -> CueSampler:
         """Load a diffusion prior (``sampler.pt``) from a directory or Hub repo."""
 
-        directory = resolve_local_dir(checkpoint, **kwargs)
-        self._sampler = load_sampler(directory, device=str(self.device))
+        source = Path(checkpoint)
+        if not source.is_dir():
+            from huggingface_hub import hf_hub_download
+
+            # Refined repos point sampler_id at the joint repo, so snapshotting it would
+            # pull that model's weights too: fetch the one file instead.
+            source = Path(hf_hub_download(str(checkpoint), SAMPLER_NAME, **kwargs))
+        self._sampler = load_sampler(source, device=str(self.device))
         return self._sampler
 
     def attach_example_pool(self, source: str | Path, *, split: str = "train") -> ExamplePool:

sampler_sample.py

@@ -229,6 +229,51 @@ def load_sampler(
     )
 
 
+"""Keys ``load_sampler`` reads, plus provenance scalars. Everything else is training state."""
+INFERENCE_KEYS = (
+    "schedule",
+    "mu",
+    "sigma",
+    "layernorm_weight",
+    "layernorm_bias",
+    "dim",
+    "config",
+    "checkpoint_hash",
+    "step",
+)
+
+
+def slim_sampler_payload(payload: dict[str, Any]) -> dict[str, Any]:
+    """Drop optimizer/scheduler/scaler state and the duplicate weight copy.
+
+    Training checkpoints hold ``model``, an EMA shadow, and optimizer moments. Inference
+    uses the EMA weights, so those are promoted into ``model`` and the rest is dropped:
+    one copy that loads identically whether or not ``use_ema`` is set.
+    """
+
+    ema = payload.get("ema")
+    if isinstance(ema, dict) and "shadow" in ema:
+        weights = ema["shadow"]
+    elif "model" in payload:
+        weights = payload["model"]
+    else:
+        raise KeyError("sampler payload has neither ema['shadow'] nor model weights")
+    slim: dict[str, Any] = {"model": weights}
+    slim.update({key: payload[key] for key in INFERENCE_KEYS if key in payload})
+    return slim
+
+
+def slim_sampler_file(source: str | Path, destination: str | Path) -> Path:
+    """Rewrite a ``sampler.pt`` with only the tensors inference needs."""
+
+    source, destination = Path(source), Path(destination)
+    if source.is_dir():
+        source = source / SAMPLER_NAME
+    payload = torch.load(source, map_location="cpu", mmap=True, weights_only=False)
+    torch.save(slim_sampler_payload(payload), destination)
+    return destination
+
+
 def as_condition_tensor(
     condition: np.ndarray | torch.Tensor | list,
     *,