Suche mit EmbeddingLists: ColBERT und ColPali

Dieses Tutorial zeigt, wie Sie mit der EmbeddingList-Suche auf StructArray-Vektor-Unterfeldern in Milvus Retrieval-Systeme im ColBERT- und ColPali-Stil erstellen können. Verwenden Sie diese Methode, wenn sowohl Ihre Abfrage als auch die gespeicherten Daten als Vektorlisten dargestellt sind und Sie eine Retrieval-Methode mit später Interaktion auf Entitätsebene unter Verwendung von „ MAX_SIM* “-Metriken wünschen.

Die Grundlagen zu StructArray, auf denen dieses Tutorial basiert, finden Sie unter „Erstellen eines StructArray-Feldes“, „Indizieren von StructArray-Feldern“ und „Grundlegende Vektorsuche mit StructArray“. Dieses Tutorial konzentriert sich auf die ColBERT- und ColPali-Workflows und nicht auf die allgemeine StructArray-Syntax.

Übersicht

Um ein Textabrufsystem aufzubauen, müssen Sie Dokumente möglicherweise in Blöcke aufteilen und jeden Block zusammen mit seinen Einbettungen als Entität in einer Vektordatenbank speichern, um Präzision und Genauigkeit zu gewährleisten – insbesondere bei langen Dokumenten, bei denen Volltext-Einbettungen die semantische Spezifität verwässern oder die Eingabegrenzen des Modells überschreiten könnten.

Die Speicherung von Daten in Blöcken führt jedoch zu blockweisen Suchergebnissen, was bedeutet, dass bei der Suche zunächst relevante Segmente und nicht zusammenhängende Dokumente identifiziert werden. Um dies zu beheben, sollten Sie eine zusätzliche Nachbearbeitung der Suchergebnisse durchführen.

ColBERT (arXiv: 2004.12832) ist ein Text-Text-Retrieval-System, das durch kontextbezogene späte Interaktionen über BERT eine effiziente und effektive Passagensuche ermöglicht. Es ermöglicht eine unabhängige tokenweise Kodierung von Suchanfragen und Dokumenten und berechnet deren Ähnlichkeit.

Token-basierte Kodierung

Bei der Datenaufnahme in ColBERT wird jedes Dokument in Token zerlegt, die anschließend vektorisiert und als Einbettungsliste gespeichert werden, wie in $∈Rn×dd \rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ E d = [ $e$ d1 $,{}^{}e$ d2 $,\dots ,e$ dn $]$ $\in$ R ${}^{n\times d}$. Wenn eine Abfrage eingeht, wird sie ebenfalls tokenisiert, vektorisiert und als Einbettungsliste gespeichert, wie in $∈Rm×dq \rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ $\to$ ${}^E$ q $=$ [e q1 $,{}^e$ q2 $,\dots ,{}^e$ qm $]\in$ ${}^{Rm\times d}$.

In den obigen Formeln gilt:

• $dd$: ein Dokument

• $qq$: die Abfrage

• $EdE_d$: die Einbettungsliste, die das Dokument repräsentiert.

• $EqE_q$: die Einbettungsliste, die die Abfrage repräsentiert.

• $) [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$$,e_{}$,…,unddn$]\in$ R${}^{n\times d}$: Die Anzahl der Vektoreinbettungen in der Einbettungsliste, die das Dokument repräsentieren, liegt im Bereich von $\R^{n×d}$.

• $[e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$q1,eq2,…,eqm$]\in$ R${}^{m\times d}$: Die Anzahl der Vektor-Einbettungen in der Einbettungsliste, die die Abfrage repräsentieren, liegt im Bereich von $\R^{m×d}$.

Späte Interaktion

Sobald die Vektorisierung abgeschlossen ist, wird die Abfrage-Einbettungsliste Token für Token mit jeder Dokument-Einbettungsliste verglichen, um den endgültigen Ähnlichkeitswert zu ermitteln.

Wie in der obigen Abbildung dargestellt, enthält die Abfrage zwei Token, nämlich „ machine “ und „ learning “, und das Dokument im Fenster enthält vier Token: „ neural “, „ network “, „ python “ und „ tutorial “. Sobald diese Token vektorisiert sind, werden die Vektor-Einbettungen jedes Abfrage-Tokens mit denen im Dokument verglichen, um eine Liste von Ähnlichkeitswerten zu erhalten. Anschließend werden die höchsten Werte aus jeder Werteliste addiert, um den Endwert zu ermitteln. Der Prozess zur Ermittlung des Endwerts eines Dokuments wird als „maximale Ähnlichkeit“ (MAX_SIM) bezeichnet. Weitere Informationen zur maximalen Ähnlichkeit finden Sie unter „Maximale Ähnlichkeit“.

ColPali-Erweiterung

Basierend auf ColBERT schlägt ColPali (arXiv: 2407.01449) einen neuartigen Ansatz für die visuell reichhaltige Dokumentensuche vor, der Vision-Language-Modelle (VLMs) nutzt. Bei der Datenaufnahme wird jede Dokumentseite in ein hochauflösendes Bild gerendert und anschließend in Patches aufgeteilt, anstatt tokenisiert zu werden. So kann beispielsweise ein Dokumentseitenbild mit einer Größe von 448 × 448 Pixel 1.024 Patches erzeugen, die jeweils 14 × 14 Pixel groß sind.

Diese Methode bewahrt nicht-textuelle Informationen wie das Dokumentlayout, Bilder und Tabellenstrukturen, die bei der Verwendung von reinen Textsuchsystemen verloren gehen.

Das in ColPali verwendete VLM heißt PaliGemma (arXiv: 2407.07726) und besteht aus einem Bild-Encoder (SigLIP-400M), einem reinen Decoder-Sprachmodell (Gemma2-2B) sowie eine lineare Schicht, die die Ausgabe des Bild-Encoders in den Vektorraum des Sprachmodells projiziert, wie in der obigen Abbildung dargestellt.

${\mathrm{Bei}}_{}$ der Datenaufnahme wird eine Dokumentseite, die als Rohbild dargestellt wird, in mehrere visuelle Bereiche unterteilt, von denen jeder eingebettet wird, um eine Liste von Vektoreinbettungen zu erzeugen. Anschließend werden diese in den Vektorraum des Sprachmodells projiziert, um die endgültige Einbettungsliste zu erhalten, wie in $\rightarrow E_d = [e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$ E d = [ $e$ d1 $,e$ d2 $,\dots ,e$ dn $]$ $\in$ R ${}^{n\times d}$. Wenn eine Anfrage eingeht, wird sie tokenisiert, und jedes Token wird eingebettet, um eine Liste von Vektor-Einbettungen zu erzeugen, wie in $∈Rm×dq \rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ $\to$ E q $=$ [e q1 $,e$ q2 $,\dots ,e$ qm $]\in$ R ${}^{m\times d}$. Anschließend wurde MAX_SIM angewendet, um die beiden Einbettungslisten zu vergleichen und die endgültige Punktzahl zwischen der Suchanfrage und der Dokumentseite zu ermitteln.

ColBERT-Textabrufsystem

In diesem Abschnitt richten wir ein ColBERT-Textabrufsystem unter Verwendung von StructArray ein. Richten Sie zuvor eine Milvus v2.6.x-Instanz ein und beschaffen Sie sich ein Cohere-Zugriffstoken.

Schritt 1: Installieren der Abhängigkeiten

Führen Sie den folgenden Befehl aus, um die Abhängigkeiten zu installieren.

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

Schritt 2: Laden des Cohere-Datensatzes

In diesem Beispiel verwenden wir den Wikipedia-Datensatz von Cohere und rufen die ersten 10.000 Datensätze ab. Informationen zu diesem Datensatz finden Sie auf dieser Seite.

from datasets import load_dataset

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

Durch Ausführen der oben genannten Skripte wird der Datensatz heruntergeladen, falls er lokal nicht verfügbar ist. Jeder Datensatz im Datensatz ist ein Absatz aus einer Wikipedia-Seite. Die folgende Tabelle zeigt die Struktur dieses Datensatzes.

Spaltenname	Beschreibung
_id	A Datensatz-ID
url	Die URL des aktuellen Datensatzes.
title	Der Titel des Quelldokuments.
text	Ein Absatz aus dem Quelldokument.
emb	Einbettungen des Textes aus dem Quelldokument.

Schritt 3: Absätze nach Titel gruppieren

Um nach Dokumenten statt nach Absätzen zu suchen, sollten wir die Absätze nach Titel gruppieren.

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

In diesem Code speichern wir die gruppierten Absätze als Dokumente und nehmen sie in die Liste „ data “ auf. Jedes Dokument verfügt über einen Schlüssel „ paragraphs “, bei dem es sich um eine Liste von Absätzen handelt; jedes Absatzobjekt enthält die Schlüssel „ text “ und „ emb “.

Schritt 4: Erstellen einer Sammlung für den Cohere-Datensatz

Sobald die Daten bereitstehen, erstellen wir eine Sammlung. In der Sammlung ist „ paragraphs “ ein „StructArray“-Feld. Eine allgemeine Erläuterung zu „StructArray“-Schemas finden Sie unter „Erstellen eines StructArray-Feldes“.

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
)

Schritt 5: Cohere-Datensatz in die Sammlung einfügen

Nun können wir die vorbereiteten Daten in die oben erstellte Sammlung einfügen.

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

Schritt 6: Suche innerhalb des Cohere-Datensatzes

Entsprechend dem Design von ColBERT sollte der Suchtext tokenisiert und anschließend in eine „EmbeddingList“ eingebettet werden. In diesem Schritt verwenden wir dasselbe Modell, mit dem Cohere die Einbettungen für die Absätze im Wikipedia-Datensatz generiert hat.

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

Im Code werden die Suchtexte in „ query_inputs “ in Token zerlegt und in eine Liste von Float-Vektoren eingebettet. Anschließend können Sie die „EmbeddingList“ von Milvus verwenden, um wie folgt eine Ähnlichkeitssuche durchzuführen.

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]",
    limit=10,
    output_fields=["title"]
)

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

Die Ausgabe des obigen Codes sieht in etwa wie folgt aus:

# 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

Jeder Wert für die paarweise Kosinus-Ähnlichkeit liegt im Bereich von -1 bis 1. Der endgültige „ MAX_SIM_COSINE “-Wert kann größer sein als 1, da er mehrere maximale Ähnlichkeitswerte auf Token-Ebene aggregiert.

ColPali-Dokumentensuchsystem

In diesem Abschnitt richten wir mithilfe von StructArray ein ColPali-basiertes Dokumentensuchsystem ein. Richten Sie zuvor eine Milvus-Instanz der Version 2.6.x ein.

Schritt 1: Installieren Sie die Abhängigkeiten

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

Schritt 2: Laden des Vidore-Datensatzes

In diesem Abschnitt verwenden wir einen Vidore-Datensatz namens „vidore_v2_finance_en“. Dieser Datensatz ist ein Korpus aus Jahresberichten des Bankensektors, der für Aufgaben zum Verständnis langer Dokumente vorgesehen ist. Er ist einer der 10 Korpora, aus denen sich der ViDoRe v3-Benchmark zusammensetzt. Details zu diesem Datensatz finden Sie auf dieser Seite.

from datasets import load_dataset

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

Durch Ausführen der oben genannten Skripte wird der Datensatz heruntergeladen, falls er lokal nicht verfügbar ist. Jeder Datensatz im Datensatz entspricht einer Seite aus einem Finanzbericht. Die folgende Tabelle zeigt die Struktur dieses Datensatzes.

Spaltenname	Beschreibung
corpus_id	Ein Datensatz im Korpus
image	Das Seitenbild in Byte.
doc_id	Die beschreibende Dokument-ID.
page_number_in_doc	Die Seitenzahl der aktuellen Seite im Dokument.

Schritt 3: Einbettungen für die Seitenbilder generieren

Wie im Abschnitt „Übersicht“ dargestellt, handelt es sich beim ColPali-Modell um ein VLM, das Bilder in den Vektorraum eines Textmodells projiziert. In diesem Schritt verwenden wir das neueste ColPali-Modell vidore/colpali-v1.3. Details zu diesem Modell finden Sie auf dieser Seite.

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)

Sobald das Modell bereit ist, können Sie wie folgt versuchen, Patches für ein bestimmtes Bild zu generieren.

from PIL import Image
from io import BytesIO

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

# Decode the image bytes and generate patch embeddings.
images = [Image.open(BytesIO(row["image"]["bytes"]))]
batch_images = processor.process_images(images).to(model.device)

with torch.no_grad():
    patches_embeddings = model(**batch_images)[0]

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

# [1031, 128]

Im obigen Code skaliert das ColPali-Modell das Bild auf 448 × 448 Pixel und teilt es anschließend in Patches mit einer Größe von jeweils 14 × 14 Pixel auf. Zuletzt werden diese Patches in 1.031 Embeddings mit jeweils 128 Dimensionen eingebettet.

Sie können Einbettungen für alle Bilder mithilfe einer Schleife wie folgt generieren:

data = []

for _, row in df.iterrows():
    corpus_id = row["corpus_id"]
    images = [Image.open(BytesIO(row["image"]["bytes"]))]
    batch_images = processor.process_images(images).to(model.device)

    with torch.no_grad():
        patches = model(**batch_images)[0]

    doc_id = row["doc_id"]
    page_number_in_doc = row["page_number_in_doc"]

    data.append({
        "corpus_id": corpus_id,
        "patches": [
            {"emb": emb.float().cpu().tolist()}
            for emb in patches
        ],
        "doc_id": doc_id,
        "page_number_in_doc": page_number_in_doc,
    })

Schritt 4: Erstellen einer Sammlung für den Datensatz der Finanzberichte

Sobald die Daten bereit sind, erstellen wir eine Sammlung. In der Sammlung ist „ patches “ ein StructArray-Feld. Jedes Struct-Element speichert eine Patch-Einbettung. Informationen zu den Indizierungsanforderungen für StructArray-Vektor-Unterfelder finden Sie unter „StructArray-Felder indizieren“.

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
)

Schritt 5: Einfügen der Finanzberichte in die Sammlung

Nun können wir die vorbereiteten Finanzberichte in die Sammlung einfügen.

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

An der Ausgabe können Sie erkennen, dass alle Seiten aus dem Vidore-Datensatz eingefügt wurden.

Schritt 6: Suche in den Finanzberichten

Sobald die Daten bereit sind, können wir wie folgt Suchvorgänge in den Daten der Sammlung durchführen:

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].float().cpu().tolist())

results = client.search(
    collection_name="financial_reports",
    data=[query_emb_list],
    anns_field="patches[emb]",
    limit=10,
    output_fields=["doc_id", "page_number_in_doc"]
)