Views
No views yet

colpali-engine==0.3.11.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}