Recuperación de documentos visuales

Los documentos pueden contener datos multimodales si incluyen gráficos, tablas y elementos visuales además de texto. Recuperar información de estos documentos es un desafío, ya que los modelos de recuperación de texto por sí solos no pueden manejar datos visuales y los modelos de recuperación de imágenes carecen de la granularidad y las capacidades de procesamiento de documentos necesarias.

La recuperación de documentos visuales puede ayudar a obtener información de todo tipo de documentos, incluida la generación aumentada por recuperación multimodal (RAG). Estos modelos aceptan documentos (como imágenes) y textos, y calculan las puntuaciones de similitud entre ellos.

Esta guía muestra cómo indexar y recuperar documentos con ColPali.

Para casos de uso a gran escala, es posible que desee indexar y recuperar documentos con una base de datos vectorial.

Asegúrate de que Transformers y Datasets estén instalados.

pip install -q datasets transformers

Indexaremos un conjunto de datos de documentos relacionados con avistamientos de OVNIs. Filtramos los ejemplos en los que falta nuestra columna de interés. El conjunto de datos contiene varias columnas; nos interesa la columna specific_detail_query, que contiene un resumen breve del documento, y la columna image, que contiene nuestros 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
})

Carguemos el modelo y el tokenizador.

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

Pasa la consulta de texto al procesador y devuelve los embeddings de texto indexados del modelo. Para la búsqueda de imagen a texto, reemplaza el parámetro text en ColPaliProcessor con el parámetro images para pasar imágenes.

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

Indexa las imágenes de forma offline y, durante la inferencia, devuelve los embeddings del texto de consulta para obtener los embeddings de imagen más cercanos.

Guarde la imagen y los embeddings de la imagen escribiéndolos en el conjunto de datos con map, como se muestra a continuación. Añada una columna embeddings que contenga los embeddings indexados. Los embeddings de ColPali ocupan mucho almacenamiento, por lo que elimínelos del acelerador y guárdelos en la CPU como vectores 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 la inferencia en línea, cree una función para buscar los embeddings de imágenes por lotes y recuperar las k imágenes más relevantes. La función siguiente devuelve los índices en el conjunto de datos y sus puntuaciones para un conjunto de datos indexado dado, embeddings de texto, número de resultados principales y tamaño del 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

Genere los embeddings de texto y páselos a la función anterior para devolver los índices del conjunto de datos y las puntuaciones.

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

Muestre las imágenes para ver los documentos relacionados con Marte.

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