compute_trainer_pair_encoding_matrix

compute_trainer_pair_encoding_matrix

compute_trainer_pair_encoding_matrix(trainers, *, split='test', max_size=None, use_cache=True, metric='positive_mean', full_eval=None, encoder_eval='auto', allow_incomplete=False, seed_selection=None, seed_aggregate='mean', dispersion=None, positive_class_by_trainer=None)

Compare how trainers encode each other's binary target-pair data.

Rows are source trainers. Columns are target class pairs. For every row trainer, the function obtains encoder outputs on a full evaluation set, subsets those outputs to the column trainer's two target classes, and aggregates the requested metric. Binary pairs must either use a non_/non- negative class naming convention or set TrainingArguments.positive_class explicitly.

Parameters:

Name	Type	Description	Default
trainers	Dict[str, object]	Mapping from display/trainer id to trainer object. Each trainer must expose exactly two target_classes.	required
split	str	Data split for encoder evaluation and cache lookup.	'test'
max_size	Optional[int]	Optional maximum number of examples used by encoder evaluation/cache keys.	None
use_cache	bool	Whether cached encoder outputs may be loaded in encoder_eval="auto" mode.	True
metric	str	Cell metric. Supported values are "positive_mean", "negative_mean", and "positive_minus_negative".	'positive_mean'
full_eval	Optional[bool]	Controls whether encoder evaluation includes other classes. None defaults to True for split="test" and False otherwise. This becomes include_other_classes in trainer.evaluate_encoder.	None
encoder_eval	Literal['auto', 'cached', 'recompute']	Encoder evaluation policy. "auto" loads cached encoder outputs when available and computes missing rows; "cached" requires existing cached encoder outputs; "recompute" ignores encoder caches and recomputes rows.	'auto'
allow_incomplete	bool	If True, missing row data or empty row/column subsets produce NaN cells instead of raising.	False
seed_selection	Optional[str]	"best" compares the selected best model/cache. "all_convergent" evaluates saved convergent seed runs and aggregates their per-seed scores.	None
seed_aggregate	str	Aggregation for multi-seed cells. Supported values are "mean", "median", "min", and "max".	'mean'
dispersion	Optional[str]	Optional dispersion metadata for multi-seed cells. Supported values are "none", "std", "range", and "minmax".	None
positive_class_by_trainer	Optional[Dict[str, str]]	Optional explicit positive class by trainer id. This is used by positive trainer suites whose positive/negative pairs do not follow the non_/non- naming heuristic.	None

Returns:

Type	Description
Dict[str, Any]	A comparison payload with measure, model_ids, rows,
Dict[str, Any]	columns, matrix, split, max_size, metric,
Dict[str, Any]	positive_class_by_column, negative_class_by_column,
Dict[str, Any]	available_mask, full_eval, encoder_eval,
Dict[str, Any]	allow_incomplete, seed_selection, seed_aggregate,
Dict[str, Any]	dispersion, n_matrix, cell_stats, and multi_seed.
Dict[str, Any]	global_n or global_n_range is included when seed counts are
Dict[str, Any]	available.

Raises:

Type	Description
ValueError	If fewer than two trainers are passed, a trainer is not a binary target-pair trainer, the positive class cannot be inferred, an argument value is unsupported, required encoder data is missing, or a row/column subset is empty while allow_incomplete=False.

Source code in gradiend/comparison/trainer_pair_encoding.py

def compute_trainer_pair_encoding_matrix(
    trainers: Dict[str, object],
    *,
    split: str = "test",
    max_size: Optional[int] = None,
    use_cache: bool = True,
    metric: str = "positive_mean",
    full_eval: Optional[bool] = None,
    encoder_eval: Literal["auto", "cached", "recompute"] = "auto",
    allow_incomplete: bool = False,
    seed_selection: Optional[str] = None,
    seed_aggregate: str = "mean",
    dispersion: Optional[str] = None,
    positive_class_by_trainer: Optional[Dict[str, str]] = None,
) -> Dict[str, Any]:
    """Compare how trainers encode each other's binary target-pair data.

    Rows are source trainers. Columns are target class pairs. For every row
    trainer, the function obtains encoder outputs on a full evaluation set,
    subsets those outputs to the column trainer's two target classes, and
    aggregates the requested metric. Binary pairs must either use a
    ``non_``/``non-`` negative class naming convention or set
    ``TrainingArguments.positive_class`` explicitly.

    Args:
        trainers: Mapping from display/trainer id to trainer object. Each
            trainer must expose exactly two ``target_classes``.
        split: Data split for encoder evaluation and cache lookup.
        max_size: Optional maximum number of examples used by encoder
            evaluation/cache keys.
        use_cache: Whether cached encoder outputs may be loaded in
            ``encoder_eval="auto"`` mode.
        metric: Cell metric. Supported values are ``"positive_mean"``,
            ``"negative_mean"``, and ``"positive_minus_negative"``.
        full_eval: Controls whether encoder evaluation includes other classes.
            ``None`` defaults to ``True`` for ``split="test"`` and ``False``
            otherwise. This becomes ``include_other_classes`` in
            ``trainer.evaluate_encoder``.
        encoder_eval: Encoder evaluation policy. ``"auto"`` loads cached
            encoder outputs when available and computes missing rows;
            ``"cached"`` requires existing cached encoder outputs;
            ``"recompute"`` ignores encoder caches and recomputes rows.
        allow_incomplete: If ``True``, missing row data or empty row/column
            subsets produce ``NaN`` cells instead of raising.
        seed_selection: ``"best"`` compares the selected best model/cache.
            ``"all_convergent"`` evaluates saved convergent seed runs and
            aggregates their per-seed scores.
        seed_aggregate: Aggregation for multi-seed cells. Supported values are
            ``"mean"``, ``"median"``, ``"min"``, and ``"max"``.
        dispersion: Optional dispersion metadata for multi-seed cells.
            Supported values are ``"none"``, ``"std"``, ``"range"``, and
            ``"minmax"``.
        positive_class_by_trainer: Optional explicit positive class by trainer
            id. This is used by positive trainer suites whose positive/negative
            pairs do not follow the ``non_``/``non-`` naming heuristic.

    Returns:
        A comparison payload with ``measure``, ``model_ids``, ``rows``,
        ``columns``, ``matrix``, ``split``, ``max_size``, ``metric``,
        ``positive_class_by_column``, ``negative_class_by_column``,
        ``available_mask``, ``full_eval``, ``encoder_eval``,
        ``allow_incomplete``, ``seed_selection``, ``seed_aggregate``,
        ``dispersion``, ``n_matrix``, ``cell_stats``, and ``multi_seed``.
        ``global_n`` or ``global_n_range`` is included when seed counts are
        available.

    Raises:
        ValueError: If fewer than two trainers are passed, a trainer is not a
            binary target-pair trainer, the positive class cannot be inferred,
            an argument value is unsupported, required encoder data is missing,
            or a row/column subset is empty while ``allow_incomplete=False``.
    """
    if not isinstance(trainers, dict) or len(trainers) < 2:
        raise ValueError("trainers must be a dict with at least 2 trainers")
    if metric not in {"positive_mean", "negative_mean", "positive_minus_negative"}:
        raise ValueError(
            "Currently supported cross-encoding metrics are "
            "'positive_mean', 'negative_mean', and 'positive_minus_negative'"
        )
    encoder_eval = str(encoder_eval).strip().lower()  # type: ignore[assignment]
    if encoder_eval not in {"auto", "cached", "recompute"}:
        raise ValueError("encoder_eval must be one of: 'auto', 'cached', 'recompute'")
    seed_selection = resolve_seed_selection_for_trainers(trainers, seed_selection)
    if dispersion is None:
        dispersion = resolve_dispersion_for_trainers(trainers, None)
    _validate_aggregate_dispersion_combo(seed_aggregate, dispersion)
    include_other_classes = _resolve_full_eval(full_eval, split)
    ids = list(trainers.keys())
    explicit_positive_by_trainer = {
        str(key): str(value)
        for key, value in (positive_class_by_trainer or {}).items()
        if value is not None
    }
    positive_by_col: Dict[str, str] = {}
    negative_by_col: Dict[str, str] = {}
    pair_by_col: Dict[str, List[str]] = {}
    row_encoder_dfs: Dict[str, Any] = {}
    for trainer_id in ids:
        trainer = trainers[trainer_id]
        target_classes = [str(c) for c in (trainer.target_classes or [])]
        training_args = getattr(trainer, "training_args", None) or getattr(trainer, "_training_args", None)
        explicit_positive = explicit_positive_by_trainer.get(str(trainer_id))
        if explicit_positive is None and training_args is not None:
            explicit_positive = getattr(training_args, "positive_class", None)
        positive_by_col[trainer_id] = _infer_positive_class_from_pair(target_classes, explicit_positive)
        negative_by_col[trainer_id] = next(
            cls for cls in target_classes if str(cls) != str(positive_by_col[trainer_id])
        )
        pair_by_col[trainer_id] = target_classes
        if seed_selection == "best":
            encoder_df = (
                _load_cached_encoder_df(trainer, split=split, max_size=max_size)
                if encoder_eval in {"auto", "cached"} and use_cache
                else None
            )
            if encoder_df is None and encoder_eval in {"auto", "recompute"}:
                eval_model = _load_eval_model_for_trainer(trainer)
                try:
                    eval_result = trainer.evaluate_encoder(
                        model_with_gradiend=eval_model,
                        split=split,
                        max_size=max_size,
                        use_cache=bool(use_cache and encoder_eval == "auto"),
                        return_df=True,
                        plot=False,
                        include_other_classes=include_other_classes,
                    )
                finally:
                    del eval_model
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                encoder_df = eval_result.get("encoder_df") if isinstance(eval_result, dict) else None
            row_encoder_dfs[trainer_id] = encoder_df
        else:
            seed_paths = [
                seed_path
                for _, seed_path in resolve_seed_run_entries(trainer, seed_selection)
            ]
            if not seed_paths:
                seed_paths = [trainer.model_path]
            encoder_dfs: List[Any] = []
            best_seed_path = trainer.get_best_seed_run_path() if hasattr(trainer, "get_best_seed_run_path") else None
            reused_best_cache = False
            if encoder_eval in {"auto", "cached"} and use_cache and best_seed_path is not None:
                cached_best_df = _load_cached_encoder_df(trainer, split=split, max_size=max_size)
                if cached_best_df is not None:
                    encoder_dfs.append(cached_best_df)
                    reused_best_cache = True
            if encoder_eval == "cached":
                row_encoder_dfs[trainer_id] = encoder_dfs
                value = row_encoder_dfs[trainer_id]
                if value is None or (isinstance(value, list) and len(value) == 0):
                    if allow_incomplete:
                        row_encoder_dfs[trainer_id] = None
                        continue
                    raise ValueError(
                        f"Trainer {trainer_id!r} has no cached encoder data for cross-encoding. "
                        "Run suite.evaluate_encoder(split='test', full_eval=True, return_df=True) first, "
                        "or use encoder_eval='auto'."
                    )
                continue
            for seed_path in seed_paths:
                if not isinstance(seed_path, str) or not os.path.isdir(seed_path):
                    continue
                if reused_best_cache and best_seed_path is not None and os.path.normcase(seed_path) == os.path.normcase(best_seed_path):
                    continue
                eval_model = _load_eval_model_for_trainer(trainer, load_directory=seed_path)
                training_args = getattr(trainer, "training_args", None) or getattr(trainer, "_training_args", None)
                original_experiment_dir = getattr(training_args, "experiment_dir", None) if training_args is not None else None
                try:
                    if training_args is not None:
                        training_args.experiment_dir = None
                    eval_result = trainer.evaluate_encoder(
                        model_with_gradiend=eval_model,
                        split=split,
                        max_size=max_size,
                        use_cache=False,
                        return_df=True,
                        plot=False,
                        include_other_classes=include_other_classes,
                    )
                finally:
                    if training_args is not None:
                        training_args.experiment_dir = original_experiment_dir
                    del eval_model
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                encoder_df = eval_result.get("encoder_df") if isinstance(eval_result, dict) else None
                if encoder_df is not None:
                    encoder_dfs.append(encoder_df)
            row_encoder_dfs[trainer_id] = encoder_dfs
        value = row_encoder_dfs[trainer_id]
        if value is None or (isinstance(value, list) and len(value) == 0):
            if allow_incomplete:
                row_encoder_dfs[trainer_id] = None
                continue
            raise ValueError(
                f"Trainer {trainer_id!r} has no full encoder data for cross-encoding. "
                "Run suite.evaluate_encoder(split='test', full_eval=True, return_df=True) first, "
                "or use encoder_eval='auto'."
            )
    matrix = [[float("nan")] * len(ids) for _ in range(len(ids))]
    available_mask = [[False] * len(ids) for _ in range(len(ids))]
    n_matrix = [[0] * len(ids) for _ in range(len(ids))]
    cell_stats: List[List[Dict[str, Any]]] = []
    all_n: List[int] = []
    for i, row_id in enumerate(ids):
        row_df_value = row_encoder_dfs[row_id]
        row_dfs = row_df_value if isinstance(row_df_value, list) else [row_df_value]
        if len(row_dfs) == 1 and row_dfs[0] is None:
            continue
        stats_row: List[Dict[str, Any]] = []
        for j, col_id in enumerate(ids):
            scores: List[float] = []
            for row_df in row_dfs:
                subset_df = _subset_encoder_df_for_target_classes(row_df, pair_by_col[col_id])
                if len(subset_df) == 0:
                    continue
                positive_mean = _extract_positive_mean_from_df(subset_df, positive_by_col[col_id])
                if metric == "positive_mean":
                    scores.append(positive_mean)
                    continue
                negative_mean = _extract_negative_mean_from_df(subset_df, negative_by_col[col_id])
                if metric == "negative_mean":
                    scores.append(negative_mean)
                else:
                    scores.append(positive_mean - negative_mean)
            if not scores:
                if allow_incomplete:
                    stats = {"aggregate": float("nan"), "n": 0}
                    stats_row.append(stats)
                    continue
                raise ValueError(
                    f"Cross-encoding found no rows for column pair {pair_by_col[col_id]!r} in row cache {row_id!r}. "
                    "This likely means the encoder cache was not created with include_other_classes=True."
                )
            stats = _aggregate_seed_scores(scores, seed_aggregate=seed_aggregate, dispersion=dispersion)
            matrix[i][j] = float(stats["aggregate"])
            available_mask[i][j] = True
            n_matrix[i][j] = int(stats["n"])
            stats_row.append(stats)
            all_n.append(int(stats["n"]))
        cell_stats.append(stats_row)
    payload: Dict[str, Any] = {
        "measure": f"cross_encoding_{metric}",
        "model_ids": ids,
        "matrix": matrix,
        "rows": ids,
        "columns": ids,
        "split": split,
        "max_size": max_size,
        "metric": metric,
        "positive_class_by_column": positive_by_col,
        "negative_class_by_column": negative_by_col,
        "available_mask": available_mask,
        "full_eval": include_other_classes,
        "encoder_eval": encoder_eval,
        "allow_incomplete": allow_incomplete,
        "seed_selection": seed_selection,
        "seed_aggregate": seed_aggregate,
        "dispersion": dispersion,
        "explicit_positive_class_by_trainer": explicit_positive_by_trainer,
        "n_matrix": n_matrix,
        "cell_stats": cell_stats,
        "multi_seed": seed_selection != "best",
    }
    if all_n:
        if min(all_n) == max(all_n):
            payload["global_n"] = int(all_n[0])
        else:
            payload["global_n_range"] = [int(min(all_n)), int(max(all_n))]
    return payload