Xuanxin/otalign-wiki-supervised

OTAlign: Supervised Monolingual Word Aligner

This repository provides a trained checkpoint for OTAlign, based on:

Yuki Arase, Han Bao, and Sho Yokoi. Unbalanced Optimal Transport for Unbalanced Word Alignment. ACL 2023. https://aclanthology.org/2023.acl-long.219/

The original implementation is available at https://github.com/yukiar/OTAlign.

This is an unofficial checkpoint and is not maintained by the original authors.

Usage

git clone https://github.com/yukiar/OTAlign
cd OTAlign/src
pip install torch transformers pytorch_lightning POT huggingface_hub

import numpy as np
import torch
from torch.utils.data import DataLoader
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download

from model import NeuralAligner
from util import AlignmentDataset, alignment_collate, get_token_word_mapping, send_batch_to_gpu

model = NeuralAligner.load_from_checkpoint(
    hf_hub_download("Xuanxin/otalign-wiki-supervised", "model.ckpt"),
    strict=False,
)

model.eval()
model.freeze()
model.to("cuda")

tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
model.tokenizer = tokenizer
model.align_hypr = 0.4242

src = "You're good at selling yourself, thanks to Dr. Bettelheim.".lower().split()
tgt = "You're good at promoting yourself, thanks to Dr. Bettelheim.".lower().split()

enc = tokenizer(
    text=[src],
    text_pair=[tgt],
    padding=True,
    truncation=True,
    is_split_into_words=True,
    return_offsets_mapping=True,
    return_tensors="pt",
)

sub2word, tgt_start = get_token_word_mapping(enc, tokenizer)

dataset = AlignmentDataset([src], [tgt], enc, [""], sub2word, tgt_start)
batch = next(iter(DataLoader(dataset, batch_size=1, collate_fn=alignment_collate)))

with torch.no_grad():
    P = model.bidirectional_align(send_batch_to_gpu(batch, model.device))[0]

print(sorted(
    (int(i), int(j))
    for i, j in np.transpose(np.nonzero(P > model.align_hypr))
))
# [(0, 0), (1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6), (7, 7), (8, 8)]

Training

Setting	Value
Encoder	bert-base-uncased
Training data	MultiMWA-Wiki, sure_and_possible
OT variant	uot
Weights	uniform
Distance	cos
ot_hyp	0.88235 (estimated in the unsupervised stage)
distortion	0.1
Learning rate	9.0e-5
Batch size	64
Early-stopping patience	5 (on validation F1)
Seed	42
Alignment threshold	0.4242 (selected on the validation set)

The training data is not redistributed. See https://github.com/yukiar/MultiMWA.

License and Citation

MIT, following the original OTAlign repository.

If you use this checkpoint, please cite:

@inproceedings{arase-etal-2023-unbalanced,
    title = "Unbalanced Optimal Transport for Unbalanced Word Alignment",
    author = "Arase, Yuki and Bao, Han and Yokoi, Sho",
    booktitle = "Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers)",
    year = "2023",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2023.acl-long.219",
    doi = "10.18653/v1/2023.acl-long.219",
    pages = "3966--3986",
}