SymmetricTrainerSuite

Bases: TrainerSuite

TrainerSuite for symmetric pair semantics such as variable contrasts.

Source code in gradiend/trainer/suite/base.py

def __init__(
    self,
    trainer_cls: Type[Trainer],
    *trainer_args: Any,
    target_classes: Optional[Sequence[str]] = None,
    target_pairs: Optional[Sequence[Sequence[str]]] = None,
    pair_definitions: Optional[Sequence[Any]] = None,
    run_id: Optional[str] = None,
    pair_filter: Optional[Callable[[Tuple[str, str]], bool]] = None,
    pair_id_fn: Optional[Callable[[Tuple[str, str]], str]] = None,
    pair_label_fn: Optional[Callable[[Tuple[str, str]], str]] = None,
    retain_models_in_memory: bool = True,
    model_device: str = "cpu",
    **trainer_kwargs: Any,
) -> None:
    classes = _infer_target_classes_from_inputs(
        target_classes=target_classes,
        trainer_kwargs=trainer_kwargs,
        trainer_cls=trainer_cls,
        trainer_args=trainer_args,
        pair_definitions=pair_definitions,
        target_pairs=target_pairs,
    )
    generated_incomplete_classes = _load_generated_incomplete_classes_from_kwargs(dict(trainer_kwargs))
    if generated_incomplete_classes:
        incomplete_set = set(generated_incomplete_classes)
        removed = [cls for cls in classes if cls in incomplete_set]
        if removed:
            logger.warning(
                "Excluding generated incomplete classes from TrainerSuite target classes: %s. "
                "Pairs containing these classes will not be built.",
                removed,
            )
            classes = [cls for cls in classes if cls not in incomplete_set]
    if len(classes) < 2:
        raise ValueError("TrainerSuite requires at least 2 target classes")
    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 model_device not in {"cpu", "cuda"}:
        raise ValueError("model_device must be 'cpu' or 'cuda'")
    if not issubclass(trainer_cls, Trainer):
        raise TypeError("trainer_cls must be a Trainer subclass")

    self.trainer_cls = trainer_cls
    self._trainer_args = trainer_args
    self._trainer_kwargs = dict(trainer_kwargs)
    self.target_classes = classes
    self.target_pairs = _normalize_target_pairs(target_pairs)
    if generated_incomplete_classes and self.target_pairs is not None:
        incomplete_set = set(generated_incomplete_classes)
        skipped = [pair for pair in self.target_pairs if pair[0] in incomplete_set or pair[1] in incomplete_set]
        if skipped:
            logger.warning("Skipping TrainerSuite target pairs with generated incomplete classes: %s", skipped)
        self.target_pairs = [
            pair for pair in self.target_pairs
            if pair[0] not in incomplete_set and pair[1] not in incomplete_set
        ]
    self._input_pair_definitions = _normalize_pair_definitions(
        pair_definitions,
        pair_id_fn=pair_id_fn or _default_child_id,
        pair_label_fn=pair_label_fn or _default_child_label,
    )
    self.run_id = run_id
    self.pair_filter = pair_filter
    self.pair_id_fn = pair_id_fn or _default_child_id
    self.pair_label_fn = pair_label_fn or _default_child_label
    self.retain_models_in_memory = bool(retain_models_in_memory)
    self.model_device = model_device

    self.pairs: List[Tuple[str, str]] = []
    self.pair_by_id: Dict[str, Tuple[str, str]] = {}
    self.pair_definitions: Dict[str, SuitePairDefinition] = {}
    self.trainers: Dict[str, Trainer] = {}
    self.label_mapping: Dict[str, str] = {}
    self._models: Dict[str, Any] = {}
    self._shared_base_model: Optional[Any] = None
    self._shared_tokenizer: Optional[Any] = None
    self._shared_model_key: Optional[str] = None

    candidate_pair_definitions = self._resolve_pair_definitions()
    if generated_incomplete_classes:
        incomplete_set = set(generated_incomplete_classes)
        filtered_definitions: List[SuitePairDefinition] = []
        skipped = 0
        for definition in candidate_pair_definitions:
            filtered = _filter_generated_incomplete_pair_definition(definition, incomplete_set)
            if filtered is None:
                skipped += 1
                continue
            filtered_definitions.append(filtered)
        if skipped:
            logger.warning(
                "Skipped %s TrainerSuite pair definitions because generated incomplete classes are excluded.",
                skipped,
            )
        candidate_pair_definitions = filtered_definitions
    if not candidate_pair_definitions:
        raise ValueError(f"{self.__class__.__name__} did not resolve any pair definitions")
    _validate_pair_definitions_against_data(
        trainer_kwargs=self._trainer_kwargs,
        pair_definitions=candidate_pair_definitions,
        trainer_cls=self.trainer_cls,
        trainer_args=self._trainer_args,
    )

    for definition in candidate_pair_definitions:
        pair = (str(definition.target_classes[0]), str(definition.target_classes[1]))
        if self.pair_filter is not None and not bool(self.pair_filter(pair)):
            continue
        child_id = str(definition.child_id) if definition.child_id is not None else self.pair_id_fn(pair)
        if child_id in self.trainers:
            raise ValueError(f"Duplicate child id generated for pair {pair}: {child_id!r}")
        self.pairs.append(pair)
        self.pair_by_id[child_id] = pair
        self.pair_definitions[child_id] = definition
        self.label_mapping[child_id] = str(definition.label) if definition.label is not None else self.pair_label_fn(pair)
        self.trainers[child_id] = self._build_child_trainer(definition=definition, child_id=child_id)
    self._validate_shared_model_compatibility()

_resolve_pair_definitions

_resolve_pair_definitions()

Source code in gradiend/trainer/suite/symmetric.py

def _resolve_pair_definitions(self) -> List[SuitePairDefinition]:
    if self._input_pair_definitions is not None:
        return list(self._input_pair_definitions)
    candidate_pairs = self.target_pairs or list(combinations(self.target_classes, 2))
    definitions: List[SuitePairDefinition] = []
    for pair in candidate_pairs:
        canon = _canonical_symmetric_pair((str(pair[0]), str(pair[1])))
        definitions.append(
            SuitePairDefinition(
                target_classes=canon,
                child_id=self.pair_id_fn(canon),
                label=self.pair_label_fn(canon),
            )
        )
    return definitions

compute_anchor_aligned_encoding_matrix

compute_anchor_aligned_encoding_matrix(feature_classes, *, encoder_summary=None, split='test', max_size=None, use_cache=True, full_eval=True, aggregate='mean', alignment='factual', column_ids=None)

Feature-class cross-encoding for symmetric pairs.

Rows are anchor feature classes (aggregated across GRADIENDs whose pair contains that class, with automatic sign alignment). Columns are evaluated feature classes.

Parameters:

Name	Type	Description	Default
feature_classes	Sequence[str]	Ordered feature classes used as matrix rows/columns.	required
encoder_summary	Optional[Dict[str, Any]]	Optional precomputed suite encoder result.	None
split	str	Encoder split used when running evaluation.	'test'
max_size	Optional[int]	Optional encoder-evaluation cap.	None
use_cache	bool	Whether to use cached encoder results.	True
full_eval	bool	Whether encoder evaluation includes all transitions.	True
aggregate	str	Aggregate used when multiple pair models cover one anchor.	'mean'
alignment	str	Column alignment mode.	'factual'
column_ids	Optional[Sequence[str]]	Optional explicit output columns.	None

Source code in gradiend/trainer/suite/symmetric.py

def compute_anchor_aligned_encoding_matrix(
    self,
    feature_classes: Sequence[str],
    *,
    encoder_summary: Optional[Dict[str, Any]] = None,
    split: str = "test",
    max_size: Optional[int] = None,
    use_cache: bool = True,
    full_eval: bool = True,
    aggregate: str = "mean",
    alignment: str = "factual",
    column_ids: Optional[Sequence[str]] = None,
) -> Dict[str, Any]:
    """
    Feature-class cross-encoding for symmetric pairs.

    Rows are anchor feature classes (aggregated across GRADIENDs whose pair
    contains that class, with automatic sign alignment). Columns are evaluated
    feature classes.

    Args:
        feature_classes: Ordered feature classes used as matrix rows/columns.
        encoder_summary: Optional precomputed suite encoder result.
        split: Encoder split used when running evaluation.
        max_size: Optional encoder-evaluation cap.
        use_cache: Whether to use cached encoder results.
        full_eval: Whether encoder evaluation includes all transitions.
        aggregate: Aggregate used when multiple pair models cover one anchor.
        alignment: Column alignment mode.
        column_ids: Optional explicit output columns.
    """
    if encoder_summary is None:
        encoder_summary = self.evaluate_encoder(
            split=split,
            max_size=max_size,
            use_cache=use_cache,
            plot=False,
            return_df=True,
            full_eval=full_eval,
        )
    return compute_anchor_aligned_encoding_matrix(
        pair_by_id=self.pair_by_id,
        encoder_summary=encoder_summary,
        feature_classes=feature_classes,
        aggregate=aggregate,
        alignment=alignment,
        column_ids=column_ids,
        source_by_id=source_by_id_from_trainers(self.trainers),
    )

plot_cross_encoding_heatmap

plot_cross_encoding_heatmap(feature_classes, *, alignment='factual', column_ids=None, encoder_summary=None, split='test', max_size=None, use_cache=True, full_eval=True, aggregate='mean', order='input', cluster=False, pretty_groups=None, **plot_kwargs)

Plot oriented cross-encoding heatmap for symmetric pairwise GRADIENDs.

Parameters:

Name	Type	Description	Default
feature_classes	Sequence[str]	Ordered feature classes used as matrix row anchors.	required
alignment	str	Column alignment mode (factual, counterfactual, or transition).	'factual'
column_ids	Optional[Sequence[str]]	Optional explicit output columns.	None
encoder_summary	Optional[Dict[str, Any]]	Optional precomputed suite encoder result.	None
split	str	Encoder split used when running evaluation.	'test'
max_size	Optional[int]	Optional encoder-evaluation cap.	None
use_cache	bool	Whether to use cached encoder results.	True
full_eval	bool	Whether encoder evaluation includes all transitions.	True
aggregate	str	Aggregate used when multiple pair models cover one anchor.	'mean'
order	Any	Heatmap ordering strategy or explicit order.	'input'
cluster	bool	If True, cluster heatmap rows/columns.	False
pretty_groups	Optional[Dict[str, List[str]]]	Optional display groups.	None
**plot_kwargs	Any	Forwarded to comparison heatmap plotting.	{}

Source code in gradiend/trainer/suite/symmetric.py

def plot_cross_encoding_heatmap(
    self,
    feature_classes: Sequence[str],
    *,
    alignment: str = "factual",
    column_ids: Optional[Sequence[str]] = None,
    encoder_summary: Optional[Dict[str, Any]] = None,
    split: str = "test",
    max_size: Optional[int] = None,
    use_cache: bool = True,
    full_eval: bool = True,
    aggregate: str = "mean",
    order: Any = "input",
    cluster: bool = False,
    pretty_groups: Optional[Dict[str, List[str]]] = None,
    **plot_kwargs: Any,
) -> Dict[str, Any]:
    """Plot oriented cross-encoding heatmap for symmetric pairwise GRADIENDs.

    Args:
        feature_classes: Ordered feature classes used as matrix row anchors.
        alignment: Column alignment mode (``factual``, ``counterfactual``,
            or ``transition``).
        column_ids: Optional explicit output columns.
        encoder_summary: Optional precomputed suite encoder result.
        split: Encoder split used when running evaluation.
        max_size: Optional encoder-evaluation cap.
        use_cache: Whether to use cached encoder results.
        full_eval: Whether encoder evaluation includes all transitions.
        aggregate: Aggregate used when multiple pair models cover one anchor.
        order: Heatmap ordering strategy or explicit order.
        cluster: If True, cluster heatmap rows/columns.
        pretty_groups: Optional display groups.
        **plot_kwargs: Forwarded to comparison heatmap plotting.
    """
    from gradiend.visualizer.heatmaps.encoding import plot_cross_encoding_heatmap

    if encoder_summary is None:
        encoder_summary = self.evaluate_encoder(
            split=split,
            max_size=max_size,
            use_cache=use_cache,
            plot=False,
            return_df=True,
            full_eval=full_eval,
        )
    return plot_cross_encoding_heatmap(
        self.trainers,
        feature_classes,
        alignment=alignment,
        column_ids=column_ids,
        encoder_summary=encoder_summary,
        split=split,
        max_size=max_size,
        use_cache=use_cache,
        full_eval=full_eval,
        aggregate=aggregate,
        order=order,
        cluster=cluster,
        pretty_groups=pretty_groups,
        **plot_kwargs,
    )