Skip to content

pelinker.model_selection_checkpoint

Structured checkpoint I/O for pelinker.model_selection runs.

FailureRecord dataclass

One recorded failure for a combination.

Source code in pelinker/model_selection_checkpoint.py
@dataclass
class FailureRecord:
    """One recorded failure for a combination."""

    combination_key: str
    error: str
    at: str

ModelSelectionCheckpoint dataclass

On-disk checkpoint for resumable model selection runs.

Source code in pelinker/model_selection_checkpoint.py
@dataclass
class ModelSelectionCheckpoint:
    """On-disk checkpoint for resumable model selection runs."""

    version: int = CHECKPOINT_VERSION
    run_fingerprint: str = ""
    created_at: str = ""
    updated_at: str = ""
    completed_combinations: list[str] = field(default_factory=list)
    summaries_by_key: dict[str, dict[str, str | float | None]] = field(
        default_factory=dict
    )
    singleton_scores_by_key: dict[str, float] = field(default_factory=dict)
    stages: dict[str, StageState] = field(default_factory=dict)
    failures: list[FailureRecord] = field(default_factory=list)
    checkpoint_fusion_pairs: int | None = None
    checkpoint_fusion_triples: int | None = None

    def to_json_dict(self) -> dict[str, Any]:
        return {
            "version": self.version,
            "run_fingerprint": self.run_fingerprint,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "checkpoint_fusion_pairs": self.checkpoint_fusion_pairs,
            "checkpoint_fusion_triples": self.checkpoint_fusion_triples,
            "completed_combinations": sorted(self.completed_combinations),
            "summaries_by_key": {
                k: dict(v) for k, v in sorted(self.summaries_by_key.items())
            },
            "singleton_scores_by_key": {
                k: float(v) for k, v in sorted(self.singleton_scores_by_key.items())
            },
            "stages": dict(sorted(self.stages.items())),
            "failures": [
                {"combination_key": f.combination_key, "error": f.error, "at": f.at}
                for f in self.failures
            ],
        }

    @staticmethod
    def from_json_dict(data: dict[str, Any]) -> ModelSelectionCheckpoint:
        if int(data.get("version", -1)) != CHECKPOINT_VERSION:
            raise ValueError(
                f"Unsupported checkpoint version: {data.get('version')!r}; "
                f"expected {CHECKPOINT_VERSION}"
            )
        failures_raw = data.get("failures") or []
        failures = [
            FailureRecord(
                combination_key=str(fr["combination_key"]),
                error=str(fr["error"]),
                at=str(fr["at"]),
            )
            for fr in failures_raw
        ]
        cfp = data.get("checkpoint_fusion_pairs")
        cft = data.get("checkpoint_fusion_triples")
        return ModelSelectionCheckpoint(
            version=int(data["version"]),
            run_fingerprint=str(data["run_fingerprint"]),
            created_at=str(data.get("created_at", "")),
            updated_at=str(data.get("updated_at", "")),
            completed_combinations=list(data.get("completed_combinations") or []),
            summaries_by_key=dict(data.get("summaries_by_key") or {}),
            singleton_scores_by_key={
                str(k): float(v)
                for k, v in (data.get("singleton_scores_by_key") or {}).items()
            },
            stages=dict(data.get("stages") or {}),
            failures=failures,
            checkpoint_fusion_pairs=int(cfp) if cfp is not None else None,
            checkpoint_fusion_triples=int(cft) if cft is not None else None,
        )