DecoderEvaluator

Decoder evaluation: grid search over (feature_factor, lr), cache, and compute best selection summary. Uses trainer for eval dataframe and model evaluation.

compute_metric_summaries

compute_metric_summaries(trainer, results, *, selector, metrics=None, extractor=default_extract_candidates, feature_factor_from_id=lambda cid: cid[0], lr_from_id=lambda cid: cid[1], empty_default_id='base', class_to_ff=None)

Build a per-metric (i.e., target classes) summary for a decoder grid evaluation.

This is a thin wrapper around the module-level compute_metric_summaries that resolves default metrics from the trainer and passes through selector and extraction behavior.

Parameters:

Name	Type	Description	Default
trainer	Any	Trainer-like object used to resolve default metrics when metrics is None (via get_target_feature_classes).	required
results	Mapping[Any, Mapping[str, Any]]	Mapping from candidate id to evaluation result entries. The expected structure is the same as produced by decoder evaluation.	required
selector	SelectionPolicy	Policy used to choose the best candidate per metric (e.g., LMS thresholding (LMSThresholdPolicy), or metric*lms (LMSTimesMetricPolicy).	required
metrics	Optional[Sequence[str]]	Optional list of metric names to summarize. If None, uses the trainer's target feature classes.	None
extractor	Callable[[Mapping[Any, Mapping[str, Any]]], Tuple[List[Candidate], BaseContext]]	Function that converts the raw results mapping into candidates and a base context.	default_extract_candidates
feature_factor_from_id	Callable[[CandidateId], float]	Function to extract feature_factor from a candidate id.	lambda cid: cid[0]
lr_from_id	Callable[[CandidateId], float]	Function to extract learning_rate from a candidate id.	lambda cid: cid[1]
empty_default_id	str	Fallback id used when no candidate is selected (for comparison with base). When the selector returns None, we first try the candidate with learning_rate != 0 and smallest absolute value; only if none exists do we use this default (representing the base model).	'base'
class_to_ff	Optional[Mapping[str, float]]	Optional mapping from class id to the feature factor that strengthens it. Used for strengthen and *_weaken metrics so each summary only considers candidates with the correct rewrite orientation for that class.	None

Returns:

Type	Description
Dict[str, Dict[str, Any]]	A dict keyed by metric name with values containing selected metric
Dict[str, Dict[str, Any]]	value, feature_factor, learning_rate, and id.

Source code in gradiend/evaluator/decoder.py

def compute_metric_summaries(
        self,
        trainer: Any,
        results: Mapping[Any, Mapping[str, Any]],
        *,
        selector: SelectionPolicy,
        metrics: Optional[Sequence[str]] = None,
        extractor: Callable[
            [Mapping[Any, Mapping[str, Any]]], Tuple[List[Candidate], BaseContext]] = default_extract_candidates,
        feature_factor_from_id: Callable[[CandidateId], float] = lambda cid: cid[0],
        lr_from_id: Callable[[CandidateId], float] = lambda cid: cid[1],
        empty_default_id: str = "base",
        class_to_ff: Optional[Mapping[str, float]] = None,
) -> Dict[str, Dict[str, Any]]:
    """
    Build a per-metric (i.e., target classes) summary for a decoder grid evaluation.

    This is a thin wrapper around the module-level `compute_metric_summaries` that
    resolves default metrics from the trainer and passes through selector and
    extraction behavior.

    Args:
        trainer: Trainer-like object used to resolve default metrics when
            `metrics` is None (via `get_target_feature_classes`).
        results: Mapping from candidate id to evaluation result entries. The
            expected structure is the same as produced by decoder evaluation.
        selector: Policy used to choose the best candidate per metric (e.g., LMS thresholding (LMSThresholdPolicy),
            or metric*lms (LMSTimesMetricPolicy).
        metrics: Optional list of metric names to summarize. If None, uses the
            trainer's target feature classes.
        extractor: Function that converts the raw `results` mapping into
            candidates and a base context.
        feature_factor_from_id: Function to extract feature_factor from a
            candidate id.
        lr_from_id: Function to extract learning_rate from a candidate id.
        empty_default_id: Fallback id used when no candidate is selected (for comparison with base).
            When the selector returns None, we first try the candidate with learning_rate != 0 and smallest
            absolute value; only if none exists do we use this default (representing the base model).
        class_to_ff: Optional mapping from class id to the feature factor
            that strengthens it. Used for strengthen and ``*_weaken`` metrics
            so each summary only considers candidates with the correct
            rewrite orientation for that class.

    Returns:
        A dict keyed by metric name with values containing selected metric
        `value`, `feature_factor`, `learning_rate`, and `id`.
    """
    metrics = list(metrics) if metrics is not None else trainer.get_target_feature_classes()
    return compute_metric_summaries(
        results,
        metrics=metrics,
        selector=selector,
        extractor=extractor,
        feature_factor_from_id=feature_factor_from_id,
        lr_from_id=lr_from_id,
        empty_default_id=empty_default_id,
        class_to_ff=class_to_ff,
    )

evaluate_decoder

evaluate_decoder(trainer, model_with_gradiend=None, feature_factors=None, lrs=None, part='decoder', output_path=None, selector=None, summary_extractor=default_extract_candidates, summary_feature_factor_from_id=lambda cid: cid['feature_factor'] if isinstance(cid, dict) else cid[0], summary_lr_from_id=lambda cid: cid['learning_rate'] if isinstance(cid, dict) else cid[1], summary_empty_default_id='base', use_cache=None, split='test', max_size=None, max_size_training_like=None, max_size_neutral=None, eval_batch_size=None, training_like_df=None, neutral_df=None, summary_metrics=None, target_class=None, increase_target_probabilities=True, plot=False, show=None, plot_kwargs=None)

Run decoder grid evaluation and return summary + grid for one direction (strengthen or weaken).

Only the dataset and feature-factor combinations required for the requested direction are computed. Use increase_target_probabilities=True (default) for strengthen, False for weaken.

Parameters:

Name	Type	Description	Default
trainer	Any	Trainer (protocol) with get_model, _model_for_decoder_eval, _get_decoder_eval_dataframe, and evaluate_base_model.	required
model_with_gradiend	Any	ModelWithGradiend instance or path. If None, uses trainer.get_model().	None
feature_factors	Optional[List[float]]	List of feature factors to test. If None, derived from direction and target classes.	None
lrs	Optional[List[float]]	List of learning rates to test.	None
part	str	which part of GRADIEND is used to derive GRADIEND-modified models (options: 'encoder-weight' | 'decoder-weight' | 'decoder-bias' | 'decoder-sum' | 'decoder'). All options besides decoder are independent of the feature factor (e.g., using the encoder weights as update direction), while decoder computes the update direction via dec(feature_factor) (and is the default).	'decoder'
output_path	Optional[str]	Optional explicit cache path. Overrides experiment_dir-based cache path.	None
selector	Optional[SelectionPolicy]	SelectionPolicy, e.g. LMSThresholdPolicy(ratio=0.99) or LMSTimesMetricPolicy().	None
summary_extractor	Callable[[Mapping[Any, Mapping[str, Any]]], Tuple[List[Candidate], BaseContext]]	Candidate extractor for summary computation. Use a custom extractor to add derived metrics (e.g. bpi, fpi, mpi) to candidates; then pass summary_metrics so they are summarized.	default_extract_candidates
summary_feature_factor_from_id	Callable[[CandidateId], float]	Function to extract feature_factor from candidate id.	lambda cid: cid['feature_factor'] if isinstance(cid, dict) else cid[0]
summary_lr_from_id	Callable[[CandidateId], float]	Function to extract lr from candidate id.	lambda cid: cid['learning_rate'] if isinstance(cid, dict) else cid[1]
summary_empty_default_id	str	Fallback id used when no candidate is selected (for comparison with base). When the selector returns None, we first try the candidate with learning_rate != 0 and smallest absolute value; only if none exists do we use this default (representing the base model).	'base'
use_cache	Optional[bool]	If True, use cached results when available; if False, recompute.	None
split	Optional[Any]	Dataset split used for training-like decoder evaluation rows. Defaults to "test".	'test'
max_size	Optional[int]	Shared evaluation-size alias. If set and explicit decoder caps are omitted, caps both training-like decoder rows and neutral/LMS rows.	None
max_size_training_like	Optional[int]	Maximum size for generated training-like eval data.	None
max_size_neutral	Optional[int]	Maximum size for generated neutral eval data (and LMS text cap).	None
eval_batch_size	Optional[int]	Common eval batch size used for LMS.	None
training_like_df	Optional[Any]	Optional explicit training-like DataFrame for probability scoring.	None
neutral_df	Optional[Any]	Optional explicit neutral DataFrame for LMS scoring.	None
summary_metrics	Optional[Sequence[str]]	Optional list of metric names to summarize. If None, uses direction and target classes (see increase_target_probabilities).	None
target_class	Optional[Union[str, List[str]]]	If set, evaluate only for this target class (or list of classes). Restricts feature factors and datasets to those needed for the given class(es) for efficiency. When None, evaluates for all trainer target classes.	None
increase_target_probabilities	bool	If True (default), compute strengthen summaries only (keys e.g. "3SG", "3PL"). If False, compute weaken summaries only (keys e.g. "3SG_weaken", "3PL_weaken"). Only the dataset–feature-factor combinations required for the chosen direction are evaluated.	True
plot	bool	If True, after selection run any missing dataset evaluations needed for plotting, update cache incrementally, then call the trainer's plot_probability_shifts.	False
show	Optional[bool]	If True, display the plot (e.g. plt.show()). If False, only save to file. When None and plot=True, defaults to True (same as evaluate_encoder: plot implies show).	None
plot_kwargs	Optional[Dict[str, Any]]	Optional dict of options forwarded to plot_probability_shifts when plot=True. E.g. plot_kwargs=dict(figsize=(5, 3), show=False). The evaluate_decoder show argument overrides plot_kwargs["show"] when set.	None

Returns:

Type

Description

Dict[str, Any]

Flat dict with: For strengthen (increase_target_probabilities=True): one entry per target class (e.g. dec_result['3SG']). For weaken (increase_target_probabilities=False): one entry per target class with "_weaken" suffix (e.g. dec_result['3SG_weaken']). Each summary entry contains selected metric value, feature_factor, learning_rate, id, a strengthen flag, and LMS fields (lms, base_lms). 'grid': candidate id -> full evaluation results. When plot=True, also 'plot_paths' and 'plot_path'.

Source code in gradiend/evaluator/decoder.py

def evaluate_decoder(
    self,
    trainer: Any,
    model_with_gradiend: Any = None,
    feature_factors: Optional[List[float]] = None,
    lrs: Optional[List[float]] = None,
    part: str = "decoder",
    output_path: Optional[str] = None,
    selector: Optional[SelectionPolicy] = None,
    summary_extractor: Callable[[Mapping[Any, Mapping[str, Any]]], Tuple[List[Candidate], BaseContext]] = default_extract_candidates,
    summary_feature_factor_from_id: Callable[[CandidateId], float] = lambda cid: cid['feature_factor'] if isinstance(cid, dict) else cid[0],
    summary_lr_from_id: Callable[[CandidateId], float] = lambda cid: cid['learning_rate'] if isinstance(cid, dict) else cid[1],
    summary_empty_default_id: str = "base",
    use_cache: Optional[bool] = None,
    split: Optional[Any] = "test",
    max_size: Optional[int] = None,
    max_size_training_like: Optional[int] = None,
    max_size_neutral: Optional[int] = None,
    eval_batch_size: Optional[int] = None,
    training_like_df: Optional[Any] = None,
    neutral_df: Optional[Any] = None,
    summary_metrics: Optional[Sequence[str]] = None,
    target_class: Optional[Union[str, List[str]]] = None,
    increase_target_probabilities: bool = True,
    plot: bool = False,
    show: Optional[bool] = None,
    plot_kwargs: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Run decoder grid evaluation and return summary + grid for one direction (strengthen or weaken).

    Only the dataset and feature-factor combinations required for the requested direction are
    computed. Use increase_target_probabilities=True (default) for strengthen, False for weaken.

    Args:
        trainer: Trainer (protocol) with get_model, _model_for_decoder_eval, _get_decoder_eval_dataframe,
                 and evaluate_base_model.
        model_with_gradiend: ModelWithGradiend instance or path. If None, uses trainer.get_model().
        feature_factors: List of feature factors to test. If None, derived from direction and target classes.
        lrs: List of learning rates to test.
        part: which part of GRADIEND is used to derive GRADIEND-modified models (options:  'encoder-weight' |
            'decoder-weight' | 'decoder-bias' | 'decoder-sum' | 'decoder'). All options besides `decoder` are
            independent of the feature factor (e.g., using the encoder weights as update direction), while `decoder`
            computes the update direction via dec(feature_factor) (and is the default).
        output_path: Optional explicit cache path. Overrides experiment_dir-based cache path.
        selector: SelectionPolicy, e.g. LMSThresholdPolicy(ratio=0.99) or LMSTimesMetricPolicy().
        summary_extractor: Candidate extractor for summary computation. Use a custom extractor to add
            derived metrics (e.g. bpi, fpi, mpi) to candidates; then pass summary_metrics so they are summarized.
        summary_feature_factor_from_id: Function to extract feature_factor from candidate id.
        summary_lr_from_id: Function to extract lr from candidate id.
        summary_empty_default_id: Fallback id used when no candidate is selected (for comparison with base).
            When the selector returns None, we first try the candidate with learning_rate != 0 and smallest
            absolute value; only if none exists do we use this default (representing the base model).
        use_cache: If True, use cached results when available; if False, recompute.
        split: Dataset split used for training-like decoder evaluation rows.
            Defaults to ``"test"``.
        max_size: Shared evaluation-size alias. If set and
            explicit decoder caps are omitted, caps both training-like decoder
            rows and neutral/LMS rows.
        max_size_training_like: Maximum size for generated training-like eval data.
        max_size_neutral: Maximum size for generated neutral eval data (and LMS text cap).
        eval_batch_size: Common eval batch size used for LMS.
        training_like_df: Optional explicit training-like DataFrame for probability scoring.
        neutral_df: Optional explicit neutral DataFrame for LMS scoring.
        summary_metrics: Optional list of metric names to summarize. If None, uses direction and target classes
            (see increase_target_probabilities).
        target_class: If set, evaluate only for this target class (or list of classes). Restricts
            feature factors and datasets to those needed for the given class(es) for efficiency.
            When None, evaluates for all trainer target classes.
        increase_target_probabilities: If True (default), compute **strengthen** summaries only (keys e.g. "3SG", "3PL").
            If False, compute **weaken** summaries only (keys e.g. "3SG_weaken", "3PL_weaken"). Only the
            dataset–feature-factor combinations required for the chosen direction are evaluated.
        plot: If True, after selection run any missing dataset evaluations needed for plotting,
            update cache incrementally, then call the trainer's plot_probability_shifts.
        show: If True, display the plot (e.g. plt.show()). If False, only save to file. When None
            and plot=True, defaults to True (same as evaluate_encoder: plot implies show).
        plot_kwargs: Optional dict of options forwarded to plot_probability_shifts when plot=True.
            E.g. plot_kwargs=dict(figsize=(5, 3), show=False). The evaluate_decoder ``show`` argument
            overrides plot_kwargs[\"show\"] when set.

    Returns:
        Flat dict with:

          - For strengthen (increase_target_probabilities=True): one entry per target class (e.g. dec_result['3SG']).
          - For weaken (increase_target_probabilities=False): one entry per target class with \"_weaken\" suffix

            (e.g. dec_result['3SG_weaken']).

          - Each summary entry contains selected metric `value`, `feature_factor`, `learning_rate`, `id`,

            a `strengthen` flag, and LMS fields (`lms`, `base_lms`).

          - 'grid': candidate id -> full evaluation results.
          - When plot=True, also 'plot_paths' and 'plot_path'.
    """
    logger.info(f"Starting decoder evaluation with part={part}")
    use_cache = trainer._resolve_artifact_use_cache(use_cache, fallback=False)
    if max_size_training_like is None:
        max_size_training_like = max_size
    if max_size_neutral is None:
        max_size_neutral = max_size
    trainer_config = getattr(trainer, "config", None)
    training_args = getattr(trainer, "_training_args", None)
    if max_size_training_like is None:
        if hasattr(trainer, "_default_from_training_args"):
            max_size_training_like = trainer._default_from_training_args(
                max_size_training_like, "decoder_eval_max_size_training_like"
            )
        elif training_args is not None:
            max_size_training_like = getattr(training_args, "decoder_eval_max_size_training_like", None)
    if max_size_training_like is None and trainer_config is not None:
        max_size_training_like = getattr(trainer_config, "decoder_eval_lms_max_samples", None)
    if max_size_neutral is None:
        if hasattr(trainer, "_default_from_training_args"):
            max_size_neutral = trainer._default_from_training_args(
                max_size_neutral, "decoder_eval_max_size_neutral"
            )
        elif training_args is not None:
            max_size_neutral = getattr(training_args, "decoder_eval_max_size_neutral", None)
    if max_size_neutral is None and trainer_config is not None:
        max_size_neutral = getattr(trainer_config, "decoder_eval_lms_max_samples", None)

    if selector is None:
        selector = LMSThresholdPolicy(ratio=0.99)

    raw_model = model_with_gradiend or trainer.get_model()
    if isinstance(raw_model, str):
        trust_remote_code = getattr(getattr(trainer, "_training_args", None), "trust_remote_code", False)
        raw_model = ModelWithGradiend.from_pretrained(raw_model, trust_remote_code=trust_remote_code)

    target_classes = trainer.get_target_feature_classes()
    if target_class is not None:
        classes_to_eval: List[str] = (
            [target_class] if isinstance(target_class, str) else list(target_class)
        )
        for c in classes_to_eval:
            if c not in (target_classes or []):
                raise ValueError(
                    f"target_class={target_class!r} must be one of trainer target classes {target_classes}. "
                    f"Got {c!r}."
                )
    else:
        classes_to_eval = target_classes or []

    tcs = set(target_classes or [])
    base_metrics = summary_metrics if summary_metrics is not None else classes_to_eval
    metrics_for_summary = list(base_metrics) if base_metrics else list(classes_to_eval or [])
    # Strengthen: we maximize P(target_class) on the *other* class's dataset (e.g. P(3SG) on 3PL data).
    # The trainer keys probs by dataset class (which data we ran on), so the raw result key is the
    # other class (e.g. "3PL"). We expose the result under the target class the user asked for
    # (e.g. "3SG") and hide the internal key so the API matches user intent.
    summary_key_aliases: Optional[Dict[str, str]] = None
    prob_on_other_class = getattr(
        getattr(trainer, "config", None), "decoder_eval_prob_on_other_class", True
    )
    _, decoder_row_wise = trainer._resolve_decoder_eval_targets(training_like_df=None)
    if summary_metrics is None and classes_to_eval:
        if increase_target_probabilities:
            if prob_on_other_class and not decoder_row_wise:
                # probs[target] is P(target) on the other class's factual dataset.
                metrics_for_summary = list(classes_to_eval)
            else:
                required_for_strengthen = {d for c in classes_to_eval for d in (tcs - {c})}
                metrics_for_summary = list(required_for_strengthen)
                summary_key_aliases = {
                    c: (tcs - {c}).pop() for c in classes_to_eval if len(tcs - {c}) == 1
                }
        else:
            # Weaken only: keys "3SG_weaken", "3PL_weaken"
            metrics_for_summary = [f"{c}_weaken" for c in classes_to_eval]

    # Per-class strengthen ff (see gradiend.model._source_target module docstring).
    class_to_ff: Optional[Dict[str, float]] = None
    if target_classes:
        class_to_ff = {
            cls: derive_feature_factor_for_class(trainer, raw_model, cls)
            for cls in target_classes
        }

    if feature_factors is None:
        if increase_target_probabilities and classes_to_eval and class_to_ff:
            feature_factors = list(dict.fromkeys(class_to_ff[c] for c in classes_to_eval))
        elif not increase_target_probabilities and classes_to_eval and class_to_ff:
            feature_factors = []
            for c in classes_to_eval:
                for o in target_classes:
                    if o == c:
                        continue
                    ff = class_to_ff[o]
                    if ff not in feature_factors:
                        feature_factors.append(ff)
            if not feature_factors:
                feature_factors = default_decoder_feature_factors(trainer, raw_model)
        else:
            feature_factors = default_decoder_feature_factors(trainer, raw_model)

    model_with_gradiend = trainer._model_for_decoder_eval(raw_model)
    path = model_with_gradiend.name_or_path
    base_model = model_with_gradiend.base_model
    tokenizer = model_with_gradiend.tokenizer
    run_id = getattr(trainer, "run_id", None)
    model_id = os.path.basename(path) if path and str(path).startswith("results/models") else path

    if lrs is None:
        lrs = [
            m * 10 ** e
            for e in range(2, -4, -1)
            for m in [5, 2, 1]
            if m * 10 ** e <= 100
        ]

    experiment_dir = trainer.experiment_dir
    cache_file = resolve_decoder_grid_cache_path(experiment_dir, explicit_path=output_path)
    if use_cache and not cache_file:
        raise ValueError(
            "evaluate_decoder(use_cache=True) requires experiment_dir on the trainer or output_path. "
            "Set experiment_dir on TrainingArguments or pass output_path= to specify the cache location."
        )
    if cache_file:
        os.makedirs(os.path.dirname(cache_file), exist_ok=True)
    split_cache_key = _decoder_split_cache_key(split)

    pairs = [(ff, lr) for ff in feature_factors for lr in lrs]
    pairs = sorted(pairs)  # deterministic order for reproducible decoder evaluation
    expected_results = len(pairs) + 1
    logger.info(
        "Decoder eval feature_factors=%s class_to_ff=%s source=%s",
        list(feature_factors),
        class_to_ff,
        getattr(raw_model, "source", None),
    )

    relevant_results: Dict[Any, Dict[str, Any]] = {}
    all_results: Dict[Any, Dict[str, Any]] = {}

    if use_cache and cache_file and os.path.isfile(cache_file):
        try:
            with open(cache_file, "r") as f:
                payload = json.load(f)
            cached_part = payload.get("part")
            cached_feature_factors = payload.get("feature_factors")
            cached_lrs = payload.get("lrs")
            cached_split = payload.get("split", "test")
            cached_max_size_training_like = payload.get("max_size_training_like")
            cached_max_size_neutral = payload.get("max_size_neutral")
            cache_matches = (
                cached_part == part
                and cached_feature_factors == list(feature_factors)
                and cached_lrs == list(lrs)
                and cached_split == split_cache_key
                and cached_max_size_training_like == max_size_training_like
                and cached_max_size_neutral == max_size_neutral
            )
            if cache_matches:
                relevant_results = convert_results_to_dict(payload.get("results", []))
                if len(relevant_results) >= expected_results:
                    logger.info("Using cached decoder grid from %s", cache_file)
                    summary = self.compute_metric_summaries(
                        trainer,
                        relevant_results,
                        selector=selector,
                        metrics=metrics_for_summary,
                        extractor=summary_extractor,
                        feature_factor_from_id=summary_feature_factor_from_id,
                        lr_from_id=summary_lr_from_id,
                        empty_default_id=summary_empty_default_id,
                        class_to_ff=class_to_ff,
                    )
                    if summary_key_aliases:
                        # Copy first so we don't overwrite when target_key and result_key cross (e.g. commutative <-> non-commutative)
                        alias_updates = {
                            target_key: summary[result_key]
                            for target_key, result_key in summary_key_aliases.items()
                            if result_key in summary
                        }
                        for target_key, val in alias_updates.items():
                            summary[target_key] = val
                        # Only remove internal result keys; keep keys that are also target class keys (multi-class case)
                        for k in set(summary_key_aliases.values()):
                            if k not in summary_key_aliases:
                                summary.pop(k, None)
                    if not plot:
                        return {**summary, "grid": relevant_results}
                    # plot=True: get full df and run fill-in + plot
                    training_like_df, neutral_df = trainer._get_decoder_eval_dataframe(
                        tokenizer,
                        max_size_training_like=max_size_training_like,
                        max_size_neutral=max_size_neutral,
                        split=split,
                        cached_training_like_df=None,
                        cached_neutral_df=None,
                    )
                    full_training_like_df = training_like_df
                    dataset_class_col = "label_class" if "label_class" in getattr(training_like_df, "columns", []) else "factual_id"
                    if full_training_like_df is not None and dataset_class_col in getattr(full_training_like_df, "columns", []):
                        _refresh_probs_by_dataset_for_plotting(
                            trainer,
                            model_with_gradiend=model_with_gradiend,
                            base_model=base_model,
                            tokenizer=tokenizer,
                            relevant_results=relevant_results,
                            training_like_df=full_training_like_df,
                            neutral_df=neutral_df,
                            part=part,
                            max_size_training_like=max_size_training_like,
                            max_size_neutral=max_size_neutral,
                            eval_batch_size=eval_batch_size,
                        )
                    if cache_file:
                        try:
                            payload_update = {
                                "part": part,
                                "split": split_cache_key,
                                "max_size_training_like": max_size_training_like,
                                "max_size_neutral": max_size_neutral,
                                "feature_factors": list(feature_factors),
                                "lrs": list(lrs),
                                "results": convert_results_to_list(relevant_results),
                            }
                            with open(cache_file, "w") as f:
                                json.dump(payload_update, f, indent=2)
                        except Exception as e:
                            logger.warning("Error writing decoder cache %s: %s", cache_file, e)
                    plot_paths: List[str] = []
                    if hasattr(trainer, "plot_probability_shifts"):
                        try:
                            plot_keys_override = list(summary_key_aliases.keys()) if summary_key_aliases else None
                            plot_paths = _plot_all_target_classes(
                                trainer,
                                summary,
                                relevant_results,
                                experiment_dir=getattr(trainer, "experiment_dir", None),
                                run_id=run_id,
                                show=show,
                                increase_target_probabilities=increase_target_probabilities,
                                plot_keys_override=plot_keys_override,
                                plot_kwargs=plot_kwargs,
                            )
                        except ImportError as e:
                            logger.warning("Skipping decoder probability-shift plots: %s", e)
                    out = {**summary, "grid": relevant_results}
                    if plot_paths:
                        out["plot_paths"] = plot_paths
                        out["plot_path"] = plot_paths[0] if len(plot_paths) == 1 else None
                    return out
            else:
                logger.info("Decoder cache mismatch (part/split/size/feature_factors/lrs); recomputing.")
        except Exception as e:
            logger.warning("Error loading cached decoder results: %s", e)

    if max_size_training_like is None:
        if training_args is not None:
            max_size_training_like = getattr(training_args, "decoder_eval_max_size_training_like", None)
        if max_size_training_like is None and trainer_config is not None:
            max_size_training_like = getattr(trainer_config, "decoder_eval_lms_max_samples", None)
    if max_size_neutral is None:
        if training_args is not None:
            max_size_neutral = getattr(training_args, "decoder_eval_max_size_neutral", None)
        if max_size_neutral is None and trainer_config is not None:
            max_size_neutral = getattr(trainer_config, "decoder_eval_lms_max_samples", None)

    if training_like_df is None or neutral_df is None:
        training_like_df, neutral_df = trainer._get_decoder_eval_dataframe(
            tokenizer,
            max_size_training_like=max_size_training_like,
            max_size_neutral=max_size_neutral,
            split=split,
            cached_training_like_df=training_like_df,
            cached_neutral_df=neutral_df,
        )

    # Restrict to datasets required for the requested direction (efficiency).
    # Strengthen: maximize P(target) on *other* class's dataset → need other's data; probs key = dataset class.
    # Weaken: maximize (1 - P(class) on class's data) → need class's data.
    full_training_like_df = training_like_df
    dataset_class_col = "label_class" if "label_class" in getattr(training_like_df, "columns", []) else "factual_id"
    if summary_metrics is not None:
        required_datasets = tcs
    elif increase_target_probabilities:
        required_datasets = {d for c in classes_to_eval for d in (tcs - {c})}
    else:
        required_datasets = set(classes_to_eval)
    if dataset_class_col in getattr(training_like_df, "columns", []) and required_datasets:
        training_like_df_selection = full_training_like_df[
            full_training_like_df[dataset_class_col].astype(str).isin(required_datasets)
        ]
        if len(training_like_df_selection) > 0:
            training_like_df = training_like_df_selection

    _LARGE_DATASET = 10000
    if max_size_training_like is None and len(training_like_df) > _LARGE_DATASET:
        logger.warning(
            "decoder eval: max_size_training_like is not set and training data has %d rows. "
            "Computation may be slow. Consider setting decoder_eval_max_size_training_like or max_size_training_like to cap.",
            len(training_like_df),
        )
    if max_size_neutral is None and len(neutral_df) > _LARGE_DATASET:
        logger.warning(
            "decoder eval: max_size_neutral is not set and neutral data has %d rows. "
            "LMS computation may be slow. Set decoder_eval_max_size_neutral or max_size_neutral to cap.",
            len(neutral_df),
        )

    if "base" not in relevant_results or not use_cache:
        logger.debug("Evaluating base model...")
        base_results = trainer.evaluate_base_model(
            base_model,
            tokenizer,
            use_cache=use_cache,
            cache_folder=f"base_split_{split_cache_key}_max_{max_size_training_like}_neutral_{max_size_neutral}",
            training_like_df=training_like_df,
            neutral_df=neutral_df,
            max_size_training_like=max_size_training_like,
            max_size_neutral=max_size_neutral,
            eval_batch_size=eval_batch_size,
        )
        if isinstance(base_results, dict):
            base_results["id"] = "base"
        all_results["base"] = base_results
        relevant_results["base"] = base_results

    for feature_factor, lr in gradiend_tqdm(
        pairs,
        desc=f"Evaluate GRADIEND {run_id or ''}",
        total=len(pairs),
        position=0,
    ):
        id_key = (feature_factor, lr)
        if id_key in relevant_results and use_cache:
            continue

        modified_model = model_with_gradiend.rewrite_base_model(
            learning_rate=lr,
            feature_factor=feature_factor,
            part=part,
        )
        modified_results = trainer.evaluate_base_model(
            modified_model,
            tokenizer,
            use_cache=use_cache,
            cache_folder=f"split_{split_cache_key}_max_{max_size_training_like}_neutral_{max_size_neutral}_{feature_factor}_{lr}",
            model_id=model_id,
            training_like_df=training_like_df,
            neutral_df=neutral_df,
            max_size_training_like=max_size_training_like,
            max_size_neutral=max_size_neutral,
            eval_batch_size=eval_batch_size,
        )
        if isinstance(modified_results, dict):
            modified_results["id"] = {"feature_factor": feature_factor, "learning_rate": lr}
        all_results[id_key] = modified_results
        relevant_results[id_key] = modified_results

        del modified_model
        torch.cuda.empty_cache()

    summary = self.compute_metric_summaries(
        trainer,
        relevant_results,
        selector=selector,
        metrics=metrics_for_summary,
        extractor=summary_extractor,
        feature_factor_from_id=summary_feature_factor_from_id,
        lr_from_id=summary_lr_from_id,
        empty_default_id=summary_empty_default_id,
        class_to_ff=class_to_ff,
    )
    if summary_key_aliases:
        # Copy first so we don't overwrite when target_key and result_key cross (e.g. commutative <-> non-commutative)
        alias_updates = {
            target_key: summary[result_key]
            for target_key, result_key in summary_key_aliases.items()
            if result_key in summary
        }
        for target_key, val in alias_updates.items():
            summary[target_key] = val
        # Only remove internal result keys; keep keys that are also target class keys (multi-class case)
        for k in set(summary_key_aliases.values()):
            if k not in summary_key_aliases:
                summary.pop(k, None)

    plot_paths: List[str] = []
    if plot and full_training_like_df is not None and dataset_class_col in getattr(full_training_like_df, "columns", []):
        _refresh_probs_by_dataset_for_plotting(
            trainer,
            model_with_gradiend=model_with_gradiend,
            base_model=base_model,
            tokenizer=tokenizer,
            relevant_results=relevant_results,
            training_like_df=full_training_like_df,
            neutral_df=neutral_df,
            part=part,
            max_size_training_like=max_size_training_like,
            max_size_neutral=max_size_neutral,
            eval_batch_size=eval_batch_size,
        )
        if cache_file:
            try:
                payload = {
                    "part": part,
                    "split": split_cache_key,
                    "max_size_training_like": max_size_training_like,
                    "max_size_neutral": max_size_neutral,
                    "feature_factors": list(feature_factors),
                    "lrs": list(lrs),
                    "results": convert_results_to_list(relevant_results),
                }
                with open(cache_file, "w") as f:
                    json.dump(payload, f, indent=2)
            except Exception as e:
                logger.warning("Error writing decoder cache %s: %s", cache_file, e)
        if hasattr(trainer, "plot_probability_shifts"):
            try:
                plot_keys_override = list(summary_key_aliases.keys()) if summary_key_aliases else None
                plot_paths = _plot_all_target_classes(
                    trainer,
                    summary,
                    relevant_results,
                    experiment_dir=getattr(trainer, "experiment_dir", None),
                    run_id=run_id,
                    show=show,
                    increase_target_probabilities=increase_target_probabilities,
                    plot_keys_override=plot_keys_override,
                    plot_kwargs=plot_kwargs,
                )
            except ImportError as e:
                logger.warning("Skipping decoder probability-shift plots: %s", e)

    if cache_file and not plot:
        try:
            payload = {
                "part": part,
                "split": split_cache_key,
                "max_size_training_like": max_size_training_like,
                "max_size_neutral": max_size_neutral,
                "feature_factors": list(feature_factors),
                "lrs": list(lrs),
                "results": convert_results_to_list(relevant_results),
            }
            with open(cache_file, "w") as f:
                json.dump(payload, f, indent=2)
        except Exception as e:
            logger.warning("Error writing decoder cache %s: %s", cache_file, e)

    out = {**summary, "grid": relevant_results}
    if plot and plot_paths:
        out["plot_paths"] = plot_paths
        out["plot_path"] = plot_paths[0] if len(plot_paths) == 1 else None
    return out