Chunking tardio para RAG: Implementação com a Jina AI

Naive Chunking e suas limitações no RAG

Em um pipeline RAG, os documentos são divididos em partes menores antes de serem incorporados e armazenados em um banco de dados vetorial. Cada bloco é processado de forma independente e usado para recuperação quando são feitas consultas. No entanto, essa abordagem de "fragmentação ingênua" geralmente perde importantes contextos de longa distância.

O problema surge porque o chunking tradicional divide os documentos sem considerar como as informações estão conectadas. Por exemplo, em um documento sobre Paris, a frase "a cidade" pode acabar em um trecho diferente daquele em que "Paris" é mencionado. Sem o contexto completo, o modelo de recuperação pode ter dificuldades para vincular essas referências, o que leva a resultados menos precisos. Esse problema é ainda pior em documentos longos em que o contexto principal está espalhado em várias seções.

Chunking tardio: Preservação do contexto na divisão de documentos

A fragmentação tardia resolve o problema alterando o momento em que você divide o documento. Em vez de dividir o documento em partes primeiro, o chunking tardio incorpora o documento inteiro usando um modelo de contexto longo. Somente depois disso é que ele divide o documento em partes menores.

Esses são os principais benefícios do late chunking:

• Mantém o contexto: O chunking tardio garante que cada pedaço mantenha o contexto geral ao incorporar o documento inteiro primeiro. Dessa forma, as referências e conexões no texto permanecem intactas nas incorporações de pedaços.
• Melhor recuperação: As incorporações de pedaços criadas por meio de chunking tardio são mais ricas e precisas, melhorando os resultados de recuperação em sistemas RAG porque o modelo entende melhor o documento.
• Lida com textos longos: É excelente para documentos muito longos que os modelos tradicionais não conseguem processar de uma só vez devido aos limites de tokens.

Usando modelos de contexto longo como o jinaai/jina-embeddings-v2-base-en da Jina, que suporta até 8192 tokens, o chunking tardio permite incorporar seções de texto grandes de forma eficaz antes de dividi-las em pedaços.

Implementação do Late Chunking

Aqui está um guia passo a passo para ajudar você a implementar a fragmentação tardia usando o modelo de incorporação de contexto longo da Jina. Você pode obter a chave de API da Jina gratuitamente aquie usaremos o seguinte texto de entrada como demonstração:

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

Etapa 1: Obter blocos e anotações de extensão

Primeiro, use sua chave de API do Jina e a função auxiliar abaixo para dividir o texto de entrada em partes. Esses blocos vêm com anotações de intervalo que ajudam a dividir a incorporação do documento posteriormente. A API da Jina usa limites naturais, como quebras de parágrafo ou de frase, para garantir que os blocos façam sentido e mantenham seu significado.

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

Etapa 2: Tokenize o texto e gere embeddings de documentos em nível de token

Primeiro, use um tokenizador compatível com modelos de contexto longo, como o embeddings-v2-base-en da Jina, para dividir o documento inteiro em tokens. Em seguida, use um modelo de transformador de contexto longo para criar embeddings para cada token. Isso significa que cada palavra ou token em seu documento recebe uma incorporação exclusiva que captura seu significado.

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

Etapa 3: Separação tardia

Quando você tiver as incorporações de tokens para todo o documento, estará pronto para o chunking tardio. Use as anotações de abrangência da etapa 1 para dividir esses tokens embeddings em partes menores. Em seguida, aplique o pooling médio para calcular a média dos embeddings em cada bloco, criando um único embedding para cada bloco. Agora, temos incorporação de partes com fortes informações contextuais sobre todo o documento.

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

Etapa 4: Resultados de chunking tardio vs. chunking tradicional

Para que você entenda os benefícios do chunking tardio, vamos compará-lo com o chunking tradicional:

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

Como você pode ver no segundo e terceiro blocos, a fragmentação tradicional mostra pontuações de similaridade de 70 a 75% quando comparada à palavra "Berlin". No entanto, com a fragmentação tardia, que mantém o contexto de todo o documento, essas pontuações aumentam para 82% a 84%. Isso mostra que a fragmentação tardia faz um trabalho melhor na preservação do contexto e na criação de incorporações mais significativas, resultando em resultados de pesquisa mais precisos.

Conclusão

A fragmentação tardia é uma melhoria importante para os sistemas de recuperação de documentos, especialmente nos pipelines RAG. Ao esperar para dividir o documento até que ele esteja totalmente incorporado, o chunking tardio mantém o contexto completo em cada pedaço. Isso resulta em incorporações mais precisas e significativas.