Recherche à l'aide de listes d'incorporation

Cette page explique comment mettre en place un système de recherche de texte ColBERT et un système de recherche de texte ColPali à l'aide du tableau de structures de Milvus, qui vous permet de stocker un document avec ses morceaux vectorisés dans des listes d'intégration.

Vue d'ensemble

Pour construire un système de recherche de texte, il peut être nécessaire de diviser les documents en morceaux et de stocker chaque morceau avec ses encastrements en tant qu'entité dans une base de données vectorielle afin de garantir la précision et l'exactitude, en particulier pour les longs documents où les encastrements en texte intégral pourraient diluer la spécificité sémantique ou dépasser les limites d'entrée du modèle.

Cependant, le stockage des données en morceaux conduit à des résultats de recherche par morceaux, ce qui signifie que la recherche identifie d'abord des segments pertinents plutôt que des documents cohérents. Pour remédier à ce problème, il convient d'effectuer un traitement supplémentaire après la recherche.

ColBERT (arXiv : 2004.12832) est un système de recherche texte-texte qui offre une recherche de passage efficace et effective par le biais d'interactions tardives contextualisées sur BERT. Il permet un encodage indépendant par jeton des requêtes et des documents et calcule leur similarité.

Encodage par jeton

Lors de l'ingestion des données dans ColBERT, chaque document est divisé en jetons, qui sont ensuite vectorisés et stockés sous la forme d'une liste d'intégration, comme dans $\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}$. Lorsqu'une requête arrive, elle est également symbolisée, vectorisée et stockée sous la forme d'une liste d'intégration, comme dans $\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}$.

Dans les formules ci-dessus,

• $dd$ d : un document

• $qq$ q : la requête

• $EdE_d$ Ed: la liste d'intégration qui représente le document.

• $EqE_q$ Eq: la liste d'intégration qui représente la requête.

• $[e_{d1}, e_{d2}, \dots, e_{dn}] ∈ \R^{n×d}$d1, ed2,$...$, edn]$\in$ R${}^{n\times d}$: le nombre de vecteurs d'intégration dans la liste d'intégration représentant le document est compris entre $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}$: le nombre de vecteurs d'intégration dans la liste d'intégration représentant la requête est compris entre $Rm×d\R^{m×d}$ R ${}^{m\times d}$.

Interaction tardive

Une fois la vectorisation terminée, la liste d'intégration de la requête est comparée à la liste d'intégration de chaque document, élément par élément, afin de déterminer le score de similarité final.

Interaction tardive

Comme le montre le diagramme ci-dessus, la requête contient deux tokens, à savoir machine et learning, et le document dans la fenêtre contient quatre tokens : neural network , python et tutorial. Une fois que ces jetons sont vectorisés, l'intégration vectorielle de chaque jeton de la requête est comparée à celle du document afin d'obtenir une liste de scores de similarité. Les scores les plus élevés de chaque liste sont ensuite additionnés pour obtenir le score final. Le processus de détermination du score final d'un document est connu sous le nom de similarité maximale(MAX_SIM). Pour plus de détails sur la similitude maximale, voir Similitude maximale.

Extension de ColPali

Basé sur ColBERT, ColPali (arXiv : 2407.01449) propose une nouvelle approche pour l'extraction de documents visuellement riches qui s'appuie sur les modèles vision-langage (VLM). Lors de l'ingestion des données, chaque page de document est convertie en une image haute résolution, puis divisée en patchs, plutôt qu'en jetons. Par exemple, une image de page de document de 448 x 448 pixels peut produire 1 024 parcelles, chacune mesurant 14 x 14 pixels.

Cette méthode permet de préserver les informations non textuelles, telles que la mise en page du document, les images et les structures des tableaux, qui sont perdues lors de l'utilisation de systèmes de recherche textuelle.

Extension de Copali

Le VLM utilisé dans ColPali est appelé PaliGemma (arXiv : 2407.07726), qui comprend un encodeur d'images(SigLIP-400M), un modèle de langage décodeur uniquement(Gemma2-2B) et une couche linéaire qui projette la sortie de l'encodeur d'images dans l'espace vectoriel du modèle de langage, comme le montre le diagramme ci-dessus.

Lors de l'ingestion des données, une page de document, représentée sous la forme d'une image brute, est divisée en plusieurs patchs visuels, chacun d'entre eux étant intégré pour générer une liste d'intégrations vectorielles. Ils sont ensuite projetés dans l'espace vectoriel du modèle de langage pour obtenir la liste d'intégration finale, comme dans $\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}$. Lorsqu'une requête arrive, elle est transformée en jeton, et chaque jeton est intégré pour générer une liste d'intégrations vectorielles, comme dans $\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}$. Ensuite, MAX_SIM a été appliqué pour comparer les deux listes d'intégration et obtenir le score final entre la requête et la page du document.

Système de recherche de texte ColBERT

Dans cette section, nous allons mettre en place un système de recherche de texte ColBERT à l'aide de la matrice de structures de Milvus. Avant cela, configurez une instance Milvus v2.6.x, un cluster Zilliz Cloud compatible avec Milvus v2.6.x et obtenez un jeton d'accès Cohere.

Etape 1 : Installer les dépendances

Exécuter la commande suivante pour installer les dépendances.

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

Étape 2 : Charger le jeu de données Cohere

Dans cet exemple, nous allons utiliser l'ensemble de données Wikipédia de Cohere et récupérer les 10 000 premiers enregistrements. Vous pouvez trouver des informations sur ce jeu de données sur cette page.

from datasets import load_dataset

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

L'exécution des scripts ci-dessus permettra de télécharger le jeu de données s'il n'est pas disponible localement. Chaque enregistrement de l'ensemble de données est un paragraphe d'une page Wikipédia. Le tableau suivant montre la structure de cet ensemble de données.

Colonne Nom	Description de la colonne
_id	ID de l'enregistrement
url	L'URL de l'enregistrement actuel.
title	Le titre du document source.
text	Un paragraphe du document source.
emb	Embeddings du texte du document source.

Étape 3 : Regrouper les paragraphes par titre

Pour rechercher des documents plutôt que des paragraphes, nous devons regrouper les paragraphes par titre.

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

Dans ce code, nous stockons les paragraphes regroupés en tant que documents et les incluons dans la liste data. Chaque document possède une clé paragraphs, qui est une liste de paragraphes ; chaque objet paragraphe contient les clés text et emb.

Étape 4 : Créer une collection pour l'ensemble de données Cohere

Une fois que les données sont prêtes, nous allons créer une collection. Dans la collection, il y a un champ nommé paragraphs, qui est un tableau de structures.

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
)

Étape 5 : Insérer l'ensemble de données Cohere dans la collection

Nous pouvons maintenant insérer les données préparées dans la collection que nous avons créée ci-dessus.

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

Étape 6 : Recherche dans l'ensemble de données Cohere

Selon la conception de ColBERT, le texte de la requête doit être tokenisé et ensuite incorporé dans une EmbeddingList. Dans cette étape, nous utiliserons le même modèle que Cohere a utilisé pour générer des embeddings pour les paragraphes de l'ensemble de données Wikipedia.

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

Dans le code, les textes de la requête sont organisés en tokens dans query_inputs et intégrés dans une liste de vecteurs flottants. Vous pouvez ensuite utiliser EmbeddingList de Milvus pour effectuer une recherche de similarité comme suit.

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

La sortie du code ci-dessus est similaire à ce qui suit :

# 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

Le score de similarité cosinus va de -1 à 1, et les scores de similarité dans la sortie ci-dessus montrent clairement la somme de plusieurs scores de similarité au niveau des jetons.

Système de recherche de texte ColPali

Dans cette section, nous allons mettre en place un système d'extraction de texte basé sur ColPali à l'aide de la matrice de structures de Milvus. Avant cela, configurez une instance Milvus v2.6.x et un cluster Zilliz Cloud compatible avec Milvus v2.6.x.

Etape 1 : Installer les dépendances

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

Étape 2 : Charger l'ensemble de données Vidore

Dans cette section, nous allons utiliser un jeu de données Vidore nommé vidore_v2_finance_fr. Ce jeu de données est un corpus de rapports annuels du secteur bancaire, destiné à des tâches de compréhension de documents longs. Il fait partie des 10 corpus composant le Benchmark ViDoRe v3. Vous pouvez trouver des détails sur ce jeu de données sur cette page.

from datasets import load_dataset

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

L'exécution des scripts ci-dessus téléchargera le jeu de données s'il n'est pas disponible localement. Chaque enregistrement de l'ensemble de données est une page d'un rapport financier. Le tableau suivant montre la structure de ce jeu de données.

Nom de la colonne	Description de la colonne
corpus_id	Un enregistrement dans le corpus
image	L'image de la page en octets.
doc_id	L'identifiant descriptif du document.
page_number_in_doc	Le numéro de page de la page actuelle du document.

Étape 3 : Génération d'enchâssements pour les images de la page

Comme illustré dans la section Vue d'ensemble, le modèle ColPali est un VLM qui projette des images dans l'espace vectoriel d'un modèle de texte. Dans cette étape, nous utiliserons le dernier modèle ColPali vidore/colpali-v1.3. Vous pouvez trouver des détails sur ce modèle sur cette page.

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)

Une fois le modèle prêt, vous pouvez essayer de générer des patchs pour une image spécifique comme suit.

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]

Dans le code ci-dessus, le modèle ColPali redimensionne l'image à 448 x 448 pixels, puis la divise en patchs, chacun mesurant 14 x 14 pixels. Enfin, ces patchs sont intégrés dans 1 031 embeddings, chacun ayant 128 dimensions.

Vous pouvez générer des embeddings pour toutes les images à l'aide d'une boucle comme suit :

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

Étape 4 : Création d'une collection pour l'ensemble de données des rapports financiers

Une fois que les données sont prêtes, nous allons créer une collection. Dans la collection, un champ nommé patches est un tableau de structures.

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
)

Étape 5 : Insérer les rapports financiers dans la collection

Nous pouvons maintenant insérer les rapports financiers préparés dans la collection.

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

Le résultat montre que toutes les pages de l'ensemble de données Vidore ont été insérées.

Étape 6 : Recherche dans les rapports financiers

Une fois que les données sont prêtes, nous pouvons effectuer des recherches dans les données de la collection comme suit :

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