Trainer

Bases: FeatureLearningDefinition

Abstract base trainer for GRADIEND models with HuggingFace-like API.

The Trainer class is an abstract base class that provides the main interface for training, evaluating, and working with GRADIEND models. It cannot be instantiated directly; you must use a concrete subclass such as TextPredictionTrainer that implements the required abstract methods from FeatureLearningDefinition.

Abstract Methods:

Subclasses must implement the following abstract methods from FeatureLearningDefinition: - create_training_data(): Create training dataset without gradient computation - create_gradient_training_dataset(): Create training dataset with gradient computation - _get_decoder_eval_dataframe(): Get DataFrames for decoder evaluation - _get_decoder_eval_targets(): Get target tokens for decoder evaluation - evaluate_base_model(): Evaluate a single model for decoder evaluation - _analyze_encoder(): Analyze encoder by encoding gradients from training data

Class Management:

The trainer uses target_classes to specify which classes are used for training. The pair property is automatically inferred from target_classes when exactly two target classes are specified (i.e., pair = (target_classes[0], target_classes[1])). The all_classes property includes all classes in the dataset (including non-target classes) and can be inferred from data if not explicitly set.

• target_classes: Classes used for training (required)
• all_classes: All classes in dataset (optional, inferred from data if not set)
• pair: Automatically computed from target_classes when len(target_classes) == 2

Key Features:

• Model Management: Stores model at construction time with lazy loading and caching
• Training: Full training pipeline with support for pre-pruning, training, and post-pruning
• Multi-Seed Training: Automatic multi-seed training with convergence tracking and best seed selection
• Evaluation: Integrated encoder and decoder evaluation with caching
• Visualization: Delegates plotting to Evaluator/Visualizer
• Device Management: Easy model device movement (CPU/CUDA)

Basic Usage:

from gradiend.trainer.text.prediction.trainer import TextPredictionTrainer
from gradiend.trainer.core.arguments import TrainingArguments
import pandas as pd

# Initialize trainer with model and training arguments
# Note: Use TextPredictionTrainer (or another concrete subclass), not the abstract Trainer directly
args = TrainingArguments(
    experiment_dir="./results",
    train_batch_size=32,
    num_epochs=10,
    learning_rate=1e-3,
)
trainer = TextPredictionTrainer(
    model="gpt2",
    args=args,
    run_id="runs/experiment_gpt2",
    data=your_dataframe,  # Modality-specific data (could also be a HF dataset id)
    target_classes=["class1", "class2"],  # Target classes the GRADIEND -> these gets encoded as +-1
)

# Train the model
trainer.train()

# Evaluate encoder and decoder
enc_results = trainer.evaluate_encoder()
dec_results = trainer.evaluate_decoder()

# Plot results
trainer.plot_encoder_distributions()
trainer.plot_training_convergence()

Multi-Seed Training:

When TrainingArguments.max_seeds > 1, the trainer automatically runs multiple training runs with different random seeds. It tracks convergence metrics, selects the best seed based on selection scores, and writes a comprehensive seed report.

args = TrainingArguments(
    max_seeds=5,
    min_convergent_seeds=2,
    convergent_metric="correlation",
    convergent_score_threshold=0.5,
)
trainer = Trainer(model="gpt2", args=args)
trainer.train()  # Runs 5 seeds, selects best

Pruning:

The trainer supports both pre-pruning (before training) and post-pruning (after training):

from gradiend.trainer.core.pruning import PrePruneConfig, PostPruneConfig

# Pre-prune: gradient-based pruning before training
pre_cfg = PrePruneConfig(n_samples=1000, topk=0.5, source="diff")
args.pre_prune_config = pre_cfg

# Post-prune: weight-based pruning after training
post_cfg = PostPruneConfig(topk=0.3, part="decoder-weight")
args.post_prune_config = post_cfg

trainer.train()  # Automatically applies pre-prune and post-prune

Evaluation:

The trainer provides convenient methods for encoder and decoder evaluation:

# Encoder evaluation: analyze gradient encodings
enc_results = trainer.evaluate_encoder(split="test", max_size=1000)
# Returns: correlation, encoded values, mean_by_class, etc.

# Decoder evaluation: grid search over feature_factor and learning_rate
dec_results = trainer.evaluate_decoder(use_cache=True)
# Returns: summary (best configs per metric) and grid (all results)

# Combined evaluation
results = trainer.evaluate()
# Returns: {"encoder": enc_results, "decoder": dec_results}

Model Access:

# Get the trained model (cached after first load)
model = trainer.get_model()

# Load a specific checkpoint
model = trainer.get_model(load_directory="./results/model")

# Move model to device
trainer.cuda(device=0)  # or trainer.cpu()

Architecture:

The Trainer subclasses FeatureLearningDefinition and adds: - Model storage and lazy loading (get_model()) - Training arguments management (training_args property) - Lazy Evaluator initialization (evaluator property) - Experiment directory resolution (experiment_dir property)

Training logic lives in _train(); subclasses can override this method to customize behavior.

Args: model: Model identifier (string path) or ModelWithGradiend instance. If string, the model is loaded lazily on first access via get_model(). target_classes: Optional list of target class names for training. If None, subclasses can determine target_classes from data by setting self._target_classes during initialization or data loading. Default: None args: Optional TrainingArguments instance. Can also be passed as kwargs to train(). run_id: Optional run identifier. When set, creates subdirectory under experiment_dir. n_features: Number of latent features (default: 1). evaluator_class: Optional Evaluator class. Defaults to Evaluator. **kwargs: Additional attributes to set on the trainer instance.

Attributes: training_args: TrainingArguments instance (if provided). experiment_dir: Resolved experiment directory (experiment_dir/run_id if run_id set). model_path: Current model path (initial model or path after training). evaluator: Lazy-initialized Evaluator instance.

Methods: train(): Train GRADIEND model with optional pre/post-pruning. evaluate_encoder(): Analyze encoder performance (correlation, encodings). evaluate_decoder(): Grid search decoder configurations. evaluate(): Run both encoder and decoder evaluation. get_model(): Get the trainer's ModelWithGradiend instance (cached). load_model(): Load a ModelWithGradiend instance from a specific directory. pre_prune(): Run pre-pruning before training. post_prune(): Run post-pruning after training. plot_encoder_distributions(): Plot encoder distribution visualizations. plot_training_convergence(): Plot training convergence metrics. rewrite_base_model(): Rewrite base model(s) using decoder evaluation results, optionally save to disk.

See Also: - FeatureLearningDefinition: Abstract base class providing data creation and evaluation protocols - TextPredictionTrainer: Concrete implementation for text-based models (MLM/CLM) - TrainingArguments: Configuration for training behavior - Evaluator: Evaluation and visualization orchestration - PrePruneConfig, PostPruneConfig: Pruning configuration

Note: This class is abstract and cannot be instantiated directly. Use a concrete subclass such as TextPredictionTrainer that implements the required abstract methods.

Source code in gradiend/trainer/trainer.py

def __init__(
    self,
    model: Union[str, Any],
    target_classes: Optional[List[str]] = None,
    args: Optional[Any] = None,
    run_id: Optional[str] = None,
    n_features: int = 1,
    evaluator_class: Optional[Type] = None,
):
    if run_id is not None and not isinstance(run_id, str):
        raise TypeError(f"run_id must be str or None, got {type(run_id).__name__}")
    if not isinstance(n_features, int):
        raise TypeError(f"n_features must be int, got {type(n_features).__name__}")
    if n_features < 1:
        raise ValueError(f"n_features must be >= 1, got {n_features}")

    super().__init__(target_classes, run_id, n_features)
    self._model_arg = model
    self._base_model_arg = model  # Original model; never changes (model_path does after train)
    self._model_instance: Optional[Any] = None
    self._training_args = args
    self._evaluator_class = evaluator_class if evaluator_class is not None else Evaluator
    self._evaluator: Optional[Any] = None
    self._model_with_gradiend_cls: Optional[Type[Any]] = None

_base_model_arg instance-attribute

_base_model_arg = model

_evaluator instance-attribute

_evaluator = None

_evaluator_class instance-attribute

_evaluator_class = evaluator_class if evaluator_class is not None else Evaluator

_model_arg instance-attribute

_model_arg = model

_model_instance instance-attribute

_model_instance = None

_model_with_gradiend_cls instance-attribute

_model_with_gradiend_cls = None

_training_args instance-attribute

_training_args = args

base_model_path property

base_model_path

Original model passed at construction (base model id or path, e.g. 'bert-base-cased').

evaluator property

evaluator

Lazy-init Evaluator(trainer).

experiment_dir property

experiment_dir

Experiment directory for this trainer.

If training_args.experiment_dir is set, returns that (with run_id subdir if run_id is set).

model_path property

model_path

Current model path: base model before training, GRADIEND output dir after train().

training_args property

training_args

__str__

__str__()

Source code in gradiend/trainer/trainer.py

def __str__(self) -> str:
    return (
        f"Trainer(model={self._model_arg!r}, run_id={self.run_id!r}, "
        f"pair={self.pair}, target_classes={self._target_classes}, "
        f"experiment_dir={self.experiment_dir!r})"
    )

_experiment_dir

_experiment_dir()

Root directory for this experiment (experiment_dir, or experiment_dir/run_id when run_id is set).

Source code in gradiend/trainer/trainer.py

def _experiment_dir(self) -> Optional[str]:
    """Root directory for this experiment (experiment_dir, or experiment_dir/run_id when run_id is set)."""
    if self.training_args is None:
        return None
    exp = self.training_args.experiment_dir
    if not exp:
        return None
    exp = exp.rstrip("/\\")
    if self.run_id and str(self.run_id).strip():
        return os.path.join(exp, str(self.run_id).strip().strip("/\\"))
    return exp

_train

_train(output_dir, args, model, model_with_gradiend_cls, callbacks)

Run GRADIEND training (cache check, model creation, data, core loop, save). Override in subclasses to customize behavior. Returns path to saved model.

Parameters:

Name	Type	Description	Default
output_dir	str	Directory to save the trained model.	required
args	Any	TrainingArguments instance.	required
model	Any	Model identifier (string path) or ModelWithGradiend instance.	required
model_with_gradiend_cls	Any	ModelWithGradiend subclass to use when creating model from string path.	required
callbacks	Any	Optional list of TrainingCallback instances.	required

Returns:

Type	Description
str	Path to saved model directory.

Source code in gradiend/trainer/trainer.py

def _train(
    self,
    output_dir: str,
    args: Any,
    model: Any,
    model_with_gradiend_cls: Any,
    callbacks: Any,
) -> str:
    """
    Run GRADIEND training (cache check, model creation, data, core loop, save).
    Override in subclasses to customize behavior. Returns path to saved model.

    Args:
        output_dir: Directory to save the trained model.
        args: TrainingArguments instance.
        model: Model identifier (string path) or ModelWithGradiend instance.
        model_with_gradiend_cls: ModelWithGradiend subclass to use when creating model from string path.
        callbacks: Optional list of TrainingCallback instances.

    Returns:
        Path to saved model directory.
    """
    args.output_dir = output_dir
    config = args
    # Supervised_decoder: correlation is meaningless; skip evaluation and use loss for best checkpoint
    if getattr(config, "supervised_decoder", False):
        config.do_eval = False

    if config.use_cache and has_saved_model(output_dir):
        logger.info(
            f"GRADIEND model already exists at {output_dir}, skipping training. Use use_cache=False to retrain."
        )
        # Ensure data is loaded so target_classes, pair, etc. are available for evaluate_encoder / evaluate_decoder
        getattr(self, "_ensure_data_for_training", lambda: None)()
        return output_dir
    if has_saved_model(output_dir):
        invalidate_experiment_caches(self.experiment_dir)

    if isinstance(model, str):
        if model_with_gradiend_cls is None:
            raise ValueError(
                "model_with_gradiend_cls is required when model is a string. "
                "For text models, use: model_with_gradiend_cls=TextModelWithGradiend"
            )
        # Ensure data (and thus pair/target_classes) is loaded before creating the model
        # so from_pretrained can set feature_class_encoding_direction on the model
        getattr(self, "_ensure_data_for_training", lambda: None)()
        # Resolve to custom prediction head (e.g. decoder MLM head) when it exists
        load_path = self.resolve_model_path(model)
        model_with_gradiend = create_model_with_gradiend(
            load_path,
            feature_definition=self,
            model_class=model_with_gradiend_cls,
            training_args=config,
            trust_remote_code=getattr(config, "trust_remote_code", False),
        )
    else:
        model_with_gradiend = model

    # Store model instance so get_model() returns the training model during training
    self._model_instance = model_with_gradiend

    training_data = self.create_training_data(
        model_with_gradiend,
        batch_size=config.train_batch_size,
        max_size=config.train_max_size,
    )

    if len(training_data) == 0:
        raise ValueError("Training data is empty; cannot train. Check create_training_data implementation and train_max_size.")

    gradient_dataset = self.create_gradient_training_dataset(
        training_data,
        model_with_gradiend,
        cache_dir=None,
        use_cached_gradients=config.use_cached_gradients,
        dtype=model_with_gradiend.gradiend.torch_dtype,
        device=model_with_gradiend.gradiend.device_encoder,
    )
    dl_kwargs = {"batch_size": 1, "shuffle": False}
    if getattr(config, "seed", None) is not None:
        g = torch.Generator()
        g.manual_seed(int(config.seed))
        dl_kwargs["generator"] = g
    dataloader = DataLoader(gradient_dataset, **dl_kwargs)

    eval_dataset = None
    if config.do_eval and config.eval_steps > 0:
        include_other = (
            getattr(getattr(self, "training_args", None), "add_identity_for_other_classes", False)
            and getattr(self, "classes", None)
            and len(getattr(self, "classes", [])) > 2
        )
        # Use dedicated train-time encoder eval cap when set, otherwise fall back to encoder_eval_max_size
        train_eval_max_size = getattr(config, "encoder_eval_train_max_size", None)
        if train_eval_max_size is None:
            train_eval_max_size = getattr(config, "encoder_eval_max_size", None)
        eval_dataset = self.create_eval_data(
            model_with_gradiend,
            split="val",
            source=config.source,
            max_size=train_eval_max_size,
            include_other_classes=include_other,
        )
        if len(eval_dataset) == 0:
            logger.warning("Evaluation dataset is empty; no in-training evaluation will be performed.")
            eval_dataset = None

    if (
        config.evaluate_fn is None
        and eval_dataset is not None
        and config.do_eval
        and config.eval_steps > 0
    ):
        def _default_evaluate(config_dict=None, training_stats=None, **eval_kwargs):
            # Use the same evaluation logic as post-training evaluation
            return self.evaluator.evaluate_encoder(
                eval_data=eval_dataset,
                use_cache=False,
            )
        config.evaluate_fn = _default_evaluate
        logger.debug(
            "Using default in-training evaluation (evaluator.evaluate_encoder on val data); "
            "correlation and mean_by_class will be tracked."
        )

    core_train(
        model_with_gradiend,
        dataloader,
        training_args=config,
        callbacks=callbacks,
    )
    # core_train handles best-checkpoint selection. Do not overwrite with the
    # in-memory final-step model when a selected model already exists at output_dir.
    if getattr(config, "save_only_best", False) and has_saved_model(output_dir):
        logger.info("Using best checkpoint selected during training at %s", output_dir)
    else:
        model_with_gradiend.save_pretrained(output_dir)
        logger.info(f"Saved trained model to {output_dir}")

    # Select model path for post-training usage:
    # - save_only_best=True: best has been moved to output_dir by core_train
    # - save_only_best=False: keep output_dir as final, but use output_dir_best when available
    selected_output_dir = output_dir
    best_output_dir = f"{output_dir}_best"
    if not getattr(config, "save_only_best", False) and has_saved_model(best_output_dir):
        selected_output_dir = best_output_dir
        logger.info("Selected best checkpoint for post-training model state: %s", selected_output_dir)
    # Release GPU memory before next seed in multi-seed training
    if getattr(config, "max_seeds", 1) > 1:
        self._model_instance = None
        del model_with_gradiend
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    return selected_output_dir

cpu

cpu()

Move loaded model to CPU. No-op if model not loaded.

Source code in gradiend/trainer/trainer.py

def cpu(self) -> "Trainer":
    """Move loaded model to CPU. No-op if model not loaded."""
    return self.to("cpu")

cuda

cuda(device=None)

Move loaded model to CUDA. device: None (default cuda), int (cuda:N), or str/torch.device.

Source code in gradiend/trainer/trainer.py

def cuda(self, device: Any = None) -> "Trainer":
    """Move loaded model to CUDA. device: None (default cuda), int (cuda:N), or str/torch.device."""
    if device is None:
        return self.to("cuda")
    if isinstance(device, int):
        return self.to(f"cuda:{device}")
    return self.to(device)

encode

encode(**kwargs)

Encode eval data; return list of encoded values.

Source code in gradiend/trainer/trainer.py

def encode(self, **kwargs: Any) -> Any:
    """Encode eval data; return list of encoded values."""
    model = self.get_model()
    eval_data = self.create_eval_data(model, **kwargs)
    return [model.encode(entry["source"], return_float=True) for entry in eval_data]

evaluate

evaluate(*, kwargs_encoder=None, kwargs_decoder=None, **kwargs)

Run encoder and decoder evaluation; return combined dict.

Source code in gradiend/trainer/trainer.py

def evaluate(self, *, kwargs_encoder: dict = None, kwargs_decoder: dict = None, **kwargs: Any) -> Dict[str, Any]:
    """Run encoder and decoder evaluation; return combined dict."""
    if kwargs_encoder is not None and not isinstance(kwargs_encoder, dict):
        raise TypeError(f"kwargs_encoder must be dict or None, got {type(kwargs_encoder).__name__}")
    if kwargs_decoder is not None and not isinstance(kwargs_decoder, dict):
        raise TypeError(f"kwargs_decoder must be dict or None, got {type(kwargs_decoder).__name__}")
    return self.evaluator.evaluate(kwargs_encoder=kwargs_encoder, kwargs_decoder=kwargs_decoder, **kwargs)

evaluate_decoder

evaluate_decoder(lrs=None, feature_factors=None, use_cache=None, max_size_training_like=None, max_size_neutral=None, eval_batch_size=None, training_like_df=None, neutral_df=None, selector=None, summary_extractor=None, summary_metrics=None, target_class=None, increase_target_probabilities=True, plot=False, show=None)

Run decoder grid evaluation for one direction (strengthen or weaken).

Delegates to evaluator.evaluate_decoder. Only the datasets and feature-factor combinations required for the chosen direction are computed. When use_cache=True and experiment_dir is set, cached grid results are reused when available.

Parameters:

Name	Type	Description	Default
lrs	Optional[Sequence[float]]	Optional sequence of learning rates to evaluate. If None, defaults are taken from TrainingArguments.decoder_eval_lrs.	None
feature_factors	Optional[Sequence[float]]	Optional sequence of feature factors to evaluate. If None, defaults are taken from TrainingArguments.decoder_eval_feature_factors.	None
use_cache	Optional[bool]	If True, reuse cached decoder grid results when available under experiment_dir. If None, defaults are taken from TrainingArguments.use_cache (default: False).	None
max_size_training_like	Optional[int]	Maximum number of samples per variant for training-like decoder evaluation data. If None, defaults are taken from TrainingArguments.decoder_eval_max_size_training_like.	None
max_size_neutral	Optional[int]	Maximum number of samples per variant for neutral decoder evaluation data (and LMS text cap). If None, defaults are taken from TrainingArguments.decoder_eval_max_size_neutral.	None
eval_batch_size	Optional[int]	Optional batch size used during decoder evaluation (e.g. for LMS calls). If None, an appropriate default is chosen by the evaluator.	None
training_like_df	Optional[DataFrame]	Optional pre-computed training-like DataFrame. When provided, this is used instead of creating training-like evaluation data inside the evaluator.	None
neutral_df	Optional[DataFrame]	Optional pre-computed neutral DataFrame. When provided, this is used instead of creating neutral evaluation data inside the evaluator.	None
selector	Optional[Any]	Optional selection policy (e.g. SelectionPolicy) controlling how the best candidate per metric is chosen.	None
summary_extractor	Optional[Any]	Optional callable that post-processes raw decoder results and attaches derived metrics (e.g. bpi, fpi, mpi) before summarization.	None
summary_metrics	Optional[Sequence[str]]	Optional sequence of metric names to summarize (e.g. ["bpi", "fpi", "mpi"]).	None
target_class	Optional[Union[str, List[str]]]	Optional target class id (or list of ids) to evaluate. When set (e.g. "3SG"), restricts evaluation to that class (or classes) for efficiency. When None, all target classes are evaluated.	None
increase_target_probabilities	bool	If True (default), compute strengthen summaries only (keys like "3SG"). If False, compute weaken summaries only (keys like "3SG_weaken").	True
plot	bool	If True, after selection run any missing evaluations needed for plotting, update cache, then create decoder plots.	False
show	Optional[bool]	Controls whether plots are shown when plot=True. If True, display plots; if False, only save them. When None and plot=True, defaults to True.	None

Returns:

Type	Description
Dict[str, Any]	Dict with flattened decoder summaries. For strengthen, keys like dec["3SG"]; for weaken,
Dict[str, Any]	keys like dec["3SG_weaken"]. Each summary entry includes value, feature_factor, learning_rate,
Dict[str, Any]	id, strengthen, lms, base_lms. The dict always includes "grid", and when plot=True also
Dict[str, Any]	"plot_paths" and/or "plot_path".

Source code in gradiend/trainer/trainer.py

def evaluate_decoder(
    self,
    lrs: Optional[Sequence[float]] = None,
    feature_factors: Optional[Sequence[float]] = None,
    use_cache: Optional[bool] = None,
    max_size_training_like: Optional[int] = None,
    max_size_neutral: Optional[int] = None,
    eval_batch_size: Optional[int] = None,
    training_like_df: Optional[pd.DataFrame] = None,
    neutral_df: Optional[pd.DataFrame] = None,
    selector: Optional[Any] = None,
    summary_extractor: 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,
) -> Dict[str, Any]:
    """
    Run decoder grid evaluation for one direction (strengthen or weaken).

    Delegates to `evaluator.evaluate_decoder`. Only the datasets and feature-factor combinations
    required for the chosen direction are computed. When use_cache=True and experiment_dir is set,
    cached grid results are reused when available.

    Args:
        lrs: Optional sequence of learning rates to evaluate. If None, defaults are taken from
            `TrainingArguments.decoder_eval_lrs`.
        feature_factors: Optional sequence of feature factors to evaluate. If None, defaults are taken
            from `TrainingArguments.decoder_eval_feature_factors`.
        use_cache: If True, reuse cached decoder grid results when available under experiment_dir.
            If None, defaults are taken from `TrainingArguments.use_cache` (default: False).
        max_size_training_like: Maximum number of samples per variant for training-like decoder
            evaluation data. If None, defaults are taken from
            `TrainingArguments.decoder_eval_max_size_training_like`.
        max_size_neutral: Maximum number of samples per variant for neutral decoder evaluation data
            (and LMS text cap). If None, defaults are taken from
            `TrainingArguments.decoder_eval_max_size_neutral`.
        eval_batch_size: Optional batch size used during decoder evaluation (e.g. for LMS calls).
            If None, an appropriate default is chosen by the evaluator.
        training_like_df: Optional pre-computed training-like DataFrame. When provided, this is used
            instead of creating training-like evaluation data inside the evaluator.
        neutral_df: Optional pre-computed neutral DataFrame. When provided, this is used instead of
            creating neutral evaluation data inside the evaluator.
        selector: Optional selection policy (e.g. `SelectionPolicy`) controlling how the best candidate
            per metric is chosen.
        summary_extractor: Optional callable that post-processes raw decoder results and attaches
            derived metrics (e.g. bpi, fpi, mpi) before summarization.
        summary_metrics: Optional sequence of metric names to summarize (e.g. ["bpi", "fpi", "mpi"]).
        target_class: Optional target class id (or list of ids) to evaluate. When set (e.g. "3SG"),
            restricts evaluation to that class (or classes) for efficiency. When None, all
            target classes are evaluated.
        increase_target_probabilities: If True (default), compute **strengthen** summaries only
            (keys like "3SG"). If False, compute **weaken** summaries only (keys like "3SG_weaken").
        plot: If True, after selection run any missing evaluations needed for plotting, update cache,
            then create decoder plots.
        show: Controls whether plots are shown when plot=True. If True, display plots; if False,
            only save them. When None and plot=True, defaults to True.

    Returns:
        Dict with flattened decoder summaries. For strengthen, keys like dec["3SG"]; for weaken,
        keys like dec["3SG_weaken"]. Each summary entry includes value, feature_factor, learning_rate,
        id, strengthen, lms, base_lms. The dict always includes "grid", and when plot=True also
        "plot_paths" and/or "plot_path".
    """
    if not isinstance(increase_target_probabilities, bool):
        raise TypeError(f"increase_target_probabilities must be bool, got {type(increase_target_probabilities).__name__}")
    if not isinstance(plot, bool):
        raise TypeError(f"plot must be bool, got {type(plot).__name__}")
    if show is not None and not isinstance(show, bool):
        raise TypeError(f"show must be bool or None, got {type(show).__name__}")
    if max_size_training_like is not None and not isinstance(max_size_training_like, int):
        raise TypeError(f"max_size_training_like must be int or None, got {type(max_size_training_like).__name__}")
    if max_size_training_like is not None and max_size_training_like < 0:
        raise ValueError(f"max_size_training_like must be >= 0, got {max_size_training_like}")
    if max_size_neutral is not None and not isinstance(max_size_neutral, int):
        raise TypeError(f"max_size_neutral must be int or None, got {type(max_size_neutral).__name__}")
    if max_size_neutral is not None and max_size_neutral < 0:
        raise ValueError(f"max_size_neutral must be >= 0, got {max_size_neutral}")
    if eval_batch_size is not None and not isinstance(eval_batch_size, int):
        raise TypeError(f"eval_batch_size must be int or None, got {type(eval_batch_size).__name__}")
    if eval_batch_size is not None and eval_batch_size < 1:
        raise ValueError(f"eval_batch_size must be >= 1, got {eval_batch_size}")

    use_cache = self._default_from_training_args(use_cache, "use_cache", fallback=False)
    max_size_training_like = self._default_from_training_args(
        max_size_training_like, "decoder_eval_max_size_training_like"
    )
    max_size_neutral = self._default_from_training_args(
        max_size_neutral, "decoder_eval_max_size_neutral"
    )
    lrs = self._default_from_training_args(lrs, "decoder_eval_lrs")
    feature_factors = self._default_from_training_args(feature_factors, "decoder_eval_feature_factors")
    return self.evaluator.evaluate_decoder(
        lrs=lrs,
        feature_factors=feature_factors,
        use_cache=use_cache,
        max_size_training_like=max_size_training_like,
        max_size_neutral=max_size_neutral,
        eval_batch_size=eval_batch_size,
        training_like_df=training_like_df,
        neutral_df=neutral_df,
        selector=selector,
        summary_extractor=summary_extractor,
        summary_metrics=summary_metrics,
        target_class=target_class,
        increase_target_probabilities=increase_target_probabilities,
        plot=plot,
        show=show,
    )

evaluate_encoder

evaluate_encoder(model_with_gradiend=None, encoder_df=None, split='test', source=None, max_size=None, neutral_data_df=None, use_cache=None, return_df=False, plot=False, plot_kwargs=None, is_decoder_only_model=None, pre_load_gradients=None, **kwargs)

Run encoder analysis and return correlation metrics; default model to current trainer model.

When encoder_df is provided (DataFrame or dict with "encoder_df" key), skips encoding and computes metrics from that DataFrame. Otherwise uses _analyze_encoder to produce the DataFrame (training + neutral variants), then delegates to EncoderEvaluator.

Parameters:

Name	Type	Description	Default
model_with_gradiend	Optional[Any]	Optional ModelWithGradiend instance to evaluate. If None, uses the trainer's current model (via get_model()).	None
encoder_df	Optional[Union[DataFrame, Dict[str, Any]]]	Optional DataFrame or dict with "encoder_df" key. If provided, skips encoding and computes metrics from this data. Use evaluate_encoder(return_df=True) to get such a dict, or pass a pre-computed DataFrame directly.	None
split	str	Dataset split to use for evaluation. Default: "test".	'test'
source	Optional[str]	Source type for gradient creation. If None, uses default from training args or "factual". Options: "factual", "counterfactual", etc.	None
max_size	Optional[int]	Maximum number of samples per variant to encode. If None, uses encoder_eval_max_size from training args.	None
neutral_data_df	Optional[DataFrame]	Optional DataFrame with neutral examples (neutral_dataset variant). If provided, these will be encoded in addition to training data.	None
use_cache	Optional[bool]	If True, use cached encoder evaluation when available. If None, uses use_cache from training args (default: False).	None
return_df	bool	If True, include encoder_df (full DataFrame with type column) in result.	False
plot	bool	If True, create encoder distribution plot from analyzed data.	False
plot_kwargs	Optional[Dict[str, Any]]	Optional dict of options forwarded to plot_encoder_distributions when plot=True. E.g. plot_kwargs=dict(target_and_neutral_only=True, show=False). Any argument accepted by plot_encoder_distributions can be passed here.	None
is_decoder_only_model	Optional[bool]	Whether the model is decoder-only (causal LM). If None, inferred from the model.	None
pre_load_gradients	Optional[bool]	If True, pre-load cached gradients when available. If None, uses use_cached_gradients from training args (default: False).	None
**kwargs	Any	Additional arguments passed to _analyze_encoder and create_eval_data.	{}

Returns:

Type	Description
Dict[str, Any]	Dict with unified encoder metrics: n_samples, all_data, training_only,
Dict[str, Any]	target_classes_only, correlation, mean_by_class, mean_by_type, boundaries;
Dict[str, Any]	optionally neutral_mean_by_type, mean_by_feature_class, label_value_to_class_name.
Dict[str, Any]	If return_df=True, includes "encoder_df" key.

Source code in gradiend/trainer/trainer.py

def evaluate_encoder(
    self,
    model_with_gradiend: Optional[Any] = None,
    encoder_df: Optional[Union[pd.DataFrame, Dict[str, Any]]] = None,
    split: str = "test",
    source: Optional[str] = None,
    max_size: Optional[int] = None,
    neutral_data_df: Optional[pd.DataFrame] = None,
    use_cache: Optional[bool] = None,
    return_df: bool = False,
    plot: bool = False,
    plot_kwargs: Optional[Dict[str, Any]] = None,
    is_decoder_only_model: Optional[bool] = None,
    pre_load_gradients: Optional[bool] = None,
    **kwargs: Any,
) -> Dict[str, Any]:
    """
    Run encoder analysis and return correlation metrics; default model to current trainer model.

    When encoder_df is provided (DataFrame or dict with "encoder_df" key), skips encoding
    and computes metrics from that DataFrame. Otherwise uses _analyze_encoder to produce
    the DataFrame (training + neutral variants), then delegates to EncoderEvaluator.

    Args:
        model_with_gradiend: Optional ModelWithGradiend instance to evaluate.
            If None, uses the trainer's current model (via get_model()).
        encoder_df: Optional DataFrame or dict with "encoder_df" key. If provided, skips
            encoding and computes metrics from this data. Use evaluate_encoder(return_df=True)
            to get such a dict, or pass a pre-computed DataFrame directly.
        split: Dataset split to use for evaluation. Default: "test".
        source: Source type for gradient creation. If None, uses default from training args
            or "factual". Options: "factual", "counterfactual", etc.
        max_size: Maximum number of samples per variant to encode. If None, uses
            encoder_eval_max_size from training args.
        neutral_data_df: Optional DataFrame with neutral examples (neutral_dataset variant).
            If provided, these will be encoded in addition to training data.
        use_cache: If True, use cached encoder evaluation when available.
            If None, uses use_cache from training args (default: False).
        return_df: If True, include encoder_df (full DataFrame with type column) in result.
        plot: If True, create encoder distribution plot from analyzed data.
        plot_kwargs: Optional dict of options forwarded to plot_encoder_distributions when
            plot=True. E.g. plot_kwargs=dict(target_and_neutral_only=True, show=False).
            Any argument accepted by plot_encoder_distributions can be passed here.
        is_decoder_only_model: Whether the model is decoder-only (causal LM).
            If None, inferred from the model.
        pre_load_gradients: If True, pre-load cached gradients when available.
            If None, uses use_cached_gradients from training args (default: False).
        **kwargs: Additional arguments passed to _analyze_encoder and create_eval_data.

    Returns:
        Dict with unified encoder metrics: n_samples, all_data, training_only,
        target_classes_only, correlation, mean_by_class, mean_by_type, boundaries;
        optionally neutral_mean_by_type, mean_by_feature_class, label_value_to_class_name.
        If return_df=True, includes "encoder_df" key.
    """
    if not isinstance(split, str):
        raise TypeError(f"split must be str, got {type(split).__name__}")
    if max_size is not None and not isinstance(max_size, int):
        raise TypeError(f"max_size must be int or None, got {type(max_size).__name__}")
    if max_size is not None and max_size < 0:
        raise ValueError(f"max_size must be >= 0, got {max_size}")
    if not isinstance(return_df, bool):
        raise TypeError(f"return_df must be bool, got {type(return_df).__name__}")
    if not isinstance(plot, bool):
        raise TypeError(f"plot must be bool, got {type(plot).__name__}")
    if use_cache is not None and not isinstance(use_cache, bool):
        raise TypeError(f"use_cache must be bool or None, got {type(use_cache).__name__}")

    resolved_encoder_df = _resolve_encoder_df(encoder_df)
    if resolved_encoder_df is not None:
        encoder_df_out = resolved_encoder_df
    else:
        mwg = model_with_gradiend if model_with_gradiend is not None else self.get_model()
        if max_size is None:
            max_size = self._default_from_training_args(max_size, "encoder_eval_max_size")
        if source is not None:
            kwargs["source"] = source
        if is_decoder_only_model is not None:
            kwargs["is_decoder_only_model"] = is_decoder_only_model
        if pre_load_gradients is not None:
            kwargs["pre_load_gradients"] = pre_load_gradients
        encoder_df_out = self._analyze_encoder(
            mwg,
            split=split,
            max_size=max_size,
            neutral_data_df=neutral_data_df,
            use_cache=use_cache,
            plot=False,
            **kwargs,
        )
    result = self.evaluator.evaluate_encoder(
        encoder_df=encoder_df_out,
        use_cache=use_cache,
        split=split,
        max_size=max_size,
        **{k: v for k, v in kwargs.items() if k not in ("return_df", "plot", "plot_kwargs")},
    )
    if return_df:
        result["encoder_df"] = encoder_df_out
    if plot:
        plot_args = dict(plot_kwargs or {})
        if "output" not in plot_args and self.experiment_dir:
            output = resolve_encoder_plot_path(
                self.experiment_dir,
                split=split,
                max_size=max_size,
            )
            if output:
                plot_args["output"] = output
        if "show" not in plot_args:
            plot_args["show"] = True
        if hasattr(self, "config") and getattr(self.config, "img_format", None) is not None:
            plot_args.setdefault("img_format", self.config.img_format)
        self.plot_encoder_distributions(encoder_df=encoder_df_out, **plot_args)
    return result

get_encoder_metrics

get_encoder_metrics(model_path=None, encoder_df=None, **kwargs)

Get unified encoder metrics from encoder_df or from cached results.

Parameters:

Name	Type	Description	Default
model_path	Optional[str]	Path to the model; defaults to the trainer's current model path.	None
encoder_df	Optional[DataFrame]	Optional DataFrame with encoded values and labels. If provided, metrics are computed directly from this DataFrame (same format as evaluate_encoder output). Use when you already have encoder outputs, e.g. from evaluate_encoder(return_df=True).	None
**kwargs	Any	Additional arguments passed to the base implementation (e.g. split, use_cache). When using cache instead of encoder_df, pass the same kwargs you use for evaluate_encoder.	{}

Returns:

Type	Description
Any	Dict with n_samples, correlation, mean_by_class, etc., or None if encoder_df is empty.

Source code in gradiend/trainer/trainer.py

def get_encoder_metrics(
    self,
    model_path: Optional[str] = None,
    encoder_df: Optional[pd.DataFrame] = None,
    **kwargs: Any,
) -> Any:
    """
    Get unified encoder metrics from encoder_df or from cached results.

    Args:
        model_path: Path to the model; defaults to the trainer's current model path.
        encoder_df: Optional DataFrame with encoded values and labels. If provided, metrics
            are computed directly from this DataFrame (same format as evaluate_encoder output).
            Use when you already have encoder outputs, e.g. from evaluate_encoder(return_df=True).
        **kwargs: Additional arguments passed to the base implementation (e.g. split, use_cache).
            When using cache instead of encoder_df, pass the same kwargs you use for
            evaluate_encoder.

    Returns:
        Dict with n_samples, correlation, mean_by_class, etc., or None if encoder_df is empty.
    """
    path = model_path if model_path is not None else self.model_path
    return super().get_encoder_metrics(path, encoder_df=encoder_df, **kwargs)

get_encodings

get_encodings(model_path=None, **kwargs)

Get encodings; default model_path to current trainer model path.

Source code in gradiend/trainer/trainer.py

def get_encodings(self, model_path: Optional[str] = None, **kwargs: Any) -> Any:
    """Get encodings; default model_path to current trainer model path."""
    path = model_path if model_path is not None else self.model_path
    return super().get_encodings(path, **kwargs)

get_model

get_model(use_cache=None, *, load_directory=None, **kwargs)

Get the trainer's ModelWithGradiend instance.

Returns the in-memory model when set (e.g. during training), otherwise loads from load_directory or model_path. The loaded instance is always cached in memory for subsequent calls. After multi-seed training the cache is cleared so the next get_model() loads from the selected best-seed directory and then caches that instance.

To load a model from a different directory, use load_model() or pass load_directory=.

Note: use_cache (e.g. TrainingArguments.use_cache) applies only to disk/output caches (skip when files exist); the in-memory model from get_model() is always cached.

Parameters:

Name	Type	Description	Default
use_cache	Optional[bool]	Ignored; kept for API compatibility. Disk cache is controlled elsewhere.	None
load_directory	Optional[Any]	If provided, load from this path (GRADIEND checkpoint expected).	None
**kwargs	Any	Passed to model_with_gradiend_cls.from_pretrained.	{}

Returns:

Type	Description
Any	ModelWithGradiend instance.

Source code in gradiend/trainer/trainer.py

def get_model(
    self,
    use_cache: Optional[bool] = None,
    *,
    load_directory: Optional[Any] = None,
    **kwargs: Any,
) -> Any:
    """
    Get the trainer's ModelWithGradiend instance.

    Returns the in-memory model when set (e.g. during training), otherwise loads from
    load_directory or model_path. The loaded instance is always cached in memory for
    subsequent calls. After multi-seed training the cache is cleared so the next
    get_model() loads from the selected best-seed directory and then caches that instance.

    To load a model from a different directory, use load_model() or pass load_directory=.

    Note: use_cache (e.g. TrainingArguments.use_cache) applies only to disk/output
    caches (skip when files exist); the in-memory model from get_model() is always cached.

    Args:
        use_cache: Ignored; kept for API compatibility. Disk cache is controlled elsewhere.
        load_directory: If provided, load from this path (GRADIEND checkpoint expected).
        **kwargs: Passed to model_with_gradiend_cls.from_pretrained.

    Returns:
        ModelWithGradiend instance.
    """
    load_directory = load_directory if load_directory is not None else kwargs.pop("load_directory", None)
    # Return in-memory model when set (e.g. during training), unless loading from a specific path
    if load_directory is None and self._model_instance is not None:
        return self._model_instance
    # Pass definition so models get pair/classes when loading
    kwargs.setdefault("definition", self)
    if self._training_args is not None:
        kwargs.setdefault("training_args", self._training_args)
        if "trust_remote_code" not in kwargs:
            kwargs.setdefault("trust_remote_code", getattr(self._training_args, "trust_remote_code", False))
    load_directory = load_directory if load_directory is not None else self.model_path
    model = super().create_model_with_gradiend(load_directory, **kwargs)
    # Always cache in memory; use_cache elsewhere is for disk/output only
    self._model_instance = model
    return model

get_training_stats

get_training_stats(model_path=None)

Load training stats; default model_path to current trainer model path.

Source code in gradiend/trainer/trainer.py

def get_training_stats(self, model_path: Optional[str] = None) -> Optional[Dict[str, Any]]:
    """Load training stats; default model_path to current trainer model path."""
    path = model_path if model_path is not None else self.model_path
    return super().get_training_stats(path)

load_model

load_model(load_directory, model_with_gradiend_cls=None, **kwargs)

Load a ModelWithGradiend instance from a specific directory.

This method loads a model from a different checkpoint/directory than the trainer's current model_path. Use this for loading different checkpoints (e.g., for comparison) or loading specific seed runs in multi-seed training.

Parameters:

Name	Type	Description	Default
load_directory	str	Directory or path to load the model from.	required
model_with_gradiend_cls	Optional[Type[Any]]	Optional ModelWithGradiend subclass. If None, uses self.model_with_gradiend_cls or self.default_model_with_gradiend_cls.	None
**kwargs	Any	Passed to model_with_gradiend_cls.from_pretrained (e.g. feature_definition=self for text models).	{}

Returns:

Type	Description
Any	ModelWithGradiend instance loaded from the specified directory.

Source code in gradiend/trainer/trainer.py

def load_model(
    self,
    load_directory: str,
    model_with_gradiend_cls: Optional[Type[Any]] = None,
    **kwargs: Any,
) -> Any:
    """
    Load a ModelWithGradiend instance from a specific directory.

    This method loads a model from a different checkpoint/directory than the trainer's
    current model_path. Use this for loading different checkpoints (e.g., for comparison)
    or loading specific seed runs in multi-seed training.

    Args:
        load_directory: Directory or path to load the model from.
        model_with_gradiend_cls: Optional ModelWithGradiend subclass. If None, uses
            self.model_with_gradiend_cls or self.default_model_with_gradiend_cls.
        **kwargs: Passed to model_with_gradiend_cls.from_pretrained (e.g. feature_definition=self for text models).

    Returns:
        ModelWithGradiend instance loaded from the specified directory.
    """
    if not isinstance(load_directory, str):
        raise TypeError(f"load_directory must be str, got {type(load_directory).__name__}")

    # Pass feature_definition so text models get pair/classes when loading
    kwargs.setdefault("feature_definition", self)
    kwargs.setdefault("require_gradiend_model", True)  # load_model expects GRADIEND checkpoint
    if self._training_args is not None:
        kwargs.setdefault("training_args", self._training_args)
        if "trust_remote_code" not in kwargs:
            kwargs.setdefault("trust_remote_code", getattr(self._training_args, "trust_remote_code", False))
    return super().create_model_with_gradiend(load_directory, model_with_gradiend_cls=model_with_gradiend_cls, **kwargs)

plot_encoder_distributions

plot_encoder_distributions(**kwargs)

Delegate to evaluator.plot_encoder_distributions.

Source code in gradiend/trainer/trainer.py

def plot_encoder_distributions(self, **kwargs: Any) -> Any:
    """Delegate to evaluator.plot_encoder_distributions."""
    return self.evaluator.plot_encoder_distributions(**kwargs)

plot_encoder_scatter

plot_encoder_scatter(**kwargs)

Delegate to evaluator.plot_encoder_scatter (interactive Plotly scatter).

Source code in gradiend/trainer/trainer.py

def plot_encoder_scatter(self, **kwargs: Any) -> Any:
    """Delegate to evaluator.plot_encoder_scatter (interactive Plotly scatter)."""
    return self.evaluator.plot_encoder_scatter(**kwargs)

plot_training_convergence

plot_training_convergence(**kwargs)

Plot training convergence (means by class/feature_class and correlation). Delegates to Evaluator/Visualizer. If experiment_dir is set, output path is auto-resolved.

Source code in gradiend/trainer/trainer.py

def plot_training_convergence(self, **kwargs: Any) -> Any:
    """
    Plot training convergence (means by class/feature_class and correlation).
    Delegates to Evaluator/Visualizer. If experiment_dir is set, output path is auto-resolved.
    """
    return self.evaluator.plot_training_convergence(**kwargs)

post_prune

post_prune(post_cfg=None)

Run post-prune (weight-based) on the current model and keep it in memory. Subsequent evaluation (e.g. evaluate_encoder) will use the pruned model. Does not save to disk.

Uses self._training_args.post_prune_config when post_cfg is None.

Source code in gradiend/trainer/trainer.py

def post_prune(self, post_cfg: Optional[Any] = None) -> "Trainer":
    """
    Run post-prune (weight-based) on the current model and keep it in memory.
    Subsequent evaluation (e.g. evaluate_encoder) will use the pruned model. Does not save to disk.

    Uses self._training_args.post_prune_config when post_cfg is None.
    """
    from gradiend.trainer.core.pruning import post_prune as _post_prune

    cfg = post_cfg if post_cfg is not None else getattr(self._training_args, "post_prune_config", None)
    if cfg is None:
        raise ValueError("Post-prune config required: pass post_cfg or set training_args.post_prune_config.")
    if post_cfg is None and getattr(cfg, "inplace", True) is False:
        raise ValueError(
            "post_prune_config.inplace=False does not make sense when post-prune is run automatically "
            "from train(). Use inplace=True so the pruned model is kept for subsequent evaluation."
        )
    logger.info(
        "Post-pruning: loading trained model and applying weight-based prune (this may take a while) ...",
    )
    model = self.get_model()
    model = _post_prune(model, cfg)
    self._model_arg = model
    self._model_instance = model
    logger.info("Post-prune done.")
    return self

post_training

post_training(model_with_gradiend, **kwargs)

Optional post-training hook.

Subclasses can override this to perform additional evaluation, logging, or analysis after training. The default implementation is a no-op so that definitions are not required to implement it.

Source code in gradiend/trainer/trainer.py

def post_training(self, model_with_gradiend, **kwargs):
    """
    Optional post-training hook.

    Subclasses can override this to perform additional evaluation,
    logging, or analysis after training. The default implementation
    is a no-op so that definitions are not required to implement it.
    """
    return None

pre_prune

pre_prune(pre_cfg=None, *, inplace=True)

Run pre-prune (gradient-mean then prune) and keep the pruned model in memory. The next train() will use this model. Does not save to disk; save explicitly if needed.

Uses self._training_args.pre_prune_config when pre_cfg is None.

Source code in gradiend/trainer/trainer.py

def pre_prune(
    self,
    pre_cfg: Optional[Any] = None,
    *,
    inplace: bool = True,
) -> "Trainer":
    """
    Run pre-prune (gradient-mean then prune) and keep the pruned model in memory.
    The next train() will use this model. Does not save to disk; save explicitly if needed.

    Uses self._training_args.pre_prune_config when pre_cfg is None.
    """
    cfg = pre_cfg if pre_cfg is not None else getattr(self._training_args, "pre_prune_config", None)
    if cfg is None:
        raise ValueError("Pre-prune config required: pass pre_cfg or set training_args.pre_prune_config.")
    if pre_cfg is None and not inplace:
        raise ValueError(
            "pre_prune(inplace=False) does not make sense when run automatically from train() "
            "(config from TrainingArguments). Use inplace=True so the pruned model is used for training."
        )
    logger.info("Pre-pruning: loading model and preparing data ...")
    model = self.get_model()
    max_size = getattr(self._training_args, "train_max_size", None) if self._training_args else None
    training_data = self.create_training_data(model, batch_size=1, max_size=max_size)
    model = _pre_prune(model, training_data, cfg, definition=self, inplace=inplace)
    self._model_arg = model
    self._model_instance = model
    return self

rewrite_base_model

rewrite_base_model(decoder_results=None, target_class=None, increase_target_probabilities=True, output_dir=None, base_model=None, decoder_stats_metric_name=None, **decoder_stats_kwargs)

Rewrite the base model by applying GRADIEND decoder updates based on decoder evaluation results.

The decoder evaluation selects a feature factor and learning rate per target class and direction (strengthening vs weakening). This method applies the selected config: by default it strengthens the given target class(es); use increase_target_probabilities=False to apply the weakening config instead (evaluate_decoder currently only produces strengthen summaries).

Accepts/loads internally the cached decoder results when experiment_dir is set. Optionally saves the rewritten model(s) to disk if output_dir is provided.

When called on a Trainer, the trained model is used automatically (base_model not needed). When experiment_dir is set, decoder_results can be omitted and will be loaded from cache (requires evaluate_decoder to have been run with use_cache=True so the decoder stats cache exists).

Parameters:

Name	Type	Description	Default
decoder_results	Optional[Dict[str, Any]]	Output from evaluate_decoder. Optional when experiment_dir is set (then loaded from cache).	None
target_class	Optional[Union[str, List[str]]]	Target class(es) to rewrite for. Must be key(s) present in decoder_results summary (e.g. per-class ids like "3SG", "masc_nom", or "combined_score"). Pass a single string for one model, or a list of strings for one rewritten model per class. For strengthening, use the class id (e.g. "masc_nom"); for weakening, the summary key is "_weaken" but you pass the class id here and set increase_target_probabilities=False.	None
increase_target_probabilities	bool	If True (default), apply the config that strengthens the target class (higher probability for that class). If False, apply the config that weakens it (uses opposite feature factor; evaluate_decoder currently only produces strengthen summaries).	True
output_dir	Optional[str]	Optional directory where the rewritten model(s) should be saved. If provided, models are saved to disk. If None, models are returned in memory only. For a single target_class, this is the exact save directory. For multiple target_classes, experiment_dir must be set (output_dir is only used for a single target_class).	None
base_model	Optional[Any]	ModelWithGradiend to rewrite; if None, uses current trainer model.	None
decoder_stats_metric_name	Optional[str]	When loading decoder_results from cache, the summary key used to locate the cache file. Defaults to first target_class or "combined_score".	None
**decoder_stats_kwargs	Any	Used to locate the cache file (feature_factors, lrs, etc.).	{}

Returns:

Type	Description
Union[Any, List[Any], str, List[str]]	If output_dir is None: Rewritten model when target_class is a single string; list of models when target_class is a list.
Union[Any, List[Any], str, List[str]]	If output_dir is provided: Path to the saved rewritten model directory, or list of paths when target_class is a list.

Source code in gradiend/trainer/trainer.py

def rewrite_base_model(
    self,
    decoder_results: Optional[Dict[str, Any]] = None,
    target_class: Optional[Union[str, List[str]]] = None,
    increase_target_probabilities: bool = True,
    output_dir: Optional[str] = None,
    base_model: Optional[Any] = None,
    decoder_stats_metric_name: Optional[str] = None,
    **decoder_stats_kwargs: Any,
) -> Union[Any, List[Any], str, List[str]]:
    """
    Rewrite the base model by applying GRADIEND decoder updates based on decoder evaluation results.

    The decoder evaluation selects a feature factor and learning rate per target class and direction
    (strengthening vs weakening). This method applies the selected config: by default it strengthens
    the given target class(es); use increase_target_probabilities=False to apply the weakening config
    instead (evaluate_decoder currently only produces strengthen summaries).

    Accepts/loads internally the cached decoder results when experiment_dir is set. Optionally saves
    the rewritten model(s) to disk if output_dir is provided.

    When called on a Trainer, the trained model is used automatically (base_model not needed).
    When experiment_dir is set, decoder_results can be omitted and will be loaded from cache
    (requires evaluate_decoder to have been run with use_cache=True so the decoder stats cache exists).

    Args:
        decoder_results: Output from evaluate_decoder. Optional when
            experiment_dir is set (then loaded from cache).
        target_class: Target class(es) to rewrite for. Must be key(s) present in decoder_results
            summary (e.g. per-class ids like "3SG", "masc_nom", or "combined_score"). Pass a single
            string for one model, or a list of strings for one rewritten model per class. For
            strengthening, use the class id (e.g. "masc_nom"); for weakening, the summary key is
            "<class>_weaken" but you pass the class id here and set increase_target_probabilities=False.
        increase_target_probabilities: If True (default), apply the config that strengthens the
            target class (higher probability for that class). If False, apply the config that weakens
            it (uses opposite feature factor; evaluate_decoder currently only produces strengthen summaries).
        output_dir: Optional directory where the rewritten model(s) should be saved.
            If provided, models are saved to disk. If None, models are returned in memory only.
            For a single target_class, this is the exact save directory. For multiple target_classes,
            experiment_dir must be set (output_dir is only used for a single target_class).
        base_model: ModelWithGradiend to rewrite; if None, uses current trainer model.
        decoder_stats_metric_name: When loading decoder_results from cache, the summary key
            used to locate the cache file. Defaults to first target_class or "combined_score".
        **decoder_stats_kwargs: Used to locate the cache file (feature_factors, lrs, etc.).

    Returns:
        If output_dir is None:
            Rewritten model when target_class is a single string; list of models when target_class is a list.
        If output_dir is provided:
            Path to the saved rewritten model directory, or list of paths when target_class is a list.
    """
    if base_model is None:
        base_model = self.get_model()
    if base_model is None:
        raise ValueError(
            "Base model is required to rewrite base model. "
            "Provide a ModelWithGradiend instance or ensure the trainer has a valid model path."
        )

    if decoder_results is None:
        experiment_dir = self.experiment_dir
        if not experiment_dir:
            raise ValueError(
                "decoder_results is required when experiment_dir is not set. "
                "Run evaluate_decoder() first or set experiment_dir on TrainingArguments."
            )
        load_metric = (
            decoder_stats_metric_name
            or (target_class[0] if isinstance(target_class, list) and target_class else target_class)
            or "combined_score"
        )
        stats_file = resolve_decoder_stats_path(
            experiment_dir,
            metric_name=load_metric,
            feature_factors=decoder_stats_kwargs.get("feature_factors"),
            lrs=decoder_stats_kwargs.get("lrs"),
            topk=decoder_stats_kwargs.get("topk"),
            part=decoder_stats_kwargs.get("part"),
            topk_part=decoder_stats_kwargs.get("topk_part"),
        )
        if stats_file and os.path.isfile(stats_file):
            try:
                decoder_results = read_decoder_stats_file(stats_file)
            except Exception as e:
                raise ValueError(
                    f"Failed to load decoder results from cache at {stats_file}: {e}. "
                    "Run evaluate_decoder() first or provide decoder_results explicitly."
                )
        else:
            raise ValueError(
                f"No decoder results cache found at {stats_file}. "
                "Run evaluate_decoder() first or provide decoder_results explicitly."
            )

    _reserved = {"grid", "plot_path", "plot_paths"}
    summary_source = {k: v for k, v in decoder_results.items() if k not in _reserved}
    raw_keys: List[str] = (
        [target_class] if isinstance(target_class, str) else (target_class or [])
    )
    if not raw_keys:
        raise ValueError(
            "target_class must be a non-empty string or list of strings present in decoder results. "
            f"Available keys: {list(summary_source.keys())}"
        )
    keys_to_process: List[str] = [
        k if increase_target_probabilities else f"{k}_weaken" for k in raw_keys
    ]

    # When output_dir is passed but empty, user intended to save but gave no path
    if output_dir is not None and not str(output_dir).strip():
        raise ValueError(
            "Cannot save rewritten model: no output path. "
            "Set experiment_dir on TrainingArguments or pass a non-empty output_dir to rewrite_base_model."
        )

    # Check if we need to save
    should_save = output_dir is not None and str(output_dir).strip()
    if should_save:
        has_experiment_dir = bool(self.experiment_dir and str(self.experiment_dir).strip())
        has_output_dir = bool(output_dir is not None and str(output_dir).strip())
        if not has_experiment_dir and not has_output_dir:
            raise ValueError(
                "Cannot save rewritten model: no output path. "
                "Set experiment_dir on TrainingArguments or pass output_dir to rewrite_base_model."
            )
        if len(keys_to_process) > 1 and not has_experiment_dir:
            raise ValueError(
                "Cannot save multiple rewritten models without experiment_dir. "
                "Set experiment_dir on TrainingArguments (output_dir is only used for a single target_class)."
            )

    rewritten_models: List[Any] = []
    for key in keys_to_process:
        summary = summary_source.get(key)
        if not summary or "feature_factor" not in summary or "learning_rate" not in summary:
            hint = ""
            if not increase_target_probabilities and not key.endswith("_weaken"):
                hint = " evaluate_decoder currently only produces strengthen summaries."
            raise ValueError(
                f"Decoder results do not contain summary for metric '{key}'. "
                f"Available keys: {list(summary_source.keys())}.{hint}"
            )
        feature_factor = summary["feature_factor"]
        lr = summary["learning_rate"]
        rewritten = base_model.rewrite_base_model(
            learning_rate=lr,
            feature_factor=feature_factor,
        )
        rewritten_models.append(rewritten)

    # Save if output_dir was provided
    if should_save:
        saved_paths: List[str] = []
        for i, key in enumerate(keys_to_process):
            summary = summary_source.get(key)
            feature_factor = summary["feature_factor"]
            lr = summary["learning_rate"]
            explicit = output_dir if len(keys_to_process) == 1 else None
            key_output_dir = require_output_path(
                self.experiment_dir, explicit, ARTIFACT_MODEL_CHANGED, target_class=key
            )
            os.makedirs(key_output_dir, exist_ok=True)
            rewritten_models[i].save_pretrained(key_output_dir)
            logger.info(
                f"Saved rewritten model to {key_output_dir} "
                f"(feature_factor={feature_factor}, lr={lr}, metric={key})"
            )
            saved_paths.append(key_output_dir)
        return saved_paths[0] if len(saved_paths) == 1 else saved_paths

    # Return models in memory
    return rewritten_models[0] if len(rewritten_models) == 1 else rewritten_models

to

to(device)

Move loaded model to the given device (str or torch.device). No-op if model not loaded.

Source code in gradiend/trainer/trainer.py

def to(self, device: Any) -> "Trainer":
    """Move loaded model to the given device (str or torch.device). No-op if model not loaded."""
    model = self.get_model()
    if model is not None and hasattr(model, "to"):
        model.to(device)
    return self

train

train(output_dir=None, model=None, model_with_gradiend_cls=None, callbacks=None, **training_args_overrides)

Train GRADIEND using stored TrainingArguments (and optional overrides).

Parameters:

Name	Type	Description	Default
output_dir	Optional[str]	Directory to save the trained model. If None, resolved from experiment_dir (from TrainingArguments) or uses a temporary directory.	None
model	Optional[Union[str, Any]]	Model to train. If None, uses the model passed at Trainer initialization. Can be a string path or ModelWithGradiend instance.	None
model_with_gradiend_cls	Optional[Type[Any]]	ModelWithGradiend subclass to use when creating model from string path. Required if model is a string. If None, uses self.model_with_gradiend_cls or self.default_model_with_gradiend_cls (set by subclasses like TextPredictionTrainer). Examples: TextModelWithGradiend for text models.	None
callbacks	Optional[List[Any]]	Optional list of TrainingCallback instances for custom training behavior. If None, default callbacks are used (evaluation, normalization, checkpoint, logging).	None
**training_args_overrides	Any	Keyword arguments that override TrainingArguments values. These are merged with self.training_args (if set) or used to create new TrainingArguments. Examples: learning_rate=1e-3, num_epochs=10, experiment_dir="./results".	{}

Returns:

Type	Description
Union[Trainer, Dict[int, Any]]	Trainer instance (for single-seed training) or Dict[int, Any] mapping seed -> model
Union[Trainer, Dict[int, Any]]	(when keep_seed_runs=True in multi-seed training).

Multi-seed behavior (when TrainingArguments.max_seeds > 1):

• Each seed is trained from the same base model (or checkpoint path) but with a different random seed applied to PyTorch, Python's random, and NumPy.
• For each seed, training statistics are collected (including encoder correlation and best checkpoints). A training-time score ("training_score") is derived from these.
• Optionally, an additional encoder evaluation on the validation split is run via evaluate_encoder(split="val"), capped by seed_selection_eval_max_size (or encoder_eval_max_size when unset). Its correlation becomes "eval_correlation".
• The "selection_score" for each seed is:
  • eval_correlation when available,
  • otherwise training_score.
• A convergence metric (correlation or loss) and threshold are used to count how many seeds "converged" (see TrainingArguments.convergent_metric and convergent_score_threshold).

After the loop, a seed_report.json is written under /seeds containing: - top-level convergence info (metric, threshold, best seed, etc.) - a per-seed breakdown with training_score, eval_correlation, selection_score, convergence_metric_value, and convergence flags.

Source code in gradiend/trainer/trainer.py

def train(
    self,
    output_dir: Optional[str] = None,
    model: Optional[Union[str, Any]] = None,
    model_with_gradiend_cls: Optional[Type[Any]] = None,
    callbacks: Optional[List[Any]] = None,
    **training_args_overrides: Any,
) -> Union["Trainer", Dict[int, Any]]:
    """
    Train GRADIEND using stored TrainingArguments (and optional overrides).

    Args:
        output_dir: Directory to save the trained model. If None, resolved from
            experiment_dir (from TrainingArguments) or uses a temporary directory.
        model: Model to train. If None, uses the model passed at Trainer initialization.
            Can be a string path or ModelWithGradiend instance.
        model_with_gradiend_cls: ModelWithGradiend subclass to use when creating model from string path.
            Required if model is a string. If None, uses self.model_with_gradiend_cls or
            self.default_model_with_gradiend_cls (set by subclasses like TextPredictionTrainer).
            Examples: TextModelWithGradiend for text models.
        callbacks: Optional list of TrainingCallback instances for custom training behavior.
            If None, default callbacks are used (evaluation, normalization, checkpoint, logging).
        **training_args_overrides: Keyword arguments that override TrainingArguments values.
            These are merged with self.training_args (if set) or used to create new TrainingArguments.
            Examples: learning_rate=1e-3, num_epochs=10, experiment_dir="./results".

    Returns:
        Trainer instance (for single-seed training) or Dict[int, Any] mapping seed -> model
        (when keep_seed_runs=True in multi-seed training).

    Multi-seed behavior (when TrainingArguments.max_seeds > 1):

    - Each seed is trained from the same base model (or checkpoint path) but with a different
      random seed applied to PyTorch, Python's random, and NumPy.
    - For each seed, training statistics are collected (including encoder correlation
      and best checkpoints). A training-time score ("training_score") is derived from these.
    - Optionally, an additional encoder evaluation on the validation split is run via
      evaluate_encoder(split="val"), capped by seed_selection_eval_max_size (or
      encoder_eval_max_size when unset). Its correlation becomes "eval_correlation".
    - The "selection_score" for each seed is:
        * eval_correlation when available,
        * otherwise training_score.
    - A convergence metric (correlation or loss) and threshold are used to count how many
      seeds "converged" (see TrainingArguments.convergent_metric and convergent_score_threshold).

    After the loop, a seed_report.json is written under <experiment_dir>/seeds containing:
    - top-level convergence info (metric, threshold, best seed, etc.)
    - a per-seed breakdown with training_score, eval_correlation, selection_score,
      convergence_metric_value, and convergence flags.
    """
    if output_dir is not None and not isinstance(output_dir, str):
        raise TypeError(f"output_dir must be str or None, got {type(output_dir).__name__}")

    # Set env vars for reproducibility before any CUDA/BLAS use (PyTorch docs: set early)
    os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", ":4096:8")
    os.environ.setdefault("OMP_NUM_THREADS", "1")
    os.environ.setdefault("MKL_NUM_THREADS", "1")
    # Resolve model: use provided model or fall back to initialization model
    model = model or self._model_arg

    # Resolve model_with_gradiend_cls: use provided, then instance property, then default property
    if model_with_gradiend_cls is None:
        model_with_gradiend_cls = self.model_with_gradiend_cls or self.default_model_with_gradiend_cls

    # Merge TrainingArguments: start with stored args, then apply overrides
    args: Optional[TrainingArguments] = self._training_args
    if args is not None:
        args = TrainingArguments.from_dict({**args.to_dict(), **training_args_overrides})
    elif training_args_overrides:
        args = TrainingArguments.from_dict(training_args_overrides)
    else:
        args = TrainingArguments()
    args.__post_init__()  # validate e.g. not both supervised_encoder and supervised_decoder

    # Apply seed early so data loading, model init, and training are deterministic when seed is set
    if getattr(args, "seed", None) is not None:
        _apply_seed(int(args.seed))

    # Support output_dir in training_args_overrides
    if output_dir is None and "output_dir" in training_args_overrides:
        output_dir = training_args_overrides.pop("output_dir")

    # Resolve output_dir from merged args (so experiment_dir in training_args_overrides is respected; path matches cache check)
    exp_dir = args.experiment_dir
    if exp_dir and self.run_id and str(self.run_id).strip():
        exp_dir = os.path.join(exp_dir.rstrip("/\\"), str(self.run_id).strip().strip("/\\"))
    elif exp_dir:
        exp_dir = exp_dir.rstrip("/\\")
    output_dir = resolve_output_path(exp_dir, output_dir, ARTIFACT_MODEL)
    if output_dir is None:
        output_dir = tempfile.mkdtemp(prefix="gradiend_train_")
        logger.debug(
            "No output path or experiment_dir set; using temp dir %s. "
            "Model will be saved there; copy or save elsewhere if you need it after this run.",
            output_dir,
        )

    # Check cache before pre_prune/train/post_prune so we skip all of them when reusing
    if getattr(args, "use_cache", True) and has_saved_model(output_dir):
        logger.info(
            f"GRADIEND model already exists at {output_dir}, skipping training. Use use_cache=False to retrain."
        )
        self._model_arg = output_dir
        self._model_instance = None
        self._last_train_used_cache = True
        # Ensure data is loaded so target_classes, pair, etc. are available for evaluate_encoder / evaluate_decoder
        getattr(self, "_ensure_data_for_training", lambda: None)()
        return self

    self._last_train_used_cache = False
    def _set_seed(seed_value: int) -> None:
        _apply_seed(seed_value)

    def _cleanup_seed_model_files(seed_model_dir: str) -> None:
        if not os.path.isdir(seed_model_dir):
            return
        remove_files = [
            "model.safetensors",
            "pytorch_model.bin",
            "config.json",
        ]
        for fn in remove_files:
            path = os.path.join(seed_model_dir, fn)
            if os.path.exists(path):
                try:
                    os.remove(path)
                except Exception:
                    pass
        for fn in os.listdir(seed_model_dir):
            if fn.startswith("mapping_") or fn.startswith("input_index_map"):
                path = os.path.join(seed_model_dir, fn)
                if os.path.exists(path):
                    try:
                        os.remove(path)
                    except Exception:
                        pass
        model_subdir = os.path.join(seed_model_dir, "model")
        if os.path.isdir(model_subdir):
            try:
                shutil.rmtree(model_subdir)
            except Exception:
                pass

    def _best_score_from_stats(stats: Optional[Dict[str, Any]]) -> Optional[float]:
        if not stats:
            return None
        bsc = stats.get("best_score_checkpoint") or {}
        if args.supervised_decoder:
            val = bsc.get("loss")
            if isinstance(val, (int, float)):
                return -float(val)
        val = bsc.get("correlation")
        if isinstance(val, (int, float)):
            return float(val)
        ts = stats.get("training_stats") or {}
        val = ts.get("correlation")
        return float(val) if isinstance(val, (int, float)) else None

    # Multi-seed training loop
    if getattr(args, "max_seeds", 1) > 1:
        if not exp_dir:
            seed_runs_dir = args.seed_runs_dir or tempfile.mkdtemp(prefix="gradiend_seeds_")
        else:
            seed_runs_dir = args.seed_runs_dir or os.path.join(exp_dir, "seeds")
        os.makedirs(seed_runs_dir, exist_ok=True)

        if args.seed is None:
            seeds = list(range(args.max_seeds))
        else:
            seeds = [int(args.seed) + i for i in range(args.max_seeds)]

        convergent_metric = (args.convergent_metric or ("loss" if args.supervised_decoder else "correlation")).lower()
        threshold = args.convergent_score_threshold
        min_convergent = args.min_convergent_seeds
        convergent_count = 0

        # Use dedicated seed-selection eval cap when set, otherwise fall back to encoder_eval_max_size
        seed_selection_eval_max_size = getattr(args, "seed_selection_eval_max_size", None)
        if seed_selection_eval_max_size is None:
            seed_selection_eval_max_size = getattr(args, "encoder_eval_max_size", None)

        if not isinstance(model, str):
            model_path = getattr(model, "name_or_path", None)
            if not model_path:
                raise ValueError("Multi-seed training requires a model path (string) or a model with name_or_path.")
            model_for_runs = model_path
        else:
            model_for_runs = model
        # Resolve to custom prediction head (e.g. decoder MLM head) BEFORE experiment_dir override
        load_model_path = (
            self.resolve_model_path(model_for_runs) if isinstance(model_for_runs, str) else model_for_runs
        )

        best_seed = None
        best_score = None
        seed_results: Dict[int, str] = {}
        seed_report: List[Dict[str, Any]] = []
        early_stop_reason: Optional[str] = None

        for seed_value in seeds:
            _set_seed(seed_value)
            seed_dir = os.path.join(seed_runs_dir, f"seed_{seed_value}")
            seed_output_dir = seed_dir
            os.makedirs(seed_output_dir, exist_ok=True)

            used_cache = False
            trained = False
            if args.use_cache and has_saved_model(seed_output_dir):
                logger.info("Seed %s: cached model found at %s; skipping training.", seed_value, seed_output_dir)
                out_path = seed_output_dir
                used_cache = True
            else:
                logger.info("Seed %s: starting training with output_dir=%s", seed_value, seed_output_dir)
                model_instance = load_model_path
                if args.pre_prune_config is not None:
                    getattr(self, "_ensure_data_for_training", lambda: None)()
                    model_instance = create_model_with_gradiend(
                        load_model_path,
                        feature_definition=self,
                        model_class=model_with_gradiend_cls,
                        training_args=args,
                        trust_remote_code=getattr(args, "trust_remote_code", False),
                    )
                    max_size = args.train_max_size
                    training_data = self.create_training_data(model_instance, batch_size=1, max_size=max_size)
                    model_instance = _pre_prune(
                        model_instance,
                        training_data,
                        args.pre_prune_config,
                        definition=self,
                        inplace=True,
                    )
                # Temporarily override experiment_dir to seed-specific directory to avoid cache collisions
                original_experiment_dir = args.experiment_dir
                args.experiment_dir = seed_output_dir
                try:
                    out_path = self._train(
                        output_dir=seed_output_dir,
                        args=args,
                        model=model_instance,
                        model_with_gradiend_cls=model_with_gradiend_cls,
                        callbacks=callbacks,
                    )
                finally:
                    # Restore original experiment_dir
                    args.experiment_dir = original_experiment_dir
                trained = True
                if args.post_prune_config is not None:
                    logger.info("Seed %s: running post-prune after training ...", seed_value)
                    from gradiend.trainer.core.pruning import post_prune as _post_prune
                    seed_model = self.get_model(load_directory=out_path, use_cache=False)
                    seed_model = _post_prune(seed_model, args.post_prune_config)
                    seed_model.save_pretrained(out_path)
                    logger.info("Seed %s: saved post-pruned model to %s", seed_value, out_path)
                    del seed_model  # Release reference for GC

            seed_results[seed_value] = out_path
            stats = self.get_training_stats(out_path)
            score = _best_score_from_stats(stats)

            prev_model_arg = self._model_arg
            prev_model_instance = self._model_instance
            eval_corr = None
            # Optionally skip expensive full validation eval:
            # - never needed for loss-based convergence (convergent_metric == "loss")
            # - can be skipped when training_score (score) is clearly below threshold
            # - skip when we already have a convergent seed (no need to recompute encoder analysis)
            run_full_eval = convergent_metric != "loss"
            if (
                run_full_eval
                and threshold is not None
                and isinstance(score, (int, float))
                and score < threshold
            ):
                run_full_eval = False
            if run_full_eval and seed_value != seeds[0] and convergent_count >= 1:
                run_full_eval = False
            try:
                self._model_arg = out_path
                self._model_instance = None
                if run_full_eval:
                    eval_result = self.evaluate_encoder(
                        split="val",
                        max_size=seed_selection_eval_max_size,
                        use_cache=False,
                    )
                    if isinstance(eval_result, dict):
                        eval_corr = eval_result.get("correlation")
                        if not isinstance(eval_corr, (int, float)):
                            eval_corr = None
            except Exception as e:
                logger.warning("Seed %s: full validation encoder eval failed: %s", seed_value, e)
            finally:
                self._model_arg = prev_model_arg
                # Don't restore _model_instance in multi-seed: keep it cleared to free GPU memory
                if getattr(args, "max_seeds", 1) <= 1:
                    self._model_instance = prev_model_instance
                else:
                    self._model_instance = None

            selection_score = eval_corr if eval_corr is not None else score
            if selection_score is not None and (best_score is None or selection_score > best_score):
                best_score = selection_score
                best_seed = seed_value

            metric_val = None
            if stats:
                bsc = stats.get("best_score_checkpoint") or {}
                if convergent_metric == "loss":
                    metric_val = bsc.get("loss")
                    if metric_val is None:
                        metric_val = (stats.get("training_stats") or {}).get("loss")
                else:
                    metric_val = bsc.get("correlation")
                    if metric_val is None:
                        metric_val = (stats.get("training_stats") or {}).get("correlation")

            converged = False
            if isinstance(metric_val, (int, float)):
                if convergent_metric == "loss":
                    if metric_val <= threshold:
                        convergent_count += 1
                        converged = True
                else:
                    # When convergent_mean_by_class_threshold is set, require abs_mean_by_type['training'] >= that value.
                    # When it is None, or when the metric was not recorded (abs_training_mean is None), use only correlation.
                    abs_training_mean = (stats.get("abs_mean_by_type") or {}).get("training")
                    mean_ok = (
                        args.convergent_mean_by_class_threshold is None
                        or abs_training_mean is None
                        or (isinstance(abs_training_mean, (int, float)) and abs_training_mean >= args.convergent_mean_by_class_threshold)
                    )
                    if metric_val >= threshold and mean_ok:
                        convergent_count += 1
                        converged = True

            seed_report.append(
                {
                    "seed": seed_value,
                    "output_dir": out_path,
                    "trained": trained,
                    "used_cache": used_cache,
                    "training_score": score,
                    "eval_correlation": eval_corr,
                    "selection_score": selection_score,
                    "convergence_metric": convergent_metric,
                    "convergence_metric_value": metric_val,
                    "threshold": threshold,
                    "converged": converged,
                }
            )
            # Per-seed summary: path, converged, reason, convergent count / min required, max seeds
            conv_status = "Yes" if converged else "No"
            reason = ""
            if convergent_metric == "loss":
                reason = f"loss={metric_val}" if isinstance(metric_val, (int, float)) else "no metric"
            else:
                reason = f"{convergent_metric}={metric_val:.4f}" if isinstance(metric_val, (int, float)) else "no metric"
                if not converged and isinstance(metric_val, (int, float)) and threshold is not None:
                    if metric_val < threshold:
                        reason += f" (below threshold {threshold:.4f})"
                    else:
                        reason += " (mean-by-class criterion not met)"
            min_req = min_convergent if min_convergent is not None else "—"
            logger.info(
                "Finished seed %s (%s). Converged: %s (%s). Convergent: %s / min_required: %s, max_seeds: %s.",
                seed_value,
                out_path,
                conv_status,
                reason,
                convergent_count,
                min_req,
                args.max_seeds,
            )

            if min_convergent is not None and convergent_count >= min_convergent:
                logger.info(
                    "Convergence reached: %s seeds meet %s threshold %.4f.",
                    convergent_count, convergent_metric, float(threshold),
                )
                early_stop_reason = (
                    f"min_convergent_seeds reached: {convergent_count} >= {min_convergent} "
                    f"with metric={convergent_metric} threshold={threshold}"
                )
                break

            # Release GPU memory before loading next seed to avoid accumulation
            self._model_instance = None
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        if best_seed is None:
            raise RuntimeError("Multi-seed training finished, but no valid training stats were found.")

        report = {
            "convergence_metric": convergent_metric,
            "threshold": threshold,
            "min_convergent_seeds": min_convergent,
            "max_seeds": args.max_seeds,
            "seeds_tried": [r.get("seed") for r in seed_report],
            "convergent_count": convergent_count,
            "best_seed": best_seed,
            "best_selection_score": best_score,
            "early_stop_reason": early_stop_reason,
            "runs": seed_report,
        }
        try:
            report_path = os.path.join(seed_runs_dir, "seed_report.json")
            with open(report_path, "w") as f:
                json.dump(report, f, indent=2)
            if not is_under_temp_dir(report_path):
                logger.info("Wrote seed report to %s", report_path)
        except Exception as e:
            logger.warning("Failed to write seed report: %s", e)

        best_path = seed_results[best_seed]
        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
        shutil.copytree(best_path, output_dir)
        if not is_under_temp_dir(output_dir):
            logger.info("Selected best seed=%s -> %s", best_seed, output_dir)

        # Clear cached model so the next get_model() loads from output_dir (the selected best seed).
        # That first load is then cached; subsequent get_model() calls reuse the same instance.
        self._model_arg = output_dir
        self._model_instance = None

        # Check convergence and warn if non-convergent
        if convergent_count == 0 and min_convergent is not None and min_convergent > 0:
            logger.warning(
                "Multi-seed training completed but no seeds converged: "
                "convergent_count=0 (required: %s) for metric=%s threshold=%.4f. "
                "Model may not have reached the convergence threshold during training.",
                min_convergent,
                convergent_metric,
                threshold if threshold is not None else 0.0,
            )

        # Update training.json in output_dir with convergence info
        convergence_info = {
            "converged": convergent_count > 0 if min_convergent is not None and min_convergent > 0 else True,
            "convergent_count": convergent_count,
            "min_convergent_seeds": min_convergent,
            "convergence_metric": convergent_metric,
            "threshold": threshold,
        }
        try:
            from gradiend.trainer.core.stats import load_training_stats, write_training_stats
            stats = load_training_stats(output_dir)
            if stats:
                write_training_stats(
                    output_dir,
                    training_stats=stats.get("training_stats", {}),
                    best_score_checkpoint=stats.get("best_score_checkpoint", {}),
                    training_args=stats.get("training_args", {}),
                    time_stats=stats.get("time"),
                    losses=stats.get("losses"),
                    convergence_info=convergence_info,
                )
        except Exception as e:
            logger.debug("Could not update convergence_info in training.json: %s", e)

        if not args.keep_seed_runs:
            for seed_value, seed_path in seed_results.items():
                _cleanup_seed_model_files(seed_path)

        # Auto-save convergence plot when experiment_dir is set
        if exp_dir:
            try:
                path = self.plot_training_convergence(show=False, experiment_dir=exp_dir)
                if path:
                    logger.info("Saved training convergence plot: %s", path)
            except Exception as e:
                logger.warning("Failed to save training convergence plot: %s", e)

        if args.keep_seed_runs:
            seed_models: Dict[int, Any] = {}
            for seed_value, seed_path in seed_results.items():
                try:
                    seed_models[seed_value] = self.load_model(
                        seed_path,
                        device="cpu",
                        device_encoder="cpu",
                        device_decoder="cpu",
                        device_base_model="cpu",
                    )
                except Exception:
                    seed_models[seed_value] = seed_path
            return seed_models

        return self

    if args.seed is not None:
        _set_seed(int(args.seed))

    logger.info(f"Starting GRADIEND training with output_dir={output_dir}")

    if getattr(args, "pre_prune_config", None) is not None:
        self.pre_prune(inplace=True)
        model = self._model_arg

    out_path = self._train(
        output_dir=output_dir,
        args=args,
        model=model,
        model_with_gradiend_cls=model_with_gradiend_cls,
        callbacks=callbacks,
    )
    self._model_arg = out_path
    # Keep selected model in memory so immediate evaluation uses the chosen checkpoint.
    try:
        self._model_instance = self.get_model(load_directory=out_path, use_cache=False)
    except Exception as e:
        logger.warning("Could not load selected model into memory from %s: %s", out_path, e)
        self._model_instance = None

    if getattr(args, "post_prune_config", None) is not None:
        logger.info("Post-prune config set: running post-prune after training ...")
        self.post_prune()
        # Persist pruned model in place to the same output directory
        self.get_model().save_pretrained(out_path)
        logger.info("Saved post-pruned model to %s", out_path)

    # Auto-save convergence plot when experiment_dir is set
    if exp_dir:
        try:
            path = self.plot_training_convergence(show=False, experiment_dir=exp_dir)
            if path:
                logger.info("Saved training convergence plot: %s", path)
        except Exception as e:
            logger.warning("Failed to save training convergence plot: %s", e)

    return self