manu commited on
Commit
ca1e9e6
1 Parent(s): d44207d

Create README.md

Browse files
Files changed (1) hide show
  1. README.md +139 -0
README.md ADDED
@@ -0,0 +1,139 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: mit
3
+ library_name: colpali
4
+ language:
5
+ - en
6
+ tags:
7
+ - vidore
8
+ ---
9
+ # ColPali: Visual Retriever based on PaliGemma-3B with ColBERT strategy
10
+
11
+ ColPali is a model based on a novel model architecture and training strategy based on Vision Language Models (VLMs) to efficiently index documents from their visual features.
12
+ It is a [PaliGemma-3B](https://huggingface.co/google/paligemma-3b-mix-448) extension that generates [ColBERT](https://arxiv.org/abs/2004.12832)- style multi-vector representations of text and images.
13
+ It was introduced in the paper [ColPali: Efficient Document Retrieval with Vision Language Models](https://arxiv.org/abs/2407.01449) and first released in [this repository](https://github.com/ManuelFay/colpali)
14
+
15
+ ## Model Description
16
+
17
+ This model is trained with 150k samples from the Docmatix dataset (and not the original train set) - with mined hard negatives + in-batch cntrastive loss !
18
+
19
+ This model is built iteratively starting from an off-the-shelf [SigLIP](https://huggingface.co/google/siglip-so400m-patch14-384) model.
20
+ We finetuned it to create [BiSigLIP](https://huggingface.co/vidore/bisiglip) and fed the patch-embeddings output by SigLIP to an LLM, [PaliGemma-3B](https://huggingface.co/google/paligemma-3b-mix-448) to create [BiPali](https://huggingface.co/vidore/bipali).
21
+
22
+ One benefit of inputting image patch embeddings through a language model is that they are natively mapped to a latent space similar to textual input (query).
23
+ This enables leveraging the [ColBERT](https://arxiv.org/abs/2004.12832) strategy to compute interactions between text tokens and image patches, which enables a step-change improvement in performance compared to BiPali.
24
+
25
+ ## Model Training
26
+
27
+ ### Dataset
28
+ Our training dataset of 127,460 query-page pairs is comprised of train sets of openly available academic datasets (63%) and a synthetic dataset made up of pages from web-crawled PDF documents and augmented with VLM-generated (Claude-3 Sonnet) pseudo-questions (37%).
29
+ Our training set is fully English by design, enabling us to study zero-shot generalization to non-English languages. We explicitly verify no multi-page PDF document is used both [*ViDoRe*](https://huggingface.co/collections/vidore/vidore-benchmark-667173f98e70a1c0fa4db00d) and in the train set to prevent evaluation contamination.
30
+ A validation set is created with 2% of the samples to tune hyperparameters.
31
+
32
+ *Note: Multilingual data is present in the pretraining corpus of the language model (Gemma-2B) and potentially occurs during PaliGemma-3B's multimodal training.*
33
+
34
+ ### Parameters
35
+
36
+ All models are trained for 1 epoch on the train set. Unless specified otherwise, we train models in `bfloat16` format, use low-rank adapters ([LoRA](https://arxiv.org/abs/2106.09685))
37
+ with `alpha=32` and `r=32` on the transformer layers from the language model,
38
+ as well as the final randomly initialized projection layer, and use a `paged_adamw_8bit` optimizer.
39
+ We train on an 8 GPU setup with data parallelism, a learning rate of 5e-5 with linear decay with 2.5% warmup steps, and a batch size of 32.
40
+
41
+ ## Usage
42
+
43
+ ```python
44
+ import torch
45
+ import typer
46
+ from torch.utils.data import DataLoader
47
+ from tqdm import tqdm
48
+ from transformers import AutoProcessor
49
+ from PIL import Image
50
+
51
+ from colpali_engine.models.paligemma_colbert_architecture import ColPali
52
+ from colpali_engine.trainer.retrieval_evaluator import CustomEvaluator
53
+ from colpali_engine.utils.colpali_processing_utils import process_images, process_queries
54
+ from colpali_engine.utils.image_from_page_utils import load_from_dataset
55
+
56
+
57
+ def main() -> None:
58
+ """Example script to run inference with ColPali"""
59
+
60
+ # Load model
61
+ model_name = "manu/colpali-3b-mix-448-docmatix-only"
62
+ model = ColPali.from_pretrained("google/paligemma-3b-mix-448", torch_dtype=torch.bfloat16, device_map="cuda").eval()
63
+ model.load_adapter(model_name)
64
+ processor = AutoProcessor.from_pretrained(model_name)
65
+
66
+ # select images -> load_from_pdf(<pdf_path>), load_from_image_urls(["<url_1>"]), load_from_dataset(<path>)
67
+ images = load_from_dataset("vidore/docvqa_test_subsampled")
68
+ queries = ["From which university does James V. Fiorca come ?", "Who is the japanese prime minister?"]
69
+
70
+ # run inference - docs
71
+ dataloader = DataLoader(
72
+ images,
73
+ batch_size=4,
74
+ shuffle=False,
75
+ collate_fn=lambda x: process_images(processor, x),
76
+ )
77
+ ds = []
78
+ for batch_doc in tqdm(dataloader):
79
+ with torch.no_grad():
80
+ batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
81
+ embeddings_doc = model(**batch_doc)
82
+ ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
83
+
84
+ # run inference - queries
85
+ dataloader = DataLoader(
86
+ queries,
87
+ batch_size=4,
88
+ shuffle=False,
89
+ collate_fn=lambda x: process_queries(processor, x, Image.new("RGB", (448, 448), (255, 255, 255))),
90
+ )
91
+
92
+ qs = []
93
+ for batch_query in dataloader:
94
+ with torch.no_grad():
95
+ batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
96
+ embeddings_query = model(**batch_query)
97
+ qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
98
+
99
+ # run evaluation
100
+ retriever_evaluator = CustomEvaluator(is_multi_vector=True)
101
+ scores = retriever_evaluator.evaluate(qs, ds)
102
+ print(scores.argmax(axis=1))
103
+
104
+
105
+ if __name__ == "__main__":
106
+ typer.run(main)
107
+
108
+ ```
109
+
110
+ ## Limitations
111
+
112
+ - **Focus**: The model primarily focuses on PDF-type documents and high-ressources languages, potentially limiting its generalization to other document types or less represented languages.
113
+ - **Support**: The model relies on multi-vector retreiving derived from the ColBERT late interaction mechanism, which may require engineering efforts to adapt to widely used vector retrieval frameworks that lack native multi-vector support.
114
+
115
+ ## License
116
+
117
+ ColPali's vision language backbone model (PaliGemma) is under `gemma` license as specified in its [model card](https://huggingface.co/google/paligemma-3b-mix-448). The adapters attached to the model are under MIT license.
118
+
119
+ ## Contact
120
+
121
+ - Manuel Faysse: [email protected]
122
+ - Hugues Sibille: [email protected]
123
+ - Tony Wu: [email protected]
124
+
125
+ ## Citation
126
+
127
+ If you use any datasets or models from this organization in your research, please cite the original dataset as follows:
128
+
129
+ ```bibtex
130
+ @misc{faysse2024colpaliefficientdocumentretrieval,
131
+ title={ColPali: Efficient Document Retrieval with Vision Language Models},
132
+ author={Manuel Faysse and Hugues Sibille and Tony Wu and Bilel Omrani and Gautier Viaud and Céline Hudelot and Pierre Colombo},
133
+ year={2024},
134
+ eprint={2407.01449},
135
+ archivePrefix={arXiv},
136
+ primaryClass={cs.IR},
137
+ url={https://arxiv.org/abs/2407.01449},
138
+ }
139
+ ```