test: add WSD case study reproduction (Bamman & Burns Table 2)

Binary sense classification with 10-fold CV across 201 lemmas.
Validates accuracy against reference 75.4% with +/-2% tolerance.
Uses latin.sense.data from original repo. Trains a separate
BertForSequenceLabeling model per lemma for up to 100 epochs,
accumulates dev/test correct counts across all lemmas, and picks
the best dev epoch to report test accuracy.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

tests/test_wsd.py

@@ -0,0 +1,386 @@
+"""Word sense disambiguation case study -- Bamman & Burns (2020) Table 2.
+
+Reproduces the WSD experiment: for each polysemous lemma (201 total),
+train a binary classifier (sense I vs sense II) on BERT embeddings with
+10-fold cross-validation. Pick the best epoch on dev, evaluate on test.
+
+Reference results (from original logs):
+  OVERALL: epoch 9, accuracy 0.754, n=1070
+"""
+
+import random
+from pathlib import Path
+
+import numpy as np
+import pytest
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+import torch.optim as optim
+from transformers import AutoTokenizer, BertModel
+
+random.seed(1)
+torch.manual_seed(0)
+np.random.seed(0)
+
+MODEL_PATH = "/tmp/latin-bert-hub"
+DATA_PATH = (
+    Path(__file__).parent.parent
+    / ".claude/reference/latin-bert/case_studies/wsd/data/latin.sense.data"
+)
+
+REF_ACCURACY = 0.754
+TOLERANCE = 0.02  # WSD has more variance due to per-lemma training
+BATCH_SIZE = 32
+DROPOUT_RATE = 0.25
+BERT_DIM = 768
+MAX_EPOCHS = 100
+
+# Special tokens that should not go through subword encoding
+_SPECIAL_TOKENS = {"[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"}
+
+
+def _word_to_subtokens(tokenizer, word):
+    """Get subtoken strings for a single word.
+
+    Special tokens ([CLS], [SEP], etc.) are returned as-is.
+    Regular words are lowercased and tokenized through the subword pipeline,
+    matching the original LatinTokenizer.tokenize() behavior.
+    """
+    if word in _SPECIAL_TOKENS:
+        return [word]
+    return tokenizer.tokenize(word.lower())
+
+
+class BertForSequenceLabeling(nn.Module):
+    """BERT + linear classifier for sequence labeling (binary WSD)."""
+
+    def __init__(self, tokenizer, bert_model, freeze_bert=False,
+                 num_labels=2):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.num_labels = num_labels
+        self.bert = bert_model
+        self.bert.eval()
+        if freeze_bert:
+            for param in self.bert.parameters():
+                param.requires_grad = False
+        self.dropout = nn.Dropout(DROPOUT_RATE)
+        self.classifier = nn.Linear(BERT_DIM, num_labels)
+
+    def forward(self, input_ids, attention_mask=None, transforms=None,
+                labels=None):
+        device = input_ids.device
+        if attention_mask is not None:
+            attention_mask = attention_mask.to(device)
+        if transforms is not None:
+            transforms = transforms.to(device)
+        if labels is not None:
+            labels = labels.to(device)
+
+        outputs = self.bert(input_ids, attention_mask=attention_mask)
+        sequence_output = outputs[0]
+        out = torch.matmul(transforms, sequence_output)
+        logits = self.classifier(out)
+
+        if labels is not None:
+            loss_fct = CrossEntropyLoss(ignore_index=-100)
+            return loss_fct(
+                logits.view(-1, self.num_labels), labels.view(-1)
+            )
+        return logits
+
+    def get_batches(self, sentences, max_batch):
+        """Tokenize and batch with subword-to-word transform matrices."""
+        all_data, all_masks, all_labels, all_transforms = [], [], [], []
+
+        for sentence in sentences:
+            tok_ids, input_mask, labels, transform = [], [], [], []
+            all_toks = []
+            n = 0
+            for word in sentence:
+                toks = _word_to_subtokens(self.tokenizer, word[0])
+                all_toks.append(toks)
+                n += len(toks)
+
+            cur = 0
+            for idx, word in enumerate(sentence):
+                toks = all_toks[idx]
+                ind = list(np.zeros(n))
+                for j in range(cur, cur + len(toks)):
+                    ind[j] = 1.0 / len(toks)
+                cur += len(toks)
+                transform.append(ind)
+                tok_ids.extend(
+                    self.tokenizer.convert_tokens_to_ids(toks)
+                )
+                input_mask.extend(np.ones(len(toks)))
+                labels.append(int(word[1]))
+
+            all_data.append(tok_ids)
+            all_masks.append(input_mask)
+            all_labels.append(labels)
+            all_transforms.append(transform)
+
+        lengths = np.array([len(l) for l in all_data])
+        ordering = np.argsort(lengths)
+
+        ordered_data = [None] * len(all_data)
+        ordered_masks = [None] * len(all_data)
+        ordered_labels = [None] * len(all_data)
+        ordered_transforms = [None] * len(all_data)
+
+        for i, ind in enumerate(ordering):
+            ordered_data[i] = all_data[ind]
+            ordered_masks[i] = all_masks[ind]
+            ordered_labels[i] = all_labels[ind]
+            ordered_transforms[i] = all_transforms[ind]
+
+        batched_data = []
+        batched_mask = []
+        batched_labels = []
+        batched_transforms = []
+
+        i = 0
+        current_batch = max_batch
+
+        while i < len(ordered_data):
+            bd = ordered_data[i:i + current_batch]
+            bm = ordered_masks[i:i + current_batch]
+            bl = ordered_labels[i:i + current_batch]
+            bt = ordered_transforms[i:i + current_batch]
+
+            ml = max(len(s) for s in bd)
+            mlabel = max(len(l) for l in bl)
+
+            for j in range(len(bd)):
+                for _k in range(len(bd[j]), ml):
+                    bd[j].append(0)
+                    bm[j].append(0)
+                    for z in range(len(bt[j])):
+                        bt[j][z].append(0)
+                for _k in range(len(bl[j]), mlabel):
+                    bl[j].append(-100)
+                for _k in range(len(bt[j]), mlabel):
+                    bt[j].append(np.zeros(ml))
+
+            batched_data.append(torch.LongTensor(bd))
+            batched_mask.append(torch.FloatTensor(bm))
+            batched_labels.append(torch.LongTensor(bl))
+            batched_transforms.append(torch.FloatTensor(bt))
+
+            i += current_batch
+            if ml > 100:
+                current_batch = 12
+            if ml > 200:
+                current_batch = 6
+
+        return (batched_data, batched_mask, batched_labels,
+                batched_transforms, ordering)
+
+
+def _get_labs(before, target, after, label):
+    """Build a labeled sentence for WSD.
+
+    Only the target word gets a real label (0 or 1); all context words
+    get -100 (ignored by CrossEntropyLoss).
+    """
+    sent = []
+    for word in before.split(" "):
+        if word:  # skip empty strings from split on empty/whitespace
+            sent.append((word, -100))
+    sent.append((target, label))
+    for word in after.split(" "):
+        if word:
+            sent.append((word, -100))
+    return sent
+
+
+def _read_wsd_data(filename):
+    """Read WSD data file, return dict of lemma -> {0: [...], 1: [...]}."""
+    lemmas = {}
+    with open(filename) as f:
+        for line in f:
+            cols = line.split("\t")
+            lemma = cols[0]
+            label = cols[1]
+            before = cols[2]
+            target = cols[3]
+            after = cols[4].rstrip()
+            if lemma not in lemmas:
+                lemmas[lemma] = {0: [], 1: []}
+            if label == "I":
+                lemmas[lemma][0].append(
+                    _get_labs(before, target, after, 0)
+                )
+            elif label == "II":
+                lemmas[lemma][1].append(
+                    _get_labs(before, target, after, 1)
+                )
+    return lemmas
+
+
+def _get_splits(data):
+    """10-fold cross-validation splits.
+
+    For each sense (0 and 1), examples are assigned to folds by index.
+    testFold = idx % 10, devFold = testFold - 1 (wrapping to 9).
+    """
+    trains, tests, devs = [], [], []
+    for _i in range(10):
+        trains.append([])
+        tests.append([])
+        devs.append([])
+
+    for sense_idx in [0, 1]:
+        for idx, sent in enumerate(data[sense_idx]):
+            test_fold = idx % 10
+            dev_fold = test_fold - 1 if test_fold > 0 else 9
+            for i in range(10):
+                if i == test_fold:
+                    tests[i].append(sent)
+                elif i == dev_fold:
+                    devs[i].append(sent)
+                else:
+                    trains[i].append(sent)
+
+    for i in range(10):
+        random.shuffle(trains[i])
+        random.shuffle(tests[i])
+        random.shuffle(devs[i])
+
+    return trains, devs, tests
+
+
+def _evaluate(model, batched_data, batched_mask, batched_labels,
+              batched_transforms, device):
+    """Evaluate model on batched data, return (correct, total).
+
+    Mirrors the original evaluate() method which returns (cor, tot),
+    with accumulation happening outside this function.
+    """
+    model.eval()
+    cor = 0
+    tot = 0
+    with torch.no_grad():
+        for b in range(len(batched_data)):
+            logits = model(
+                batched_data[b].to(device),
+                attention_mask=batched_mask[b],
+                transforms=batched_transforms[b],
+            )
+            logits = logits.cpu()
+            size = batched_labels[b].shape
+            logits = logits.view(-1, size[1], 2)
+            preds = np.argmax(logits.numpy(), axis=2)
+            for row in range(size[0]):
+                for col in range(size[1]):
+                    if batched_labels[b][row][col] != -100:
+                        if preds[row][col] == batched_labels[b][row][col]:
+                            cor += 1
+                        tot += 1
+    return cor, tot
+
+
+@pytest.mark.slow
+def test_wsd_accuracy():
+    """Reproduce WSD case study from Bamman & Burns (2020).
+
+    Trains a separate binary classifier per lemma (201 lemmas) with
+    10-fold cross-validation. Uses fold 0 splits (train/dev/test).
+    Accumulates dev and test correct/total across all lemmas at each
+    epoch, then picks the best dev epoch and reports test accuracy.
+    """
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_PATH, trust_remote_code=True
+    )
+    data = _read_wsd_data(str(DATA_PATH))
+
+    dev_cors = [0.0] * MAX_EPOCHS
+    test_cors = [0.0] * MAX_EPOCHS
+    dev_n = [0.0] * MAX_EPOCHS
+    test_n = [0.0] * MAX_EPOCHS
+
+    for lemma_idx, lemma in enumerate(data):
+        print(f"\n[{lemma_idx + 1}/{len(data)}] {lemma}")
+
+        # Fresh model per lemma
+        bert_model = BertModel.from_pretrained(MODEL_PATH)
+        model = BertForSequenceLabeling(
+            tokenizer, bert_model, freeze_bert=False, num_labels=2
+        )
+        model.to(device)
+
+        trains, devs, tests = _get_splits(data[lemma])
+        train_data = trains[0]
+        dev_data = devs[0]
+        test_data = tests[0]
+
+        train_b, train_m, train_l, train_t, _ = model.get_batches(
+            train_data, BATCH_SIZE
+        )
+        dev_b, dev_m, dev_l, dev_t, _ = model.get_batches(
+            dev_data, BATCH_SIZE
+        )
+        test_b, test_m, test_l, test_t, _ = model.get_batches(
+            test_data, BATCH_SIZE
+        )
+
+        optimizer = optim.Adam(model.parameters(), lr=5e-5)
+
+        for epoch in range(MAX_EPOCHS):
+            model.train()
+            for b in range(len(train_b)):
+                loss = model(
+                    train_b[b].to(device),
+                    attention_mask=train_m[b],
+                    transforms=train_t[b],
+                    labels=train_l[b],
+                )
+                loss.backward()
+                optimizer.step()
+                model.zero_grad()
+
+            # Evaluate dev
+            c, t = _evaluate(
+                model, dev_b, dev_m, dev_l, dev_t, device
+            )
+            dev_cors[epoch] += c
+            dev_n[epoch] += t
+
+            # Evaluate test
+            c, t = _evaluate(
+                model, test_b, test_m, test_l, test_t, device
+            )
+            test_cors[epoch] += c
+            test_n[epoch] += t
+
+        # Print per-lemma dev accuracy summary
+        for epoch in range(MAX_EPOCHS):
+            if dev_n[epoch] > 0:
+                dev_acc = dev_cors[epoch] / dev_n[epoch]
+                print(
+                    f"  DEV: epoch={epoch} acc={dev_acc:.3f} "
+                    f"lemma={lemma} n={dev_n[epoch]}"
+                )
+
+    # Find best dev epoch, report test accuracy at that epoch
+    best_epoch = max(
+        range(MAX_EPOCHS),
+        key=lambda i: dev_cors[i] / dev_n[i] if dev_n[i] > 0 else 0,
+    )
+    test_accuracy = test_cors[best_epoch] / test_n[best_epoch]
+
+    print(
+        f"\nOVERALL: epoch={best_epoch}, "
+        f"accuracy={test_accuracy:.3f}, "
+        f"n={test_n[best_epoch]}"
+    )
+    print(f"Reference: epoch=9, accuracy={REF_ACCURACY}, n=1070")
+
+    assert abs(test_accuracy - REF_ACCURACY) < TOLERANCE, (
+        f"WSD accuracy {test_accuracy:.3f} outside tolerance "
+        f"of {REF_ACCURACY} +/- {TOLERANCE}"
+    )