inference.py

import argparse
import json
from pathlib import Path

import joblib
import pandas as pd


MODEL_PATH = Path(__file__).parent / "model.joblib"

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


def load_model():
    if not MODEL_PATH.exists():
        raise FileNotFoundError(
            f"Model artifact not found: {MODEL_PATH}"
        )

    return joblib.load(MODEL_PATH)


def validate_input(data: dict) -> None:
    missing = [feature for feature in FEATURES if feature not in data]

    if missing:
        raise ValueError(
            "Missing required features: " + ", ".join(missing)
        )


def predict_governance_risk(data: dict) -> dict:
    validate_input(data)

    model = load_model()

    frame = pd.DataFrame(
        [{feature: data[feature] for feature in FEATURES}]
    )

    prediction = model.predict(frame)[0]

    result = {
        "risk_tier": str(prediction),
    }

    if hasattr(model, "predict_proba"):
        probabilities = model.predict_proba(frame)[0]
        classes = model.classes_

        result["class_probabilities"] = {
            str(label): round(float(probability), 6)
            for label, probability in zip(classes, probabilities)
        }

    return result


def main():
    parser = argparse.ArgumentParser(
        description="Run the AIGov governance risk classifier."
    )

    parser.add_argument(
        "--json",
        required=True,
        help="Governance scenario encoded as JSON.",
    )

    args = parser.parse_args()

    data = json.loads(args.json)
    result = predict_governance_risk(data)

    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()