Add support for token pooling for vision retriever wrapper

[tonywu71]

Description

• Add support for token pooling for VisionRetriever class
• Add tests for the VisionRetriever class

Example usage:

from colpali_engine.compression import HierarchicalTokenPooler
from vidore_benchmark.retrievers import VisionRetriever

# Initialize your model and processor
# ...

# Create token pooler (reduces embedding size)
token_pooler = HierarchicalTokenPooler()

# Initialize retriever with token pooling
retriever = VisionRetriever(
    model=model,
    processor=processor,
    token_pooler=token_pooler,
    num_workers=2  # Optional: for faster processing
)

# Use as normal
query_embeddings = retriever.forward_queries(queries)
passage_embeddings = retriever.forward_passages(
    passages,
    pooling_kwargs={"pool_factor": 3},
)  # will use token pooling
scores = retriever.get_scores(query_embeddings, passage_embeddings)

Tests (local run)

[Copilot]

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Comments suppressed due to low confidence (1)

tests/retrievers/test_vision_retriever.py:53

• [nitpick] The type annotation 'ColIdefics3Retriever' is inconsistent with the fixture, which instantiates a VisionRetriever. Consider updating the type annotation to match the actual instance returned by the fixture.

def test_forward_queries(self, retriever: ColIdefics3Retriever, queries_fixture):

[QuentinJGMace]

All good for me, sorry for the delay