04-04-03 — Contextual retrieval e late chunking

O problema dos chunks sem contexto

Quando você corta um documento em chunks e embeda cada um isoladamente, informações de contexto se perdem. Um chunk pode dizer "O valor previsto aumentou 15% neste trimestre" — mas o embedding não sabe que se refere à Vale, ao contrato X, ou ao exercício fiscal 2023. Esse chunk isolado vai ter uma representação vetorial vaga demais para ser encontrado por queries específicas.

Contextual Retrieval (Anthropic)

Publicado pela Anthropic em setembro de 2024, o método é elegantemente simples: para cada chunk, use o LLM para gerar um mini-contexto de 1-2 frases que situe o chunk no documento. Esse contexto é prepended ao chunk antes de embeddar.

def generate_contextual_chunk(document: str, chunk: str, llm_client) -> str:
    """
    Gera contexto situacional para o chunk dentro do documento.
    Use um modelo barato (gpt-4o-mini, claude-haiku) para reduzir custo.
    """
    prompt = f"""Você receberá um documento e um trecho específico desse documento.
Gere uma frase curta (máximo 2 sentenças) que contextualize esse trecho
dentro do documento, mencionando: (1) o assunto geral do documento,
(2) como o trecho se relaciona com o contexto maior.
Não repita informação que já está no trecho.


{document[:3000]}  



{chunk}


Contexto situacional (máximo 2 frases):"""

    response = llm_client.chat.completions.create(
        model="gpt-4o-mini",  # modelo barato — escala para muitos chunks
        messages=[{"role": "user", "content": prompt}],
        max_tokens=100
    )
    context = response.choices[0].message.content.strip()

    # Prepend o contexto ao chunk antes de embeddar
    return f"{context}\n\n{chunk}"

# Pipeline de indexação com contextual retrieval
def index_with_context(document_text: str, chunks: list[str]) -> list[dict]:
    contextualized_chunks = []
    for chunk in chunks:
        contextual = generate_contextual_chunk(document_text, chunk, client)
        embedding = embed(contextual)  # embeda o chunk COM contexto
        contextualized_chunks.append({
            "text": chunk,               # armazena o chunk original para exibição
            "contextual_text": contextual,  # armazena o texto contextualizado
            "embedding": embedding
        })
    return contextualized_chunks

Custo de contextual retrieval

O custo extra vem de gerar contexto para cada chunk durante a indexação. Para um corpus de 10.000 chunks usando gpt-4o-mini:

• ~200 tokens por chamada (documento resumido + chunk + prompt)
• 10.000 × 200 tokens = 2M tokens de input
• Custo: ~$0.30 (gpt-4o-mini a $0.15/1M tokens input)

É um investimento único na indexação, não por query. Vale muito a pena.

Late Chunking (JinaAI)

Abordagem diferente para o mesmo problema. Em vez de adicionar contexto textualmente, a ideia é aproveitar a atenção do modelo de embedding sobre o documento completo antes de chunkar.

Processo:

1. Tokenize o documento completo
2. Rode o encoder transformer sobre o documento inteiro — cada token obtém representação com atenção ao documento todo
3. Faça mean pooling por chunk (média dos embeddings dos tokens do chunk)

Resultado: cada chunk tem embedding que "viu" o documento inteiro, não só seus próprios tokens.

from transformers import AutoTokenizer, AutoModel
import torch

model_name = "jinaai/jina-embeddings-v2-base-en"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

def late_chunking(document: str, chunk_boundaries: list[tuple[int,int]]) -> list[list[float]]:
    """
    chunk_boundaries: lista de (start_char, end_char) para cada chunk
    retorna: lista de embeddings, um por chunk
    """
    # 1. Tokenizar documento completo
    inputs = tokenizer(document, return_tensors="pt", truncation=True, max_length=8192)

    # 2. Encoder sobre o documento inteiro
    with torch.no_grad():
        outputs = model(**inputs)

    token_embeddings = outputs.last_hidden_state[0]  # shape: [seq_len, hidden_size]

    # 3. Mean pooling por chunk
    chunk_embeddings = []
    for start_char, end_char in chunk_boundaries:
        # Mapear caracteres para tokens
        start_token = inputs.char_to_token(start_char) or 0
        end_token = inputs.char_to_token(end_char - 1) or -1
        chunk_emb = token_embeddings[start_token:end_token+1].mean(dim=0)
        chunk_embeddings.append(chunk_emb.tolist())

    return chunk_embeddings

Comparativo das duas técnicas

Para a maioria dos casos, contextual retrieval é mais prático por ser independente do modelo de embedding e escalar para documentos longos.