core/metrics.py

from collections.abc import Sequence

from sklearn.metrics import mutual_info_score


def compute_mi_scores(
    cluster_ids: Sequence[int],
    reasoning_types: Sequence[str],
    model_ids: Sequence[str],
) -> dict[str, float]:
    if not cluster_ids:
        return {
            "MI(cluster, reasoning_type)": 0.0,
            "MI(cluster, model_identity)": 0.0,
        }

    return {
        "MI(cluster, reasoning_type)": float(mutual_info_score(cluster_ids, reasoning_types)),
        "MI(cluster, model_identity)": float(mutual_info_score(cluster_ids, model_ids)),
    }