cogtran2 0.2.0

CogTran2 - Deep Neural Cognate Detection

The original version of this repository is available on the branch legacy.

Install the package and dependencies:

pip install cogtran2

Install the package and dependencies from the current branch:

python -m pip install .

The no-language-token model is the new default that accepts ASJP representations without language identifiers, so it can be used with ASJP data from any source. The pretrained plain-ASJP checkpoint trained for 90% proportion is available at mahesh27/CogTran2.

Cognate prediction example

Provide IPA transcriptions for each form. If starting from ordinary spelling, convert it to IPA with a language-aware G2P tool first; this package does not perform orthography-to-IPA conversion. The example uses unequal-sized “bird” (12 forms) and “dog” (3 forms) MSAs to demonstrate batch collation. If forms are already ASJP-tokenized, pass asjp=True to align_msa_sca to skip the IPA-to-ASJP conversion.

import torch

from cogtran2 import (
    CogTranDataCollator,
    CogTranTokenizer,
    CogTranTriAttn,
    align_msa_sca,
    upgma_clusters,
)

checkpoint = "mahesh27/CogTran2"
meanings = {
    "bird": [
        "cɛːm", "sɛːm", "sɛːmpar", "ceːm", "seːp", "cim",
        "sɨm", "ci̤m", "cem", "seːm", "com", "kaˀjleːmʔ",
    ],
    "dog": ["cɔː", "kɔʔ", "cɑː"],
}

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = CogTranTokenizer.from_pretrained(checkpoint)
model = CogTranTriAttn.from_pretrained(checkpoint).to(device).eval()

features = []
for meaning, ipa_forms in meanings.items():
    msa = align_msa_sca([tuple(form) for form in ipa_forms])
    unknown = sorted({
        segment
        for row in msa
        for segment in row
        if segment != "-"
        and tokenizer.convert_tokens_to_ids(segment) == tokenizer.unk_token_id
    })
    if unknown:
        raise ValueError(f"{meaning}: unknown ASJP segments: {unknown}")
    features.append({
        "input_ids": [
            tokenizer.encode_form(row, max_length=model.config.max_columns)
            for row in msa
        ]
    })

collator = CogTranDataCollator(
    tokenizer,
    max_length=model.config.max_columns,
    max_rows=model.config.max_rows,
)
batch = {
    name: tensor.to(device)
    for name, tensor in collator(features).items()
}

with torch.inference_mode():
    logits = model(**batch).logits

for batch_index, (meaning, ipa_forms) in enumerate(meanings.items()):
    pair_scores = (
        logits[batch_index, :len(ipa_forms), :len(ipa_forms)]
        .softmax(dim=-1)[..., 1]
        .cpu()
        .tolist()
    )
    clusters = upgma_clusters(pair_scores, threshold=0.4)
    print(meaning)
    for ipa, cluster in zip(ipa_forms, clusters):
        print(f"  {ipa} -> cognate cluster {cluster}")

The collator pads the two MSAs together to [2, 12, 9] (batch, forms, aligned tokens). With the Hub checkpoint, “bird” yields clusters including cɛːm → 0, sɛːmpar → 1, and kaˀjleːmʔ → 2; “dog” includes cɔː → 0 and kɔʔ → 1. Cluster IDs are local to each meaning and have no intrinsic meaning.

Training from scratch

Download fast-cognate-detection.zip to the repository root, then rebuild the training, test, and proportion data with:

python scripts/generate_proportion_splits.py

Run the training suite with run_prop_splits.sh:

bash run_prop_splits.sh

Test scores

Comparison with language-tokens models is given as follows. Mean ± sample standard deviation over five splits where available.

This specific version implements its MSA attention and pairwise triangle-attention stack in native PyTorch rather than importing Fair-ESM or the original AlphaFold modules at runtime. These implementation differences can affect results; in particular, the language-token scores may differ from the legacy model.

Citation

@inproceedings{akavarapu-bhattacharya-2024-automated,
    title = "Automated Cognate Detection as a Supervised Link Prediction Task with Cognate Transformer",
    author = "Akavarapu, V.S.D.S.Mahesh  and
      Bhattacharya, Arnab",
    editor = "Graham, Yvette  and
      Purver, Matthew",
    booktitle = "Proceedings of the 18th Conference of the {E}uropean Chapter of the {A}ssociation for {C}omputational {L}inguistics (Volume 1: Long Papers)",
    month = mar,
    year = "2024",
    address = "St. Julian{'}s, Malta",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2024.eacl-long.58/",
    doi = "10.18653/v1/2024.eacl-long.58",
    pages = "965--975"
}