SayedShaun/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

Dataset	Loss	Start F1	End F1	Exact Match	Joint F1
ALL	0.512	0.8129	0.8170	0.7962	0.8087
semeval2007	0.517	0.8088	0.8088	0.7934	0.8037
semeval2013	0.524	0.8054	0.8096	0.7835	0.7995
semeval2015	0.611	0.7916	0.7955	0.7769	0.7880
senseval2	0.527	0.8094	0.8146	0.7927	0.8056
senseval3	0.476	0.8151	0.8232	0.8043	0.8142

Cosine Similarity

Dataset	Loss	F1	Precision	Recall	Accuracy
ALL	0.5684	0.8024	0.8024	0.8024	0.8024
semeval2007	0.5524	0.8066	0.8066	0.8066	0.8066
semeval2013	0.4821	0.8303	0.8303	0.8303	0.8303
semeval2015	0.6726	0.7965	0.7965	0.7965	0.7965
senseval2	0.5827	0.7993	0.7993	0.7993	0.7993
senseval3	0.5064	0.8000	0.8000	0.8000	0.8000

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}
}