Skip to content

GlossBERT

GlossBertLinker ports Huang et al.'s GlossBERT (Sent-CLS-WS) cross-encoder for WSD: each mention/gloss pair is scored by a single BertForSequenceClassification forward pass. Four examples cover it end to end: loading the paper's own checkpoint, training the encoder from scratch, a full-corpus run, and LLM re-ranking on top of it.

Reproducing the paper's published results

Unlike every other example on this site, this one runs no training at all -- see Verifying the training path for that side instead. GlossBertLinker loads Huang et al.'s own published checkpoint for their best-reported configuration, GlossBERT(Sent-CLS-WS), directly into GlossBertEncoder -- its BertForSequenceClassification architecture is verbatim what the checkpoint contains -- and evaluates it via GlossBertLinker.link, the exact same production code path a freshly-trained model would use. See glossbert.py's module docstring for how closely the candidate generation, gloss formatting and weak-supervision marking were ported from the reference implementation's own data-prep code, not just the paper text.

Requires the checkpoint (Sent_CLS_WS.zip, linked from the reference repo's README) at ~/Downloads/Sent_CLS_WS.zip, and UFSAC 2.1 extracted to ~/data/ufsac-public-2.1/ -- see the script's own docstring for both. Candidates are every WordNet sense of each mention's lemma (ExactMatch(top_k=50) against a WnEntitySource), matching the paper's own candidate generation exactly (no gold-POS filtering, no popularity prior).

"""Reproduces GlossBERT(Sent-CLS-WS)'s published results using the paper's
own checkpoint, through this repo's `GlossBertLinker`/`UfsacDataset` pipeline.

No training happens here -- this loads Huang et al.'s own published
checkpoint (https://arxiv.org/pdf/1908.07245.pdf, best-reported
configuration) directly into `GlossBertEncoder` (its `BertForSequenceClassification`
architecture is verbatim what the checkpoint contains) and evaluates it
via `GlossBertLinker.link()`, the exact same production code path a
freshly-trained model would use.

Setup:

1. Download the checkpoint from the link in
   https://github.com/HSLCY/GlossBERT#news (`Sent_CLS_WS.zip`) and place
   it at `~/Downloads/Sent_CLS_WS.zip` (or point `_CHECKPOINT_ZIP` below
   at wherever you put it). It's a plain
   `config.json`/`pytorch_model.bin`/`vocab.txt` triple, extracted once
   into `~/.cache/linkingtk/glossbert_checkpoint/` on first run.
2. Download UFSAC 2.1 (https://github.com/getalp/UFSAC#content-of-the-repository)
   and extract it to `~/data/ufsac-public-2.1/` (or point `_UFSAC_DIR`
   below elsewhere) -- this script reads its `raganato_*.xml` files, the
   same Raganato et al. (2017) WSD evaluation framework data GlossBERT's
   own paper reports numbers against.

Candidates are every WordNet sense of each mention's lemma
(`ExactMatch(top_k=50)` against a `WnEntitySource` -- wide enough that no
real word's sense count gets truncated), matching the paper's own
candidate generation (no gold-POS filtering, no popularity prior).

Published target (the paper's README, "this checkpoint" row -- the
checkpoint this script actually loads, not the original paper run):
SE07=72.1, SE2=77.7, SE3=75.9, SE13=76.8, SE15=79.3, ALL=77.2 (micro F1,
which reduces to precision@1 in this always-answered, one-prediction-
per-instance setting).

Run with: `uv run python examples/glossbert_reproduction.py`
"""

from __future__ import annotations

import zipfile
from pathlib import Path

import torch

from linkingtk.algorithms.wsd.glossbert import GlossBertEncoder, GlossBertLinker
from linkingtk.blocking.exact import ExactMatch
from linkingtk.datasets.ufsac import UfsacDataset
from linkingtk.eval import Evaluator

_CHECKPOINT_ZIP = Path.home() / "Downloads" / "Sent_CLS_WS.zip"
_CHECKPOINT_DIR = Path.home() / ".cache" / "linkingtk" / "glossbert_checkpoint"
_UFSAC_DIR = Path.home() / "data" / "ufsac-public-2.1"

# (UFSAC filename stem, README column label, README's "this checkpoint" score)
_EVAL_SETS = [
    ("raganato_semeval2007", "SE07", 72.1),
    ("raganato_senseval2", "SE2", 77.7),
    ("raganato_senseval3", "SE3", 75.9),
    ("raganato_semeval2013", "SE13", 76.8),
    ("raganato_semeval2015", "SE15", 79.3),
    ("raganato_ALL", "ALL", 77.2),
]


def _ensure_checkpoint_extracted() -> Path:
    if not (_CHECKPOINT_DIR / "pytorch_model.bin").exists():
        _CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
        with zipfile.ZipFile(_CHECKPOINT_ZIP) as archive:
            archive.extractall(_CHECKPOINT_DIR)
    return _CHECKPOINT_DIR


def main() -> None:
    checkpoint_dir = _ensure_checkpoint_extracted()
    device = "cuda" if torch.cuda.is_available() else "cpu"

    encoder = GlossBertEncoder.from_checkpoint(checkpoint_dir, max_length=512)
    linker = GlossBertLinker(model=encoder)
    linker.model.to(device)
    blocking = ExactMatch(top_k=50)

    print(f"{'dataset':<8} {'precision@1':>12} {'published':>10}")
    for stem, label, published in _EVAL_SETS:
        mentions, senses, ground_truth = UfsacDataset(source=str(_UFSAC_DIR / f"{stem}.xml")).load()

        results = linker.link(mentions, senses, blocking=blocking)
        predictions = [(result.source_id, result.target_id) for result in results]
        report = Evaluator.evaluate(predictions=predictions, ground_truth=ground_truth)
        score = report.metrics["precision@1"] * 100

        print(f"{label:<8} {score:>11.1f}% {published:>9.1f}%")


if __name__ == "__main__":
    main()

Run with:

uv run python examples/glossbert_reproduction.py

dataset   precision@1  published
SE07            72.5%      72.1%
SE2             77.5%      77.7%
SE3             75.7%      75.9%
SE13            76.7%      76.8%
SE15            78.7%      79.3%
ALL             76.7%      77.2%

Every dataset lands within 0.6 points of its published number. Getting here took four rounds of real diagnosis, each one a concrete, verified bug -- not a single number was accepted at face value.

Bugs found and fixed

1 -- mentions were labeled by surface form, not lemma. UfsacDataset originally set each mention's labels to its literal surface text (e.g. "caught"). Candidate generation (ExactMatch querying WnEntitySource by label) needs the dictionary lemma ("catch") -- WordNet doesn't index inflected forms. The first real run measured precision@1 = 52.6% on SE07 with recall well below precision (35.4%), the tell that a third of mentions were producing no candidates at all. Fixed by using each word's lemma attribute for labels (both UfsacDataset and SemCorDataset had this bug), keeping the literal surface form available from context's span for _gloss_text's formatting, which needs it instead (GlossBERT's gloss side is "<surface form> : <gloss>", not the lemma).

2 -- case-sensitive post-filtering dropped wn's own case-insensitive matches. ExactMatch's EntitySource branch verifies each of search()'s results actually carries the queried label, to guard against a loosely-matching search implementation. But wn.synsets("friday") is itself case-insensitive -- it finds the synset WordNet lemmatizes as "Friday" -- so the case-sensitive post-filter silently dropped it, and every other capitalized/proper-noun-like lemma. Made the post-filter case-insensitive.

3 -- single-answer scoring on multi-answer gold data. Some WSD gold-standard instances genuinely accept more than one correct sense key. SemEval-2015's multi-answer rate (18%, 184/1022) is 20x SemEval-2007's (0.9%) -- and it's exactly the dataset with the disproportionately large gap after fixes 1-2 (71.3% vs. published 79.3%). Fixed by having UfsacDataset emit one ground_truth row per valid sense key and Evaluator.evaluate accept a match against any of a source's rows (recall's denominator now the distinct-source count, not the row count -- unchanged for every single-answer caller elsewhere in this repo).

4 -- WordNet sense-key convention mismatches between UFSAC and omw-en:1.4. Found by cross-checking against GlossBERT's own official scorer (see below), which revealed our own Evaluator.evaluate numbers were precision-correct but the official recall was still short -- i.e. we were dropping real mentions from dataset1 entirely, not just scoring them wrong. Two distinct lemma-lookup mismatches, both in sources/wn.py:

  • Phrasal verbs: UFSAC's classic-WordNet convention joins multi-word lemmas with underscores ("point_out"), but omw-en:1.4/omw-en:2.0 index the same lemma space-separated ("point out") -- confirmed directly (wn.synsets("point_out", ...) returns nothing, wn.synsets("point out", ...) returns 3). ~4-5% of SemEval-2007's instances.
  • Adjective satellites: UFSAC's wn30_key layer tags "peculiar" (meaning "specific") as type 3 (plain adjective) -- "peculiar%3:00:00:specific:00" -- but omw-en:1.4 stores the identical sense as type 5 (adjective satellite) instead. ~7.5% of SensEval-2's instances.

Both fixed with a query-side fallback (retry with underscores replaced by spaces; retry with the adjective/adjective-satellite type digit swapped) in WnEntitySource.search, sensekey_to_synset_id and synset_id_to_sensekey, plus the matching normalization in ExactMatch's post-filter.

Cross-checked against GlossBERT's own official scorer

1.7 points still felt worth chasing further after fix 3, so rather than trust our own Evaluator.evaluate a fourth time, this was cross-verified against the reference implementation's own evaluation tool: Raganato et al.'s Scorer.java, bundled in the GlossBERT repo. Compiled it with a locally-extracted JDK (no root needed -- apt-get download+dpkg-deb -x), converted our predictions back to sense-key form via the new synset_id_to_sensekey helper, and ran it against the actual official gold-key files (not just this repo's own UfsacDataset-loaded ground truth).

Before fix 4, the official scorer showed precision matching our own numbers almost exactly (e.g. SE07 72.9% both ways) but recall noticeably lower (SE07 69.7%) -- independent confirmation that our Evaluator.evaluate math was already correct, and that the gap was missing coverage (mentions silently dropped from dataset1), not a scoring bug. After fix 4, every dataset scores P = R = F1 against the official scorer, matching this repo's own Evaluator.evaluate numbers to the decimal -- full agreement between two independent implementations, and 100% of every gold file's instances answered.

Verifying the training path

Reproducing the paper's published results validates GlossBertEncoder's inference path -- architecture, candidate generation, text formatting -- by loading the paper's own published checkpoint. It never touches CrossEncoderTrainer at all. This example is the counterpart: it trains a fresh bert-base-uncased GlossBertEncoder from scratch via CrossEncoderTrainer on a real slice of SemCor, and checks that real learning happens.

Small subset, not the paper's setup -- the paper trains on all ~226K sense-tagged SemCor instances for 6 epochs; this trains on the first 6 Brown Corpus documents (~3K instances) for 3, so it runs in well under a minute. The point isn't to approach the published numbers (that's what the checkpoint reproduction is for) -- it's to verify the training code itself end to end: CrossEncoderTrainer's BCE loss, hard-negative mining via blocking, optimizer/warmup/weight-decay setup, all against a real pretrained backbone and real data, not a toy synthetic model.

"""Trains GlossBertEncoder from scratch (real `bert-base-uncased`, real
gradient descent) on a real slice of SemCor, and evaluates on held-out
documents -- the training-path counterpart to
`examples/glossbert_reproduction.py`, which only exercises inference
against the paper's published checkpoint.

**Small subset, not the paper's setup.** The paper trains on all of
SemCor (~226K sense-tagged instances) for 6 epochs at `batch_size=64`,
`max_seq_length=512` -- multiple GPU-hours even on this sandbox's single
RTX 4090. This script instead trains on the first `_TRAIN_DOCS` Brown
Corpus documents (~3K instances) for 3 epochs at `max_length=128`, runs in
under a minute, and exists to verify the *training path itself* --
`CrossEncoderTrainer`'s loss, hard-negative mining, and optimizer setup
end to end on real data and a real pretrained backbone -- not to
reproduce the paper's published numbers (that's what the checkpoint-based
`glossbert_reproduction.py` is for). Don't expect these numbers to
approach the published ones; a fraction of the data and epochs will not
reach that quality, and that's not the point here.

**Regression context**: this script's own `eval_dataset2=senses` wiring
matters -- `CrossEncoderTrainer._evaluate()` defaults to deriving its
candidate pool from `eval_data`'s own (already-correct) targets alone,
which for WSD silently omits every real decoy sense and overstates
Hits@1. Measured directly during this benchmark's development: 60%
self-reported Hits@1 vs. 33% via the real `GlossBertLinker.link()` path
on the same held-out mentions -- see `CrossEncoderTrainer`'s own
docstring for the fix.

Requires UFSAC 2.1 extracted to `~/data/ufsac-public-2.1/` -- see
`glossbert_reproduction.py`'s docstring.

Run with: `uv run python examples/glossbert_benchmark.py`
"""

from __future__ import annotations

from pathlib import Path

import torch

from linkingtk.algorithms.wsd.glossbert import GlossBertLinker
from linkingtk.blocking.exact import ExactMatch
from linkingtk.core.entity import Entity, label_texts
from linkingtk.datasets.ufsac import UfsacDataset
from linkingtk.eval import Evaluator
from linkingtk.sources.wn import WnEntitySource
from linkingtk.train.arguments import TrainingArguments
from linkingtk.train.cross_encoder import CrossEncoderTrainer

_UFSAC_DIR = Path.home() / "data" / "ufsac-public-2.1"
_TRAIN_DOCS = 6
_EVAL_DOCS = 3


def _document_id(mention_id: str) -> str:
    # UfsacDataset's positional mention ids (semcor.xml has no native
    # per-word id, unlike the raganato_*.xml files) -- "ufsac:{doc}:{sent}:{index}".
    return mention_id.split(":")[1]


def _resolve_pairs(
    ground_truth: list[tuple[str, str]],
    mentions_by_id: dict[str, Entity],
    senses: WnEntitySource,
) -> list[tuple[Entity, Entity]]:
    cache: dict[str, Entity | None] = {}
    pairs = []
    for mention_id, synset_id in ground_truth:
        if synset_id not in cache:
            cache[synset_id] = senses.get(synset_id)
        sense = cache[synset_id]
        if sense is not None:
            pairs.append((mentions_by_id[mention_id], sense))
    return pairs


def main() -> None:
    mentions, senses, ground_truth = UfsacDataset(source=str(_UFSAC_DIR / "semcor.xml")).load()
    mentions_by_id = {m.id: m for m in mentions}

    doc_ids = sorted({_document_id(m.id) for m in mentions})
    train_docs = set(doc_ids[:_TRAIN_DOCS])
    eval_docs = set(doc_ids[_TRAIN_DOCS : _TRAIN_DOCS + _EVAL_DOCS])
    train_gt = [(mid, tid) for mid, tid in ground_truth if _document_id(mid) in train_docs]
    eval_gt = [(mid, tid) for mid, tid in ground_truth if _document_id(mid) in eval_docs]
    train_pairs = _resolve_pairs(train_gt, mentions_by_id, senses)
    eval_pairs = _resolve_pairs(eval_gt, mentions_by_id, senses)
    eval_mentions = list({m.id: m for m, _s in eval_pairs}.values())
    print(
        f"{len(train_pairs)} train instances ({_TRAIN_DOCS} docs) / "
        f"{len(eval_pairs)} eval instances ({_EVAL_DOCS} docs)"
    )

    blocking = ExactMatch(top_k=30)

    # Most-frequent-sense baseline: WordNet's own sense order (most-common
    # first) with no model at all -- the standard WSD sanity-check floor.
    mfs_predictions = [
        (mention.id, candidates[0].id)
        for mention in eval_mentions
        for candidates in [senses.search(label_texts(mention)[0], top_k=1)]
        if candidates
    ]
    mfs_report = Evaluator.evaluate(predictions=mfs_predictions, ground_truth=eval_gt)
    print(f"Most-frequent-sense baseline: {mfs_report.metrics}")

    linker = GlossBertLinker(model_name_or_path="bert-base-uncased", max_length=128)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    linker.model.to(device)

    before = Evaluator.evaluate(
        predictions=[
            (r.source_id, r.target_id)
            for r in linker.link(eval_mentions, senses, blocking=blocking)
        ],
        ground_truth=eval_gt,
    )
    print(f"Untrained precision@1: {before.metrics}")

    args = TrainingArguments(
        output_dir="./models/glossbert_semcor_subset",
        learning_rate=2e-5,
        negative_samples_ratio=3,
        num_epochs=3,
        batch_size=16,
        weight_decay=0.01,
        warmup_ratio=0.1,
        device=device,
    )
    trainer = CrossEncoderTrainer(
        model=linker.model,
        args=args,
        train_data=train_pairs,
        eval_data=eval_pairs,
        eval_dataset2=senses,
        blocking=blocking,
    )
    trainer.train()
    print("Per-epoch held-out Hits@1 (via CrossEncoderTrainer.eval_history):")
    for i, epoch_report in enumerate(trainer.eval_history):
        print(f"  epoch {i + 1}: {epoch_report.metrics}")

    after = Evaluator.evaluate(
        predictions=[
            (r.source_id, r.target_id)
            for r in linker.link(eval_mentions, senses, blocking=blocking)
        ],
        ground_truth=eval_gt,
    )
    print("Trained precision@1 (via GlossBertLinker.link, the real production path):")
    print(f"  {after.metrics}")


if __name__ == "__main__":
    main()

Run with:

uv run python examples/glossbert_benchmark.py

2964 train instances (6 docs) / 669 eval instances (3 docs)
Most-frequent-sense baseline: {'precision@1': 0.205, 'recall': 0.205, 'f1': 0.205}
Untrained precision@1: {'precision@1': 0.138, 'recall': 0.138, 'f1': 0.138}
Per-epoch held-out Hits@1 (via CrossEncoderTrainer.eval_history):
  epoch 1: {'Hits@1': 0.238, 'MRR': 0.423}
  epoch 2: {'Hits@1': 0.324, 'MRR': 0.490}
  epoch 3: {'Hits@1': 0.333, 'MRR': 0.496}
Trained precision@1 (via GlossBertLinker.link, the real production path):
  {'precision@1': 0.333, 'recall': 0.333, 'f1': 0.333}

Real, meaningful learning on a real pretrained backbone: 13.8% (untrained) → 33.3% (trained, 3 epochs on ~3K instances), clearing the most-frequent-sense baseline (20.5%) by a wide margin.

A real bug this caught

The first version of this benchmark reported a 60% Hits@1 from CrossEncoderTrainer.eval_history -- but the same trained model, evaluated independently via GlossBertLinker.link() on the identical held-out mentions, measured only 33%. Nearly 2x apart, on numbers that should agree almost exactly.

The cause: CrossEncoderTrainer._evaluate() (mirroring Trainer._evaluate()'s convention for EA/EL) derived its candidate pool from eval_data's own targets -- correct for EA/EL, where the eval split's own KB subset already stands in for the full target set, but silently wrong for WSD: the eval split's own gold senses are a far narrower, far less confusable candidate pool than a mention's real lemma-wide sense inventory. With only the correct answers ever offered as candidates, the reported number was measuring something closer to "can the model recognize its own training targets" than "can it disambiguate."

Fixed by adding CrossEncoderTrainer(..., eval_dataset2=...): pass the real target set (here, the same WnEntitySource link() itself queries) and _evaluate() blocks against that instead. After the fix, eval_history's last epoch (33.3%) matches the independent link() check exactly -- see CrossEncoderTrainer's own docstring, and tests/train/test_cross_encoder.py::TestEvalDataset2 for the regression coverage.

Full-corpus training run

Verifying the training path verifies the training path on a small, fast slice of SemCor. This script is the full-scale counterpart: it trains GlossBertEncoder on the entire real SemCor corpus (~230K sense-tagged instances) with near-paper-faithful hyperparameters, then evaluates against UFSAC's real Raganato et al. framework test sets -- directly comparable to Reproducing the paper's published results's checkpoint-based numbers, but from a model this repo's own CrossEncoderTrainer trained, not the paper's published weights.

Long-running -- meant to run for hours, unattended (nohup ... & disown, not a foreground uv run). See the script's own docstring for measured throughput and expected wall-clock time on a single GPU.

"""Trains GlossBertEncoder on the *full* real SemCor corpus (~230K
sense-tagged instances) with near-paper-faithful hyperparameters, then
evaluates the result against UFSAC's real Raganato et al. framework test
sets -- the full-scale counterpart to `examples/glossbert_benchmark.py`'s
small, fast verification slice.

**Long-running.** Unlike every other example in this repo, this is meant
to run for hours, unattended (`nohup ... & disown`, not a foreground
`uv run`). Measured directly on this sandbox's single RTX 4090 before
launching: ~545 rows/sec (positives + mined negatives combined) at
`batch_size=64`, `max_length=512`, so ~230K positives (expanding to
~700K-1M rows with negatives) is roughly 20-30 minutes per epoch, low
single-digit GB of peak GPU memory (WSD sentences are short despite the
paper's `max_length=512`) -- comfortably an overnight run at the paper's
own 6 epochs, with margin.

**Hyperparameters, matching the paper (`README.md`/`run_classifier_WSD_sent.py`
in the reference implementation) where GPU-safe**: `learning_rate=2e-5`,
`weight_decay=0.01`, `warmup_ratio=0.1`, `max_length=512`,
`num_epochs=6`. `batch_size=64` matches the paper's own effective batch
(no gradient accumulation needed here -- measured peak memory leaves
comfortable headroom on a 24GB GPU at this scale). Candidates are every
WordNet sense of a lemma across all POS
(`ExactMatch(top_k=50)`/`negative_samples_ratio=49`), matching the
paper's own candidate generation, not an artificially-capped subset.

**Dev/test split matches the paper's own protocol**: train on all of
SemCor, monitor on SemEval-2007 (the paper's own dev set) via
`CrossEncoderTrainer.eval_history` every epoch, and only evaluate on the
real held-out test sets (SensEval-2/3, SemEval-2013/2015, ALL) once,
after training completes -- exactly like `glossbert_reproduction.py`'s
checkpoint-based evaluation, so the two are directly comparable.

Requires UFSAC 2.1 extracted to `~/data/ufsac-public-2.1/` -- see
`glossbert_reproduction.py`'s docstring.

Run with (from the repo root), detached so it survives the shell exiting:
nohup uv run python examples/glossbert_full_training.py \ > /tmp/glossbert_full_training.log 2>&1 & disown 
"""

from __future__ import annotations

import logging
import time
from pathlib import Path

import torch

from linkingtk.algorithms.wsd.glossbert import GlossBertLinker
from linkingtk.blocking.exact import ExactMatch
from linkingtk.core.entity import Entity
from linkingtk.datasets.ufsac import UfsacDataset
from linkingtk.eval import Evaluator
from linkingtk.sources.wn import WnEntitySource
from linkingtk.train.arguments import TrainingArguments
from linkingtk.train.cross_encoder import CrossEncoderTrainer

_UFSAC_DIR = Path.home() / "data" / "ufsac-public-2.1"
_OUTPUT_DIR = Path("./models/glossbert_semcor_full")

# (UFSAC filename stem, README column label, README's published "this checkpoint" score)
_TEST_SETS = [
    ("raganato_senseval2", "SE2", 77.7),
    ("raganato_senseval3", "SE3", 75.9),
    ("raganato_semeval2013", "SE13", 76.8),
    ("raganato_semeval2015", "SE15", 79.3),
    ("raganato_ALL", "ALL", 77.2),
]


def _resolve_pairs(
    ground_truth: list[tuple[str, str]],
    mentions_by_id: dict[str, Entity],
    senses: WnEntitySource,
) -> list[tuple[Entity, Entity]]:
    cache: dict[str, Entity | None] = {}
    pairs = []
    for mention_id, synset_id in ground_truth:
        if synset_id not in cache:
            cache[synset_id] = senses.get(synset_id)
        sense = cache[synset_id]
        if sense is not None:
            pairs.append((mentions_by_id[mention_id], sense))
    return pairs


def main() -> None:
    logging.basicConfig(
        level=logging.WARNING, format="%(asctime)s %(levelname)s %(message)s", force=True
    )
    logging.getLogger("linkingtk").setLevel(logging.INFO)
    t0 = time.time()
    train_mentions, senses, train_gt = UfsacDataset(source=str(_UFSAC_DIR / "semcor.xml")).load()
    train_mentions_by_id = {m.id: m for m in train_mentions}
    train_pairs = _resolve_pairs(train_gt, train_mentions_by_id, senses)
    print(f"[{time.time() - t0:.0f}s] loaded {len(train_pairs)} SemCor training instances")

    dev_mentions, _senses, dev_gt = UfsacDataset(
        source=str(_UFSAC_DIR / "raganato_semeval2007.xml")
    ).load()
    dev_mentions_by_id = {m.id: m for m in dev_mentions}
    dev_pairs = _resolve_pairs(dev_gt, dev_mentions_by_id, senses)
    print(f"[{time.time() - t0:.0f}s] loaded {len(dev_pairs)} SemEval-2007 dev instances")

    device = "cuda" if torch.cuda.is_available() else "cpu"
    linker = GlossBertLinker(model_name_or_path="bert-base-uncased", max_length=512)
    linker.model.to(device)

    blocking = ExactMatch(top_k=50)
    args = TrainingArguments(
        output_dir=str(_OUTPUT_DIR),
        learning_rate=2e-5,
        negative_samples_ratio=49,
        num_epochs=6,
        batch_size=64,
        weight_decay=0.01,
        warmup_ratio=0.1,
        device=device,
    )
    trainer = CrossEncoderTrainer(
        model=linker.model,
        args=args,
        train_data=train_pairs,
        eval_data=dev_pairs,
        eval_dataset2=senses,
        blocking=blocking,
    )

    print(f"[{time.time() - t0:.0f}s] starting training ({args.num_epochs} epochs)")
    trainer.train()
    for i, epoch_report in enumerate(trainer.eval_history):
        print(f"[{time.time() - t0:.0f}s] epoch {i + 1} dev (SE07): {epoch_report.metrics}")
    print(f"[{time.time() - t0:.0f}s] training complete, checkpoint at {_OUTPUT_DIR}/model.pt")

    print(f"[{time.time() - t0:.0f}s] evaluating on real held-out test sets")
    print(f"{'dataset':<8} {'precision@1':>12} {'published':>10}")
    for stem, label, published in _TEST_SETS:
        mentions, _senses, ground_truth = UfsacDataset(
            source=str(_UFSAC_DIR / f"{stem}.xml")
        ).load()
        results = linker.link(mentions, senses, blocking=blocking)
        predictions = [(r.source_id, r.target_id) for r in results]
        report = Evaluator.evaluate(predictions=predictions, ground_truth=ground_truth)
        score = (report.metrics["precision@1"] or 0.0) * 100
        print(f"{label:<8} {score:>11.1f}% {published:>9.1f}%")
    print(f"[{time.time() - t0:.0f}s] done")


if __name__ == "__main__":
    main()

Run with (from the repo root), detached so it survives the shell exiting:

nohup uv run python examples/glossbert_full_training.py \
    > /tmp/glossbert_full_training.log 2>&1 & disown

LLM re-ranking over GlossBERT candidates

LlmRerankerLinker re-ranks GlossBertLinker's own top-k candidate senses with a real local LLM. Filed as #23 -- unlike ChatEA (#22), this isn't a port of a specific paper; it generalizes the same "cheap retrieval, then let an LLM re-rank only the top-k" idea to Word Sense Disambiguation, reusing this repo's already-benchmarked GlossBertLinker as the retrieval stage.

Acceptance bar is relative, not an external paper's number

There's nothing published to chase here -- #23 doesn't port a paper. The real question: does re-ranking GlossBERT's own top-k candidate senses with an LLM improve precision@1 over GlossBERT's own already-benchmarked checkpoint-reproduction baseline (ALL=76.7%, see Reproducing the paper's published results/#39)?

Setup is identical to glossbert_reproduction.py -- see Reproducing the paper's published results for the checkpoint/UFSAC download steps. Also requires a local Ollama server with the target model pulled -- ollama pull llama2:13b by default, pass --model for a different one.

"""Benchmarks LlmRerankerLinker (#23) layered on top of GlossBertLinker's own
cross-encoder scoring, using the paper's published checkpoint -- same setup
as `examples/glossbert_reproduction.py`.

**Acceptance bar is relative, not an external paper's number.** Unlike #22
(ChatEA), #23 doesn't port a specific paper -- there's nothing published to
chase. The question this script answers: does re-ranking GlossBERT's own
top-k candidate senses with an LLM improve precision@1 over GlossBERT's
own already-benchmarked checkpoint-reproduction baseline (ALL=76.7%, per
`glossbert_reproduction.py`/#39)?

First reproduces the full published table exactly like
`glossbert_reproduction.py` (no LLM calls, no cost). Then, on the "ALL"
split, samples a random subset of mentions whose true sense already falls
within GlossBertLinker's own top-k candidates (`--sample`, default 30 --
matching `examples/chatea_ea.py`'s established precedent for capping
real-LLM-call benchmarks) and re-ranks just that window with
`LlmRerankerLinker`, reporting before/after precision@1.

Setup: identical to `glossbert_reproduction.py` -- see that script's
docstring for the checkpoint/UFSAC download steps. Also requires a local
Ollama server with `ollama pull llama2:13b` (or pass `--model`).

Run with: `uv run python examples/glossbert_llm_reranker_benchmark.py`
"""

from __future__ import annotations

import argparse
import random
import zipfile
from pathlib import Path

import torch

from linkingtk.algorithms.llm_reranker import LlmRerankerLinker
from linkingtk.algorithms.wsd.glossbert import GlossBertEncoder, GlossBertLinker
from linkingtk.blocking.exact import ExactMatch
from linkingtk.datasets.ufsac import UfsacDataset
from linkingtk.eval import Evaluator
from linkingtk.llm.client import create_client

_CHECKPOINT_ZIP = Path.home() / "Downloads" / "Sent_CLS_WS.zip"
_CHECKPOINT_DIR = Path.home() / ".cache" / "linkingtk" / "glossbert_checkpoint"
_UFSAC_DIR = Path.home() / "data" / "ufsac-public-2.1"

# (UFSAC filename stem, README column label, README's "this checkpoint" score)
_EVAL_SETS = [
    ("raganato_semeval2007", "SE07", 72.1),
    ("raganato_senseval2", "SE2", 77.7),
    ("raganato_senseval3", "SE3", 75.9),
    ("raganato_semeval2013", "SE13", 76.8),
    ("raganato_semeval2015", "SE15", 79.3),
    ("raganato_ALL", "ALL", 77.2),
]


def _ensure_checkpoint_extracted() -> Path:
    if not (_CHECKPOINT_DIR / "pytorch_model.bin").exists():
        _CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
        with zipfile.ZipFile(_CHECKPOINT_ZIP) as archive:
            archive.extractall(_CHECKPOINT_DIR)
    return _CHECKPOINT_DIR


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default="ollama/llama2:13b")
    parser.add_argument("--sample", type=int, default=30)
    parser.add_argument("--top-k", type=int, default=10)
    parser.add_argument("--random-seed", type=int, default=20260827)
    args = parser.parse_args()

    checkpoint_dir = _ensure_checkpoint_extracted()
    device = "cuda" if torch.cuda.is_available() else "cpu"

    encoder = GlossBertEncoder.from_checkpoint(checkpoint_dir, max_length=512)
    linker = GlossBertLinker(model=encoder)
    linker.model.to(device)
    blocking = ExactMatch(top_k=50)

    print(f"{'dataset':<8} {'precision@1':>12} {'published':>10}")
    all_mentions = all_senses = all_ground_truth = None
    for stem, label, published in _EVAL_SETS:
        mentions, senses, ground_truth = UfsacDataset(source=str(_UFSAC_DIR / f"{stem}.xml")).load()

        results = linker.link(mentions, senses, blocking=blocking)
        predictions = [(result.source_id, result.target_id) for result in results]
        report = Evaluator.evaluate(predictions=predictions, ground_truth=ground_truth)
        score = report.metrics["precision@1"] * 100

        print(f"{label:<8} {score:>11.1f}% {published:>9.1f}%")
        if label == "ALL":
            all_mentions, all_senses, all_ground_truth = mentions, senses, ground_truth

    assert all_mentions is not None and all_senses is not None and all_ground_truth is not None
    mentions_by_id = {entity.id: entity for entity in all_mentions}
    true_target_by_source = dict(all_ground_truth)

    base_results = linker.link(all_mentions, all_senses, blocking=blocking)
    base_prediction_by_source = {result.source_id: result.target_id for result in base_results}
    base_topk_by_source = {
        result.source_id: [result.target_id, *result.alternatives][: args.top_k]
        for result in base_results
    }

    improvable_ids = [
        source_id
        for source_id, top_k_ids in base_topk_by_source.items()
        if source_id in true_target_by_source and true_target_by_source[source_id] in top_k_ids
    ]
    print(
        f"Improvable (true sense within top-{args.top_k}): "
        f"{len(improvable_ids)}/{len(base_topk_by_source)}"
    )

    rng = random.Random(args.random_seed)
    sampled_ids = set(rng.sample(improvable_ids, min(args.sample, len(improvable_ids))))
    print(f"Sampled for real LLM re-ranking: {len(sampled_ids)}")

    client = create_client(args.model)
    reranker = LlmRerankerLinker(base_linker=linker, client=client, task="wsd", top_k=args.top_k)
    sampled_mentions = [mentions_by_id[source_id] for source_id in sampled_ids]
    llm_results = reranker.link(sampled_mentions, all_senses, blocking=blocking)

    merged_prediction_by_source = dict(base_prediction_by_source)
    changed_correctness = 0
    for result in llm_results:
        source_id = result.source_id
        true_target = true_target_by_source[source_id]
        was_correct = base_prediction_by_source.get(source_id) == true_target
        merged_prediction_by_source[source_id] = result.target_id
        now_correct = result.target_id == true_target
        if was_correct != now_correct:
            changed_correctness += 1
        print(f"  {source_id}: base_correct={was_correct} -> llm_correct={now_correct}")
    print(
        f"Sampled entities whose correctness changed after LLM "
        f"re-ranking: {changed_correctness}/{len(llm_results)}"
    )

    final_report = Evaluator.evaluate(
        predictions=list(merged_prediction_by_source.items()), ground_truth=all_ground_truth
    )
    final_score = final_report.metrics["precision@1"] * 100
    print(f"After LlmRerankerLinker re-ranking (ALL): {final_score:.1f}%")


if __name__ == "__main__":
    main()

Run with:

uv run python examples/glossbert_llm_reranker_benchmark.py

dataset   precision@1  published
SE07            72.5%      72.1%
SE2             77.5%      77.7%
SE3             75.7%      75.9%
SE13            76.7%      76.8%
SE15            78.7%      79.3%
ALL             76.7%      77.2%
Improvable (true sense within top-10): 7223/7253
Sampled for real LLM re-ranking: 30
Sampled entities whose correctness changed after LLM re-ranking: 1/30
After LlmRerankerLinker re-ranking (ALL): 76.7%

A real bug was found and fixed along the way, not just a null result accepted at face value. The first run of this benchmark showed 0/30 changed -- but per-entity diagnosis (not accepting a flat number at face value, see the established practice from #22/#39) revealed why: LlmRerankerLinker's confidence shortcut defaulted to threshold=0.5, tuned with a bi-encoder's roughly-[-1, 1] cosine similarity in mind. GlossBertLinker's score is a raw, unbounded cross-encoder logit margin instead -- routinely 1-15+ in magnitude on real UFSAC instances -- so a 0.5 gap was almost always exceeded, and 28 of the 30 sampled entities never reached the LLM at all, silently falling back to GlossBERT's own base ranking every time.

This was a genuine design bug, not a WSD-specific quirk: an absolute score-gap threshold is only meaningful when the caller knows the base linker's score scale, which varies by method (bounded similarity for bi-encoders, unbounded logits for cross-encoders). Fixed by making threshold default to None (shortcut disabled) rather than a value tuned for one family of base linkers -- see LlmRerankerLinker's own docstring, and the BLINK reranker benchmark, which opts back into threshold=0.5 explicitly since BLINK's cosine-similarity score does make that a sound choice.

With the LLM actually consulted for all 30 sampled entities: 22/30 were already correct beforehand, 8/30 were wrong, and the LLM correctly fixed 1 of those 8 with zero regressions -- a small but real, honestly-measured positive signal from a 30-instance sample (0.4% of the 7,253-instance ALL split), not the "no effect" the threshold bug would have reported.

Fidelity decisions vs. #22 (ChatEA)

Documented up front in llm_reranker's module docstring -- summarized here:

  • top_k truncation happens before the LLM call -- the actual point of this class vs. LlmBaseLinker (#21), which sends an LLM every blocked candidate with no first-stage narrowing.
  • Confidence shortcut (single top1-vs-top2 gap check) is a simplified analog of ChatEALinker's same idea, without its iterative window-widening -- #23 isn't tied to a specific paper's algorithm.
  • LLM-call failure falls back to the base ranking for that source entity rather than dropping it, since a reranker always has a real base score to fall back to.

This closes out #23's WSD half.