plot_encoder_by_target

Encoder distribution by masked target token (x), grouped by feature class, colored by split.

Parameters:

Name	Type	Description	Default
trainer	Any	Optional trainer used to compute encoder_df and resolve labels/convergence.	None
encoder_df	Optional[DataFrame]	Optional precomputed encoder analysis DataFrame.	None
plot_style	PlotStyle	strip (default), box, or violin.	'strip'
target_col	str	Column containing the masked target token shown on the x-axis.	'source_token'
class_col	str	Column used to group target tokens by feature class.	'source_id'
hue_col	str	Column used as plot hue, typically data_split.	'data_split'
class_order	Optional[Sequence[str]]	Optional feature-class display order.	None
interactive	bool	If True, create a Plotly strip plot with hover instead of a static Matplotlib plot.	False
legend_loc	str	Matplotlib legend location for static plots.	'upper right'
point_size	float	Marker size for strip plot points (seaborn size).	1.5
title	Optional[str]	Optional plot title. None (default) draws no title.	None
output	Optional[str]	Explicit output path. Interactive plots are written as HTML.	None
output_dir	Optional[str]	Directory used when resolving the default output filename.	None
experiment_dir	Optional[str]	Experiment directory used when resolving the default output filename.	None
show	bool	Whether to display the plot.	True
figsize	Optional[Tuple[float, float]]	Static figure size in inches.	None
jitter	float	Jitter width for static strip plots.	0.25
dodge	bool	Whether to dodge static strip points by hue.	True
height	int	Plotly figure height for interactive plots.	520
highlight_non_convergence	Optional[bool]	When True, append a non-convergence marker to the title for non-converged runs. None uses trainer settings when available.	None
**kwargs	Any	Forwarded to trainer.analyze_encoder when encoder_df is not supplied.	{}

Source code in gradiend/visualizer/encoder_by_target.py

def plot_encoder_by_target(
    trainer: Any = None,
    encoder_df: Optional[pd.DataFrame] = None,
    *,
    plot_style: PlotStyle = "strip",
    target_col: str = "source_token",
    class_col: str = "source_id",
    hue_col: str = "data_split",
    class_order: Optional[Sequence[str]] = None,
    title: Optional[str] = None,
    output: Optional[str] = None,
    output_dir: Optional[str] = None,
    experiment_dir: Optional[str] = None,
    show: bool = True,
    figsize: Optional[Tuple[float, float]] = None,
    jitter: float = 0.25,
    dodge: bool = True,
    point_size: float = 1.5,
    interactive: bool = False,
    height: int = 520,
    legend_loc: str = "upper right",
    highlight_non_convergence: Optional[bool] = None,
    **kwargs: Any,
) -> Optional[str]:
    """
    Encoder distribution by masked target token (x), grouped by feature class, colored by split.

    Args:
        trainer: Optional trainer used to compute ``encoder_df`` and resolve labels/convergence.
        encoder_df: Optional precomputed encoder analysis DataFrame.
        plot_style: ``strip`` (default), ``box``, or ``violin``.
        target_col: Column containing the masked target token shown on the x-axis.
        class_col: Column used to group target tokens by feature class.
        hue_col: Column used as plot hue, typically ``data_split``.
        class_order: Optional feature-class display order.
        interactive: If True, create a Plotly strip plot with hover instead of a static Matplotlib plot.
        legend_loc: Matplotlib legend location for static plots.
        point_size: Marker size for strip plot points (seaborn ``size``).
        title: Optional plot title. None (default) draws no title.
        output: Explicit output path. Interactive plots are written as HTML.
        output_dir: Directory used when resolving the default output filename.
        experiment_dir: Experiment directory used when resolving the default output filename.
        show: Whether to display the plot.
        figsize: Static figure size in inches.
        jitter: Jitter width for static strip plots.
        dodge: Whether to dodge static strip points by hue.
        height: Plotly figure height for interactive plots.
        highlight_non_convergence: When True, append a non-convergence marker to the title
            for non-converged runs. ``None`` uses trainer settings when available.
        **kwargs: Forwarded to ``trainer.analyze_encoder`` when ``encoder_df`` is not supplied.
    """
    removed_label_kwargs = {
        "label_points",
        "label_indices",
        "label_col",
        "label_max_chars",
        "label_formatter",
        "label_sample_per_group",
        "adjust_labels",
        "label_fontsize",
        "outlier_method",
        "outlier_k",
    }
    passed_removed = sorted(removed_label_kwargs & set(kwargs))
    if passed_removed:
        raise TypeError(
            "plot_encoder_by_target does not accept point-label arguments: "
            + ", ".join(passed_removed)
        )
    if encoder_df is None:
        if trainer is None:
            raise ValueError("Provide encoder_df or trainer")
        encoder_df = trainer.analyze_encoder(getattr(trainer, "get_model", lambda: None)(), **kwargs)
    if encoder_df is None or encoder_df.empty:
        logger.warning("No encoder data for by-target plot")
        return None

    id2label = {}
    if trainer is not None:
        id2label = dict(getattr(trainer, "_id2label", None) or {})
        config_obj = getattr(trainer, "config", None)
        config_map = getattr(config_obj, "id2label", None) if config_obj is not None else None
        if isinstance(config_map, dict):
            id2label.update(config_map)
    if class_order is None and trainer is not None:
        pair = getattr(trainer, "pair", None)
        if pair:
            class_order = list(pair)

    plot_df, target_order, split_hue_order, target_to_class, target_to_split = (
        build_encoder_target_plot_frame(
        encoder_df,
        target_col=target_col,
        class_col=class_col,
        hue_col=hue_col,
        class_order=class_order,
        id2label=id2label or None,
    )
    )
    if plot_df.empty or not target_order:
        logger.warning("No plottable target tokens for by-target encoder plot")
        return None

    highlight = resolve_highlight_non_convergence(highlight_non_convergence, trainer=trainer)
    plot_title: Optional[str] = None
    if title is not None:
        plot_title = format_label_with_convergence(
            title,
            converged=converged_for_trainer(trainer),
            highlight_non_convergence=highlight,
        )

    class_order_resolved = class_order or sorted({target_to_class[t] for t in target_order})
    plot_df["x_group"] = plot_df["target_token"].astype(str)

    style = str(plot_style).lower()
    if style not in SUPPORTED_PLOT_STYLES:
        raise ValueError(
            f"plot_style must be one of {sorted(SUPPORTED_PLOT_STYLES)}; got {plot_style!r}"
        )

    if interactive:
        return _plot_encoder_by_target_interactive(
            plot_df,
            target_order=target_order,
            split_hue_order=split_hue_order,
            title=plot_title,
            output=output,
            show=show,
            height=height,
        )

    plt = _require_matplotlib()
    sns = _require_seaborn()
    width = max(4.0, 0.14 * len(target_order))
    _figsize = figsize if figsize is not None else (width, 2.7)
    fig, ax = plt.subplots(figsize=_figsize)
    split_colors = _split_color_map(split_hue_order)
    palette = [split_colors[split] for split in split_hue_order]

    common = dict(
        data=plot_df,
        x="target_token",
        y="encoded",
        hue="plot_hue",
        hue_order=split_hue_order,
        order=target_order,
        palette=palette,
        ax=ax,
    )
    splits_per_target = plot_df.groupby("target_token")["plot_hue"].nunique(dropna=True)
    center_single_split_targets = bool(not splits_per_target.empty and splits_per_target.max() <= 1)
    effective_dodge = bool(dodge and not center_single_split_targets)
    if style == "box":
        sns.boxplot(**common, dodge=effective_dodge, fliersize=2)
    elif style == "violin":
        sns.violinplot(**common, dodge=effective_dodge, inner="box", cut=0)
    elif style == "strip":
        sns.stripplot(**common, dodge=effective_dodge, jitter=jitter, size=point_size, alpha=0.85)

    ax.set_ylabel("Encoded value")
    ax.set_xlabel("Target")
    if plot_title:
        ax.set_title(plot_title, pad=30)
    ax.tick_params(axis="x", labelsize=8)
    plt.setp(ax.get_xticklabels(), rotation=90, ha="center")
    handles, labels = ax.get_legend_handles_labels()
    mark_colors = (
        _legend_split_colors(handles, labels)
        if handles and labels
        else split_colors
    )
    _draw_split_spine_marks(
        ax,
        target_order=target_order,
        target_to_split=target_to_split,
        target_to_class=target_to_class,
        class_order=class_order_resolved,
        split_color_map=mark_colors,
    )
    if handles and labels:
        for handle in handles:
            if hasattr(handle, "set_edgecolor"):
                handle.set_edgecolor("none")
        ax.legend(handles, labels, title="Split", loc=legend_loc, fontsize=8)
    _add_feature_class_group_brackets(
        ax,
        target_order=target_order,
        target_to_class=target_to_class,
        class_order=class_order_resolved,
    )
    _tighten_target_axis_xlim(ax, len(target_order))

    fig.tight_layout(rect=(0, 0, 1, 0.88 if plot_title else 0.92))

    out = output
    if out is None:
        run_id = getattr(trainer, "run_id", None) if trainer is not None else None
        out = resolve_output_path(
            experiment_dir or (getattr(trainer, "experiment_dir", None) if trainer is not None else None),
            output_dir,
            ARTIFACT_ENCODER_PLOT,
            run_id=run_id,
        )
        if out is not None:
            base, _ = os.path.splitext(out)
            out = f"{base}_by_target_{style}.png"
        elif output_dir:
            out = os.path.join(output_dir, f"encoder_by_target_{style}.png")
    if out:
        os.makedirs(os.path.dirname(out) or ".", exist_ok=True)
        fig.savefig(out, bbox_inches="tight")
        logger.info("Saved encoder by-target plot: %s", out)
    if show:
        plt.show()
    plt.close(fig)
    return out