word-sense-disambiguation

Word Sense Disambiguation (WSD) — Fine-tuned Weights

Fine-tuned checkpoints for the Dual Architecture pipeline for Word Sense Disambiguation (span extraction + gloss-context cosine similarity), built on top of distilbert-base-uncased and trained on the SemCor dataset.

The two architectures implemented here are based on ideas from:

• Barba et al., ESC: Redesigning WSD with Extractive Sense Comprehension (NAACL 2021) — the span extraction formulation.
• Blevins & Zettlemoyer, Moving Down the Long Tail of Word Sense Disambiguation with Gloss-Informed Bi-encoders (ACL 2020) — the gloss-context bi-encoder / cosine similarity formulation.

Training/inference code, configs, and full documentation: github.com/sayedshaun/wsd

Repository Contents

├── span/
│   └── step-28000-f1-0.8037.pt      # Span extraction architecture checkpoint
└── cosine/
    └── step-12000-f1-0.8066.pt      # Gloss-context cosine similarity architecture checkpoint

Each .pt file is a raw PyTorch state_dict (not a full pickled model), saved with torch.save(model.state_dict(), path).

Architectures

• Span (SpanExtractionModel): Encodes the sentence and predicts the start/end token span of the correct gloss, QA-style.
• Cosine (WSDModel): Encodes the sentence and each candidate gloss separately, then scores senses by cosine similarity between the [CLS] representations.

Results

Span Extraction

Cosine Similarity

Inference

These weights load into the model classes defined in the wsd repo (model.py). Clone that repo first, or vendor model.py and utils.py.

1. Download the weights

hf download SayedShaun/word-sense-disambiguation --local-dir weights

2. Load a model

import torch
from transformers import AutoTokenizer
from model import WSDModel, SpanExtractionModel  # from the wsd repo

MODEL_NAME = "distilbert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# --- Cosine architecture ---
model = WSDModel(MODEL_NAME, tokenizer=tokenizer, n_sense=5).to(device)
state_dict = torch.load(
    "weights/cosine/step-12000-f1-0.8066.pt", map_location=device
)
model.load_state_dict(state_dict)
model.eval()

# --- Span extraction architecture ---
span_model = SpanExtractionModel(MODEL_NAME, tokenizer=tokenizer).to(device)
span_state_dict = torch.load(
    "weights/span/step-28000-f1-0.8037.pt", map_location=device
)
span_model.load_state_dict(span_state_dict)
span_model.eval()

3. Run evaluation with the repo's script

python predict.py \
    --data_dir "data/Evaluation_Datasets/semeval2015" \
    --model_name "distilbert-base-uncased" \
    --weight_dir "weights/cosine" \
    --pos "ALL" \
    --seed 1234 \
    --num_sense 5 \
    --max_length 256 \
    --batch_size 32 \
    --architecture "cosine"

Pass --weight_dir weights/span --architecture span to evaluate the span extraction checkpoint instead.

Training Data

• SemCor (training)
• SemEval 2007/2013/2015, Senseval 2/3 (evaluation)

Citation

@misc{wsd,
    author       = {Md Abu Sayed Shaun},
    title        = {Dual Architecture pipeline for Word Sense Disambiguation (WSD)},
    year         = {2025},
    howpublished = {\url{https://github.com/sayedshaun/wsd}},
    note         = {GitHub repository}
}