Transformers documentation
Визуальный поиск документов
Визуальный поиск документов
Документы могут содержать мультимодальные данные, если помимо текста в них присутствуют диаграммы, таблицы и визуальные элементы. Извлечение информации из таких документов представляет собой сложную задачу, поскольку модели текстового поиска не способны обрабатывать визуальные данные, а модели поиска изображений не обладают необходимой детализацией и возможностями обработки документов.
Визуальный поиск документов позволяет извлекать информацию из документов любого типа, включая мультимодальную генерацию с расширением контекста (RAG). Эти модели принимают документы (в виде изображений) и тексты и вычисляют оценки их сходства.
В этом руководстве показано, как индексировать и извлекать документы с помощью ColPali.
Для задач крупного масштаба может потребоваться индексировать и извлекать документы с помощью векторной базы данных.
Убедитесь, что установлены Transformers и Datasets.
pip install -q datasets transformers
Мы будем индексировать набор данных, содержащий документы о наблюдениях НЛО. Мы отфильтруем примеры, в которых отсутствует интересующий нас столбец. Набор данных содержит несколько столбцов; нас интересует столбец specific_detail_query, в котором хранится краткое описание документа, и столбец image, содержащий сами документы.
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
})Загрузим модель и токенизатор.
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()Передайте текстовый запрос в процессор и получите индексированные текстовые эмбеддинги из модели. Для поиска по изображениям замените параметр text в ColPaliProcessor на параметр images, чтобы передать изображения.
inputs = processor(text="a document about Mars expedition").to(model.device)
with torch.no_grad():
text_embeds = model(**inputs, return_tensors="pt").embeddingsИндексируйте изображения офлайн, а во время инференса возвращайте эмбеддинги текстового запроса, чтобы получить ближайшие эмбеддинги изображений.
Сохраните изображения и эмбеддинги изображений, записав их в датасет с помощью map, как показано ниже. Добавьте столбец embeddings, содержащий индексированные эмбеддинги. Эмбеддинги ColPali занимают много места, поэтому удалите их с акселератора и сохраните на CPU в виде векторов 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()})Для онлайн-инференса создайте функцию, которая ищет эмбеддинги изображений партиями и извлекает k наиболее релевантных изображений. Функция ниже возвращает индексы в наборе данных и их оценки для заданного индексированного набора данных, эмбеддингов текста, количества верхних результатов и размера партии.
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Сгенерируйте текстовые эмбеддинги и передайте их в функцию выше, чтобы получить индексы датасета и оценки.
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])
Отобразите изображения, чтобы просмотреть документы, связанные с Марсом.
for i in indices:
display(dataset[i]["image"])
