Chunking tardío para RAG: Aplicación con Jina AI

El Naive Chunking y sus limitaciones en el GAR

En una canalización RAG, los documentos se dividen en trozos más pequeños antes de ser incrustados y almacenados en una base de datos vectorial. Cada trozo se procesa de forma independiente y se utiliza para la recuperación cuando se hacen consultas. Sin embargo, este enfoque de "fragmentación ingenua" a menudo pierde un contexto importante a larga distancia.

El problema surge porque el chunking tradicional divide los documentos sin tener en cuenta cómo está conectada la información. Por ejemplo, en un documento sobre París, la frase "la ciudad" puede acabar en un trozo distinto de donde se menciona "París". Sin el contexto completo, el modelo de recuperación puede tener dificultades para relacionar estas referencias, lo que conduce a resultados menos precisos. Este problema es aún peor en los documentos largos en los que el contexto clave está repartido en varias secciones.

Chunking tardío: Preservar el contexto en la división de documentos

La fragmentación tardía resuelve el problema cambiando el momento en que divides el documento. En lugar de dividir primero el documento en trozos, el chunking tardío incrusta todo el documento utilizando un modelo de contexto largo. Sólo después divide el documento en trozos más pequeños.

Éstas son las principales ventajas de la fragmentación tardía:

• Keeps context: El chunking tardío garantiza que cada trozo conserve el contexto general al incrustar primero todo el documento. De este modo, las referencias y conexiones a través del texto permanecen intactas en las incrustaciones de trozos.
• Mejor recuperación: Las incrustaciones de trozos creadas mediante el chunking tardío son más ricas y precisas, lo que mejora los resultados de recuperación en los sistemas RAG porque el modelo comprende mejor el documento.
• Maneja textos largos: Es estupendo para documentos muy largos que los modelos tradicionales no pueden manejar de una sola vez debido a los límites de tokens.

Utilizando modelos de contexto largo como el de Jina jinaai/jina-embeddings-v2-base-en, que admite hasta 8192 tokens, el chunking tardío permite incrustar eficazmente grandes secciones de texto antes de dividirlas en trozos.

Aplicación de la fragmentación tardía

Aquí tienes una guía paso a paso para ayudarte a poner en práctica el chunking tardío utilizando el modelo de incrustación de contexto largo de Jina. Puedes obtener gratuitamente la clave API de Jina aquíy utilizaremos el siguiente texto de entrada como demostración:

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

Paso 1: Obtener trozos y anotaciones span

En primer lugar, utiliza tu clave API de Jina y la siguiente función de ayuda para dividir tu texto de entrada en trozos. Estos trozos vienen con anotaciones de span que ayudan a dividir la incrustación del documento más adelante. La API de Jina utiliza límites naturales, como saltos de párrafo o de frase, para garantizar que los trozos tengan sentido y conserven su 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)]

Paso 2: Tokeniza el texto y genera incrustaciones de documentos a nivel de token

En primer lugar, utiliza un tokenizador compatible con los modelos de contexto largo, como el de Jina embeddings-v2-base-en, para descomponer todo el documento en tokens. A continuación, utiliza un modelo transformador de contexto largo para crear incrustaciones para cada token. Esto significa que cada palabra o token de tu documento recibe una incrustación única que capta su 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

Paso 3: Troceado tardío

Una vez que tengas las incrustaciones de tokens de todo el documento, estarás listo para el chunking tardío. Utiliza las anotaciones de span del paso uno para dividir estas incrustaciones de token en trozos más pequeños. A continuación, aplica la agrupación de medias para promediar las incrustaciones dentro de cada trozo, creando una única incrustación para cada trozo. Ahora tenemos incrustaciones de trozos con información contextual sólida sobre todo el 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

Paso 4: Resultados del chunking tardío frente al chunking tradicional

Para comprender las ventajas de la fragmentación tardía, comparémosla con la fragmentación 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 puedes ver en el segundo y tercer chunks, el chunking tradicional muestra puntuaciones de similitud del 70-75% cuando se compara con la palabra "Berlín". Sin embargo, con el chunking tardío, que mantiene el contexto de todo el documento, estas puntuaciones suben al 82-84%. Esto demuestra que el chunking tardío conserva mejor el contexto y crea incrustaciones más significativas, lo que da lugar a resultados de búsqueda más precisos.

Conclusión

El chunking tardío es una mejora importante para los sistemas de recuperación de documentos, especialmente en los pipelines RAG. Al esperar a dividir el documento hasta que esté totalmente incrustado, el chunking tardío mantiene el contexto completo en cada trozo. Así se obtienen incrustaciones más precisas y significativas.