def plot_encoder_distributions(
trainer: Any,
encoder_df: pd.DataFrame = None,
output: Optional[str] = None,
output_dir: Optional[str] = None,
show: bool = True,
title: Union[str, bool] = True,
violin_order: Optional[List[str]] = None,
paired_legend_labels: Optional[List[str]] = None,
class_label_mapping: Optional[Dict[Any, str]] = None,
legend_name_mapping: Optional[Dict[str, str]] = None,
legend_group_mapping: Optional[Dict[str, List[str]]] = None,
ignore_missing_legend_group_labels: bool = False,
title_fontsize: Optional[float] = None,
label_fontsize: Optional[float] = None,
axis_label_fontsize: Optional[float] = None,
legend_fontsize: Optional[float] = None,
colors: Optional[Dict[str, str]] = None,
legend_loc: str = "best",
legend_ncol: Optional[int] = None,
legend_bbox_to_anchor: Optional[Tuple[float, float]] = None,
cmap: str = "tab20",
img_format: str = "png",
dpi: Optional[int] = None,
figsize: Optional[Tuple[float, float]] = None,
target_and_neutral_only: bool = True,
split_plot_mode: str = "facet",
neutral_data_split: str = DEFAULT_NEUTRAL_DATA_SPLIT,
include_neutral: bool = False,
highlight_non_convergence: Optional[bool] = None,
return_fig_ax: bool = False,
**kwargs: Any,
) -> Any:
"""
Plot encoder distributions as grouped split violins.
Requires matplotlib and seaborn. If missing, raises ImportError with install instructions.
Args:
trainer: Trainer with id, pair, get_model(), and analyze_encoder()
when encoder_df is None.
encoder_df: Pre-computed encoder analysis DataFrame (columns: encoded, label, source_id,
target_id, type, ...). If provided, no call to trainer.analyze_encoder.
output: Explicit path for saved PDF (overrides experiment_dir / output_dir).
output_dir: Directory for saved PDF when output and experiment_dir are not set.
show: If True, call plt.show() to display the plot.
title: True (default run_id), False, or custom string for the plot title.
target_and_neutral_only: If True (default), restrict the plot to the target (training)
transition(s) and neutral data only; other transitions are excluded. Uses
trainer.pair to determine the target transition(s). Set to False to show
all transitions.
split_plot_mode: When ``data_split`` has multiple values (``split="all"`` or a list),
use ``"facet"`` to draw columns per split. Removed modes ``"overlay"``
and ``"dodge"`` are accepted as aliases for ``"facet"``.
neutral_data_split: Split bucket for neutral encoder rows in multi-split plots
(default ``"test"``). Each neutral encoder type still gets its own
x-axis group and dodge/overlay hue (e.g. ``test — training masked``).
include_neutral: When False (default), multi-split plots show training rows only.
Set True to add neutral_dataset / neutral_training_masked groups.
violin_order: Optional ordering of violin groups on the x-axis.
paired_legend_labels: Optional explicit ordering of raw legend labels to plot. If set,
the labels are paired consecutively into split violins: (0,1) -> violin 1,
(2,3) -> violin 2, etc. This bypasses the default inference of split halves
from source/target ids.
class_label_mapping: Optional dict mapping individual class ids to display labels before
transition labels are built. For example, ``{"NM": "Masc. Nom."}`` turns
``"NM -> NF"`` into ``"Masc. Nom. -> Fem. Nom."``. When not provided, a
trainer ``id2label`` mapping is used when available.
legend_name_mapping: Optional dict mapping raw legend labels to display names.
legend_group_mapping: Optional dict mapping a new legend label to a list of existing
legend labels (e.g. {"Gender swap": ["masc_nom -> fem_nom",
"fem_nom -> masc_nom"]}). Grouped transitions are downsampled
to the minimum count within the group so groups are balanced.
highlight_non_convergence: When True, append a non-convergence marker to the title for
non-converged runs. ``None`` uses ``TrainingArguments.highlight_non_convergence``.
ignore_missing_legend_group_labels: If True, silently drop labels from
``legend_group_mapping`` that are not present in the plotted
encoder data and skip groups that become empty. This is useful
for experiment scripts that reuse a broad paper-style grouping
across runs whose encoder export contains only a subset.
title_fontsize: Font size for the title.
label_fontsize: Font size for axis tick labels.
axis_label_fontsize: Font size for axis labels.
legend_fontsize: Font size for legend text.
colors: Optional dict mapping legend labels to hex colors.
legend_loc: Matplotlib legend location (default "best").
legend_ncol: Number of columns for the legend. If None (default), a
sensible value is chosen based on the number of legend entries
(aiming for at most ~4 rows and up to 5 columns when many
transitions are shown).
legend_bbox_to_anchor: (x, y) for legend when placed outside (e.g. below when >6 entries). If None and >6 entries, legend is placed below the plot.
cmap: Matplotlib colormap name for palette (default "tab20").
img_format: File extension/format used when saving the figure.
dpi: Optional Matplotlib savefig DPI.
figsize: Figure size (width, height) in inches. If None, uses (max(6, 1.5 * n_groups), 3).
return_fig_ax: If True, return ``(fig, axes)`` and leave the figure open for
caller-side customization.
**kwargs: Forwarded to ``trainer.analyze_encoder`` when ``encoder_df`` is not supplied.
Returns:
Path to saved plot PDF, or "" if nothing to plot or plot was shown only (no save path).
Raises ValueError only when show=False and no output path can be determined.
"""
plt = _require_matplotlib()
sns = _require_seaborn()
# Suppress noisy INFO logs from matplotlib's categorical units used when
# plotting string labels that are parsable as numbers or dates.
import logging
for _name in ("matplotlib.category", "matplotlib.units"):
logging.getLogger(_name).setLevel(logging.WARNING)
run_id = getattr(trainer, "run_id", None)
highlight = resolve_highlight_non_convergence(highlight_non_convergence, trainer=trainer)
title = resolve_plot_title_with_convergence(
title,
trainer=trainer,
highlight_non_convergence=highlight,
default=str(run_id) if run_id else None,
)
if encoder_df is None:
raise ValueError(
"encoder_df is required. Call analyze_encoder(...) first and pass the returned DataFrame."
)
df_all = encoder_df
model_with_gradiend = getattr(trainer, "get_model", lambda: None)()
if df_all is None or df_all.empty:
raise ValueError("Encoder plot has no data (encoder_df or analyze_encoder returned empty).")
source_type = "factual"
if model_with_gradiend is not None and hasattr(model_with_gradiend, "gradiend"):
g = getattr(model_with_gradiend.gradiend, "kwargs", None) or {}
training_config = g.get("training", {}).get("training_args", {})
source_type = training_config.get("source", "factual")
training_pair = getattr(trainer, "pair", None)
training_transitions = set()
if training_pair and len(training_pair) >= 2:
training_transitions.add(f"{training_pair[0]} -> {training_pair[1]}")
training_transitions.add(f"{training_pair[1]} -> {training_pair[0]}")
# Map numeric source_id to class name (e.g. 0 -> "positive") when trainer has id2label.
trainer_id2label = getattr(trainer, "_id2label", None)
config_obj = getattr(trainer, "config", None)
config_id2label = getattr(config_obj, "id2label", None) if config_obj is not None else None
if not isinstance(trainer_id2label, dict):
trainer_id2label = {}
if not isinstance(config_id2label, dict):
config_id2label = {}
id2label = dict(config_id2label)
id2label.update(trainer_id2label)
if class_label_mapping:
id2label.update({k: v for k, v in class_label_mapping.items()})
def _source_to_display(s):
if not id2label or (isinstance(s, float) and s != s):
return s
try:
int_s = int(s)
return id2label.get(int_s, id2label.get(str(int_s), id2label.get(s, s)))
except (ValueError, TypeError):
return id2label.get(s, id2label.get(str(s), s))
if training_pair and len(training_pair) >= 2:
display_pair = [_source_to_display(training_pair[0]), _source_to_display(training_pair[1])]
training_transitions = {
f"{display_pair[0]} -> {display_pair[1]}",
f"{display_pair[1]} -> {display_pair[0]}",
}
df_training_probe = df_all[df_all["type"] == "training"]
if (
"data_split" in df_training_probe.columns
and df_training_probe["data_split"].nunique(dropna=True) > 1
):
return _plot_encoder_distributions_by_data_split(
df_all,
trainer=trainer,
_source_to_display=_source_to_display,
split_plot_mode=split_plot_mode,
neutral_data_split=neutral_data_split,
include_neutral=include_neutral,
target_and_neutral_only=target_and_neutral_only,
training_pair=training_pair,
show=show,
title=title,
run_id=run_id,
output=output,
output_dir=output_dir,
figsize=figsize,
img_format=img_format,
dpi=dpi,
cmap=cmap,
return_fig_ax=return_fig_ax,
**kwargs,
)
plot_rows = []
df_training = df_all[df_all["type"] == "training"].copy()
if not df_training.empty:
df_training["violin_group"] = df_training["source_id"].astype(str)
df_training["hue_label"] = df_training["target_id"].astype(str)
targets_per_group = (
df_training.groupby("violin_group")["hue_label"]
.apply(lambda s: sorted(s.dropna().unique().tolist()))
.to_dict()
)
def _assign_side(row):
tgts = targets_per_group.get(row["violin_group"], [])
if len(tgts) == 2:
return "Left" if row["hue_label"] == tgts[0] else "Right"
if len(tgts) == 1:
return "Left" if row["hue_label"] == tgts[0] else None
return None
df_training["side"] = df_training.apply(_assign_side, axis=1)
unique_sources = df_training["source_id"].dropna().astype(str).unique()
has_identity = (
(df_training["source_id"].astype(str) == df_training["target_id"].astype(str)).any()
)
# Use source-only labels for binary classification (two classes, no identity).
# Decision is based only on training data; it does not depend on presence of
# neutral data or on n_plot_groups, so behavior is consistent with or without neutral.
use_simple_class_labels = (
len(unique_sources) == 2 and not has_identity and source_type == "alternative"
)
if source_type == "alternative":
if use_simple_class_labels:
df_training["split_label"] = df_training["source_id"].astype(str).map(_source_to_display)
df_training["hue_label"] = df_training["source_id"].astype(str).map(_source_to_display)
df_training["violin_group"] = df_training["split_label"].astype(str)
training_transitions = set(df_training["split_label"].dropna().astype(str).unique())
else:
df_training["split_label"] = df_training.apply(
lambda r: f"{_source_to_display(r['source_id'])} -> {_source_to_display(r['target_id'])}", axis=1
)
df_training["violin_group"] = df_training["source_id"].map(_source_to_display).astype(str)
df_training["hue_label"] = df_training["target_id"].map(_source_to_display).astype(str)
df_training["is_training_transition"] = df_training["split_label"].isin(training_transitions)
else:
df_training["violin_group"] = df_training["source_id"].map(_source_to_display).astype(str)
df_training["hue_label"] = df_training["target_id"].map(_source_to_display).astype(str)
df_training["is_training_transition"] = df_training["source_id"].isin(training_pair or [])
plot_rows.append(df_training)
df_neutral_training_masked = df_all[df_all["type"] == "neutral_training_masked"].copy()
if not df_neutral_training_masked.empty:
df_neutral_training_masked["violin_group"] = "Neutral"
df_neutral_training_masked["hue_label"] = "Training masked"
df_neutral_training_masked["side"] = "Left"
plot_rows.append(df_neutral_training_masked)
df_neutral_dataset = df_all[df_all["type"] == "neutral_dataset"].copy()
if not df_neutral_dataset.empty:
df_neutral_dataset["violin_group"] = "Neutral"
df_neutral_dataset["hue_label"] = "Neutral dataset"
# When "Neutral: Training masked" is missing, put Neutral dataset on the left so the right half stays empty.
has_training_masked = not df_neutral_training_masked.empty
df_neutral_dataset["side"] = "Right" if has_training_masked else "Left"
plot_rows.append(df_neutral_dataset)
if not plot_rows:
raise ValueError("Encoder plot has no plottable rows (no training or neutral data).")
def _legend_label(g: str, label: str) -> str:
# Pass-through when already a full legend label (e.g. from paired mode or "X -> Y").
if isinstance(label, str) and (" -> " in label or label.startswith("Neutral:")):
return label
if legend_group_mapping and isinstance(label, str) and label in legend_group_mapping:
return label
if g == "Neutral":
return f"Neutral: {label}"
# Single class label (e.g. binary classification): show only source, not transition.
if g == label:
return g
return f"{g} -> {label}"
df_plot = pd.concat(plot_rows, ignore_index=True)
df_plot["legend_label"] = df_plot.apply(
lambda r: _legend_label(r["violin_group"], r["hue_label"]), axis=1
)
if target_and_neutral_only and training_transitions:
present_training_labels = set(
df_plot.loc[df_plot["type"] == "training", "legend_label"].dropna().unique().tolist()
)
matching_training_labels = present_training_labels & set(training_transitions)
keep_training_labels = matching_training_labels or present_training_labels
keep_labels = set(keep_training_labels) | {
lbl for lbl in df_plot["legend_label"].dropna().unique().tolist()
if isinstance(lbl, str) and lbl.startswith("Neutral:")
}
df_plot = df_plot[df_plot["legend_label"].isin(keep_labels)].copy()
if legend_group_mapping:
flat_labels = [lbl for labels in legend_group_mapping.values() for lbl in labels]
present_labels = set(df_plot["legend_label"].dropna().unique().tolist())
missing_labels = [lbl for lbl in flat_labels if lbl not in present_labels]
if missing_labels:
if ignore_missing_legend_group_labels:
legend_group_mapping = {
group: [lbl for lbl in labels if lbl in present_labels]
for group, labels in legend_group_mapping.items()
}
legend_group_mapping = {
group: labels for group, labels in legend_group_mapping.items() if labels
}
flat_labels = [lbl for labels in legend_group_mapping.values() for lbl in labels]
else:
raise ValueError(
"legend_group_mapping labels not present in data: %s"
% sorted(set(missing_labels))
)
indices_to_keep: set = set()
for new_label, legend_labels in legend_group_mapping.items():
subset = df_plot[df_plot["legend_label"].isin(legend_labels)]
if subset.empty:
raise ValueError(
f"Legend group '{new_label}' has no data (none of its labels are present in the plot)."
)
by_label = subset.groupby("legend_label", group_keys=False)
min_count = by_label.size().min()
if min_count == 0:
continue
sampled = by_label.sample(n=min_count, random_state=42)
indices_to_keep.update(sampled.index.tolist())
ungrouped = df_plot[~df_plot["legend_label"].isin(flat_labels)]
indices_to_keep.update(ungrouped.index.tolist())
df_plot = df_plot.loc[list(indices_to_keep)].copy()
label_to_group: Dict[str, str] = {}
for new_label, legend_labels in legend_group_mapping.items():
for lbl in legend_labels:
label_to_group[lbl] = new_label
df_plot["legend_label"] = df_plot["legend_label"].map(
lambda v: label_to_group.get(v, v)
)
used_paired_mode = True
if paired_legend_labels:
order = list(paired_legend_labels)
else:
order = df_plot["legend_label"].dropna().unique().tolist()
if not order:
raise ValueError("Encoder plot has no data to plot (no legend labels).")
df_plot = df_plot[df_plot["legend_label"].isin(order)].copy()
has_identity_transitions = False
if {"source_id", "target_id", "type"}.issubset(df_plot.columns):
identity_mask = (
(df_plot["type"] == "training")
& (df_plot["source_id"].astype(str) == df_plot["target_id"].astype(str))
)
has_identity_transitions = bool(identity_mask.any())
label_to_pair: Dict[str, Tuple[str, str]] = {}
for i, raw_label in enumerate(order):
pair_id = str(i // 2)
side = "Left" if (i % 2 == 0) else "Right"
label_to_pair[raw_label] = (pair_id, side)
df_plot["violin_group"] = df_plot["legend_label"].map(lambda s: label_to_pair[str(s)][0])
df_plot["side"] = df_plot["legend_label"].map(lambda s: label_to_pair[str(s)][1])
df_plot["hue_label"] = df_plot["legend_label"]
if used_paired_mode:
default_group_order = sorted(df_plot["violin_group"].unique().tolist(), key=lambda s: int(str(s)))
else:
training_groups = sorted([g for g in df_plot["violin_group"].unique().tolist() if g != "Neutral"])
has_neutral = "Neutral" in df_plot["violin_group"].values
default_group_order = training_groups + (["Neutral"] if has_neutral else [])
group_side_to_label = (
df_plot.drop_duplicates(["violin_group", "side"])[["violin_group", "side", "hue_label"]]
.set_index(["violin_group", "side"])["hue_label"]
.to_dict()
)
# Only show legend entries that have data; do not add placeholder entries for missing halves.
default_half_pairs = []
for g in default_group_order:
if (g, "Left") in group_side_to_label:
default_half_pairs.append((g, "Left"))
if (g, "Right") in group_side_to_label:
default_half_pairs.append((g, "Right"))
if (not used_paired_mode) and violin_order:
name_to_half = {
_legend_label(g, group_side_to_label[(g, side)]): (g, side)
for (g, side) in default_half_pairs
}
half_pairs = []
group_order = []
for name in violin_order:
if name in name_to_half:
g, side = name_to_half[name]
half_pairs.append((g, side))
if g not in group_order:
group_order.append(g)
if not half_pairs:
half_pairs = default_half_pairs
group_order = default_group_order
else:
half_pairs = default_half_pairs
group_order = default_group_order
group_to_x_id = {g: i for i, g in enumerate(group_order)}
df_plot["x_id"] = df_plot["violin_group"].map(group_to_x_id)
df_plot = df_plot[df_plot["x_id"].notna()].copy()
df_plot["x_id"] = df_plot["x_id"].astype(int)
x_id_order = list(range(len(group_order)))
df_plot["x_cat"] = pd.Categorical(df_plot["x_id"], categories=x_id_order, ordered=True)
_use_default_figsize = figsize is None
_figsize = figsize if figsize is not None else (max(6, 1.5 * len(group_order)), 3)
fig = plt.figure(figsize=_figsize)
ax = sns.violinplot(
data=df_plot,
x="x_cat",
y="encoded",
hue="side",
hue_order=["Left", "Right"],
split=True,
inner="quartile",
density_norm="width",
linewidth=0.7,
cut=0,
zorder=5,
)
_require_matplotlib() # Ensure matplotlib is available
from matplotlib.collections import PolyCollection
if ax.legend_ is not None:
ax.legend_.remove()
legend_label_to_group: Dict[str, str] = {}
if legend_group_mapping:
for new_label, labels in legend_group_mapping.items():
for lbl in labels:
legend_label_to_group[lbl] = new_label
# Only collapse to source class when the plot has exactly two violin groups
# (one target pair). With 3+ groups, each violin gets its own color and
# the legend shows full transitions. Note: with exactly two *labels*
# ("A -> B", "B -> A") both get pair_id 0, so n_plot_groups==1 and this
# collapse never ran; that case is handled earlier via use_simple_class_labels.
n_plot_groups = len([g for g in group_order if g != "Neutral"])
def _legend_group_for_half(g: str, side: str) -> str:
label = group_side_to_label.get((g, side), "")
# If the stored label is already a single class name (no " -> "), use it as-is.
# Otherwise we'd show "0 -> positive" because g is pair_id ("0") after the violin_group remap.
if isinstance(label, str) and label and " -> " not in label and not label.startswith("Neutral:"):
return legend_label_to_group.get(label, label)
raw = _legend_label(g, label)
# When there are no identity transitions, no explicit grouping/ordering,
# and only two violin groups in the plot, collapse transition labels
# to the factual class (left-hand side of "source -> target").
if (
not has_identity_transitions
and legend_group_mapping is None
and paired_legend_labels is None
and n_plot_groups == 2
and isinstance(raw, str)
):
if raw.startswith("Neutral:"):
# Keep neutral legend entries as-is.
pass
elif " -> " in raw:
raw = raw.split("->", 1)[0].strip()
return legend_label_to_group.get(raw, raw)
display_to_idx: Dict[str, int] = {}
display_to_raw: Dict[str, str] = {}
unique_displays: List[str] = []
half_to_display_idx: Dict[Tuple[str, str], int] = {}
def _display_label(raw: str) -> str:
mapped = (legend_name_mapping or {}).get(raw)
if mapped is not None:
return format_transition_label(mapped)
return format_transition_label(raw)
for (g, side) in half_pairs:
legend_group = _legend_group_for_half(g, side)
display = _display_label(legend_group)
if display not in display_to_idx:
display_to_idx[display] = len(unique_displays)
unique_displays.append(display)
display_to_raw[display] = legend_group
half_to_display_idx[(g, side)] = display_to_idx[display]
n_legend_entries = len(unique_displays)
# Derive a good default for legend_ncol when not explicitly provided.
# For many entries (e.g. ~20 transitions), aim for up to ~4 rows and at
# most 5 columns (so ~5x4 layout for 20 entries).
_legend_ncol = legend_ncol
if _legend_ncol is None:
if n_legend_entries <= 6:
_legend_ncol = 1
else:
# At most 4 rows: columns ~= ceil(n / 4), capped at 5 to avoid
# overly wide legends.
_legend_ncol = min(5, max(1, (n_legend_entries + 3) // 4))
try:
cmap_obj = plt.get_cmap(cmap)
n_cmap = getattr(cmap_obj, "N", None)
if n_cmap is not None and n_legend_entries <= n_cmap:
palette = [cmap_obj(i / n_cmap) for i in range(max(1, n_legend_entries))]
else:
palette = [
cmap_obj(i / max(1, n_legend_entries - 1)) if n_legend_entries > 1 else cmap_obj(0.5)
for i in range(max(1, n_legend_entries))
]
except (ValueError, TypeError):
palette = sns.color_palette(cmap, n_colors=max(1, n_legend_entries))
train_pair_half_set = set()
if used_paired_mode:
for (g, side) in half_pairs:
label = group_side_to_label.get((g, side))
if label and label in training_transitions:
train_pair_half_set.add((g, side))
elif not df_training.empty and "is_training_transition" in df_training.columns:
for (g, side) in half_pairs:
if g == "Neutral":
continue
label = group_side_to_label.get((g, side))
if label is None:
continue
sub = df_training[(df_training["violin_group"] == g) & (df_training["hue_label"] == label)]
if len(sub) and sub["is_training_transition"].any():
train_pair_half_set.add((g, side))
x_centers = {i: i for i in range(len(group_order))}
for coll in ax.collections:
if not isinstance(coll, PolyCollection):
continue
paths = coll.get_paths() or []
if not paths:
continue
xs = []
for p in paths:
verts = getattr(p, "vertices", None)
if verts is None or len(verts) == 0:
continue
xs.append(sum(v[0] for v in verts) / len(verts))
if not xs:
continue
x_mean = sum(xs) / len(xs)
x_id = int(round(x_mean))
if x_id not in x_centers:
continue
g = group_order[x_id]
side = "Left" if x_mean < x_centers[x_id] else "Right"
if (g, side) not in group_side_to_label:
continue
disp_idx = half_to_display_idx.get((g, side), 0)
face = palette[disp_idx % len(palette)]
coll.set_facecolor(face)
coll.set_edgecolor("black")
coll.set_alpha(1.0)
lw = 2.5 if (g, side) in train_pair_half_set else 0.7
coll.set_linewidth(lw)
if title is False:
pass
elif isinstance(title, str):
plt.title(title, fontsize=title_fontsize)
elif run_id:
plt.title(str(run_id), fontsize=title_fontsize)
ax.set_xticklabels([])
plt.xlabel("", fontsize=axis_label_fontsize)
plt.ylabel("Encoded value", fontsize=axis_label_fontsize)
if label_fontsize is not None:
ax.tick_params(labelsize=label_fontsize)
# Import matplotlib.patches locally where it's used
_require_matplotlib() # Ensure matplotlib is available before importing patches
import matplotlib.patches as mpatches
legend_handles = []
legend_labels = []
bold_legend_idxs = set()
for (g, side) in half_pairs:
if (g, side) in train_pair_half_set:
bold_legend_idxs.add(half_to_display_idx[(g, side)])
for i, display in enumerate(unique_displays):
raw_label = display_to_raw.get(display, display)
color = (colors or {}).get(display) or (colors or {}).get(raw_label) or palette[i % len(palette)]
legend_handles.append(mpatches.Patch(facecolor=color, edgecolor="black", label=display))
legend_labels.append(display)
# When many legend entries, place legend below the plot so it does not shrink the axes
legend_below = n_legend_entries > 6
if legend_below:
leg_loc = legend_loc if legend_bbox_to_anchor is not None else "upper center"
leg_bbox = legend_bbox_to_anchor if legend_bbox_to_anchor is not None else (0.5, -0.06)
legend = ax.legend(
legend_handles, legend_labels, loc=leg_loc, ncol=_legend_ncol, fontsize=legend_fontsize,
bbox_to_anchor=leg_bbox,
)
if legend_bbox_to_anchor is None:
fig = ax.get_figure()
fig.subplots_adjust(bottom=0.22)
else:
legend = ax.legend(
legend_handles,
legend_labels,
loc=legend_loc,
ncol=_legend_ncol,
fontsize=legend_fontsize,
)
for i in bold_legend_idxs:
try:
legend.get_texts()[i].set_fontweight("bold")
except Exception:
pass
experiment_dir = getattr(trainer, "experiment_dir", None)
if callable(experiment_dir):
experiment_dir = experiment_dir()
resolved = resolve_output_path(
experiment_dir, output, ARTIFACT_ENCODER_PLOT,
run_id=run_id,
)
out_path = resolved or (
os.path.join(output_dir, f"{run_id or 'run'}_encoder_distributions.pdf")
if output_dir else None
)
if out_path and img_format:
ext = img_format if img_format.startswith(".") else f".{img_format}"
out_path = os.path.splitext(out_path)[0] + ext
if legend_below and legend_bbox_to_anchor is None:
# When using default figsize, add vertical space so the plot area is not shrunk by the legend
if _use_default_figsize:
fig = ax.get_figure()
w, h = fig.get_size_inches()
fig.set_size_inches(w, h + 2.0)
plt.tight_layout(rect=[0, 0.2, 1, 1])
else:
plt.tight_layout()
plt.grid(axis="y", alpha=0.3, zorder=0)
if out_path:
os.makedirs(os.path.dirname(out_path), exist_ok=True)
save_kwargs: Dict[str, Any] = {"bbox_inches": "tight"}
if dpi is not None:
save_kwargs["dpi"] = dpi
plt.savefig(out_path, **save_kwargs)
logger.info("Saved encoder distribution plot: %s", out_path)
elif not show and not return_fig_ax:
plt.close()
raise ValueError(
"No output path or experiment_dir set and show=False. "
"Set experiment_dir on TrainingArguments, or pass output= / output_dir= to save, or show=True to display only."
)
if show:
plt.show()
if return_fig_ax:
return fig, ax
plt.close(fig)
return out_path or ""