Views
No views yet

colpali-engine==0.3.11.MultiVectorEncoder. Queries are text, and documents can be page images or audio.pip install "sentence-transformers[image,audio,video]>=6.0.0"1from sentence_transformers import MultiVectorEncoder
2
3model = MultiVectorEncoder("vidore/colqwen-omni-v0.1")
4
5queries = [
6 "What is the variable represented on the y-axis of the graph?",
7 "Total outlay is maximum in which year?",
8]
9documents = [
10 f"https://huggingface.co/datasets/sentence-transformers/example-documents/resolve/main/doc{i}.jpg" for i in range(1, 5)
11]
12
13query_embeddings = model.encode_query(queries)
14document_embeddings = model.encode_document(documents)
15print(f"Query 0 shape: {tuple(query_embeddings[0].shape)}")
16print(f"Document 0 shape: {tuple(document_embeddings[0].shape)}")
17# Query 0 shape: (61, 128)
18# Document 0 shape: (1034, 128)
19
20# MaxSim late-interaction scoring (rows = queries, columns = documents)
21scores = model.similarity(query_embeddings, document_embeddings)
22print(scores)
23# tensor([[53.5625, 49.2036, 46.6958, 45.4949],
24# [45.6436, 53.1328, 45.0957, 45.5176]])1from datasets import Audio, load_dataset
2
3dataset = load_dataset("eustlb/dailytalk-conversations-grouped", split="train[:20]")
4dataset = dataset.cast_column("audio", Audio(sampling_rate=16_000))
5audios = [row["array"] for row in dataset["audio"]] # raw mono waveforms, float32 at 16 kHz
6
7audio_embeddings = model.encode_document(audios)
8scores = model.similarity(model.encode_query(["medicine for car nausea"]), audio_embeddings)[0]
9
10# zero-shot (trained on images only, no transcription step): the "nausea" query matches the "carsickness" recording
11top_scores, top_indices = scores.topk(3)
12for score, index in zip(top_scores.tolist(), top_indices.tolist()):
13 print(f"{score:.2f} {dataset[index]['texts'][0]}")
14# 50.88 Excuse me? Do you have anything for a carsickness?
15# 46.05 Excuse me, could you tell me where you have got that music book?
16# 46.01 Jeff, I'm going to the supermarket. Do you want to come with me?[!NOTE] Documents are tiled adaptively, so their embeddings vary in length (1034 to 1062 tokens for the four example pages). MaxSim handles that, andmodel.similaritymasks the padding for you.
colpali-engine is installed from source or with a version superior to 0.3.11.pip install git+https://github.com/illuin-tech/colpali1
2import torch
3from PIL import Image
4from transformers.utils.import_utils import is_flash_attn_2_available
5from tqdm import tqdm
6from torch.utils.data import DataLoader
7
8from colpali_engine.models import ColQwen2_5Omni, ColQwen2_5OmniProcessor
9
10model = ColQwen2_5Omni.from_pretrained(
11 "vidore/colqwen-omni-v0.1",
12 torch_dtype=torch.bfloat16,
13 device_map="cuda", # or "mps" if on Apple Silicon
14 attn_implementation="flash_attention_2" # if is_flash_attn_2_available() else None,
15).eval()
16processor = ColQwen2_5OmniProcessor.from_pretrained("vidore/colqwen-omni-v0.1")
17
18dataset = load_dataset("eustlb/dailytalk-conversations-grouped", split="train[:500]")
19audios = [x["array"] for x in dataset["audio"]]
20
21
22dataloader = DataLoader(
23 dataset=audios,
24 batch_size=2,
25 shuffle=False,
26 collate_fn=lambda x: processor.process_audios(x),
27)
28
29ds = []
30for batch_doc in tqdm(dataloader):
31 with torch.no_grad():
32 batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
33 embeddings_doc = model(**batch_doc)
34 ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
35
36def get_results(query: str, k=10):
37 batch_queries = processor.process_queries([query]).to(model.device)
38
39 # Forward pass
40 with torch.no_grad():
41 query_embeddings = model(**batch_queries)
42
43 scores = processor.score_multi_vector(query_embeddings, ds)
44 # get top-5 scores
45 return scores[0].topk(k).indices.tolist()
46
47res = get_results("A person looking for a taxi")
48
49# In colab
50display(Audio(dataset[res[0]]["audio"]["array"], autoplay=True, rate=dataset[res[0]]["audio"]["sampling_rate"]))1@misc{faysse2024colpaliefficientdocumentretrieval,
2 title={ColPali: Efficient Document Retrieval with Vision Language Models},
3 author={Manuel Faysse and Hugues Sibille and Tony Wu and Bilel Omrani and Gautier Viaud and Céline Hudelot and Pierre Colombo},
4 year={2024},
5 eprint={2407.01449},
6 archivePrefix={arXiv},
7 primaryClass={cs.IR},
8 url={https://arxiv.org/abs/2407.01449},
9}