Update README.md

README.md

@@ -18,7 +18,7 @@ It was introduced in the paper [ColPali: Efficient Document Retrieval with Visio
 
 ## Version specificity
 
-This version is trained with `colpali-engine==0.2.0`.
+This version is trained with `colpali-engine==0.2.0` but can be loaded for any version `>=0.2.0`.
 
 Compared to `colpali`, this version is trained with right padding for queries to fix unwanted tokens in the query encoding.
 It also stems from the fixed `vidore/colpaligemma-3b-pt-448-base` to guarantee deterministic projection layer initialization.
@@ -52,89 +52,53 @@ We train on an 8 GPU setup with data parallelism, a learning rate of 5e-5 with l
 
 ## Usage
 
+Install [`colpali-engine`](https://github.com/illuin-tech/colpali):
+
 ```bash
-pip install colpali-engine==0.2.0
+pip install colpali-engine>=0.3.0,<0.4.0
 ```
 
+Then run the following code:
 
 ```python
+from typing import cast
+
 import torch
-import typer
-from torch.utils.data import DataLoader
-from tqdm import tqdm
-from transformers import AutoProcessor
 from PIL import Image
 
-from colpali_engine.models.paligemma_colbert_architecture import ColPali
-from colpali_engine.trainer.retrieval_evaluator import CustomEvaluator
-from colpali_engine.utils.colpali_processing_utils import process_images, process_queries
-from colpali_engine.utils.image_from_page_utils import load_from_dataset
-
-
-def main() -> None:
-    """Example script to run inference with ColPali"""
-
-    # Load model
-    model_name = "vidore/colpali-v1.2"
-    model = ColPali.from_pretrained("vidore/colpaligemma-3b-pt-448-base", torch_dtype=torch.bfloat16, device_map="cuda").eval()
-    model.load_adapter(model_name)
-    model = model.eval()
-    processor = AutoProcessor.from_pretrained(model_name)
-
-    # select images -> load_from_pdf(<pdf_path>),  load_from_image_urls(["<url_1>"]), load_from_dataset(<path>)
-    images = load_from_dataset("vidore/docvqa_test_subsampled")
-    queries = ["From which university does James V. Fiorca come ?", "Who is the japanese prime minister?"]
-
-    # run inference - docs
-    dataloader = DataLoader(
-        images,
-        batch_size=4,
-        shuffle=False,
-        collate_fn=lambda x: process_images(processor, x),
-    )
-    ds = []
-    for batch_doc in tqdm(dataloader):
-        with torch.no_grad():
-            batch_doc = {k: v.to(model.device) for k, v in batch_doc.items()}
-            embeddings_doc = model(**batch_doc)
-        ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
-
-    # run inference - queries
-    dataloader = DataLoader(
-        queries,
-        batch_size=4,
-        shuffle=False,
-        collate_fn=lambda x: process_queries(processor, x, Image.new("RGB", (448, 448), (255, 255, 255))),
-    )
-
-    qs = []
-    for batch_query in dataloader:
-        with torch.no_grad():
-            batch_query = {k: v.to(model.device) for k, v in batch_query.items()}
-            embeddings_query = model(**batch_query)
-        qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
-
-    # run evaluation
-    retriever_evaluator = CustomEvaluator(is_multi_vector=True)
-    scores = retriever_evaluator.evaluate(qs, ds)
-    print(scores.argmax(axis=1))
-
-
-if __name__ == "__main__":
-    typer.run(main)
-
-```
-
-**Note:** If you need to further train ColPali from this adapter, you should run:
-
-```python
-lora_config = LoraConfig.from_pretrained("vidore/colpali-v1.1")
-lora_config.inference_mode = False  # force training mode for fine-tuning
-
-model = get_peft_model(model, lora_config)
-
-print("after")
-model.print_trainable_parameters()
+from colpali_engine.models import ColPali, ColPaliProcessor
+
+model = cast(
+    ColPali,
+    ColPali.from_pretrained(
+        "vidore/colpali-v1.2",
+        torch_dtype=torch.bfloat16,
+        device_map="cuda:0",  # or "mps" if on Apple Silicon
+    ),
+)
+
+processor = cast(ColPaliProcessor, ColPaliProcessor.from_pretrained("google/paligemma-3b-mix-448"))
+
+# Your inputs
+images = [
+    Image.new("RGB", (32, 32), color="white"),
+    Image.new("RGB", (16, 16), color="black"),
+]
+queries = [
+    "Is attention really all you need?",
+    "Are Benjamin, Antoine, Merve, and Jo best friends?",
+]
+
+# Process the inputs
+batch_images = processor.process_images(images).to(model.device)
+batch_queries = processor.process_queries(queries).to(model.device)
+
+# Forward pass
+with torch.no_grad():
+    image_embeddings = model(**batch_images)
+    querry_embeddings = model(**batch_queries)
+
+scores = processor.score_multi_vector(querry_embeddings, image_embeddings)
 ```
 
 ## Limitations