| """Resolve CUE checkpoint layouts (native cue-hf, joint training, decoder refinement).""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
| from typing import Any |
|
|
| import torch |
|
|
| from .configuration_cue import CueConfig |
|
|
| JOINT_CONFIG_NAME = "unified_config.json" |
| JOINT_WEIGHTS_NAME = "model.pt" |
| REFINED_DECODER_STATE = "decoder_state.pt" |
| REFINED_CONFIG_NAME = "refinement_config.json" |
| JOINT_CHECKPOINT_KEY = "joint_training_checkpoint" |
|
|
| LAYOUT_NATIVE = "native" |
| LAYOUT_JOINT = "joint" |
| LAYOUT_REFINED = "refined" |
|
|
|
|
| def resolve_local_dir( |
| path_or_repo: str | Path, |
| *, |
| revision: str | None = None, |
| token: str | bool | None = None, |
| ) -> Path: |
| """Return a local directory for a path or a Hub repo id.""" |
|
|
| candidate = Path(path_or_repo) |
| if candidate.is_dir(): |
| return candidate |
| from huggingface_hub import snapshot_download |
|
|
| return Path( |
| snapshot_download(str(path_or_repo), revision=revision, token=token) |
| ) |
|
|
|
|
| def detect_layout(directory: str | Path) -> str: |
| directory = Path(directory) |
| if _refined_config_path(directory) is not None: |
| return LAYOUT_REFINED |
| if (directory / JOINT_CONFIG_NAME).is_file(): |
| return LAYOUT_JOINT |
| if (directory / "config.json").is_file(): |
| 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}, {REFINED_CONFIG_NAME})" |
| ) |
|
|
|
|
| def _refined_config_path(directory: Path) -> Path | None: |
| path = directory / REFINED_CONFIG_NAME |
| return path if path.is_file() else None |
|
|
|
|
| def config_from_joint_dir(directory: str | Path) -> CueConfig: |
| payload = json.loads( |
| (Path(directory) / JOINT_CONFIG_NAME).read_text(encoding="utf-8") |
| ) |
| return CueConfig.from_cue_config(payload) |
|
|
|
|
| def load_joint_state( |
| directory: str | Path, *, map_location: str | torch.device = "cpu" |
| ) -> dict[str, Any]: |
| payload = torch.load( |
| Path(directory) / JOINT_WEIGHTS_NAME, map_location=map_location, weights_only=False |
| ) |
| return payload["model"] |
|
|
|
|
| def resolve_refined(directory: str | Path) -> tuple[Path, Path]: |
| """Return (joint checkpoint dir, decoder state path) for a refinement checkpoint.""" |
|
|
| directory = Path(directory) |
| config_path = _refined_config_path(directory) |
| if config_path is None: |
| raise ValueError(f"{directory} has no refinement config") |
| payload = json.loads(config_path.read_text(encoding="utf-8")) |
| joint = payload.get(JOINT_CHECKPOINT_KEY) |
| if not joint: |
| raise ValueError( |
| f"{config_path} has no {JOINT_CHECKPOINT_KEY}; cannot find the " |
| "joint training checkpoint holding the encoder" |
| ) |
| decoder_state = directory / REFINED_DECODER_STATE |
| if not decoder_state.is_file(): |
| raise FileNotFoundError(f"refinement checkpoint missing {decoder_state}") |
| return Path(str(joint)), decoder_state |
|
|
