Training

arguments

Configuration for Trainer and CrossEncoderTrainer runs.

TrainingArguments dataclass

TrainingArguments(
    output_dir: str,
    learning_rate: float = 3e-05,
    negative_samples_ratio: int = 5,
    use_peft: bool = False,
    peft_config: Any | None = None,
    num_epochs: int = 3,
    batch_size: int = 32,
    loss: Literal["infonce", "margin"] = "infonce",
    temperature: float = 0.05,
    margin: float = 1.0,
    device: str = "cpu",
    weight_decay: float = 0.01,
    warmup_ratio: float = 0.0,
)

Hyperparameters and settings for a training run.

Attributes:

Name	Type	Description
output_dir	str	Directory where the trained model's state_dict is written.
learning_rate	float	Optimizer learning rate.
negative_samples_ratio	int	Number of hard negatives sampled per positive pair during batch generation, via sample_hard_negatives.
use_peft	bool	Whether to wrap the model with LoRA adapters via PEFT. Requires peft_config to be set.
peft_config	Any | None	A peft.LoraConfig describing which of the model's layers to adapt (e.g. target_modules). Ignored unless use_peft is True, but required in that case -- there's no safe universal default, since target_modules names differ per model architecture (e.g. ["query", "value"] vs ["q_proj", "v_proj"]); a wrong guess would silently adapt the wrong (or no) layers. Typed loosely (Any) to avoid a hard import of peft at the dataclass-definition level.
num_epochs	int	Number of passes over the training data.
batch_size	int	Number of entity pairs per batch.
loss	Literal['infonce', 'margin']	Contrastive objective for Trainer (ignored by CrossEncoderTrainer, which always trains with binary cross-entropy -- there's no separately-embedded pair of vectors to contrast). "infonce" (default) treats every other pair's target in the batch, plus each anchor's own mined hard negatives, as negatives and applies a softmax-temperature cross-entropy loss. "margin" applies torch.nn.MarginRankingLoss between each anchor's positive and each of its mined hard negatives.
temperature	float	Softmax temperature for the "infonce" loss. Ignored by CrossEncoderTrainer.
margin	float	Margin for the "margin" loss. Ignored by CrossEncoderTrainer.
device	str	Torch device to train on, e.g. "cpu" (default) or "cuda". Resolved via resolve_device.
weight_decay	float	AdamW weight decay, applied to every parameter except bias and LayerNorm weights (standard BERT fine-tuning convention). Defaults to 0.01, torch.optim.AdamW's own default -- preserves Trainer's original behavior, which relied on that implicit default rather than setting it explicitly.
warmup_ratio	float	Fraction of total optimizer steps spent linearly warming up the learning rate before linearly decaying it back to zero, via transformers.get_linear_schedule_with_warmup. 0.0 (the default) skips the scheduler entirely -- a flat learning rate for the whole run, every existing caller's behavior before this field existed.

trainer

Unified training loop for trainable linking modules.

Supports bi-encoder/GNN contrastive training (InfoNCE or Margin Ranking Loss) with hard negative sampling drawn from a BlockingStrategy via sample_hard_negatives, and optional LoRA adapters via PEFT (see TrainingArguments.use_peft).

model must satisfy TrainableModel: a torch.nn.Module exposing a differentiable encode(entities) -> torch.Tensor method, one embedding row per entity, on the model's own device. No linker in this repo implements that yet -- this class is validated against a small in-repo toy encoder (see tests/train/test_trainer.py); future bi-encoder/GNN linkers implement encode() to plug into it.

TrainableModel

Bases: Protocol

Structural type for the model Trainer trains.

Any torch.nn.Module exposing a differentiable encode method satisfies this -- gradients must flow from its output back to self.parameters(), and it must return its output on the same device its parameters currently live on (Trainer moves the model via .to(device) but never touches tensors itself).

encode

encode(entities: list[Entity]) -> torch.Tensor

Return a (len(entities), dim) tensor of embeddings, one row per entity.

Source code in src/linkingtk/train/trainer.py

def encode(self, entities: list[Entity]) -> torch.Tensor:
    """Return a ``(len(entities), dim)`` tensor of embeddings, one row per entity."""
    ...

Trainer

Trainer(
    model: TrainableModel,
    args: TrainingArguments,
    train_data: list[tuple[Entity, Entity]],
    eval_data: list[tuple[Entity, Entity]] | None = None,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
)

Trains a TrainableModel via contrastive learning.

Parameters:

Name	Type	Description	Default
model	TrainableModel	The model to train.	required
args	TrainingArguments	Hyperparameters and settings, see TrainingArguments.	required
train_data	list[tuple[Entity, Entity]]	(source_entity, target_entity) positive pairs.	required
eval_data	list[tuple[Entity, Entity]] | None	Optional held-out positive pairs. If given, every epoch's end evaluates Hits@1/Hits@10/MRR over an exhaustive (not blocking-restricted) ranking of every eval entity, and appends the result to eval_history.	None
blocking	BlockingStrategy	Strategy used to mine hard negatives from train_data, via sample_hard_negatives. Defaults to ExactMatch, matching every other linker in this repo -- pass something tolerant of partial overlap (e.g. LabelOverlap) for it to find anything beyond train_data's own true pairs.	DEFAULT_BLOCKING

Attributes:

Name	Type	Description
eval_history	list[EvaluationReport]	One EvaluationReport per epoch, appended in train() when eval_data is given.

Source code in src/linkingtk/train/trainer.py

def __init__(
    self,
    model: TrainableModel,
    args: TrainingArguments,
    train_data: list[tuple[Entity, Entity]],
    eval_data: list[tuple[Entity, Entity]] | None = None,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> None:
    self.model = model
    self.args = args
    self.train_data = train_data
    self.eval_data = eval_data
    self.blocking = blocking
    self.eval_history: list[EvaluationReport] = []

train

train() -> None

Run the training loop, updating self.model in place.

Source code in src/linkingtk/train/trainer.py

def train(self) -> None:
    """Run the training loop, updating ``self.model`` in place."""
    args = self.args
    device = resolve_device(args.device)
    self.model.to(device)

    self.model = apply_peft(self.model, args)

    dataset1, dataset2, ground_truth = _entities_and_ground_truth(self.train_data)
    negatives_by_source = _group_by_source(
        sample_hard_negatives(
            dataset1, dataset2, ground_truth, self.blocking, top_k=args.negative_samples_ratio
        )
    )

    num_batches_per_epoch = -(-len(self.train_data) // args.batch_size)
    optimizer, scheduler = build_optimizer(
        self.model, args, num_batches_per_epoch * args.num_epochs
    )

    for _ in range(args.num_epochs):
        shuffled = list(self.train_data)
        random.shuffle(shuffled)
        for start in range(0, len(shuffled), args.batch_size):
            batch = shuffled[start : start + args.batch_size]
            loss = self._batch_loss(batch, negatives_by_source)
            if loss is None:
                continue
            optimizer.zero_grad()
            loss.backward()  # type: ignore[no-untyped-call]
            optimizer.step()
            if scheduler is not None:
                scheduler.step()

        if self.eval_data is not None:
            self.eval_history.append(self._evaluate())

    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    torch.save(self.model.state_dict(), Path(args.output_dir) / "model.pt")

cross_encoder

Training loop for cross-encoder linking modules.

Complements Trainer (bi-encoder/GNN models: independently encode() each side, then contrast the resulting embeddings) with the shape a true cross-encoder needs: one joint forward pass per (entity1, entity2) pair, scored directly, with no independently comparable embedding space to take a dot product over. GlossBertLinker is the first (and, as of this writing, only) consumer.

model must satisfy CrossEncoderModel: a torch.nn.Module exposing a differentiable score(pairs) -> torch.Tensor method, one logit-margin per pair (higher = more likely a true match).

CrossEncoderModel

Bases: Protocol

Structural type for the model CrossEncoderTrainer trains.

Any torch.nn.Module exposing a differentiable score method satisfies this -- gradients must flow from its output back to self.parameters(), and it must return its output on the same device its parameters currently live on (CrossEncoderTrainer moves the model via .to(device) but never touches tensors itself).

score

score(pairs: list[tuple[Entity, Entity]]) -> torch.Tensor

Return a (len(pairs),) tensor: one logit-margin per pair.

Source code in src/linkingtk/train/cross_encoder.py

def score(self, pairs: list[tuple[Entity, Entity]]) -> torch.Tensor:
    """Return a ``(len(pairs),)`` tensor: one logit-margin per pair."""
    ...

CrossEncoderTrainer

CrossEncoderTrainer(
    model: CrossEncoderModel,
    args: TrainingArguments,
    train_data: list[tuple[Entity, Entity]],
    eval_data: list[tuple[Entity, Entity]] | None = None,
    eval_dataset2: list[Entity]
    | EntitySource
    | None = None,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
)

Trains a CrossEncoderModel with binary cross-entropy.

Parameters:

Name	Type	Description	Default
model	CrossEncoderModel	The model to train.	required
args	TrainingArguments	Hyperparameters and settings, see TrainingArguments. loss/temperature/margin are ignored -- binary cross-entropy on model's own logit-margin is the only objective a cross-encoder needs (there's no separately embedded pair of vectors to contrast).	required
train_data	list[tuple[Entity, Entity]]	(source_entity, target_entity) positive pairs.	required
eval_data	list[tuple[Entity, Entity]] | None	Optional held-out positive pairs. If given, every epoch's end scores blocking.candidate_pairs(eval_dataset1, eval_dataset2) and appends a Hits@1 report to eval_history. Deliberately not exhaustive against a full target set the way Trainer._evaluate is: unlike EA/EL (where blocking is a scalability shortcut over a semantically meaningful full target set, so exhaustive ranking is the honest measure -- see issue #37/#38), WSD blocking is always exact, because ranking a mention against an unrelated lemma's senses isn't a meaningful question in the first place. A mention's own blocking-restricted candidates -- typically just its lemma's own senses -- are the right ranking scope here, matching the standard WSD "all-words" evaluation protocol (and every published WSD system, including GlossBERT itself).	None
eval_dataset2	list[Entity] | EntitySource | None	The real target set (or EntitySource, e.g. a WnEntitySource) to block eval_data's sources against. Defaults to the deduplicated targets within eval_data itself, matching Trainer._evaluate's convention -- correct for EA/EL, where the eval split's own KB subset already stands in for the full target set, but silently too easy for WSD: the eval split's own gold senses are a far narrower (and far less confusable) candidate pool than a mention's real lemma-wide sense inventory. Concretely, without this the reported Hits@1 can come out roughly 2x optimistic (measured directly: 60% self-reported vs. 33% measured via GlossBertLinker's own link(), same held-out mentions, same trained model -- always pass the real dataset2 for WSD.	None
blocking	BlockingStrategy	Strategy used to mine hard negatives from train_data (via sample_hard_negatives) and to restrict eval candidates. Defaults to ExactMatch, matching WSD's own "blocking is always exact" convention.	DEFAULT_BLOCKING

Attributes:

Name	Type	Description
eval_history	list[EvaluationReport]	One EvaluationReport per epoch, appended in train() when eval_data is given.

Source code in src/linkingtk/train/cross_encoder.py

def __init__(
    self,
    model: CrossEncoderModel,
    args: TrainingArguments,
    train_data: list[tuple[Entity, Entity]],
    eval_data: list[tuple[Entity, Entity]] | None = None,
    eval_dataset2: list[Entity] | EntitySource | None = None,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> None:
    self.model = model
    self.args = args
    self.train_data = train_data
    self.eval_data = eval_data
    self.eval_dataset2 = eval_dataset2
    self.blocking = blocking
    self.eval_history: list[EvaluationReport] = []

train

train() -> None

Run the training loop, updating self.model in place.

Source code in src/linkingtk/train/cross_encoder.py

def train(self) -> None:
    """Run the training loop, updating ``self.model`` in place."""
    args = self.args
    device = resolve_device(args.device)
    self.model.to(device)
    self.model = apply_peft(self.model, args)

    dataset1, dataset2, ground_truth = _entities_and_ground_truth(self.train_data)
    negatives = sample_hard_negatives(
        dataset1, dataset2, ground_truth, self.blocking, top_k=args.negative_samples_ratio
    )
    examples: list[tuple[Entity, Entity, float]] = [
        (entity1, entity2, 1.0) for entity1, entity2 in self.train_data
    ] + [(entity1, entity2, 0.0) for entity1, entity2 in negatives]

    num_batches_per_epoch = -(-len(examples) // args.batch_size)
    optimizer, scheduler = build_optimizer(
        self.model, args, num_batches_per_epoch * args.num_epochs
    )

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    logger.info(
        "CrossEncoderTrainer: %d examples (%d positive, %d mined negative), "
        "%d batches/epoch, %d epochs",
        len(examples),
        len(self.train_data),
        len(negatives),
        num_batches_per_epoch,
        args.num_epochs,
    )

    for epoch in range(args.num_epochs):
        shuffled = list(examples)
        random.shuffle(shuffled)
        loss_total = 0.0
        for start in range(0, len(shuffled), args.batch_size):
            batch = shuffled[start : start + args.batch_size]
            pairs = [(entity1, entity2) for entity1, entity2, _label in batch]
            labels = torch.tensor([label for _entity1, _entity2, label in batch], device=device)
            logits = self.model.score(pairs)
            loss = functional.binary_cross_entropy_with_logits(logits, labels)
            optimizer.zero_grad()
            loss.backward()  # type: ignore[no-untyped-call]
            optimizer.step()
            if scheduler is not None:
                scheduler.step()
            loss_total += float(loss.item())

        mean_loss = loss_total / max(num_batches_per_epoch, 1)

        if self.eval_data is not None:
            self.eval_history.append(self._evaluate())
            logger.info(
                "CrossEncoderTrainer: epoch %d/%d done, mean loss %.4f, eval %s",
                epoch + 1,
                args.num_epochs,
                mean_loss,
                self.eval_history[-1].metrics,
            )
        else:
            logger.info(
                "CrossEncoderTrainer: epoch %d/%d done, mean loss %.4f",
                epoch + 1,
                args.num_epochs,
                mean_loss,
            )

        # Saved after *every* epoch, not just at the end -- a long run
        # (e.g. full-corpus WSD training) can span hours unattended;
        # without this, an interruption on epoch 5 of 6 loses everything.
        torch.save(self.model.state_dict(), output_dir / "model.pt")
        logger.info("CrossEncoderTrainer: checkpoint saved to %s", output_dir / "model.pt")

ewiser_trainer

Training loop for EwiserEncoder.

A dedicated trainer, not a reuse of CrossEncoderTrainer: EWISER's real objective is full-inventory classification (one softmax cross-entropy per annotated word, over the entire SenseVocabulary), not a pairwise margin/BCE loss over (mention, candidate) pairs -- there's no independently-embedded pair of vectors to contrast, and batching is per-sentence (every annotated word in a sentence scored by one shared forward pass), not per-pair.

Confirmed directly against the reference's own training code (ewiser/fairseq_ext/criterions/weighted_cross_entropy.py): despite the name, "weighted cross-entropy" is not frequency-based class balancing -- the weight is 1.0 on every real sense and 0.0 only on the handful of reserved/special vocabulary slots (a padding mask, not class weighting, and never actually reachable as a gold class anyway). This trainer's loss is plain torch.nn.functional.cross_entropy over only the resolved (word position, gold vocabulary index) pairs -- unresolved mentions are filtered out before the loss is built (see _group_into_sentence_batches), which does the same masking job as the reference's own ignore_index without needing one.

EwiserTrainer

EwiserTrainer(
    model: EwiserEncoder,
    args: TrainingArguments,
    train_data: tuple[list[Entity], list[tuple[str, str]]],
    eval_data: tuple[list[Entity], list[tuple[str, str]]]
    | None = None,
    eval_dataset2: list[Entity]
    | EntitySource
    | None = None,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
    freeze_output_epochs: int = 0,
    output_freeze_lr: float = 0.0001,
    output_unfreeze_lr: float = 1e-05,
)

Trains an EwiserEncoder with full-inventory cross-entropy.

Parameters:

Name	Type	Description	Default
model	EwiserEncoder	The EwiserEncoder to train.	required
args	TrainingArguments	Hyperparameters and settings, see TrainingArguments. negative_samples_ratio/loss/temperature/margin/ use_peft are ignored -- there's no candidate sampling or embedding-pair objective here, and PEFT/LoRA has nothing meaningful to adapt on a model whose only trainable parameters are already the small decoder (the encoder stays frozen, see EwiserEncoder's own docstring). args.learning_rate is the single-stage (or post-freeze) learning rate; see output_freeze_lr for the freeze-stage rate.	required
train_data	tuple[list[Entity], list[tuple[str, str]]]	(mentions, ground_truth) -- mentions a list[Entity] (the shape SemCorDataset/ UfsacDataset produce as dataset1), ground_truth a list of (mention_id, synset_id) pairs. A mention whose gold synset isn't in model.vocabulary is dropped from training (logged, not an error) -- the same "unresolved -> excluded" convention UfsacDataset itself already uses.	required
eval_data	tuple[list[Entity], list[tuple[str, str]]] | None	Optional held-out (mentions, ground_truth), evaluated the same way link() would score it (via model.score + blocking), at the end of every epoch.	None
eval_dataset2	list[Entity] | EntitySource | None	The real target set (or EntitySource, e.g. a WnEntitySource) to block eval_data's mentions against. Required whenever eval_data is given -- unlike CrossEncoderTrainer, which can (if less accurately) fall back to a dataset2 derived from its own train_data's real candidate Entity objects, EwiserTrainer.train_data/eval_data carry only gold synset ids (see train_data), not full Entity objects with real labels -- a same-shaped fallback here would synthesize label-less entities that can never match anything under ExactMatch, silently produce zero candidates, and zero out every metric. Failing loudly beats that.	None
blocking	BlockingStrategy	Candidate-generation strategy for eval_data. Defaults to ExactMatch, matching WSD's "blocking is always exact" convention.	DEFAULT_BLOCKING
freeze_output_epochs	int	Number of leading epochs trained with model.decoder.logits.weight (the sense output embedding layer) frozen -- matches the reference's own "freeze-then-thaw" schedule (Section 3.4.1 of the paper), which protects an externally-pretrained (LMMS/SensEmBERT) initialization from noisy early gradients. Defaults to 0 (single-stage): freezing a randomly initialized output layer for a few epochs has no comparable benefit -- pass a nonzero value only when model was built with a real pretrained output_embedding_init (see load_synset_centroid_vectors/ build_synset_centroid_vectors_from_lmms), e.g. when continuing training from one of the published checkpoints or from-scratch with LMMS/SensEmBERT vectors -- this package doesn't bundle those vector files themselves (same as it doesn't bundle the published checkpoints), only the loaders that consume them.	0
output_freeze_lr	float	Learning rate for the leading freeze_output_epochs epochs. Only relevant when freeze_output_epochs > 0.	0.0001
output_unfreeze_lr	float	Learning rate from freeze_output_epochs onward (or for the entire run, when freeze_output_epochs=0 -- args.learning_rate is used instead in that single-stage case; see args).	1e-05

Attributes:

Name	Type	Description
eval_history	list[EvaluationReport]	One EvaluationReport per epoch, appended in train() when eval_data is given.

Source code in src/linkingtk/train/ewiser_trainer.py

def __init__(
    self,
    model: EwiserEncoder,
    args: TrainingArguments,
    train_data: tuple[list[Entity], list[tuple[str, str]]],
    eval_data: tuple[list[Entity], list[tuple[str, str]]] | None = None,
    eval_dataset2: list[Entity] | EntitySource | None = None,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
    freeze_output_epochs: int = 0,
    output_freeze_lr: float = 1e-4,
    output_unfreeze_lr: float = 1e-5,
) -> None:
    if eval_data is not None and eval_dataset2 is None:
        raise LinkingTKError(
            "EwiserTrainer requires eval_dataset2 whenever eval_data is given -- "
            "see eval_dataset2's own docstring for why no fallback is offered."
        )
    self.model = model
    self.args = args
    self.train_data = train_data
    self.eval_data = eval_data
    self.eval_dataset2 = eval_dataset2
    self.blocking = blocking
    self.freeze_output_epochs = freeze_output_epochs
    self.output_freeze_lr = output_freeze_lr
    self.output_unfreeze_lr = output_unfreeze_lr
    self.eval_history: list[EvaluationReport] = []

train

train() -> None

Run the training loop, updating self.model in place.

Source code in src/linkingtk/train/ewiser_trainer.py

def train(self) -> None:
    """Run the training loop, updating ``self.model`` in place."""
    device = resolve_device(self.args.device)
    self.model.to(device)

    batches = _group_into_sentence_batches(
        *self.train_data, self.model.vocabulary, self.args.batch_size
    )
    logger.info(
        "EwiserTrainer: %d sentence-batches/epoch, %d epochs",
        len(batches),
        self.args.num_epochs,
    )

    stage_lr = (
        self.output_freeze_lr if self.freeze_output_epochs > 0 else self.args.learning_rate
    )
    if self.freeze_output_epochs > 0:
        self.model.decoder.logits.weight.requires_grad = False
    optimizer, scheduler = build_optimizer(
        self.model,
        dataclasses.replace(self.args, learning_rate=stage_lr),
        len(batches) * self.args.num_epochs,
    )

    output_dir = Path(self.args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    for epoch in range(self.args.num_epochs):
        if epoch == self.freeze_output_epochs and self.freeze_output_epochs > 0:
            self.model.decoder.logits.weight.requires_grad = True
            optimizer, scheduler = build_optimizer(
                self.model,
                dataclasses.replace(self.args, learning_rate=self.output_unfreeze_lr),
                len(batches) * (self.args.num_epochs - epoch),
            )
            logger.info(
                "EwiserTrainer: epoch %d unfreezing decoder.logits.weight, lr -> %s",
                epoch + 1,
                self.output_unfreeze_lr,
            )

        shuffled = list(batches)
        random.shuffle(shuffled)
        loss_total = 0.0
        for sentences, targets in shuffled:
            logits_per_sentence = self.model._encode_chunk(sentences)
            losses = []
            for logits, (word_indices, gold_indices) in zip(
                logits_per_sentence, targets, strict=True
            ):
                if not word_indices:
                    continue
                word_index_tensor = torch.tensor(word_indices, device=device)
                gold_index_tensor = torch.tensor(gold_indices, device=device)
                losses.append(
                    functional.cross_entropy(
                        logits[word_index_tensor], gold_index_tensor, reduction="sum"
                    )
                )
            if not losses:
                continue
            loss = torch.stack(losses).sum() / sum(len(w) for w, _g in targets)
            optimizer.zero_grad()
            loss.backward()  # type: ignore[no-untyped-call]
            optimizer.step()
            if scheduler is not None:
                scheduler.step()
            loss_total += float(loss.item())

        mean_loss = loss_total / max(len(batches), 1)

        if self.eval_data is not None:
            self.eval_history.append(self._evaluate())
            logger.info(
                "EwiserTrainer: epoch %d/%d done, mean loss %.4f, eval %s",
                epoch + 1,
                self.args.num_epochs,
                mean_loss,
                self.eval_history[-1].metrics,
            )
        else:
            logger.info(
                "EwiserTrainer: epoch %d/%d done, mean loss %.4f",
                epoch + 1,
                self.args.num_epochs,
                mean_loss,
            )

        # Saved after *every* epoch, not just at the end -- a long run
        # can span hours unattended; without this, an interruption
        # loses everything (same rationale as CrossEncoderTrainer).
        torch.save(self.model.state_dict(), output_dir / "model.pt")
        logger.info("EwiserTrainer: checkpoint saved to %s", output_dir / "model.pt")

esc_trainer

Training loop for EscEncoder.

A dedicated trainer, not a reuse of CrossEncoderTrainer or EwiserTrainer: ESC's real objective is extractive-QA span-position cross-entropy (one joint sequence per mention, containing every one of that mention's candidates' glosses at once, with the gold candidate's token span as the target start/end position) -- not a pairwise margin/BCE loss over independent (mention, candidate) pairs (CrossEncoderTrainer), and not per-sentence full-inventory classification (EwiserTrainer). The loss itself needs no hand-rolled cross-entropy: EscEncoder.forward(..., start_positions=..., end_positions=...) already returns HF's own (start_loss + end_loss) / 2 via .loss (see EscEncoder's own module docstring).

Unlike EwiserTrainer (whose training never needs a candidate-generation step -- gold indices are resolved directly against its fixed output vocabulary), ESC's training is candidate generation: each mention's training example requires the same real, blocking-restricted candidate senses (with real gloss text) that EscLinker.link() would score at inference time. So EscTrainer takes a senses source unconditionally (not only when evaluating, like EwiserTrainer.eval_dataset2) -- used to build both training examples and (if eval_data is given) evaluation candidates.

EscTrainer

EscTrainer(
    model: EscEncoder,
    args: TrainingArguments,
    train_data: tuple[list[Entity], list[tuple[str, str]]],
    senses: list[Entity] | EntitySource,
    eval_data: tuple[list[Entity], list[tuple[str, str]]]
    | None = None,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
)

Trains an EscEncoder with extractive-QA span-position cross-entropy.

Parameters:

Name	Type	Description	Default
model	EscEncoder	The EscEncoder to train.	required
args	TrainingArguments	Hyperparameters and settings, see TrainingArguments. negative_samples_ratio/loss/temperature/margin are ignored -- there's no candidate sampling or embedding-pair objective here, the loss comes straight from the QA head (see the module docstring).	required
train_data	tuple[list[Entity], list[tuple[str, str]]]	(mentions, ground_truth) -- mentions a list[Entity], ground_truth a list of (mention_id, sense_id) pairs, same shape EwiserTrainer takes.	required
senses	list[Entity] | EntitySource	The real target sense set (or EntitySource, e.g. a WnEntitySource) blocking resolves each mention's candidates against -- required unconditionally (unlike EwiserTrainer.eval_dataset2, which is only needed for evaluation): ESC's training examples are themselves built from real candidate glosses, not a fixed output vocabulary index.	required
eval_data	tuple[list[Entity], list[tuple[str, str]]] | None	Optional held-out (mentions, ground_truth), evaluated the same way EscLinker.link() would score it (via model.score + blocking against senses), at the end of every epoch.	None
blocking	BlockingStrategy	Candidate-generation strategy for both training examples and eval_data. Defaults to ExactMatch, matching WSD's "blocking is always exact" convention (see CrossEncoderTrainer's own docstring) -- pass a POS-restricted wrapper (as examples/esc_reproduction.py/ esc_benchmark.py do) to match the reference's own candidate generation.	DEFAULT_BLOCKING

Attributes:

Name	Type	Description
eval_history	list[EvaluationReport]	One EvaluationReport per epoch, appended in train() when eval_data is given.
loss_history	list[float]	Mean training loss per epoch, always appended in train() -- the most direct signal that the QA span-position objective is actually decreasing, independent of eval_data (which, on a very small eval set, can already sit at a ceiling precision@1 with no room to visibly "improve").

Source code in src/linkingtk/train/esc_trainer.py

def __init__(
    self,
    model: EscEncoder,
    args: TrainingArguments,
    train_data: tuple[list[Entity], list[tuple[str, str]]],
    senses: list[Entity] | EntitySource,
    eval_data: tuple[list[Entity], list[tuple[str, str]]] | None = None,
    blocking: BlockingStrategy = DEFAULT_BLOCKING,
) -> None:
    self.model = model
    self.args = args
    self.train_data = train_data
    self.senses = senses
    self.eval_data = eval_data
    self.blocking = blocking
    self.eval_history: list[EvaluationReport] = []
    self.loss_history: list[float] = []

train

train() -> None

Run the training loop, updating self.model in place.

Source code in src/linkingtk/train/esc_trainer.py

def train(self) -> None:
    """Run the training loop, updating ``self.model`` in place."""
    args = self.args
    device = resolve_device(args.device)
    self.model.to(device)

    mentions, ground_truth = self.train_data
    examples = self._build_examples(mentions, ground_truth)
    num_batches_per_epoch = -(-len(examples) // args.batch_size)
    optimizer, scheduler = build_optimizer(
        self.model, args, num_batches_per_epoch * args.num_epochs
    )

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    logger.info(
        "EscTrainer: %d training examples (%d mentions given), %d batches/epoch, %d epochs",
        len(examples),
        len(mentions),
        num_batches_per_epoch,
        args.num_epochs,
    )

    for epoch in range(args.num_epochs):
        shuffled = list(examples)
        random.shuffle(shuffled)
        loss_total = 0.0
        num_batches = 0
        for start in range(0, len(shuffled), args.batch_size):
            batch = shuffled[start : start + args.batch_size]
            loss = self._train_step(batch, device)
            optimizer.zero_grad()
            loss.backward()  # type: ignore[no-untyped-call]
            optimizer.step()
            if scheduler is not None:
                scheduler.step()
            loss_total += float(loss.item())
            num_batches += 1

        mean_loss = loss_total / max(num_batches, 1)
        self.loss_history.append(mean_loss)

        if self.eval_data is not None:
            self.eval_history.append(self._evaluate())
            logger.info(
                "EscTrainer: epoch %d/%d done, mean loss %.4f, eval %s",
                epoch + 1,
                args.num_epochs,
                mean_loss,
                self.eval_history[-1].metrics,
            )
        else:
            logger.info(
                "EscTrainer: epoch %d/%d done, mean loss %.4f",
                epoch + 1,
                args.num_epochs,
                mean_loss,
            )

        # Saved after *every* epoch, not just at the end -- a long run
        # can span hours unattended; without this, an interruption
        # loses everything (same rationale as EwiserTrainer/CrossEncoderTrainer).
        torch.save(self.model.state_dict(), output_dir / "model.pt")
        logger.info("EscTrainer: checkpoint saved to %s", output_dir / "model.pt")