ibm-granite

/

granite-vision-3.3-2b-embedding

like 29

Follow

IBM Granite 5.14k

Hello!

I recently released the new MultiVectorEncoder class in Sentence Transformers v6.0, and I think this model would be a great fit for it. You can read more about the release here: https://huggingface.co/blog/multi-vector-encoder. I would also love to feature this model in the accompanying documentation.

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 custom modeling and processing code for transformers 5.x (the repository currently cannot load on transformers >= 5.0 at all).
• Integrate ibm-granite/granite-vision-3.3-2b-embedding with Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever via MultiVectorEncoder.

Details

The larger half of this PR is transformers 5.x compatibility for the existing custom code, independent of Sentence Transformers. On transformers 5.15 the repository fails to load for a cascade of reasons, each addressed here: the processor __init__ used a *args, **kwargs signature that 5.x can no longer introspect, so the components are now declared explicitly. The language model moved under model.model.language_model in the 5.x layout, so attribute paths and the embedding accessors are updated (the custom tie_weights is no longer needed and removed). The checkpoint on the Hub uses the 4.x key layout, so the model class now carries a _checkpoint_conversion_mapping plus a from_pretrained hook that applies it. The legacy chat_template.json is converted to chat_template.jinja, since the legacy file conflicts with newer template handling. Finally the processor gains a standard __call__ that shares the smart resize and padding with process_images and resolves URL or path image inputs the way stock processors do, so apply_chat_template and generic pipelines produce inputs identical to the dedicated methods. All of these are mechanical (attribute paths, key layout, processor plumbing) and do not touch the math or the weights. I tested on transformers 5.15.0. One caveat: the from_pretrained hook passes key_mapping, which recent 4.x versions accept, but older 4.x versions would need a version guard there. Happy to add one if you support those.

On top of that, the Sentence Transformers integration is config-only. The pipeline drives the repository's own embedding forward directly, so the exact embeddings the modeling code produces are used as is, followed by a small repo-local module (st_gather_image_mask.py) that reproduces the fixed 729 image-token gather of forward for documents, and a mask module that drops padding. A named chat template in additional_chat_templates/sentence_transformers.jinja reproduces the reference prompt formats, including the doubled query text and the ten <|end_of_text|> augmentation tokens.

Since config.json declares no dtype, a plain load runs in fp32. The snippet below passes model_kwargs={"dtype": "float16"} to match the model card's own usage. I did not change the config default, that call is yours to make.

Verified against the repository's own pipeline (with the fixes above) on text queries and image documents: the per-token embeddings match with cosine similarity 1.0000 on every query and document, and the MaxSim scores differ by at most 0.0015. For reference, on a 400-query NanoViDoRe v3 subset this integration scores 0.4087 nDCG@10.

pip install "sentence-transformers[image]>=6.0.0"

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder(
    "ibm-granite/granite-vision-3.3-2b-embedding",
    trust_remote_code=True,
    model_kwargs={"dtype": "float16"},
)

queries = [
    "What is the variable represented on the y-axis of the graph?",
    "Total outlay is maximum in which year?",
]
images = [
    "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc1.jpg",
    "https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc2.jpg",
]

query_embeddings = model.encode_query(queries)
image_embeddings = model.encode_document(images)
print(query_embeddings[0].shape, image_embeddings[0].shape)
# torch.Size([46, 128]) torch.Size([729, 128])

# Diagonal should have higher scores
scores = model.similarity(query_embeddings, image_embeddings)
print(scores)
# tensor([[34.0132, 25.3076],
#         [19.6362, 23.4976]], device='cuda:0')

To try this before merging, pass revision="refs/pr/14" to MultiVectorEncoder.

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

• Tom Aarsen