Skip to content

vllm.entrypoints.pooling.late_chunking

Functions:

attach_late_chunking_metadata(output, metadata)

Attach ranges only to successful, complete chunk-vector results.

Source code in vllm/entrypoints/pooling/late_chunking.py 
79
80
81
82
83
84
85
86
87
88
89
90
91
92
def attach_late_chunking_metadata(
    output: PoolingRequestOutput, metadata: LateChunkingMetadata
) -> None:
    """Attach ranges only to successful, complete chunk-vector results."""
    if not output.finished or output.error is not None:
        return
    data = output.outputs.data
    if (
        data.ndim != 2
        or data.shape[0] != len(metadata.chunks)
        or len(output.prompt_token_ids) != metadata.input_tokens
    ):
        raise ValueError("Late-chunking output does not match its token/source ranges")
    output.late_chunking = metadata

build_late_chunking_metadata(text, num_tokens, offsets, chunk_size)

Map chunks to original text, keeping special-token-only chunks.

Source code in vllm/entrypoints/pooling/late_chunking.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def build_late_chunking_metadata(
    text: str,
    num_tokens: int,
    offsets: Sequence[tuple[int, int]] | None,
    chunk_size: int,
) -> LateChunkingMetadata:
    """Map chunks to original text, keeping special-token-only chunks."""
    if num_tokens <= 0 or offsets is None or len(offsets) != num_tokens:
        raise VLLMValidationError(
            "Late chunking requires a fast tokenizer with offsets aligned to "
            "all input tokens"
        )
    previous_start = previous_end = 0
    for start, end in offsets:
        if not 0 <= start <= end <= len(text):
            raise VLLMValidationError("Token offsets are not aligned to the input text")
        if start < end:
            if start < previous_start or end < previous_end:
                raise VLLMValidationError("Token offsets are not in source-text order")
            previous_start, previous_end = start, end

    chunks = []
    for start in range(0, num_tokens, chunk_size):
        end = min(start + chunk_size, num_tokens)
        source = [(a, b) for a, b in offsets[start:end] if a < b]
        char_range = (
            (min(a for a, _ in source), max(b for _, b in source)) if source else None
        )
        chunks.append(LateChunk(token_range=(start, end), char_range=char_range))
    return LateChunkingMetadata(
        chunk_size=chunk_size, input_tokens=num_tokens, chunks=chunks
    )

prepare_late_chunking_input(config, render_params)

Validate full-document text input before requesting tokenizer offsets.

Source code in vllm/entrypoints/pooling/late_chunking.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def prepare_late_chunking_input(
    config: VllmConfig, render_params: AnyRenderParam
) -> tuple[str, TokenizeParams]:
    """Validate full-document text input before requesting tokenizer offsets."""
    params = render_params["params"]
    params.verify(config.model_config)
    if "prompts" not in render_params:
        raise VLLMValidationError("Late chunking requires plain-text input")
    prompt = cast(EncodeCMPLRenderParams, render_params)["prompts"]
    if (
        not isinstance(prompt, dict)
        or set(prompt) - {"prompt", "cache_salt"}
        or not isinstance(text := prompt.get("prompt"), str)
        or not text
    ):
        raise VLLMValidationError("Late chunking requires nonempty plain-text input")
    tok_params = render_params["tok_params"]
    if (
        tok_params.pad_prompt_tokens is not None
        or tok_params.truncate_prompt_tokens is not None
        or tok_params.do_lower_case
    ):
        raise VLLMValidationError(
            "Late chunking does not support input padding, truncation or "
            "renderer text normalization"
        )
    return text, replace(tok_params, return_token_offsets=True)