Le découpage tardif pour le RAG : Mise en œuvre avec Jina AI

Le découpage naïf et ses limites dans RAG

Dans un pipeline RAG, les documents sont divisés en petits morceaux avant d'être intégrés et stockés dans une base de données vectorielle. Chaque morceau est traité indépendamment et utilisé pour la recherche lorsque des requêtes sont effectuées. Cependant, cette approche "naïve" perd souvent le contexte important de la longue distance.

Le problème vient du fait que le découpage traditionnel divise les documents sans tenir compte de la manière dont les informations sont reliées entre elles. Par exemple, dans un document sur Paris, l'expression "la ville" peut se retrouver dans un morceau différent de celui où est mentionné "Paris". Sans le contexte complet, le modèle de recherche peut avoir du mal à relier ces références, ce qui conduit à des résultats moins précis. Ce problème est encore plus grave dans les documents longs où le contexte clé est réparti sur plusieurs sections.

Le découpage tardif : Préserver le contexte dans le découpage des documents

Le découpage tardif résout le problème en modifiant le moment où vous divisez le document. Au lieu de découper d'abord le document en morceaux, le découpage tardif incorpore l'ensemble du document à l'aide d'un modèle de contexte long. Ce n'est qu'après cela qu'il divise le document en petits morceaux.

Tels sont les principaux avantages du découpage tardif :

• Conserve le contexte : Le découpage tardif permet de s'assurer que chaque morceau conserve le contexte global en intégrant d'abord l'ensemble du document. De cette manière, les références et les connexions dans le texte restent intactes dans les morceaux intégrés.
• Meilleure récupération : Les enchâssements de morceaux créés grâce au découpage tardif sont plus riches et plus précis, ce qui améliore les résultats de recherche dans les systèmes RAG car le modèle comprend mieux le document.
• Traite les textes longs : Il est idéal pour les documents très longs que les modèles traditionnels ne peuvent pas traiter en une seule fois en raison des limites de jetons.

En utilisant des modèles à contexte long comme celui de Jina( jinaai/jina-embeddings-v2-base-en), qui prend en charge jusqu'à 8192 tokens, le découpage tardif permet d'intégrer efficacement de grandes sections de texte avant de les diviser en morceaux.

Mise en œuvre du découpage tardif

Voici un guide étape par étape pour vous aider à mettre en œuvre le découpage tardif à l'aide du modèle d'intégration du contexte long de Jina. Vous pouvez obtenir gratuitement la clé API de Jina iciNous utiliserons le texte d'entrée suivant pour la démonstration :

input_text = """Berlin is the capital and largest city of Germany, both by area and by population.
Its more than 3.85 million inhabitants make it the European Union's most populous city, as measured by population within city limits.
The city is also one of the states of Germany, and is the third smallest state in the country in terms of area."""

Étape 1 : Obtenir des morceaux et des annotations de portée

Tout d'abord, utilisez votre clé API Jina et la fonction d'aide ci-dessous pour diviser votre texte d'entrée en morceaux. Ces morceaux sont accompagnés d'annotations de portée qui aident à diviser l'intégration du document ultérieurement. L'API de Jina utilise des limites naturelles, telles que des coupures de paragraphes ou de phrases, pour s'assurer que les morceaux ont un sens et conservent leur signification.

import json
import requests

def custom_tokenize_jina_api(input_text: str):
    url = '<https://segment.jina.ai/>'
    headers = {
        'Content-Type': 'application/json',
        'Authorization': 'Bearer ENTER_YOUR_JINA_API_KEY'
    }
    data = {
        "content": input_text,
        "tokenizer": "o200k_base",
        "return_tokens": "true",
        "return_chunks": "true",
        "max_chunk_length": "1000"
    }
    # Make the API request
    response = requests.post(url, headers=headers, json=data)
    response_data = response.json()
    chunks = response_data.get("chunks", [])
    i = 1
    j = 1
    span_annotations = []
    for x in response_data['tokens']:
        if j == 1:
            j = len(x)
        else:
            j = len(x) + i
        span_annotations.append((i, j))
        i = j
    return chunks, span_annotations
chunks, span_annotations = custom_tokenize_jina_api(input_text)

print(chunks)
print(span_annotations)

['Berlin is the capital and largest city of Germany, both by area and by population.\\n\\n', "Its more than 3.85 million inhabitants make it the European Union's most populous city, as measured by population within city limits.\\n\\n", 'The city is also one of the states of Germany, and is the third smallest state in the country in terms of area.']
[(1, 17), (17, 44), (44, 69)]

Étape 2 : Tokeniser le texte et générer des encastrements de documents au niveau du token

Tout d'abord, utilisez un tokenizer compatible avec les modèles à contexte long, tel que Jina's embeddings-v2-base-en, pour décomposer l'ensemble du document en tokens. Ensuite, nous utilisons un modèle de transformation en contexte long pour créer des enchâssements pour chaque mot clé. Cela signifie que chaque mot ou token de votre document reçoit une intégration unique qui capture sa signification.

from transformers import AutoModel
from transformers import AutoTokenizer

# load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
model = AutoModel.from_pretrained('jinaai/jina-embeddings-v2-base-en', trust_remote_code=True)
inputs = tokenizer(input_text, return_tensors='pt')
model_output = model(**inputs)
model_output[0].shape

torch.Size([1, 71, 768]) # 71 represents number of tokens in the entire document

Étape 3 : Chunking tardif

Une fois que vous avez les enchâssements de jetons pour l'ensemble du document, vous êtes prêt à procéder à un découpage tardif. Utilisez les annotations de portée de l'étape 1 pour diviser ces enchâssements de jetons en plus petits morceaux. Ensuite, appliquez la méthode du pooling moyen pour faire la moyenne des intégrations dans chaque morceau, en créant une seule intégration pour chaque morceau. Nous disposons à présent d'encastrements de morceaux contenant des informations contextuelles solides sur l'ensemble du document.

def late_chunking(
    model_output: 'BatchEncoding', span_annotation: list, max_length=None
):
    token_embeddings = model_output[0]
    outputs = []
    for embeddings, annotations in zip(token_embeddings, span_annotation):
        if (
            max_length is not None
        ):  # remove annotations which go bejond the max-length of the model
            annotations = [
                (start, min(end, max_length - 1))
                for (start, end) in annotations
                if start < (max_length - 1)
            ]
        pooled_embeddings = [
            embeddings[start:end].sum(dim=0) / (end - start)
            for start, end in annotations
            if (end - start) >= 1
        ]
        pooled_embeddings = [
            embedding.detach().cpu().numpy() for embedding in pooled_embeddings
        ]
        outputs.append(pooled_embeddings)
    return outputs

embeddings = late_chunking(model_output, [span_annotations])[0]
len(embeddings)

3 # matches number of chunks in Step 1

Étape 4 : Résultats du découpage tardif par rapport au découpage traditionnel

Pour comprendre les avantages du découpage tardif, comparons-le au découpage traditionnel :

embeddings_traditional_chunking = model.encode(chunks)

import numpy as np

cos_sim = lambda x, y: np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y))
q = "Berlin"
berlin_embedding = model.encode(q)

print(q)
print('\\n')
for chunk, new_embedding, trad_embeddings in zip(chunks, embeddings, embeddings_traditional_chunking):
  print(chunk.strip())
  print(f'Late chunking:', cos_sim(berlin_embedding, new_embedding))
  print(f'Traditional chunking:', cos_sim(berlin_embedding, trad_embeddings))
  print("------------------------------------------------------------------")

Berlin
Berlin is the capital and largest city of Germany, both by area and by population.
Late chunking: 0.84954596
Traditional chunking: 0.84862185
------------------------------------------------------------------
Its more than 3.85 million inhabitants make it the European Union's most populous city, as measured by population within city limits.
Late chunking: 0.82489026
Traditional chunking: 0.70843375
------------------------------------------------------------------
The city is also one of the states of Germany, and is the third smallest state in the country in terms of area.
Late chunking: 0.84980094
Traditional chunking: 0.7534553
------------------------------------------------------------------

Comme vous pouvez le voir dans les deuxième et troisième morceaux, le découpage traditionnel montre des scores de similarité de 70 à 75 % par rapport au mot "Berlin". Cependant, avec le découpage tardif, qui maintient le contexte du document entier, ces scores passent à 82-84%. Cela montre que le découpage tardif permet de mieux préserver le contexte et de créer des liens plus significatifs, ce qui se traduit par des résultats de recherche plus précis.

Conclusion

Le découpage tardif est une amélioration majeure pour les systèmes de recherche documentaire, en particulier dans les pipelines RAG. En attendant que le document soit entièrement intégré, le découpage tardif permet de conserver l'intégralité du contexte dans chaque morceau. Cela permet d'obtenir des encastrements plus précis et plus significatifs.