Tìm kiếm tài liệu trực quan

Tài liệu có thể chứa dữ liệu đa phương thức nếu chúng bao gồm biểu đồ, bảng và hình ảnh trực quan bên cạnh văn bản. Việc truy xuất thông tin từ những tài liệu này là một thách thức vì các mô hình truy xuất văn bản đơn thuần không thể xử lý dữ liệu trực quan, trong khi các mô hình truy xuất hình ảnh lại thiếu độ chi tiết và khả năng xử lý tài liệu.

Truy xuất tài liệu trực quan có thể giúp lấy thông tin từ mọi loại tài liệu, bao gồm cả truy xuất tăng cường sinh đa phương thức (RAG). Các mô hình này nhận tài liệu (dưới dạng hình ảnh) và văn bản, sau đó tính điểm tương tự giữa chúng.

Hướng dẫn này minh họa cách lập chỉ mục và truy xuất tài liệu bằng ColPali.

Đối với các trường hợp sử dụng quy mô lớn, bạn có thể muốn lập chỉ mục và truy xuất tài liệu bằng cơ sở dữ liệu vector.

Hãy đảm bảo rằng Transformers và Datasets đã được cài đặt.

pip install -q datasets transformers

Chúng ta sẽ lập chỉ mục cho một tập dữ liệu các tài liệu liên quan đến việc nhìn thấy UFO. Chúng ta lọc các ví dụ mà ở đó cột quan tâm của chúng ta bị thiếu. Tập dữ liệu này chứa một số cột, chúng ta quan tâm đến cột specific_detail_query chứa tóm tắt ngắn gọn của tài liệu, và cột image chứa các tài liệu của chúng ta.

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

Hãy tải mô hình và bộ mã hóa.

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

Truyền truy vấn văn bản vào bộ xử lý và trả về các nhúng văn bản có chỉ số từ mô hình. Đối với tìm kiếm từ ảnh sang văn bản, thay thế tham số text trong ColPaliProcessor bằng tham số images để truyền các hình ảnh.

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

Chỉ mục hóa các hình ảnh ngoại tuyến, và trong quá trình suy luận, trả về các vector nhúng văn bản truy vấn để lấy các vector nhúng hình ảnh gần nhất.

Lưu trữ ảnh và các embedding của ảnh bằng cách ghi chúng vào dataset bằng map như minh họa bên dưới. Thêm một cột embeddings chứa các embedding đã được chỉ mục. Embedding của ColPali chiếm nhiều dung lượng lưu trữ, vì vậy hãy gỡ chúng khỏi accelerator và lưu trữ trên CPU dưới dạng vector 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()})

Đối với suy luận trực tuyến, hãy tạo một hàm để tìm kiếm các nhúng ảnh theo lô và lấy ra k ảnh liên quan nhất. Hàm bên dưới trả về các chỉ số trong bộ dữ liệu và điểm số của chúng cho một bộ dữ liệu đã lập chỉ mục, các nhúng văn bản, số lượng kết quả hàng đầu và kích thước lô.

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

Tạo các vector nhúng văn bản và truyền chúng vào hàm ở trên để trả về các chỉ số tập dữ liệu và điểm số.

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

Hiển thị các hình ảnh để xem các tài liệu liên quan đến Sao Hỏa.

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

Update on GitHub