Suche mit eingebetteten Listen

Diese Seite erklärt, wie man ein ColBERT-Textrecherchesystem und ein ColPali-Textrecherchesystem einrichtet, indem man das Array of Structs in Milvus verwendet, das es ermöglicht, ein Dokument zusammen mit seinen vektorisierten Chunks in Einbettungslisten zu speichern.

Überblick

Um ein Text-Retrieval-System aufzubauen, müssen Sie möglicherweise Dokumente in Chunks aufteilen und jeden Chunk 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 Grenzen der Modelleingabe überschreiten könnten.

Die Speicherung von Daten in Chunks führt jedoch zu chunk-artigen Suchergebnissen, d. h. die Suche identifiziert zunächst relevante Segmente und nicht zusammenhängende Dokumente. Um dies zu beheben, sollten Sie eine zusätzliche Verarbeitung nach der Suche durchführen.

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

Token-weise Kodierung

Während der Dateneingabe in ColBERT wird jedes Dokument in Token aufgeteilt, die dann vektorisiert und als Einbettungsliste gespeichert werden, wie in $\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}$. Wenn eine Anfrage eintrifft, wird sie auch tokenisiert, vektorisiert und als Einbettungsliste gespeichert, wie in $\rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ q $\to$ E q = $[e$ q1, e q2, $...$, e qm ]∈ R ${}^{m\times d}$.

In den obigen Formeln,

• $dd$ d: ein Dokument

• $qq$ q: die Abfrage

• $EdE_d$ Ed: die Einbettungsliste, die das Dokument darstellt.

• $EqE_q$ Eq: die Einbettungsliste, die die Abfrage darstellt.

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

• $[e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$q1, eq2,$...$, eqm]∈ R${}^{m\times d}$: Die Anzahl der Vektoreinbettungen in der Einbettungsliste, die die Anfrage repräsentiert, liegt im Bereich von $Rm×d\R^{m×d}$ R ${}^{m\times d}$.

Späte Interaktion

Nach Abschluss der Vektorisierung wird die Einbettungsliste der Abfrage mit jeder Einbettungsliste der Dokumente verglichen, Token für Token, um den endgültigen Ähnlichkeitswert zu ermitteln.

Späte Interaktion

Wie im obigen Diagramm dargestellt, enthält die Abfrage zwei Token, nämlich machine und learning, und das Dokument im Fenster hat vier Token: neural, network, python und tutorial. Sobald diese Token vektorisiert sind, werden die Vektoreinbettungen jedes Abfrage-Tokens mit denen im Dokument verglichen, um eine Liste von Ähnlichkeitswerten zu erhalten. Anschließend werden die höchsten Punktzahlen aus jeder Liste addiert, um die endgültige Punktzahl zu ermitteln. Das Verfahren zur Ermittlung der Endpunktzahl eines Dokuments wird als maximale Ähnlichkeit(MAX_SIM) bezeichnet. Einzelheiten zur maximalen Ähnlichkeit finden Sie unter Maximale Ähnlichkeit.

ColPali-Erweiterung

Auf der Grundlage von ColBERT schlägt ColPali (arXiv: 2407.01449) einen neuartigen Ansatz für die Suche nach visuell reichhaltigen Dokumenten vor, der auf Vision-Language-Modellen (VLMs) basiert. Während der Dateneingabe wird jede Dokumentseite in ein hochauflösendes Bild gerendert und dann in Patches aufgeteilt, anstatt sie in Token zu zerlegen. Ein Bild einer Dokumentseite mit 448 x 448 Pixeln kann beispielsweise 1.024 Patches mit einer Größe von jeweils 14 x 14 Pixeln erzeugen.

Bei dieser Methode bleiben nicht-textliche Informationen wie das Dokumentenlayout, Bilder und Tabellenstrukturen erhalten, die bei der Verwendung reiner Textsuchsysteme verloren gehen.

Copali-Erweiterung

Das in ColPali verwendete VLM heißt PaliGemma (arXiv: 2407.07726) und besteht aus einem Bildcodierer(SigLIP-400M), einem reinen Decoder-Sprachmodell(Gemma2-2B) und einer linearen Schicht, die die Ausgabe des Bildcodierers in den Vektorraum des Sprachmodells projiziert, wie im obigen Diagramm dargestellt.

Während der Dateneingabe wird eine Dokumentenseite, 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 sie 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}$ d $\to$ E d = $[e$ d1, e d2, $...$, e dn] $\in$ R ${}^{n\times d}$. Wenn eine Anfrage eintrifft, wird sie in Token umgewandelt, und jedes Token wird eingebettet, um eine Liste von Vektoreinbettungen zu erzeugen, wie in $\rightarrow E_q = [e_{q1}, e_{q2}, \dots, e_{qm}] ∈ \R^{m×d}$ q $\to$ E q = $[e$ q1, e q2, $...$, e qm ]∈ R ${}^{m\times d}$. Dann wurde MAX_SIM angewendet, um die beiden Einbettungslisten zu vergleichen und die endgültige Punktzahl zwischen der Abfrage und der Dokumentenseite zu erhalten.

ColBERT-Textabfragesystem

In diesem Abschnitt werden wir ein ColBERT Textsuchsystem mit Milvus' Array of Structs einrichten. Zuvor müssen Sie eine Milvus v2.6.x-Instanz einrichten, einen mit Milvus v2.6.x kompatiblen Zilliz-Cloud-Cluster einrichten und ein Cohere-Zugriffstoken erhalten.

Schritt 1: Installieren Sie die 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 Sie den Cohere-Datensatz

In diesem Beispiel werden wir den Wikipedia-Datensatz von Cohere verwenden und die ersten 10.000 Datensätze abrufen. 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]"
)

Wenn Sie die obigen Skripte ausführen, wird der Datensatz heruntergeladen, wenn er nicht lokal verfügbar ist. Jeder Eintrag im Datensatz ist ein Absatz einer Wikipedia-Seite. Die folgende Tabelle zeigt die Struktur dieses Datensatzes.

Spalte Name	Beschreibung
_id	Eine 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 und nicht nach Absätzen zu suchen, sollten wir die Absätze nach dem 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 hat einen Schlüssel paragraphs, der eine Liste von Absätzen ist; jedes Absatzobjekt enthält einen Schlüssel text und emb.

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

Sobald die Daten fertig sind, erstellen wir eine Sammlung. In der Sammlung gibt es ein Feld namens paragraphs, das ein Array von Structs ist.

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: Einfügen des Cohere-Datensatzes in die Sammlung

Jetzt 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

Gemäß dem Design von ColBERT sollte der Abfragetext tokenisiert und dann in eine EmbeddingList eingebettet werden. In diesem Schritt verwenden wir dasselbe Modell, das Cohere zur Generierung von Einbettungen für die Absätze im Wikipedia-Datensatz verwendet.

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 Abfragetexte in query_inputs in Token organisiert und in eine Liste von Float-Vektoren eingebettet. Dann können Sie Milvus' EmbeddingList verwenden, um eine Ähnlichkeitssuche wie folgt 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]",
    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}")

Die Ausgabe des obigen Codes ist ähnlich wie die folgende:

# 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

Der Kosinus-Ähnlichkeitswert reicht von -1 bis 1, und die Ähnlichkeitswerte in der obigen Ausgabe zeigen deutlich die Summe mehrerer Ähnlichkeitswerte auf Token-Ebene.

ColPali Text Retrieval System

In diesem Abschnitt werden wir ein ColPali-basiertes Textsuchsystem einrichten, das Milvus' Array of Structs verwendet. Zuvor müssen Sie eine Instanz von Milvus v2.6.x einrichten. Der Zilliz-Cluster ist mit Milvus v2.6.x kompatibel.

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 Sie den Vidore-Datensatz

In diesem Abschnitt werden wir einen Vidore-Datensatz namens vidore_v2_finance_de verwenden. Bei diesem Datensatz handelt es sich um ein Korpus von Geschäftsberichten aus dem Bankensektor, das für Aufgaben zum Verstehen langer Dokumente gedacht ist. Es ist eines der 10 Korpora, die den ViDoRe v3 Benchmark bilden. 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()

Wenn Sie die oben genannten Skripte ausführen, wird der Datensatz heruntergeladen, wenn er nicht lokal verfügbar ist. Jeder Datensatz im Dataset ist eine Seite aus einem Finanzbericht. Die folgende Tabelle zeigt die Struktur dieses Datasets.

Spalte Name	Beschreibung
corpus_id	Ein Datensatz im Korpus
image	Das Seitenbild in Bytes.
doc_id	Die beschreibende Dokument-ID.
page_number_in_doc	Die Seitennummer der aktuellen Seite im Dokument.

Schritt 3: Erzeugen von Einbettungen für die Seitenbilder

Wie im Abschnitt Überblick dargestellt, ist das ColPali-Modell ein VLM, das Bilder in den Vektorraum eines Textmodells projiziert. In diesem Schritt verwenden wir das neueste ColPali-Modell vidore/colpali-v1.3. Einzelheiten 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 fertig ist, können Sie wie folgt versuchen, Patches für ein bestimmtes Bild zu erzeugen.

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]

Im obigen Code ändert das ColPali-Modell die Größe des Bildes auf 448 x 448 Pixel und unterteilt es dann in Patches, die jeweils 14 x 14 Pixel groß sind. Schließlich werden diese Flecken in 1.031 Einbettungen mit jeweils 128 Dimensionen eingebettet.

Sie können die Einbettungen für alle Bilder mit Hilfe einer Schleife wie folgt erzeugen:

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

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

Sobald die Daten fertig sind, erstellen wir eine Sammlung. In der Sammlung ist ein Feld mit dem Namen patches ein Array of Structs.

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 Kollektion

Jetzt können wir die vorbereiteten Finanzberichte in die Kollektion einfügen.

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

In der Ausgabe sehen Sie, dass alle Seiten aus dem Vidore-Datensatz eingefügt wurden.

Schritt 6: Suche in den Finanzberichten

Sobald die Daten fertig sind, können wir die Daten in der Sammlung wie folgt durchsuchen:

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