Recuperação de documentos visuais

Documentos podem conter dados multimodais se incluírem gráficos, tabelas e elementos visuais, além de texto. Recuperar informações desses documentos é desafiador, pois modelos de recuperação de texto sozinhos não conseguem lidar com dados visuais, e modelos de recuperação de imagem carecem de granularidade e capacidades de processamento de documentos.

A recuperação visual de documentos pode ajudar a recuperar informações de todos os tipos de documentos, incluindo a geração aumentada por recuperação multimodal (RAG). Esses modelos aceitam documentos (como imagens) e textos e calculam as pontuações de similaridade entre eles.

Este guia demonstra como indexar e recuperar documentos com ColPali.

Para casos de uso em larga escala, você pode querer indexar e recuperar documentos com um banco de dados vetorial.

Verifique se Transformers e Datasets estão instalados.

pip install -q datasets transformers

Vamos indexar um conjunto de dados de documentos relacionados a avistamentos de OVNIs. Filtramos os exemplos em que a coluna de interesse está ausente. Ele contém várias colunas; estamos interessados na coluna specific_detail_query, que contém um resumo curto do documento, e na coluna image, que contém nossos documentos.

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
})

Vamos carregar o modelo e o 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()

Passe a consulta de texto para o processador e retorne os embeddings de texto indexados do modelo. Para busca de imagem para texto, substitua o parâmetro text em ColPaliProcessor pelo parâmetro images para passar as imagens.

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

Indexe as imagens offline e, durante a inferência, retorne os embeddings do texto da consulta para obter os embeddings de imagem mais próximos.

Armazene a imagem e os embeddings da imagem gravando-os no dataset com map, conforme mostrado abaixo. Adicione uma coluna embeddings que contenha os embeddings indexados. Os embeddings do ColPali ocupam bastante armazenamento, portanto, remova-os do acelerador e armazene-os na CPU como vetores 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()})

Para inferência online, crie uma função para buscar os embeddings de imagem em lotes e recuperar as k imagens mais relevantes. A função abaixo retorna os índices no conjunto de dados e suas pontuações para um conjunto de dados indexado dado, embeddings de texto, número de resultados principais e o tamanho do lote.

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

Gere os embeddings de texto e passe-os para a função acima para retornar os índices do conjunto de dados e as pontuações.

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])

Exiba as imagens para visualizar os documentos relacionados a Marte.

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

Update on GitHub