answerdotai

/

JaColBERTv2.5

like 23

Follow

Answer.AI 611

Hello!

Heads up, this PR was AI-generated and human-reviewed. Here's a summary of the changes as reported by my agent:

Pull Request overview

• Fix the architectures field in config.json so this checkpoint is recognised as a Stanford ColBERT model again.
• Add a Sentence Transformers usage section to the model card, plus the multi-vector and sentence-transformers tags.

Details

This repository ships the Stanford ColBERT layout: the 128-dim projection lives at the root of model.safetensors as linear.weight, and artifact.metadata carries the training configuration. Both Sentence Transformers v6.0.0 and PyLate detect that layout from config.json's architectures field, which here reads BertModel rather than the HF_ColBERT that bclavie/JaColBERTv2 and colbert-ir/colbertv2.0 ship.

The practical effect is that neither library currently loads this model correctly. linear.weight is reported as an unexpected key and discarded, a fresh projection is randomly initialised in its place, and the [unused0] / [unused1] markers, the 32-token query expansion and the punctuation masking are all skipped. PyLate additionally resizes the vocabulary from 32768 to 32770 to add its own marker tokens. The embeddings that come out of either path are not the ones this checkpoint was trained to produce.

Changing that one field to HF_ColBERT is enough. Everything else is then read from artifact.metadata that is already in the repo: the [unused0] and [unused1] prefixes, query_maxlen 32 for the [MASK] query expansion with attend_to_mask_tokens false, doc_maxlen 300, mask_punctuation, and the 128-dim projection loaded from linear.weight. No new configuration files, no new weights, and no trust_remote_code.

The field is metadata only for colbert-ai and RAGatouille, which resolve the backbone from model_type instead. That was verified rather than assumed: running colbert-ai against this repo before and after the change produces bit-identical token embeddings and identical MaxSim scores (max difference 0.0). The existing RAGatouille usage is unaffected.

For verification, colbert-ai was run on transformers==4.42.4, the version recorded in config.json, and used as the reference. Sentence Transformers reproduces it with a per-token cosine similarity of 1.000000 on the query and on all four documents, identical token counts, and MaxSim scores matching to within 1.9e-06. PyLate, once the field is fixed, matches the same reference exactly.

Modified files:

• config.json: architectures changed from ["BertModel"] to ["HF_ColBERT"], a one-line change.
• README.md: adds the Sentence Transformers usage section and the multi-vector and sentence-transformers tags.

pip install "sentence-transformers>=6.0.0" fugashi unidic-lite

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("answerdotai/JaColBERTv2.5", revision="refs/pr/5")

query = "日本で一番高い山は何ですか？"
documents = [
    "富士山は日本で最も高い山で、標高は3776メートルです。",
    "東京は日本の首都で、世界最大の都市圏の一つです。",
    "北岳は南アルプスにある山で、日本で二番目に高い山です。",
    "富士山は静岡県と山梨県にまたがる活火山です。",
]

query_embeddings = model.encode_query(query)
document_embeddings = model.encode_document(documents)
print(query_embeddings.shape, document_embeddings[0].shape)
# torch.Size([32, 128]) torch.Size([21, 128])

# MaxSim late-interaction scoring (higher is more relevant)
scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[30.6709, 28.9073, 29.6979, 29.0785]], device='cuda:0')

Note that the weights are untouched and no existing behaviour changes. The colbert-ai and RAGatouille paths produce exactly the same numbers as before, so this only repairs the two loaders that were reading the field and adds a familiar way to run the model.

I'd love to feature this model in the Sentence Transformers documentation, especially once it loads without the revision pin (that is, once this PR is merged).

Happy to tweak anything you'd like changed. Please let me know if you have any questions or feedback!

• Tom Aarsen