src/train.py

import json
from pathlib import Path

import joblib
import pandas as pd
import sklearn

from datasets import load_dataset
from sklearn.compose import ColumnTransformer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import (
    accuracy_score,
    classification_report,
    confusion_matrix,
)
from sklearn.model_selection import StratifiedKFold, cross_val_predict
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder


DATASET_ID = "aigovdev/ai-governance-scenarios"

MODEL_PATH = Path("artifacts/model.joblib")
METRICS_PATH = Path("artifacts/metrics.json")


FEATURES = [
    "sector",
    "impact",
    "decision_autonomy",
    "human_oversight",
    "monitoring",
    "traceability",
    "technical_documentation",
]


def map_risk(label: str) -> str:
    if label in {"low", "limited"}:
        return "lower"
    if label == "high":
        return "high"
    if label == "unacceptable":
        return "unacceptable"

    raise ValueError(f"Unexpected governance_risk label: {label}")


def build_pipeline() -> Pipeline:
    preprocessor = ColumnTransformer(
        transformers=[
            (
                "categorical",
                OneHotEncoder(
                    handle_unknown="ignore",
                ),
                FEATURES,
            )
        ],
        remainder="drop",
    )

    classifier = LogisticRegression(
        max_iter=2000,
        class_weight="balanced",
        random_state=42,
    )

    return Pipeline(
        [
            ("preprocessor", preprocessor),
            ("classifier", classifier),
        ]
    )


def main():
    MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)

    dataset = load_dataset(
        DATASET_ID,
        split="train",
    )

    df = dataset.to_pandas()

    df["risk_tier"] = df["governance_risk"].map(map_risk)

    X = df[FEATURES].copy()
    y = df["risk_tier"].copy()

    print("Dataset:", DATASET_ID)
    print("Examples:", len(df))
    print()
    print("Target distribution:")
    print(y.value_counts().sort_index())
    print()

    cv = StratifiedKFold(
        n_splits=3,
        shuffle=True,
        random_state=42,
    )

    pipeline = build_pipeline()

    predictions = cross_val_predict(
        pipeline,
        X,
        y,
        cv=cv,
    )

    labels = [
        "lower",
        "high",
        "unacceptable",
    ]

    accuracy = accuracy_score(
        y,
        predictions,
    )

    report = classification_report(
        y,
        predictions,
        labels=labels,
        output_dict=True,
        zero_division=0,
    )

    matrix = confusion_matrix(
        y,
        predictions,
        labels=labels,
    )

    print("Stratified 3-Fold Cross-Validation")
    print("=" * 42)

    print(
        classification_report(
            y,
            predictions,
            labels=labels,
            digits=3,
            zero_division=0,
        )
    )

    print("Confusion matrix")
    print("Labels:", labels)
    print(matrix)
    print()

    pipeline.fit(
        X,
        y,
    )

    joblib.dump(
        pipeline,
        MODEL_PATH,
    )

    metrics = {
        "dataset": DATASET_ID,
        "examples": int(len(df)),
        "target_distribution": {
            key: int(value)
            for key, value in y.value_counts().to_dict().items()
        },
        "evaluation": {
            "method": "stratified_3_fold_cross_validation",
            "accuracy": float(accuracy),
            "macro_precision": float(
                report["macro avg"]["precision"]
            ),
            "macro_recall": float(
                report["macro avg"]["recall"]
            ),
            "macro_f1": float(
                report["macro avg"]["f1-score"]
            ),
            "confusion_matrix_labels": labels,
            "confusion_matrix": matrix.tolist(),
        },
        "features": FEATURES,
        "risk_mapping": {
            "low": "lower",
            "limited": "lower",
            "high": "high",
            "unacceptable": "unacceptable",
        },
        "runtime": {
            "scikit_learn": sklearn.__version__,
        },
        "limitations": [
            "Dataset contains only 12 synthetic scenarios.",
            "Evaluation is illustrative and not a production benchmark.",
            "Risk tiers are engineering labels, not legal classifications.",
        ],
    }

    METRICS_PATH.write_text(
        json.dumps(
            metrics,
            indent=2,
        ),
        encoding="utf-8",
    )

    print("Saved model:", MODEL_PATH.resolve())
    print("Saved metrics:", METRICS_PATH.resolve())


if __name__ == "__main__":
    main()