core.model_io¶

def generate_run_id() -> str:
    """Produce a unique run directory name: ``YYYY-MM-DD_HHMMSS_run-XXXX``.

    Returns:
        A string like ``2026-02-19_013000_run-0042``.
    """
    now = datetime.now(UTC)
    suffix = f"{random.randint(0, 9999):04d}"
    return f"{now.strftime('%Y-%m-%d_%H%M%S')}_run-{suffix}"

def save_model_bundle(
    model: lgb.Booster,
    metadata: ModelMetadata,
    metrics: dict,
    confusion_df: pd.DataFrame,
    base_dir: Path,
    cat_encoders: dict | None = None,
) -> Path:
    """Persist a complete model bundle into ``base_dir/<run_id>/``.

    Writes the core files per the Model Bundle Contract plus an optional
    ``categorical_encoders.json`` mapping each categorical column to its
    sorted vocabulary list.

    Args:
        model: Trained LightGBM booster.
        metadata: Provenance record (schema hash, label set, params, etc.).
        metrics: Evaluation dict (as returned by
            :func:`~taskclf.core.metrics.compute_metrics`).
        confusion_df: Labelled confusion matrix for CSV export.
        base_dir: Parent directory (e.g. ``Path("models")``).
            A new ``<run_id>/`` subdirectory is created inside it.
        cat_encoders: Optional dict mapping categorical column names to
            fitted ``LabelEncoder`` instances.  Persisted as JSON
            vocabulary lists so inference can reconstruct them.

    Returns:
        Path to the newly created run directory.

    Raises:
        FileExistsError: If the generated run directory already exists.
    """
    run_id = generate_run_id()
    run_dir = base_dir / run_id
    if run_dir.exists():
        raise FileExistsError(f"Run directory already exists: {run_dir}")
    run_dir.mkdir(parents=True)

    model.save_model(str(run_dir / "model.txt"))

    (run_dir / "metadata.json").write_text(json.dumps(metadata.model_dump(), indent=2))

    (run_dir / "metrics.json").write_text(json.dumps(metrics, indent=2))

    confusion_df.to_csv(run_dir / "confusion_matrix.csv")

    if cat_encoders:
        vocab = {col: list(le.classes_) for col, le in cat_encoders.items()}
        (run_dir / "categorical_encoders.json").write_text(json.dumps(vocab, indent=2))

    return run_dir

def load_model_bundle(
    run_dir: Path,
    *,
    validate_schema: bool = True,
    validate_labels: bool = True,
) -> tuple[lgb.Booster, ModelMetadata, dict[str, Any]]:
    """Load a model bundle and optionally validate schema hash and label set.

    Schema validation accepts v1, v2, and v3 bundles: the bundle's
    ``schema_version`` is looked up in the known schema registry and its
    ``schema_hash`` is checked against the corresponding expected hash.

    Args:
        run_dir: Path to an existing run directory (e.g.
            ``models/2026-02-19_013000_run-0042/``).
        validate_schema: When ``True`` (the default), raise if the
            bundle's schema hash does not match the expected hash for
            its declared schema version.
        validate_labels: When ``True`` (the default), raise if the
            bundle's label set differs from the current ``LABEL_SET_V1``.

    Returns:
        A ``(model, metadata, cat_encoders)`` tuple where *cat_encoders*
        is a dict mapping column names to fitted ``LabelEncoder``
        instances. Returns an empty dict when no encoder file exists.

    Raises:
        ValueError: If validation is enabled and the schema hash or label
            set recorded in the bundle does not match the running code.
    """
    from sklearn.preprocessing import LabelEncoder

    model = lgb.Booster(model_file=str(run_dir / "model.txt"))

    raw = json.loads((run_dir / "metadata.json").read_text())
    metadata = ModelMetadata.model_validate(raw)

    if validate_schema:
        expected_hash = _SCHEMA_HASHES.get(metadata.schema_version)
        if expected_hash is None:
            raise ValueError(
                f"Unknown schema version in bundle: {metadata.schema_version!r}"
            )
        if metadata.schema_hash != expected_hash:
            raise ValueError(
                f"Schema hash mismatch: bundle has {metadata.schema_hash!r}, "
                f"expected {expected_hash!r} for schema {metadata.schema_version!r}"
            )

    if validate_labels and sorted(metadata.label_set) != sorted(LABEL_SET_V1):
        raise ValueError(
            f"Label set mismatch: bundle has {sorted(metadata.label_set)!r}, "
            f"current label set is {sorted(LABEL_SET_V1)!r}"
        )

    cat_encoders: dict[str, LabelEncoder] = {}
    enc_path = run_dir / "categorical_encoders.json"
    if enc_path.exists():
        vocab = json.loads(enc_path.read_text())
        for col, classes in vocab.items():
            le = LabelEncoder()
            le.fit(classes)
            cat_encoders[col] = le

    return model, metadata, cat_encoders

def build_metadata(
    label_set: list[str],
    train_date_from: date,
    train_date_to: date,
    params: dict[str, Any],
    *,
    dataset_hash: str,
    reject_threshold: float | None = None,
    data_provenance: Literal["real", "synthetic", "mixed"] = "real",
    unknown_category_freq_threshold: int | None = None,
    unknown_category_mask_rate: float | None = None,
    schema_version: str = LATEST_FEATURE_SCHEMA_VERSION,
) -> ModelMetadata:
    """Convenience builder that fills in schema info and git commit.

    Args:
        label_set: Task-type labels used during training.
        train_date_from: First date of the training range.
        train_date_to: Last date (inclusive) of the training range.
        params: LightGBM (or other model) hyperparameters dict.
        dataset_hash: Deterministic SHA-256 hash of the training dataset
            used for reproducibility auditing.
        reject_threshold: Reject threshold used during evaluation.
        data_provenance: Origin of the training data
            (``"real"``, ``"synthetic"``, or ``"mixed"``).
        unknown_category_freq_threshold: Minimum category frequency
            used during training (categories below this are ``__unknown__``).
        unknown_category_mask_rate: Fraction of known categories randomly
            masked to ``__unknown__`` during training.
        schema_version: ``"v1"``, ``"v2"``, or ``"v3"``.

    Returns:
        A populated ``ModelMetadata`` instance.

    Raises:
        ValueError: If *schema_version* is not recognised.
    """
    schema_hash = _SCHEMA_HASHES.get(schema_version)
    if schema_hash is None:
        raise ValueError(f"Unknown schema version: {schema_version!r}")

    return ModelMetadata(
        schema_version=schema_version,
        schema_hash=schema_hash,
        label_set=sorted(label_set),
        train_date_from=train_date_from.isoformat(),
        train_date_to=train_date_to.isoformat(),
        params=params,
        git_commit=_current_git_commit(),
        dataset_hash=dataset_hash,
        reject_threshold=reject_threshold,
        data_provenance=data_provenance,
        unknown_category_freq_threshold=unknown_category_freq_threshold,
        unknown_category_mask_rate=unknown_category_mask_rate,
    )