Late chunking voor RAG: implementatie met Jina AI

Naïeve chunking en de beperkingen ervan in RAG

In een RAG-pijplijn worden documenten opgesplitst in kleinere chunks voordat ze worden ge-embed en opgeslagen in een vectordatabank. Elke chunk wordt onafhankelijk verwerkt en gebruikt voor retrieval wanneer er queries worden uitgevoerd. Deze aanpak van "naïeve chunking" verliest echter vaak belangrijke langafstandcontext.

Het probleem ontstaat doordat traditionele chunking documenten splitst zonder rekening te houden met hoe informatie samenhangt. In een document over Parijs kan bijvoorbeeld de frase "de stad" in een andere chunk terechtkomen dan waar "Parijs" wordt genoemd. Zonder de volledige context kan het retrievalmodel moeite hebben om deze verwijzingen te koppelen, wat leidt tot minder nauwkeurige resultaten. Dit speelt nog sterker bij lange documenten waarin cruciale context over meerdere secties verspreid is.

Late chunking: context behouden bij het splitsen van documenten

Late chunking lost dit op door het moment van splitsen te veranderen. In plaats van eerst het document in stukken te verdelen, embedt late chunking het hele document met een long-contextmodel. Pas daarna wordt het document opgesplitst in kleinere chunks.

Dit zijn de belangrijkste voordelen van late chunking:

• Context blijft behouden: Late chunking zorgt ervoor dat elke chunk de overkoepelende context behoudt door eerst het hele document te embedden. Zo blijven verwijzingen en verbanden in de tekst intact in de chunk-embeddings.
• Betere retrieval: De chunk-embeddings die via late chunking ontstaan, zijn rijker en nauwkeuriger, wat de retrievalresultaten in RAG-systemen verbetert omdat het model het document beter begrijpt.
• Geschikt voor lange teksten: Ideaal voor zeer lange documenten die traditionele modellen door tokenlimieten niet in één keer aankunnen.

Met long-contextmodellen zoals Jina’s jinaai/jina-embeddings-v2-base-en, dat tot 8192 tokens ondersteunt, maakt late chunking het mogelijk om grote tekstgedeelten effectief te embedden voordat je ze in chunks opsplitst.

Late chunking implementeren

Hier is een stapsgewijze gids om late chunking te implementeren met Jina’s long-context embeddingmodel. Je kunt Jina’s API-sleutel gratis verkrijgen hier, en we gebruiken de volgende invoertekst als demo:

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."""

Stap 1: Haal chunks en span-annotaties op

Gebruik eerst je Jina API-sleutel en de onderstaande helperfunctie om je invoertekst in chunks te splitsen. Deze chunks bevatten span-annotaties die later helpen om de documentembedding te splitsen. Jina’s API gebruikt natuurlijke grenzen zoals alinea- of zinsafbrekingen om ervoor te zorgen dat de chunks logisch zijn en hun betekenis behouden.

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

Stap 2: Tokenizeer de tekst en genereer documentembeddings op tokenniveau

Gebruik eerst een tokenizer die compatibel is met long-contextmodellen, zoals Jina’s embeddings-v2-base-en, om het hele document in tokens op te delen. Gebruik vervolgens een long-context transformermodel om voor elk token embeddings te maken. Dit betekent dat elk woord of token in je document een eigen embedding krijgt die de betekenis vastlegt.

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

Stap 3: Late chunking

Zodra je de token-embeddings voor het hele document hebt, kun je late chunking toepassen. Gebruik de span-annotaties uit stap één om deze token-embeddings in kleinere chunks te splitsen. Pas daarna mean pooling toe om de embeddings binnen elke chunk te middelen en zo één embedding per chunk te maken. We hebben nu chunk-embeddings met sterke contextuele informatie over het hele 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

Stap 4: Resultaten van late chunking vs. traditionele chunking

Om de voordelen van late chunking te begrijpen, vergelijken we het met traditionele chunking:

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
------------------------------------------------------------------

Zoals je in de tweede en derde chunk ziet, geeft traditionele chunking similariteitsscores van 70–75% vergeleken met het woord "Berlin". Met late chunking, dat de context van het hele document behoudt, stijgen deze scores echter naar 82–84%. Dit laat zien dat late chunking beter is in het behouden van context en het creëren van betekenisvollere embeddings, wat resulteert in nauwkeurigere zoekresultaten.

Conclusie

Late chunking is een grote verbetering voor documentretrievalsystemen, vooral in RAG-pijplijnen. Door te wachten met splitsen totdat het document volledig is ge-embed, behoudt late chunking de volledige context in elke chunk. Dit resulteert in nauwkeurigere en betekenisvollere embeddings.