from __future__ import annotations
from .base import SuperClass
from neer_match.similarity_map import SimilarityMap
from neer_match.matching_model import DLMatchingModel
from neer_match_utilities.baseline_io import ModelBaseline
from neer_match_utilities.model import Model, EpochEndSaver, ValidationMetrics, validation_scores
from neer_match_utilities.thresholds import TUNE_METRICS
import pandas as pd
from datetime import datetime
from pathlib import Path
import shutil
import dill
import os
import numpy as np
import tensorflow.keras.backend as K
import tensorflow as tf
import math
import pickle
import warnings
[docs]
class Training(SuperClass):
"""
A class for managing and evaluating training processes, including
reordering matches, evaluating performance metrics, and exporting models.
Inherits:
---------
SuperClass : Base class providing shared attributes and methods.
"""
[docs]
def matches_reorder(self, matches: pd.DataFrame, matches_id_left: str, matches_id_right: str):
"""
Reorders a matches DataFrame to include indices from the left and
right DataFrames instead of their original IDs.
Parameters
----------
matches : pd.DataFrame
DataFrame containing matching pairs.
matches_id_left : str
Column name in the `matches` DataFrame corresponding to the left IDs.
matches_id_right : str
Column name in the `matches` DataFrame corresponding to the right IDs.
Returns
-------
pd.DataFrame
A DataFrame with columns `left` and `right`, representing the indices
of matching pairs in the left and right DataFrames.
"""
# Create local copies of the original dataframes
df_left = self.df_left.copy()
df_right = self.df_right.copy()
# Add custom indices
df_left['index_left'] = self.df_left.index
df_right['index_right'] = self.df_right.index
# Combine the datasets into one
df = pd.merge(
df_left,
matches,
left_on=self.id_left,
right_on=matches_id_left,
how='right',
validate='1:m',
suffixes=('_l', '_r')
)
df = pd.merge(
df,
df_right,
left_on=matches_id_right,
right_on=self.id_right,
how='left',
validate='m:1',
suffixes=('_l', '_r')
)
# Extract and rename index columns
matches = df[['index_left', 'index_right']].rename(
columns={
'index_left': 'left',
'index_right': 'right'
}
).reset_index(drop=True)
matches = matches.sort_values(by='left', ascending=True).reset_index(drop=True)
return matches
[docs]
def evaluate_dataframe(self, evaluation_test: dict, evaluation_train: dict):
"""
Combines and evaluates test and training performance metrics.
Parameters
----------
evaluation_test : dict
Dictionary containing performance metrics for the test dataset.
evaluation_train : dict
Dictionary containing performance metrics for the training dataset.
Returns
-------
pd.DataFrame
A DataFrame with accuracy, precision, recall, F-score, and a timestamp
for both test and training datasets.
"""
# Create DataFrames for test and training metrics
df_test = pd.DataFrame([evaluation_test])
df_test.insert(0, 'data', ['test'])
df_train = pd.DataFrame([evaluation_train])
df_train.insert(0, 'data', ['train'])
# Concatenate and calculate metrics
df = pd.concat([df_test, df_train], axis=0, ignore_index=True)
df['timestamp'] = datetime.now()
return df
[docs]
def focal_loss(alpha=0.99, gamma=1.5):
"""
Focal Loss function for binary classification tasks.
Focal Loss is designed to address class imbalance by assigning higher weights
to the minority class and focusing the model's learning on hard-to-classify examples.
It reduces the loss contribution from well-classified examples, making it
particularly effective for imbalanced datasets.
Parameters
----------
alpha : float, optional, default=0.99
Weighting factor for the positive class (minority class).
- Must be in the range [0, 1].
- A higher value increases the loss contribution from the positive class
(underrepresented class) relative to the negative class (overrepresented class).
gamma : float, optional, default=1.5
Focusing parameter that reduces the loss contribution from easy examples.
- ``gamma = 0``: No focusing, equivalent to Weighted Binary Cross-Entropy Loss (if alpha is set to 0.5).
- ``gamma > 0``: Focuses more on hard-to-classify examples.
- Larger values emphasize harder examples more strongly.
Returns
-------
loss : callable
A loss function that computes the focal loss given the true labels (`y_true`)
and predicted probabilities (`y_pred`).
Raises
------
ValueError
If `alpha` is not in the range [0, 1].
Notes
-----
- The positive class (minority or underrepresented class) is weighted by `alpha`.
- The negative class (majority or overrepresented class) is automatically weighted
by ``1 - alpha``.
- Ensure `alpha` is set appropriately to reflect the level of imbalance in the dataset.
References
----------
Lin, T.-Y., Goyal, P., Girshick, R., He, K., & Dollár, P. (2017).
Focal Loss for Dense Object Detection. In ICCV.
Explanation of Key Terms
-------------------------
- **Positive Class (Underrepresented):**
- Refers to the class with fewer examples in the dataset.
- Typically weighted by `alpha`, which should be greater than 0.5 in highly imbalanced datasets.
- **Negative Class (Overrepresented):**
- Refers to the class with more examples in the dataset.
- Its weight is automatically ``1 - alpha``.
"""
if not (0 <= alpha <= 1):
raise ValueError("Parameter `alpha` must be in the range [0, 1].")
def loss(y_true, y_pred):
# One value per pair. neer-match's evaluate passes labels of shape (n,)
# and predictions of shape (n, 1), which would otherwise broadcast to
# all n * n label-prediction combinations; Keras aligns them in fit.
y_true = tf.reshape(tf.cast(y_true, tf.float32), (-1,))
y_pred = tf.reshape(tf.cast(y_pred, tf.float32), (-1,))
# numerical safety
eps = K.epsilon()
y_pred = K.clip(y_pred, eps, 1.0 - eps)
# per-example alpha: alpha for positive, (1-alpha) for negative
alpha_t = y_true * alpha + (1.0 - y_true) * (1.0 - alpha)
# p_t is the prob of the true class
p_t = y_true * y_pred + (1.0 - y_true) * (1.0 - y_pred)
# BCE equals -log(p_t) when reduced per-example
bce = -K.log(p_t)
# focal modulation and weighting
fl = alpha_t * K.pow(1.0 - p_t, gamma) * bce
return K.mean(fl)
return loss
[docs]
def soft_f1_loss(epsilon: float = 1e-7):
"""
Soft F1 Loss for imbalanced binary classification tasks.
Soft F1 Loss provides a differentiable approximation of the F1 score,
combining precision and recall into a single metric. By optimizing
this loss, models are encouraged to balance false positives and false
negatives, which is especially useful when classes are imbalanced.
Parameters
----------
epsilon : float, optional, default=1e-7
Small constant added to numerator and denominator to avoid division
by zero and stabilize training. Must be > 0.
Returns
-------
loss : callable
A loss function that takes true labels (`y_true`) and predicted
probabilities (`y_pred`) and returns `1 - soft_f1`, so that
minimizing this loss maximizes the soft F1 score.
Raises
------
ValueError
If `epsilon` is not strictly positive.
Notes
-----
- True positives (TP), false positives (FP), and false negatives (FN)
are computed in a “soft” (differentiable) manner by summing over
probabilities rather than thresholded predictions.
- Soft F1 = (2·TP + ε) / (2·TP + FP + FN + ε).
- Loss = 1 − Soft F1, which ranges from 0 (perfect) to 1 (worst).
References
----------
- Lipton, Z. C., Elkan, C., & Naryanaswamy, B. (2014). Optimal Thresholding
of Classifiers to Maximize F1 Measure. *ECML PKDD 2014*.
Explanation of Key Terms
------------------------
- **True Positives (TP):** Sum of predicted probabilities for actual
positive examples.
- **False Positives (FP):** Sum of predicted probabilities assigned to
negative examples.
- **False Negatives (FN):** Sum of (1 − predicted probability) for
positive examples.
- **ε (epsilon):** Stabilizer to prevent division by zero when TP, FP,
and FN are all zero.
Examples
--------
```python
loss_fn = soft_f1_loss(epsilon=1e-6)
y_true = tf.constant([[1, 0, 1]], dtype=tf.float32)
y_pred = tf.constant([[0.9, 0.2, 0.7]], dtype=tf.float32)
loss_value = loss_fn(y_true, y_pred)
print(loss_value.numpy()) # e.g. 0.1…
```
"""
if epsilon <= 0:
raise ValueError("`epsilon` must be strictly positive.")
def loss(y_true, y_pred):
# One value per pair, as in focal_loss: (n,) and (n, 1) would broadcast
y_true = tf.reshape(tf.cast(y_true, tf.float32), (-1,))
y_pred = tf.reshape(tf.cast(y_pred, tf.float32), (-1,))
y_pred = tf.clip_by_value(y_pred, epsilon, 1.0 - epsilon)
# Soft counts
tp = tf.reduce_sum(y_pred * y_true)
fp = tf.reduce_sum(y_pred * (1 - y_true))
fn = tf.reduce_sum((1 - y_pred) * y_true)
# Denominator
denom = 2 * tp + fp + fn + epsilon
# Avoid NaNs from 0/0
soft_f1 = tf.where(denom > 0, (2 * tp + epsilon) / denom, tf.constant(0.0))
loss_val = 1.0 - soft_f1
return tf.where(tf.math.is_finite(loss_val), loss_val, tf.constant(1.0))
return loss
[docs]
def combined_loss(
weight_f1: float = 0.5,
epsilon: float = 1e-7,
alpha: float = 0.99,
gamma: float = 1.5
):
"""
Combined loss: weighted sum of Soft F1 loss and Focal Loss for imbalanced binary classification.
This loss blends the advantages of a differentiable F1-based objective (which balances
precision and recall) with the sample-focusing property of Focal Loss (which down-weights
easy examples). By tuning ``weight_f1``, you can interpolate between solely optimizing
for F1 score (when ``weight_f1 = 1.0``) and solely focusing on hard examples via focal loss
(when ``weight_f1 = 0.0``).
Parameters
----------
weight_f1 : float, default=0.5
Mixing coefficient in ``[0, 1]``.
- ``weight_f1 = 1.0``: optimize only Soft F1 loss.
- ``weight_f1 = 0.0``: optimize only Focal Loss.
- Intermediate values blend the two objectives proportionally.
epsilon : float, default=1e-7
Small stabilizer for Soft F1 calculation. Must be ``> 0``.
alpha : float, default=0.99
Balancing factor for Focal Loss, weighting the positive (minority) class.
Must lie in ``[0, 1]``.
gamma : float, default=1.5
Focusing parameter for Focal Loss.
- ``gamma = 0`` reduces to weighted BCE.
- Larger ``gamma`` emphasizes harder (misclassified) examples.
Returns
-------
callable
A function ``loss(y_true, y_pred)`` that computes
.. math::
\\text{CombinedLoss}
= \\text{weight\\_f1} \\cdot \\text{SoftF1}(y, \\hat{y};\\,\\varepsilon)
+ (1 - \\text{weight\\_f1}) \\cdot \\text{FocalLoss}(y, \\hat{y};\\,\\alpha, \\gamma).
Minimizing this combined loss encourages both a high F1 score
and focus on hard-to-classify samples.
Raises
------
ValueError
If ``weight_f1`` is not in ``[0, 1]``, or if ``epsilon <= 0``, or if ``alpha`` is not
in ``[0, 1]``, or if ``gamma < 0``.
Notes
-----
- **Soft F1 loss**: ``1 - \\text{SoftF1}``, where
.. math::
\\text{SoftF1} = \\frac{2\\,TP + \\varepsilon}{2\\,TP + FP + FN + \\varepsilon}.
Here ``TP``, ``FP``, and ``FN`` are *soft* counts computed from probabilities.
- **Focal Loss** down-weights well-classified examples to focus learning on difficult cases.
References
----------
- Lin, T.-Y., Goyal, P., Girshick, R., He, K., & Dollár, P. (2017).
Focal Loss for Dense Object Detection. *ICCV*.
- Lipton, Z. C., Elkan, C., & Naryanaswamy, B. (2014). Optimal Thresholding
of Classifiers to Maximize F1 Measure. *ECML PKDD 2014*.
Examples
--------
.. code-block:: python
import tensorflow as tf
loss_fn = combined_loss(weight_f1=0.5, epsilon=1e-6, alpha=0.25, gamma=2.0)
y_true = tf.constant([[1, 0, 1]], dtype=tf.float32)
y_pred = tf.constant([[0.9, 0.2, 0.7]], dtype=tf.float32)
value = loss_fn(y_true, y_pred)
print("Combined loss:", float(value.numpy()))
"""
# Validate hyper-parameters
if not (0.0 <= weight_f1 <= 1.0):
raise ValueError("`weight_f1` must be in [0, 1].")
if epsilon <= 0:
raise ValueError("`epsilon` must be strictly positive.")
if not (0.0 <= alpha <= 1.0):
raise ValueError("`alpha` must be in [0, 1].")
if gamma < 0:
raise ValueError("`gamma` must be non-negative.")
# Instantiate the individual losses
f1_fn = soft_f1_loss(epsilon)
focal_fn = focal_loss(alpha=alpha, gamma=gamma)
def loss(y_true, y_pred):
# Weighted combination
return (weight_f1 * f1_fn(y_true, y_pred)
+ (1.0 - weight_f1) * focal_fn(y_true, y_pred))
return loss
[docs]
def alpha_balanced(left, right, matches, mismatch_share:float=1.0, max_alpha:float=.95) -> float:
"""
Compute α so that α*N_pos = (1-α)*mismatch_share*N_neg, capped at max_alpha.
N_pos is the number of matches and N_neg = len(left)*len(right) - N_pos the
number of non-matching pairs, so that the positive pairs and the sampled
negative pairs carry equal total weight.
Parameters
----------
left, right : pandas.DataFrame
matches : pandas.DataFrame
mismatch_share : float, default=1.0
Share of the non-matching pairs sampled during training.
max_alpha : float, default=0.95
Upper bound on the returned α.
Returns
-------
float
α in [0,1], at most max_alpha, for focal loss (positive-class weight).
Raises
------
ValueError
If there are no pairs, or neither matches nor sampled non-matches.
"""
N_pos = len(matches)
N_neg = len(left)*len(right)-len(matches)
# Validate before dividing, so these raise instead of a ZeroDivisionError
N_total = len(left) * len(right)
if N_total <= 0:
raise ValueError("Total number of pairs is zero.")
if mismatch_share * N_neg + N_pos <= 0:
raise ValueError(
f"There are no matches and no sampled non-matches "
f"(mismatch_share={mismatch_share}), so there is nothing to balance."
)
alpha = (mismatch_share * N_neg) / (mismatch_share * N_neg + N_pos)
alpha = min(alpha, max_alpha)
return alpha
[docs]
class TrainingPipe:
"""
Orchestrates the full training and evaluation process of a deep learning
record-linkage model using a user-supplied similarity map and preprocessed data.
The class handles both training phases (soft-F1 pretraining and focal-loss fine-tuning),
dynamic learning-rate scheduling, and automatic weight-decay adaptation. It also exports
checkpoints, final models, and evaluation statistics for reproducibility.
Parameters:
-----------
model_name : str
Name assigned to the trained model. A corresponding subdirectory is created
under the project directory to store checkpoints and exports.
training_data : tuple or dict
Preprocessed training data in one of the following formats:
- Tuple: (left_train, right_train, matches_train)
- Dict: {"left": left_train, "right": right_train, "matches": matches_train}
testing_data : tuple or dict
Preprocessed testing data in one of the following formats:
- Tuple: (left_test, right_test, matches_test)
- Dict: {"left": left_test, "right": right_test, "matches": matches_test}
validation_data : tuple or dict, optional
Preprocessed validation data, in the same formats as `training_data`.
If given, the model is scored on all validation record pairs after every
epoch of both stages. The scores are logged with a `val_` prefix (e.g.
`val_F1`, `val_MCC`) in the history and in `epoch_overview.csv`, and each
stage keeps its best epoch by them rather than by the training F1:
with `tune_threshold`, by `val_tuned_MCC` (or `val_tuned_F1`), the
score at the threshold that maximises it in that epoch; otherwise by
`val_F1`. Stage 2 also counts the model it starts from, stage 1's best,
as its epoch 0 and keeps it if no epoch of its own does better.
Defaults to `None`.
early_stopping_patience : int, optional
Requires `validation_data`. If given, a stage stops once the score its
best epoch is chosen by has not improved for this many epochs.
Defaults to `None` (no early stopping).
threshold : float, optional
Probability from which a pair counts as a match, wherever none is
tuned: in training, and in the final evaluation without
`validation_data` or with `tune_threshold=False`. Defaults to 0.5.
tune_threshold : bool, optional
Whether to choose the threshold on the validation pairs, where
`validation_data` is given: in every epoch, to rank the epochs, and for
the final model, as `GradientBoostingModel` does. The final model's
threshold is saved with it as `best_threshold_`, used in its final
evaluation and written to the `threshold` column of `performance.csv`;
`Model.load` restores it. Defaults to `True`.
tune_metric : {'mcc', 'f1'}, optional
The validation score the threshold maximises. Defaults to `'mcc'`.
tune_match_share : 'train', float or None, optional
The share of matches the tuning weights the validation pairs to. Paired
record by record, a split of n records holds a share of about 1/n
matches, so that a small validation split holds far more than the
training or test pairs, and a threshold tuned on it as it is lets
through too many false positives. With `'train'` (the default), its
non-matches count as if the validation pairs had the training pairs'
share of matches; a float sets that share; None weights nothing, as up
to version 1.2.5. It applies to the `val_tuned_` scores, and so to the
choice of epochs, as well as to the final threshold.
evaluation_batch_size : int, optional
Record pairs per batch wherever the pipeline scores pairs without
training on them: the validation pairs after each epoch, the check after
stage 1, and the final evaluation of all training and test pairs. The
confusion counts and the rates computed from them (F1, MCC, ...) do not
depend on it. `Loss` does: it is neer-match's loss summed over these
batches. Defaults to 1024; neer-match's own default of 16 makes the
final evaluation take many times longer.
similarity_map : dict
User-defined similarity configuration mapping variable names to similarity measures.
Must follow the format accepted by `SimilarityMap`.
id_left_col : str, optional
Name of the unique identifier column in the left DataFrames
(`left_train` and `left_test`).
The ID column is used internally to index entities and to align
training labels. Defaults to `"id_unique"`.
id_right_col : str, optional
Name of the unique identifier column in the right DataFrames
(`right_train` and `right_test`).
The ID column is used internally to index entities and to align
training labels. Defaults to `"id_unique"`.
no_tm_pbatch : int
Target number of positive (matching) pairs per batch, in both stages and in
the warm-up pass. Used to adapt batch size dynamically via the
`required_batch_size` heuristic. Required whichever stages run.
save_architecture: bool, optional
Whether to save the model architecture in an image alongside weights when exporting.
Requires binaries of of graphviz to be installed. Otherwise, the code breaks.
Defaults to `False`.
overwrite : bool, optional
Whether to replace a model already saved under `model_name`. If False,
`execute()` raises `FileExistsError` before training starts when
``model_name/model`` exists in the working directory. Defaults to `False`.
stage_1 : bool, optional
Whether to run the first training stage (soft-F1 pretraining).
If False, the arguments `epochs_1` and `mismatch_share_1` are not required
and will be ignored.
stage_2 : bool, optional
Whether to run the second training stage (focal-loss fine-tuning).
If False, the arguments `epochs_2`, `mismatch_share_2`, `gamma`, and
`max_alpha` are not required and will be ignored.
epochs_1 : int, optional
Number of training epochs during the first phase (soft-F1 pretraining).
Required only when `stage_1=True`.
mismatch_share_1 : float, optional
Fraction of all possible negative (non-matching) pairs used during Round 1.
Required only when `stage_1=True`.
stage1_loss : str or callable, optional
Loss function used during Stage 1 (pretraining).
By default, this is ``soft_f1_loss()``, reproducing the original NeerMatch
behavior.
The argument accepts either:
- A **string** specifying a built-in or predefined loss:
``"soft_f1"`` (default) or ``"binary_crossentropy"``
- A **callable loss function**, allowing full customization:
``soft_f1_loss()``, ``focal_loss(alpha=0.25, gamma=2.0)``,
or any user-defined loss function of signature ``loss(y_true, y_pred)``
epochs_2 : int, optional
Number of training epochs during the second phase (focal-loss fine-tuning).
Required only when `stage_2=True`.
mismatch_share_2 : float, optional
Fraction of sampled negative pairs used during Round 2.
Required only when `stage_2=True`.
gamma : float, optional
Focusing parameter of the focal loss (Round 2).
Required only when `stage_2=True`.
max_alpha : float, optional
Maximum weighting factor of the positive class for focal loss (Round 2).
Required only when `stage_2=True`.
Returns:
--------
None
Notes
-----
- The pipeline assumes that the data have already been preprocessed, formatted, and tokenized.
- Round 1 (soft-F1 phase) initializes the model and emphasizes balanced learning across classes.
- Round 2 (focal-loss phase) refines the model to focus on hard-to-classify examples.
- Dynamic heuristics are used to automatically infer batch size, peak learning rate,
and weight decay.
- Model checkpoints, histories, and evaluation reports are stored in subdirectories named
after the provided ``model_name``.
- The final model, similarity map, and performance metrics are exported to disk using the
``Training.performance_statistics_export`` method for reproducibility.
- Each training stage can be enabled or disabled independently through
the ``stage_1`` and ``stage_2`` flags.
- If a stage is disabled, its hyperparameters are not required and will be ignored.
- When only one stage is active, the warm-up pass automatically adapts
to the active stage's mismatch sampling configuration.
"""
# Share of the non-matches the warm-up pass samples: small, to keep it short
WARMUP_MISMATCH_SHARE = 0.05
# ---------- Built-in helpers ----------
@staticmethod
def required_batch_size(num_matches: int,
num_left: int,
num_right: int,
mismatch_share: float,
desired_pos_per_batch: int = 8,
eps: float = 1e-12) -> tuple[int, float]:
total_mismatches = max(num_left * num_right - num_matches, 0)
sampled_negatives = mismatch_share * total_mismatches
denom = num_matches + sampled_negatives
if denom <= eps or num_matches == 0:
return max(desired_pos_per_batch, 1), 0.0
q = num_matches / denom
q = max(min(q, 1.0), eps)
batch_size = math.ceil(desired_pos_per_batch / q)
return batch_size, batch_size * q
@staticmethod
def count_trainable_params(keras_model) -> int:
return int(sum(v.shape.num_elements() for v in keras_model.trainable_variables))
@staticmethod
def suggest_peak_lr_adamw(batch_size: int,
positives_per_batch: float | None = None,
param_count: int | None = None,
base_batch: int = 256,
base_lr: float = 1e-3,
min_lr: float = 3e-4,
max_lr: float = 3e-3) -> float:
lr = base_lr * (batch_size / float(base_batch))
if positives_per_batch is not None:
x = max(min(positives_per_batch, 16.0), 1.0)
shrink = 0.6 + 0.4 * ((x - 1.0) / 15.0) # in [0.6, 1.0]
lr *= shrink
if param_count is not None:
scale = math.log10(max(param_count, 1)) - math.log10(10_000_000)
lr *= (0.8 ** max(scale, 0))
return float(max(min(lr, max_lr), min_lr))
@staticmethod
def _suggest_weight_decay_adamw(batch_size: int,
param_count: int,
learning_rate: float,
base_batch: int = 256,
base_params: int = 10_000_000,
base_lr: float = 1e-3,
base_wd: float = 5e-4,
wd_min: float = 1e-5,
wd_max: float = 5e-3) -> float:
batch_scale = base_batch / float(batch_size)
batch_scale = min(max(batch_scale, 0.5), 2.0)
param_scale = (base_params / float(param_count)) ** 0.2
param_scale = min(max(param_scale, 0.5), 2.0)
lr_scale = (learning_rate / base_lr) ** 0.5
lr_scale = min(max(lr_scale, 0.5), 2.0)
wd = base_wd * batch_scale * param_scale * lr_scale
return float(max(min(wd, wd_max), wd_min))
# ---------- Inner: WarmupCosine schedule ----------
[docs]
@tf.keras.utils.register_keras_serializable(package="custom")
class WarmupCosine(tf.keras.optimizers.schedules.LearningRateSchedule):
[docs]
def __init__(self, peak_lr, warmup_steps, total_steps, min_lr_ratio=0.1, name=None):
super().__init__()
self.peak_lr = float(peak_lr)
self.warmup_steps = int(warmup_steps)
self.total_steps = int(total_steps)
self.min_lr_ratio = float(min_lr_ratio)
self.name = name
def __call__(self, step):
step = tf.cast(step, tf.float32)
peak = tf.cast(self.peak_lr, tf.float32)
warm_steps = tf.cast(self.warmup_steps, tf.float32)
tot_steps = tf.cast(self.total_steps, tf.float32)
min_ratio = tf.cast(self.min_lr_ratio, tf.float32)
warm = peak * (step / tf.maximum(warm_steps, 1.0))
cos = peak * (
min_ratio
+ 0.5 * (1.0 - min_ratio)
* (1.0 + tf.cos(math.pi * (step - warm_steps) /
tf.maximum(tot_steps - warm_steps, 1.0)))
)
return tf.where(step < warm_steps, warm, cos)
def get_config(self):
return {
"peak_lr": self.peak_lr,
"warmup_steps": self.warmup_steps,
"total_steps": self.total_steps,
"min_lr_ratio": self.min_lr_ratio,
"name": self.name,
}
[docs]
@classmethod
def from_config(cls, config):
return cls(**config)
# ---------- Pipeline ----------
[docs]
def __init__(
self,
model_name: str,
training_data, # (left_train, right_train, matches_train) OR dict with keys
testing_data, # (left_test, right_test, matches_test) OR dict with keys
similarity_map: dict,
*,
validation_data=None, # same formats as training_data
early_stopping_patience: int | None = None,
threshold: float = 0.5,
tune_threshold: bool = True,
tune_metric: str = "mcc",
tune_match_share: float | str | None = "train",
evaluation_batch_size: int = 1024,
initial_feature_width_scales: int = 10,
feature_depths: int = 2,
initial_record_width_scale: int = 10,
record_depth: int = 4,
id_left_col: str = "id_unique",
id_right_col: str = "id_unique",
no_tm_pbatch: int | None = None,
save_architecture: bool = False,
overwrite: bool = False,
stage_1: bool = True,
stage_2: bool = True,
epochs_1: int | None = None,
mismatch_share_1: float | None = None,
stage1_loss: str | callable = soft_f1_loss(),
epochs_2: int | None = None,
mismatch_share_2: float | None = None,
gamma: float | None = None,
max_alpha: float | None = None,
):
if similarity_map is None:
raise ValueError("similarity_map is required and must not be None.")
self.similarity_map = similarity_map
self.model_name = model_name
self.initial_feature_width_scales = initial_feature_width_scales
self.feature_depths = feature_depths
self.initial_record_width_scale = initial_record_width_scale
self.record_depth = record_depth
self.stage_1 = bool(stage_1)
self.stage_2 = bool(stage_2)
if not self.stage_1 and not self.stage_2:
raise ValueError("At least one of stage_1 or stage_2 must be True.")
# ---- Stage 1 hyperparameters ----
if self.stage_1:
missing_stage1 = {
name: val
for name, val in [
("epochs_1", epochs_1),
("mismatch_share_1", mismatch_share_1),
("no_tm_pbatch", no_tm_pbatch),
]
if val is None
}
if missing_stage1:
raise ValueError(
"The following arguments must be provided when stage_1=True: "
+ ", ".join(missing_stage1.keys())
)
self.epochs_1 = int(epochs_1)
self.mismatch_share_1 = float(mismatch_share_1)
self.no_tm_pbatch = int(no_tm_pbatch)
# ---- Stage 1 loss configuration ----
if not (isinstance(stage1_loss, str) or callable(stage1_loss)):
raise ValueError("stage1_loss must be a loss name (str) or a callable.")
self.stage1_loss = stage1_loss
# ---- Stage 2 hyperparameters ----
if self.stage_2:
missing = {
name: val
for name, val in [
("epochs_2", epochs_2),
("mismatch_share_2", mismatch_share_2),
("no_tm_pbatch", no_tm_pbatch),
("gamma", gamma),
("max_alpha", max_alpha),
]
if val is None
}
if missing:
missing_keys = ", ".join(missing.keys())
raise ValueError(
f"The following arguments must be provided when stage_2=True: {missing_keys}."
)
self.epochs_2 = int(epochs_2)
self.mismatch_share_2 = float(mismatch_share_2)
self.no_tm_pbatch = int(no_tm_pbatch)
self.gamma = float(gamma)
self.max_alpha = float(max_alpha)
else:
# dummy values; not used when stage_2=False. no_tm_pbatch is not
# among them: stage 1 and the warm-up pass use it too.
self.epochs_2 = 0
self.mismatch_share_2 = 0.0
self.gamma = 0.0
self.max_alpha = 0.0
self.id_left_col = id_left_col
self.id_right_col = id_right_col
self.save_architecture = save_architecture
self.overwrite = bool(overwrite)
# Unpack user-supplied data
self.left_train, self.right_train, self.matches_train = self._unpack_split(training_data)
self.left_test, self.right_test, self.matches_test = self._unpack_split(testing_data)
self.validation_data = (
None if validation_data is None else self._unpack_split(validation_data)
)
if early_stopping_patience is not None and self.validation_data is None:
raise ValueError("early_stopping_patience requires validation_data.")
self.early_stopping_patience = early_stopping_patience
if (
not isinstance(evaluation_batch_size, int)
or isinstance(evaluation_batch_size, bool)
or evaluation_batch_size < 1
):
raise ValueError("evaluation_batch_size must be a positive integer.")
self.evaluation_batch_size = evaluation_batch_size
if not 0.0 < float(threshold) < 1.0:
raise ValueError("threshold must lie strictly between 0 and 1.")
self.threshold = float(threshold)
if tune_metric not in TUNE_METRICS:
raise ValueError(f"tune_metric must be one of {TUNE_METRICS}, not {tune_metric!r}.")
self.tune_metric = tune_metric
self.tune_threshold = bool(tune_threshold)
# Tuning needs validation pairs; without them the threshold stays fixed
self._tuning = self.tune_threshold and self.validation_data is not None
if tune_match_share == "train":
pairs = len(self.left_train) * len(self.right_train)
self._tune_match_share = len(self.matches_train) / pairs if pairs else None
elif tune_match_share is None:
self._tune_match_share = None
elif isinstance(tune_match_share, (int, float)) and not isinstance(tune_match_share, bool) \
and 0.0 < float(tune_match_share) < 1.0:
self._tune_match_share = float(tune_match_share)
else:
raise ValueError("tune_match_share must be 'train', None or a share between 0 and 1.")
self.tune_match_share = tune_match_share
# Basic sanity checks
frames = [
("left_train", self.left_train, self.id_left_col),
("right_train", self.right_train, self.id_right_col),
("left_test", self.left_test, self.id_left_col),
("right_test", self.right_test, self.id_right_col),
]
if self.validation_data is not None:
frames += [
("left_validation", self.validation_data[0], self.id_left_col),
("right_validation", self.validation_data[1], self.id_right_col),
]
for name, df, col in frames:
if col not in df.columns:
raise ValueError(f"{name} must include column '{col}'.")
self.base_dir = Path.cwd()
self.model: DLMatchingModel | None = None
# Minimal Training helper for exporting stats at the end
self.training_util = Training(
similarity_map=self.similarity_map,
df_left=self.left_train,
df_right=self.right_train,
id_left=self.id_left_col,
id_right=self.id_right_col,
)
@staticmethod
def _unpack_split(obj):
if isinstance(obj, dict):
return obj["left"], obj["right"], obj["matches"]
left, right, matches = obj
return left, right, matches
# ---------- Public entry point ----------
def execute(self):
# Fail now rather than at the final save, after all epochs have run
model_dir = self.base_dir / self.model_name / "model"
if model_dir.exists() and not self.overwrite:
raise FileExistsError(
f"Directory '{model_dir}' already exists. Pass overwrite=True to "
"replace the saved model, or choose a different model_name."
)
# Build model with provided similarity map
smap = SimilarityMap(self.similarity_map)
self.model = DLMatchingModel(
similarity_map=smap,
initial_feature_width_scales = self.initial_feature_width_scales,
feature_depths = self.feature_depths,
initial_record_width_scale = self.initial_record_width_scale,
record_depth = self.record_depth,
)
# Warmup pass to initialize shapes/BN, etc. Its batches are sized for the
# share of non-matches it samples, so they hold no_tm_pbatch matches as
# the stages' batches do.
bsz_warm, _ = self.required_batch_size(
len(self.matches_train),
len(self.left_train),
len(self.right_train),
self.WARMUP_MISMATCH_SHARE,
desired_pos_per_batch=self.no_tm_pbatch,
)
self.model.compile(threshold=self.threshold)
self.model.fit(
self.left_train,
self.right_train,
self.matches_train,
epochs=1,
batch_size=bsz_warm,
mismatch_share=self.WARMUP_MISMATCH_SHARE,
shuffle=True,
)
# Count params
P = self.count_trainable_params(self.model)
# Set by stage 1 for stage 2: the validation scores and the checkpoint
# of its chosen model
self._stage1_scores, self._stage1_best = None, None
# Round 1 (optional)
if self.stage_1:
self._round1(P)
# Round 2 (optional)
if self.stage_2:
self._round2(P)
# Save, evaluate, export
self._finalize_and_report()
@property
def _selection_metric(self) -> str:
"""The epoch metric each stage maximizes to pick its best epoch."""
if self.validation_data is None:
return "f1"
return f"val_tuned_{self.tune_metric.upper()}" if self._tuning else "val_F1"
def _validation_scores(self) -> tuple[dict, np.ndarray]:
"""The current model's scores on all validation pairs, and its predictions."""
return validation_scores(
self.model, *self.validation_data,
batch_size=self.evaluation_batch_size,
# always tuned, as the collapse check needs the best threshold's MCC
tune_metric=self.tune_metric if self._tuning else "mcc",
match_share=self._tune_match_share,
)
def _warn_if_collapsed(self, stage: str, scores: dict | None = None, history: pd.DataFrame | None = None):
"""
Warn when a stage's chosen model does not tell matches from non-matches:
no threshold gives the validation pairs a positive MCC, or, without
validation data, no epoch predicted a single training match.
"""
if scores is not None:
collapsed = not scores["val_tuned_MCC"] > 0
what = "no threshold separates any validation matches from the non-matches"
else:
collapsed = not history["f1"].max() > 0
what = f"no epoch predicted a single training match at threshold {self.threshold}"
if collapsed:
warnings.warn(
f"[TrainingPipe] {stage}: {what}. The network has probably collapsed to "
"predicting about the same probability for every pair, as networks whose "
"hidden layers are only a few units wide often do; wider layers "
"(initial_feature_width_scales, initial_record_width_scale) prevent it.",
RuntimeWarning,
stacklevel=3,
)
def _stage_callbacks(self) -> list:
"""
Callbacks for one stage's fit. Validation scoring comes first, so that
early stopping and the history see the val_ metrics it logs.
"""
callbacks = []
if self.validation_data is not None:
callbacks.append(
ValidationMetrics(
*self.validation_data,
batch_size=self.evaluation_batch_size,
tune_metric=self.tune_metric if self._tuning else None,
match_share=self._tune_match_share,
)
)
callbacks.append(EpochEndSaver(base_dir=self.base_dir, model_name=self.model_name))
if self.early_stopping_patience is not None:
callbacks.append(
tf.keras.callbacks.EarlyStopping(
monitor=self._selection_metric, mode="max", patience=self.early_stopping_patience
)
)
return callbacks
# ---------- Rounds ----------
def _round1(self, P: int):
assert self.model is not None
batch_size, expected_pos = self.required_batch_size(
len(self.matches_train), len(self.left_train), len(self.right_train),
self.mismatch_share_1, desired_pos_per_batch=self.no_tm_pbatch
)
peak_lr = self.suggest_peak_lr_adamw(batch_size, expected_pos, P)
no_obs_1 = self.mismatch_share_1 * (len(self.left_train) * len(self.right_train) - len(self.matches_train)) + len(self.matches_train)
steps_per_epoch_1 = int(round(no_obs_1 / batch_size, 0))
total_steps_1 = steps_per_epoch_1 * self.epochs_1
warmup_steps_1 = int(0.1 * total_steps_1)
lr_sched_1 = self.WarmupCosine(peak_lr, warmup_steps_1, total_steps_1, min_lr_ratio=0.001)
wd_1 = self._suggest_weight_decay_adamw(batch_size, P, peak_lr)
opt_1 = tf.keras.optimizers.AdamW(learning_rate=lr_sched_1, weight_decay=wd_1, clipnorm=1.0)
loss_1 = self._get_stage1_loss()
self.model.compile(loss=loss_1, optimizer=opt_1, threshold=self.threshold)
# Start from an empty checkpoints/, so that epochs left by an interrupted
# run are not renamed into checkpoints_stage_1/ along with this run's
checkpoints = self.base_dir / self.model_name / "checkpoints"
if checkpoints.exists():
shutil.rmtree(checkpoints)
history = self.model.fit(
self.left_train, self.right_train, self.matches_train,
epochs=self.epochs_1, batch_size=batch_size,
mismatch_share=self.mismatch_share_1, shuffle=True, callbacks=self._stage_callbacks()
)
df_epoch = pd.DataFrame(history.history)
df_epoch["epoch"] = pd.Series(history.epoch, dtype=int) + 1
(self.base_dir / self.model_name / "checkpoints").mkdir(parents=True, exist_ok=True)
df_epoch.to_csv(self.base_dir / self.model_name / "checkpoints" / "epoch_overview.csv", index=False)
best_row = df_epoch.loc[df_epoch[self._selection_metric].idxmax()]
best_epoch = int(best_row["epoch"])
best_epoch_str = f"{best_epoch:02d}"
# Load best R1 model
self.model = Model.load(self.base_dir / self.model_name / "checkpoints" / f"epoch_{best_epoch_str}")
print(f"[R1] best epoch by {self._selection_metric}:", best_epoch, "path:", self.base_dir / self.model_name / "checkpoints_stage_1" / f"epoch_{best_epoch_str}")
# Compiled with stage 1's loss, so that it can be evaluated (dummy optimizer)
probe_loss = self._get_stage1_loss()
self.model.compile(
loss=probe_loss, optimizer=tf.keras.optimizers.Adam(), threshold=self.threshold
)
if self.validation_data is not None:
# The loaded model's validation scores: a check that loading worked,
# and stage 2's epoch 0, the model it starts from
self._stage1_scores, _ = self._validation_scores()
print("[R1->R2] validation scores after load:", self._stage1_scores)
self._warn_if_collapsed("Stage 1", scores=self._stage1_scores)
else:
# Sanity-evaluate the loaded model before Round 2 training
eval_probe = self.model.evaluate(
self.left_train,
self.right_train,
self.matches_train,
mismatch_share=self.mismatch_share_1,
batch_size=self.evaluation_batch_size,
verbose=0,
)
print("[R1->R2] probe metrics after load:", eval_probe)
self._warn_if_collapsed("Stage 1", history=df_epoch)
# Preserve Round 1 checkpoints
old_path = self.base_dir / self.model_name / "checkpoints"
new_path = self.base_dir / self.model_name / "checkpoints_stage_1"
if old_path.exists():
if new_path.exists():
shutil.rmtree(new_path) # left by an earlier run
old_path.rename(new_path)
self._stage1_best = new_path / f"epoch_{best_epoch_str}"
def _get_stage1_loss(self):
"""
Resolve the configured Stage 1 loss into a *callable* loss function.
- If self.stage1_loss is "soft_f1": return soft_f1_loss()
(this preserves the original behavior)
- If self.stage1_loss is "binary_crossentropy": return a wrapped BCE
that reshapes labels/preds to be compatible with NeerMatch's evaluate loop.
- If self.stage1_loss is a callable (e.g. soft_f1_loss(), focal_loss(...)):
return it as-is.
"""
# Case 1: user passed a ready-made callable, e.g. soft_f1_loss() or focal_loss(...)
if callable(self.stage1_loss):
return self.stage1_loss
# Case 2: string selector
if self.stage1_loss == "soft_f1":
# Default behavior: identical to original code
return soft_f1_loss()
if self.stage1_loss == "binary_crossentropy":
# Wrap BCE to be robust to whatever shapes NeerMatch passes in
def bce_wrapped(y_true, y_pred):
y_true = tf.cast(y_true, tf.float32)
y_pred = tf.cast(y_pred, tf.float32)
# Flatten to 1D so shapes always match
y_true = tf.reshape(y_true, (-1,))
y_pred = tf.reshape(y_pred, (-1,))
return tf.keras.losses.binary_crossentropy(y_true, y_pred)
return bce_wrapped
# You can add more named options here if you like
raise ValueError(f"Unknown stage1_loss string: {self.stage1_loss}")
def _round2(self, P: int):
assert self.model is not None
# alpha bounded by your alpha_balanced (expects max_alpha argument in your implementation)
alpha = alpha_balanced(
left=self.left_train, right=self.right_train, matches=self.matches_train,
mismatch_share=self.mismatch_share_2, max_alpha=self.max_alpha
)
batch_size, expected_pos = self.required_batch_size(
len(self.matches_train), len(self.left_train), len(self.right_train),
self.mismatch_share_2, desired_pos_per_batch=self.no_tm_pbatch
)
no_obs_2 = self.mismatch_share_2 * (len(self.left_train) * len(self.right_train) - len(self.matches_train)) + len(self.matches_train)
steps_per_epoch_2 = int(round(no_obs_2 / batch_size, 0))
total_steps_2 = steps_per_epoch_2 * self.epochs_2
warmup_steps_2 = int(0.1 * total_steps_2)
peak_lr_2 = self.suggest_peak_lr_adamw(batch_size, expected_pos, P)
peak_lr_2 = min(peak_lr_2, 1e-4) # your cap
lr_sched_2 = self.WarmupCosine(peak_lr_2, warmup_steps_2, total_steps_2, min_lr_ratio=0.1)
wd_2 = self._suggest_weight_decay_adamw(batch_size, P, peak_lr_2)
opt_2 = tf.keras.optimizers.AdamW(learning_rate=lr_sched_2, weight_decay=wd_2, clipnorm=1.0)
self.model.compile(
loss=focal_loss(alpha=alpha, gamma=self.gamma), optimizer=opt_2, threshold=self.threshold
)
# Start from an empty checkpoints/, so that epochs left by an interrupted
# run are not renamed into checkpoints_stage_2/ along with this run's
checkpoints = self.base_dir / self.model_name / "checkpoints"
if checkpoints.exists():
shutil.rmtree(checkpoints)
history = self.model.fit(
self.left_train, self.right_train, self.matches_train,
epochs=self.epochs_2, batch_size=batch_size,
mismatch_share=self.mismatch_share_2, shuffle=True, callbacks=self._stage_callbacks()
)
df_epoch = pd.DataFrame(history.history)
df_epoch["epoch"] = pd.Series(history.epoch, dtype=int) + 1
# Stage 1's best model, the one this stage started from, is its epoch 0:
# stage 2 keeps it unless one of its own epochs scores higher
stage1_scores = self._stage1_scores if self.stage_1 else None
if stage1_scores is not None:
df_epoch = pd.concat(
[pd.DataFrame([{**stage1_scores, "epoch": 0}]), df_epoch], ignore_index=True
)
(self.base_dir / self.model_name / "checkpoints").mkdir(parents=True, exist_ok=True)
df_epoch.to_csv(self.base_dir / self.model_name / "checkpoints" / "epoch_overview.csv", index=False)
best_row = df_epoch.loc[df_epoch[self._selection_metric].idxmax()]
best_epoch = int(best_row["epoch"])
best_epoch_str = f"{best_epoch:02d}"
if best_epoch == 0:
self.model = Model.load(self._stage1_best)
print(f"[R2] best epoch by {self._selection_metric}:", best_epoch, "path:", self._stage1_best, "(stage 1's model)")
else:
self.model = Model.load(self.base_dir / self.model_name / "checkpoints" / f"epoch_{best_epoch_str}")
print(f"[R2] best epoch by {self._selection_metric}:", best_epoch, "path:", self.base_dir / self.model_name / "checkpoints_stage_2" / f"epoch_{best_epoch_str}")
# Re-compile for export/eval
self.model.compile(
loss=focal_loss(alpha=alpha, gamma=self.gamma), optimizer=opt_2, threshold=self.threshold
)
# Re-name checkpoints folder
old_path = self.base_dir / self.model_name / "checkpoints"
new_path = self.base_dir / self.model_name / "checkpoints_stage_2"
if old_path.exists():
if new_path.exists():
shutil.rmtree(new_path) # optional: only needed if rerunning
old_path.rename(new_path)
# ---------- Save / evaluate ----------
def _finalize_and_report(self):
assert self.model is not None and self.training_util is not None
# The final model's threshold: tuned on the validation pairs, as each
# epoch's was, or the fixed one
threshold, self.model.best_threshold_ = self.threshold, None
if self.validation_data is not None:
scores, _ = self._validation_scores()
if self._tuning:
threshold = self.model.best_threshold_ = scores["val_threshold"]
metric = self.tune_metric.upper()
print(
f"[TrainingPipe] threshold tuned on the validation pairs by {metric}: "
f"{threshold:.2f} (validation {metric} {scores[f'val_tuned_{metric}']:.4f} "
f"at it, {scores[f'val_{metric}']:.4f} at {self.threshold})"
)
if not self.stage_1 or self.stage_2:
self._warn_if_collapsed("Final model", scores=scores)
self.model.compile(
loss=self.model.loss, optimizer=self.model.optimizer, threshold=threshold
)
# Save the final bundle, with the threshold
Model.save(
model=self.model, target_directory=self.base_dir, name=self.model_name,
save_architecture=self.save_architecture, overwrite=self.overwrite,
)
# Remove checkpoints of a stage this run skipped. They were left by an
# earlier run under the same name and would otherwise sit next to a model
# they did not produce. This waits for the save above, so that a run that
# fails earlier leaves the previous run's files intact.
for stage_dir, stage_ran in (
("checkpoints_stage_1", self.stage_1),
("checkpoints_stage_2", self.stage_2),
):
stale = self.base_dir / self.model_name / stage_dir
if not stage_ran and stale.exists():
shutil.rmtree(stale)
# Evaluate on supplied splits. neer-match scores the batches one by one,
# in a loop that is not compiled, so the number of batches, not of
# pairs, sets how long this takes.
perf_train = self.model.evaluate(
self.left_train, self.right_train, self.matches_train,
mismatch_share=1.0, batch_size=self.evaluation_batch_size,
)
perf_test = self.model.evaluate(
self.left_test, self.right_test, self.matches_test,
mismatch_share=1.0, batch_size=self.evaluation_batch_size,
)
perf_train["threshold"] = perf_test["threshold"] = threshold
# Export stats
self.training_util.performance_statistics_export(
model=self.model, model_name=self.model_name, target_directory=self.base_dir,
evaluation_train=perf_train, evaluation_test=perf_test
)