Perolehan dokumen visual

Dokumen dapat berisi data multimodal jika mencakup grafik, tabel, dan visual selain teks. Mengambil informasi dari dokumen-dokumen ini menantang karena model pencarian teks saja tidak dapat menangani data visual, dan model pencarian gambar tidak memiliki granularitas serta kemampuan pemrosesan dokumen yang memadai.

Visual document retrieval dapat membantu mengambil informasi dari berbagai jenis dokumen, termasuk multimodal retrieval augmented generation (RAG). Model-model ini menerima dokumen (sebagai gambar) dan teks, serta menghitung skor kemiripan di antara keduanya.

Panduan ini menunjukkan cara mengindeks dan mengambil dokumen dengan ColPali.

Untuk kasus penggunaan berskala besar, Anda mungkin ingin mengindeks dan mengambil dokumen dengan database vektor.

Pastikan Transformers dan Datasets sudah terpasang.

pip install -q datasets transformers

Kami akan mengindeks dataset dokumen yang berkaitan dengan penampakan UFO. Kami memfilter contoh di mana kolom yang menjadi perhatian kami tidak ada. Dataset ini berisi beberapa kolom; kami tertarik pada kolom specific_detail_query yang berisi ringkasan singkat dari dokumen, serta kolom image yang berisi dokumen kami.

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

Mari kita muat model dan 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()

Lewatkan kueri teks ke processor dan kembalikan embedding teks terindeks dari model. Untuk pencarian dari gambar ke teks, ganti parameter text di ColPaliProcessor dengan parameter images untuk melewatkan gambar.

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

Indeks gambar secara offline, dan selama inferensi, kembalikan embedding teks query untuk mendapatkan embedding gambar terdekatnya.

Simpan gambar dan embedding gambar dengan menuliskannya ke dataset menggunakan map seperti yang ditunjukkan di bawah ini. Tambahkan kolom embeddings yang berisi embedding terindeks. Embedding ColPali membutuhkan banyak penyimpanan, jadi pindahkan dari accelerator dan simpan di CPU sebagai vektor 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()})

Untuk inferensi online, buatlah fungsi untuk mencari embedding gambar dalam batch dan mengambil k gambar paling relevan. Fungsi di bawah ini mengembalikan indeks dalam dataset beserta skor mereka untuk dataset terindeks tertentu, embedding teks, jumlah hasil teratas, dan ukuran batch.

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

Buat *text embeddings* dan teruskan ke fungsi di atas untuk mengembalikan indeks dataset dan skor.

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

Tampilkan gambar untuk melihat dokumen yang terkait dengan Mars.

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

Update on GitHub