Récupération de documents visuels

Les documents peuvent contenir des données multimodales s’ils incluent des graphiques, des tableaux et des visuels en plus du texte. La récupération d’informations à partir de ces documents est complexe, car les modèles de recherche de texte seuls ne peuvent pas traiter les données visuelles, tandis que les modèles de recherche d’images manquent de granularité et de capacités de traitement de documents.

La récupération de documents visuels permet de récupérer des informations à partir de tous les types de documents, y compris la génération augmentée par récupération multimodale (RAG). Ces modèles acceptent des documents (sous forme d’images) et des textes, et calculent les scores de similarité entre eux.

Ce guide montre comment indexer et récupérer des documents avec ColPali.

Pour les cas d’usage à grande échelle, vous pouvez indexer et récupérer des documents à l’aide d’une base de données vectorielle.

Vérifiez que Transformers et Datasets sont installés.

pip install -q datasets transformers

Nous indexerons un jeu de données de documents liés à des observations de soucoupes volantes. Nous filtrons les exemples pour lesquels notre colonne d’intérêt est manquante. Il contient plusieurs colonnes ; nous nous intéressons à la colonne specific_detail_query, qui contient un résumé court du document, et à la colonne image, qui contient nos documents.

from datasets import load_dataset

dataset = load_dataset("davanstrien/ufo-ColPali")
dataset = dataset["train"]
dataset = dataset.filter(lambda example: example["specific_detail_query"] is not None)
print(dataset)

Dataset({
    features: ['image', 'raw_queries', 'broad_topical_query', 'broad_topical_explanation', 'specific_detail_query', 'specific_detail_explanation', 'visual_element_query', 'visual_element_explanation', 'parsed_into_json'],
    num_rows: 2172
})

Chargons le modèle et le tokenizer.

import torch
from transformers import ColPaliForRetrieval, ColPaliProcessor
from accelerate import Accelerator

device = Accelerator().device

model_name = "vidore/colpali-v1.2-hf"

processor = ColPaliProcessor.from_pretrained(model_name)

model = ColPaliForRetrieval.from_pretrained(
    model_name,
    dtype=torch.bfloat16,
    device_map="auto",
).eval()

Pass the text query to the processor and return the indexed text embeddings from the model. For image-to-text search, replace the text parameter in ColPaliProcessor with the images parameter to pass images.

inputs = processor(text="a document about Mars expedition").to(model.device)
with torch.no_grad():
  text_embeds = model(**inputs, return_tensors="pt").embeddings

Indexez les images hors ligne, puis, lors de l'inférence, renvoyez les embeddings du texte de requête pour obtenir les embeddings d'images les plus proches.

Enregistrez l’image et les embeddings d’image en les écrivant dans le jeu de données à l’aide de map, comme indiqué ci-dessous. Ajoutez une colonne embeddings contenant les embeddings indexés. Les embeddings ColPali occupent beaucoup d’espace de stockage, il est donc recommandé de les retirer de l’accélérateur et de les stocker sur le CPU sous forme de vecteurs NumPy.

ds_with_embeddings = dataset.map(lambda example: {'embeddings': model(**processor(images=example["image"]).to(device), return_tensors="pt").embeddings.to(torch.float32).detach().cpu().numpy()})

Pour l'inférence en ligne, créez une fonction permettant de rechercher les embeddings d'images par lots et de récupérer les k images les plus pertinentes. La fonction ci-dessous renvoie les indices dans le jeu de données ainsi que leurs scores, pour un jeu de données indexé donné, des embeddings de texte, un nombre de résultats supérieurs et une taille de lot.

def find_top_k_indices_batched(dataset, text_embedding, processor, k=10, batch_size=4):
    scores_and_indices = []

    for start_idx in range(0, len(dataset), batch_size):

        end_idx = min(start_idx + batch_size, len(dataset))
        batch = dataset[start_idx:end_idx]        
        batch_embeddings = [torch.tensor(emb[0], dtype=torch.float32) for emb in batch["embeddings"]]
        scores = processor.score_retrieval(text_embedding.to("cpu").to(torch.float32), batch_embeddings)

        if hasattr(scores, "tolist"):
            scores = scores.tolist()[0]

        for i, score in enumerate(scores):
            scores_and_indices.append((score, start_idx + i))

    sorted_results = sorted(scores_and_indices, key=lambda x: -x[0])

    topk = sorted_results[:k]
    indices = [idx for _, idx in topk]
    scores = [score for score, _ in topk]

    return indices, scores

Générez les embeddings de texte et transmettez-les à la fonction ci-dessus pour obtenir les indices du jeu de données et les scores.

with torch.no_grad():
  text_embeds = model(**processor(text="a document about Mars expedition").to(model.device), return_tensors="pt").embeddings
indices, scores = find_top_k_indices_batched(ds_with_embeddings, text_embeds, processor, k=3, batch_size=4)
print(indices, scores)

([440, 442, 443],
 [14.370786666870117,
  13.675487518310547,
  12.9899320602417])

Affichez les images pour consulter les documents liés à Mars.

for i in indices:
  display(dataset[i]["image"])

Update on GitHub