Update README.md

README.md

@@ -24,65 +24,95 @@ pip install . -e
 Then, you can run the following inference code: 
 
 ```python
+import pprint
+from typing import List, cast
+
 import torch
-import typer
+from datasets import Dataset, load_dataset
 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
+from colpali_engine.models import ColFlor
+from colpali_engine.models import ColFlorProcessor
+from colpali_engine.utils.processing_utils import BaseVisualRetrieverProcessor
+from colpali_engine.utils.torch_utils import ListDataset, get_torch_device
+
 
+def main():
+    """
+    Example script to run inference with ColFlor.
+    """
 
-def main() -> None:
-    """Example script to run inference with ColPali"""
+    device = get_torch_device("auto")
+    print(f"Device used: {device}")
+
+    # Model name
+    model_name = "ahmed-masry/ColFlor"
 
     # Load model
-    model_name = "vidore/colpali"
-    model = ColPali.from_pretrained("vidore/colpaligemma-3b-mix-448-base", torch_dtype=torch.bfloat16, device_map="cuda").eval()
-    model.load_adapter(model_name)
-    processor = AutoProcessor.from_pretrained(model_name)
+    model = ColFlor.from_pretrained(
+        model_name,
+        #torch_dtype=torch.bfloat16,
+        device_map=device,
+    ).eval()
+
+    # Load processor
+    processor = cast(ColFlorProcessor, ColFlorProcessor.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?"]
+    if not isinstance(processor, BaseVisualRetrieverProcessor):
+        raise ValueError("Processor should be a BaseVisualRetrieverProcessor")
 
-    # run inference - docs
+    # NOTE: Only the first 16 images are used for demonstration purposes
+    dataset = cast(Dataset, load_dataset("vidore/docvqa_test_subsampled", split="test[:16]"))
+    images = dataset["image"]
+
+    # Select a few queries for demonstration purposes
+    query_indices = [12, 15]
+    queries = [dataset[idx]["query"] for idx in query_indices]
+    print("Selected queries:")
+    pprint.pprint(dict(zip(query_indices, queries)))
+
+    # Run inference - docs
     dataloader = DataLoader(
-        images,
+        dataset=ListDataset[str](images),
         batch_size=4,
         shuffle=False,
-        collate_fn=lambda x: process_images(processor, x),
+        collate_fn=lambda x: processor.process_images(x),
     )
-    ds = []
+    ds: List[torch.Tensor] = []
     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
+    # Run inference - queries
     dataloader = DataLoader(
-        queries,
+        dataset=ListDataset[str](queries),
         batch_size=4,
         shuffle=False,
-        collate_fn=lambda x: process_queries(processor, x, Image.new("RGB", (448, 448), (255, 255, 255))),
+        collate_fn=lambda x: processor.process_queries(x),
     )
 
-    qs = []
+    qs: List[torch.Tensor] = []
     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))
+    # Run scoring
+    scores = processor.score(qs, ds).cpu().numpy()
+    idx_top_1 = scores.argmax(axis=1)
+    print("Indices of the top-1 retrieved documents for each query:", idx_top_1)
+
+    # Sanity check
+    if idx_top_1.tolist() == query_indices:
+        print("The top-1 retrieved documents are correct.")
+    else:
+        print("The top-1 retrieved documents are incorrect.")
+
+    return
 
 
 if __name__ == "__main__":