Поиск с помощью вложенных списков

На этой странице объясняется, как настроить систему поиска текста ColBERT и систему поиска текста ColPali с помощью массива структур в Milvus, который позволяет хранить документ вместе с его векторизованными фрагментами в списках встраивания.

Обзор

Для создания системы текстового поиска вам может понадобиться разбить документы на фрагменты и хранить каждый фрагмент вместе с его вкраплениями как сущность в векторной базе данных, чтобы обеспечить точность и достоверность, особенно для длинных документов, где полнотекстовые вкрапления могут размыть семантическую специфику или превысить ограничения на входные данные модели.

Однако хранение данных в виде фрагментов приводит к появлению результатов поиска по фрагментам, то есть поиск изначально определяет релевантные сегменты, а не цельные документы. Чтобы решить эту проблему, необходимо выполнить дополнительную обработку после поиска.

ColBERT (arXiv: 2004.12832) - это система текстового поиска, которая предлагает эффективный и результативный поиск отрывков через контекстуализированное позднее взаимодействие через BERT. Она позволяет независимо кодировать запросы и документы с помощью токенов и вычислять их сходство.

Кодирование с помощью токенов

При вводе данных в ColBERT каждый документ разбивается на токены, которые затем векторизуются и хранятся в виде списка вложений, как в $\rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ d $\to$ E d = $[e$ d1, e d2, $...$, e dn] $\in$ R ${}^{n\times d}$. Когда поступает запрос, он также токенизируется, векторизуется и хранится в виде списка вложений, как в $\rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ q $\to$ E q = $[e$ q1, e q2, $...$, e qm] $\in$ R ${}^{m\times d}$.

В приведенных выше формулах,

• $dd$ d: документ

• $qq$ q: запрос

• $EdE_d$ Ed: список вкраплений, представляющий документ.

• $EqE_q$ Eq: список вкраплений, представляющий запрос.

• $[e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$d1, ed2,$...$, edn]$\in$ R${}^{n\times d}$: количество векторных вкраплений в списке вкраплений, представляющих документ, находится в диапазоне $Rn×d\R^{n×d}$ R ${}^{n\times d}$.

• $[e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$q1, eq2,$...$, eqm]$\in$ R${}^{m\times d}$: количество векторных вкраплений в списке вкраплений, представляющих запрос, находится в диапазоне $Rm×d\R^{m×d}$ R ${}^{m\times d}$.

Позднее взаимодействие

После завершения векторизации список вкраплений запроса сравнивается со списком вкраплений каждого документа, токен за токеном, чтобы определить окончательный балл сходства.

Позднее взаимодействие

Как показано на диаграмме выше, запрос содержит две лексемы, а именно machine и learning, а документ в окне имеет четыре лексемы: neural, network, python, и tutorial. После того как эти лексемы векторизованы, векторные вложения каждой лексемы запроса сравниваются с лексемами документа, чтобы получить список баллов сходства. Затем самые высокие баллы из каждого списка суммируются, чтобы получить итоговый балл. Процесс определения итоговой оценки документа известен как максимальное сходство(MAX_SIM). Подробнее о максимальном сходстве см. в разделе Максимальное сходство.

Расширение ColPali

Основанный на ColBERT, проект ColPali (arXiv: 2407.01449) предлагает новый подход к поиску визуально насыщенных документов, который использует модели видения-языка (VLM). При получении данных каждая страница документа преобразуется в изображение высокого разрешения, затем разбивается на участки, а не токенизируется. Например, изображение страницы документа размером 448 x 448 пикселей может создать 1024 патча, каждый размером 14 x 14 пикселей.

Этот метод позволяет сохранить нетекстовую информацию, такую как макет документа, изображения и структуры таблиц, которые теряются при использовании только текстовых поисковых систем.

Расширение Copali

VLM, используемый в ColPali, называется PaliGemma (arXiv: 2407.07726) и состоит из кодера изображений(SigLIP-400M), языковой модели только декодера(Gemma2-2B) и линейного слоя, который проецирует выход кодера изображений в векторное пространство языковой модели, как показано на рисунке выше.

При вводе данных страница документа, представленная в виде необработанного изображения, делится на несколько визуальных патчей, каждый из которых встраивается для создания списка векторных вложений. Затем они проецируются в векторное пространство языковой модели для получения окончательного списка вкраплений, как в $\rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ d $\to$ E d = $[e$ d1, e d2, $...$, e dn] $\in$ R ${}^{n\times d}$. Когда поступает запрос, он токенизируется, и каждый токен встраивается для создания списка векторных вкраплений, как в $\rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ q $\to$ E q = $[e$ q1, e q2, $...$, e qm] $\in$ R ${}^{m\times d}$. Затем был применен MAX_SIM для сравнения двух списков вкраплений и получения итоговой оценки между запросом и страницей документа.

Система поиска текста ColBERT

В этом разделе мы создадим систему поиска текстов ColBERT с использованием массива структур Milvus. Перед этим настройте экземпляр Milvus v2.6.xZilliz Cloud cluster, совместимый с Milvus v2.6.x, получите токен доступа Cohere.

Шаг 1: Установите зависимости

Выполните следующую команду для установки зависимостей.

pip install --upgrade huggingface-hub transformers datasets pymilvus cohere

Шаг 2: Загрузите набор данных Cohere

В этом примере мы будем использовать набор данных Cohere для Википедии и извлекать первые 10 000 записей. Информацию об этом наборе данных вы можете найти на этой странице.

from datasets import load_dataset

lang = "simple"
docs = load_dataset(
    "Cohere/wikipedia-2023-11-embed-multilingual-v3", 
    lang, 
    split="train[:10000]"
)

Запуск приведенных выше скриптов позволит загрузить набор данных, если он не доступен локально. Каждая запись в наборе данных представляет собой абзац со страницы Википедии. В следующей таблице показана структура этого набора данных.

Имя столбца	Описание
_id	Идентификатор записи
url	URL-адрес текущей записи.
title	Заголовок исходного документа.
text	Абзац из исходного документа.
emb	Вкрапления текста из исходного документа.

Шаг 3: Группировка абзацев по названию

Чтобы искать документы, а не абзацы, мы должны сгруппировать абзацы по заголовку.

df = docs.to_pandas()
groups = df.groupby('title')

data = []

for title, group in groups:
  data.append({
      "title": title,
      "paragraphs": [{
          "text": row['text'],
          'emb': row['emb']
      } for _, row in group.iterrows()]
  })

В этом коде мы храним сгруппированные абзацы как документы и включаем их в список data. Каждый документ имеет ключ paragraphs, который представляет собой список параграфов; каждый объект параграфа содержит ключи text и emb.

Шаг 4: Создание коллекции для набора данных Cohere

Когда данные будут готовы, мы создадим коллекцию. В коллекции есть поле paragraphs, которое представляет собой массив структур.

from pymilvus import MilvusClient, DataType

client = MilvusClient(
    uri="http://localhost:19530",
    token="root:Milvus"
)

# Create collection schema
schema = client.create_schema()

schema.add_field('id', DataType.INT64, is_primary=True, auto_id=True)
schema.add_field('title', DataType.VARCHAR, max_length=512)

# Create struct schema
struct_schema = client.create_struct_field_schema()
struct_schema.add_field('text', DataType.VARCHAR, max_length=65535)
struct_schema.add_field('emb', DataType.FLOAT_VECTOR, dim=512)

schema.add_field('paragraphs', DataType.ARRAY,
                 element_type=DataType.STRUCT,
                 struct_schema=struct_schema, max_capacity=200)

# Create index parameters
index_params = client.prepare_index_params()
index_params.add_index(
    field_name="paragraphs[emb]",
    index_type="AUTOINDEX",
    metric_type="MAX_SIM_COSINE"
)

# Create a collection
client.create_collection(
    collection_name='wiki_documents', 
    schema=schema, 
    index_params=index_params
)

Шаг 5: Вставка набора данных Cohere в коллекцию

Теперь мы можем вставить подготовленные данные в коллекцию, которую мы создали выше.

client.insert(
    collection_name='wiki_documents', 
    data=data
)

Шаг 6: Поиск в наборе данных Cohere

В соответствии с дизайном ColBERT, текст запроса должен быть токенизирован, а затем встроен в EmbeddingList. На этом шаге мы используем ту же модель, которую Cohere использовал для генерации вкраплений для абзацев в наборе данных Википедии.

import cohere

co = cohere.ClientV2("COHERE_API_KEY")

query_inputs = [
    {
        'content': [
            {'type': 'text', 'text': 'Adobe'},
        ]
    },
    {
        'content': [
            {'type': 'text', 'text': 'software'}
        ]
    }
]

embeddings = co.embed(
    inputs=query_inputs,
    model='embed-multilingual-v3.0',
    input_type="classification",
    embedding_types=["float"],
)

В коде тексты запросов организуются в лексемы в query_inputs и встраиваются в список векторов с плавающей точкой. Затем вы можете использовать Milvus's EmbeddingList для поиска сходства следующим образом.

from pymilvus.client.embedding_list import EmbeddingList

query_emb_list = EmbeddingList()

if (embeddings.embeddings.float):
  query_emb_list.add_batch(embeddings.embeddings.float)

results = client.search(
    collection_name="wiki_documents",
    data=[query_emb_list],
    anns_field="paragraphs[emb]",
    search_params={
        "metric_type": "MAX_SIM_COSINE"
    },
    limit=10,
    output_fields=["title"]
)

for hit in results[0]:
  print(f"Document {hit['entity']['title']}: {hit['distance']:.4f}")

Результат работы приведенного выше кода выглядит следующим образом:

# Document Software: 2.3035
# Document Application: 2.1875
# Document Adobe Illustrator: 2.1167
# Document Open source: 2.0542
# Document Computer: 1.9811
# Document Microsoft: 1.9784
# Document Web browser: 1.9655
# Document Program: 1.9627
# Document Website: 1.9594
# Document Computer science: 1.9460

Косинусная оценка сходства варьируется от -1 до 1, а оценки сходства в приведенном выше выводе наглядно демонстрируют сумму нескольких оценок сходства на уровне токенов.

Система поиска текстов ColPali

В этом разделе мы создадим систему поиска текстов на основе ColPali с использованием массива структур Milvus. Перед этим настройте экземпляр Milvus v2.6.x на облачном кластереZilliz Cloud, совместимом с Milvus v2.6.x.

Шаг 1: Установите зависимости

pip install --upgrade huggingface-hub transformers datasets pymilvus 'colpali-engine>=0.3.0,<0.4.0'

Шаг 2: Загрузка набора данных Vidore

В этом разделе мы будем использовать набор данных Vidore под названием vidore_v2_finance_en. Этот набор представляет собой корпус годовых отчетов из банковского сектора, предназначенный для задач понимания длинных документов. Он является одним из 10 корпусов, входящих в ViDoRe v3 Benchmark. Подробности об этом наборе данных можно найти на этой странице.

from datasets import load_dataset

ds = load_dataset("vidore/vidore_v3_finance_en", "corpus")
df = ds['test'].to_pandas()

Запуск приведенных выше скриптов загрузит набор данных, если он не доступен локально. Каждая запись в наборе данных представляет собой страницу из финансового отчета. В следующей таблице показана структура этого набора данных.

Имя столбца	Описание
corpus_id	Запись в корпусе
image	Изображение страницы в байтах.
doc_id	Описательный идентификатор документа.
page_number_in_doc	Номер текущей страницы в документе.

Шаг 3: Генерация вкраплений для изображений страниц

Как показано в разделе "Обзор", модель ColPali представляет собой VLM, которая проецирует изображения в векторное пространство текстовой модели. В этом шаге мы будем использовать последнюю версию модели ColPali vidore/colpali-v1.3. Подробности об этой модели вы можете найти на этой странице.

import torch
from typing import cast
from colpali_engine.models import ColPali, ColPaliProcessor

model_name = "vidore/colpali-v1.3"

model = ColPali.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="cuda:0",  # or "mps" if on Apple Silicon
).eval()

processor = ColPaliProcessor.from_pretrained(model_name)

Когда модель готова, можно попробовать сгенерировать патчи для конкретного изображения следующим образом.

from PIL import Image
from io import BytesIO

# Use the iterrow() generator to get the first row
row = next(df.iterrows())[1]

# Include the image in the above row in a list
images = [ Image.open(row['image']['bytes'] ]
patches = processor.process_images(images).to(model.device)
patches_embeddings = model(**patches_in_pixels)[0]

# Check the shape of the embeddings generated for the patches
print(patches_embeddings.shape)

# [1031, 128]

В приведенном выше коде модель ColPali изменяет размер изображения до 448 x 448 пикселей, затем делит его на патчи, каждый размером 14 x 14 пикселей. Наконец, эти патчи встраиваются в 1 031 вкрапление, каждое из которых имеет 128 размеров.

Вы можете сгенерировать вкрапления для всех изображений с помощью цикла следующим образом:

data = []

for index, row in df.iterrows():
  row = next(df.iterrows())[1]
  corpus_id = row['corpus_id']
  
  images = [Image.open(BytesIO(row['image']['bytes']))]
  batch_images = processor.process_images(images).to(model.device)
  patches = model(**batch_images)[0]

  doc_id = row['doc_id']
  markdown = row['markdown']
  page_number_in_doc = row['page_number_in_doc']

  data.append({
      "corpus_id": corpus_id,
      "patches": [ {"emb": emb} for emb in patches ],
      "doc_id": markdown,
      "page_number_in_doc": row['page_number_in_doc']
  })

Шаг 4: Создание коллекции для набора данных финансовых отчетов

Когда данные будут готовы, мы создадим коллекцию. В коллекции поле с именем patches представляет собой массив структур.

from pymilvus import MilvusClient, DataType

client = MilvusClient(
    uri=YOUR_CLUSTER_ENDPOINT,
    token=YOUR_API_KEY
)

schema = client.create_schema()

schema.add_field(
    field_name="corpus_id",
    datatype=DataType.INT64,
    is_primary=True
)

patch_schema = client.create_struct_field_schema()

patch_schema.add_field(
    field_name="emb",
    datatype=DataType.FLOAT_VECTOR,
    dim=128
)

schema.add_field(
    field_name="patches",
    datatype=DataType.ARRAY,
    element_type=DataType.STRUCT,
    struct_schema=patch_schema,
    max_capacity=1031
)

schema.add_field(
    field_name="doc_id",
    datatype=DataType.VARCHAR,
    max_length=512
)

schema.add_field(
    field_name="page_number_in_doc",
    datatype=DataType.INT64
)

index_params = client.prepare_index_params()

index_params.add_index(
    field_name="patches[emb]",
    index_type="AUTOINDEX",
    metric_type="MAX_SIM_COSINE"
)

client.create_collection(
    collection_name="financial_reports",
    schema=schema,
    index_params=index_params
)

Шаг 5: Вставка финансовых отчетов в коллекцию

Теперь мы можем вставить подготовленные финансовые отчеты в коллекцию.

client.insert(
    collection_name="financial_reports",
    data=data
)

Из вывода видно, что все страницы из набора данных Vidore вставлены.

Шаг 6: Поиск в финансовых отчетах

Когда данные готовы, мы можем выполнить поиск по данным в коллекции следующим образом:

from pymilvus.client.embedding_list import EmbeddingList

queries = [
    "quarterly revenue growth chart"
]

batch_queries = processor.process_queries(queries).to(model.device)

with torch.no_grad():
  query_embeddings = model(**batch_queries)

query_emb_list = EmbeddingList()
query_emb_list.add_batch(query_embeddings[0].cpu())

results = client.search(
    collection_name="financial_reports",
    data=[query_emb_list],
    anns_field="patches[emb]",
    search_params={
        "metric_type": "MAX_SIM_COSINE"
    },
    limit=10,
    output_fields=["doc_id", "page_number_in_doc"]
)