gritlm.py 7.92 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Set
4

5
import numpy as np
6
7
import torch

8
from vllm.config import ModelConfig, VllmConfig
9
from vllm.logger import init_logger
10
11
12
13
from vllm.model_executor.layers.pooler import (
    DispatchPooler,
    PoolingParamsUpdate,
)
14
15
16
17
18
19
20
21
from vllm.model_executor.layers.pooler.activations import PoolerNormalize
from vllm.model_executor.layers.pooler.seqwise import (
    SequencePooler,
    SequencePoolerHeadOutput,
    SequencePoolingMethod,
    SequencePoolingMethodOutput,
)
from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
22
from vllm.model_executor.models.llama import LlamaForCausalLM
23
from vllm.tasks import PoolingTask
24
from vllm.tokenizers import cached_tokenizer_from_config
25
from vllm.v1.pool.metadata import PoolingMetadata
26

27
from .interfaces_base import default_pooling_type
28

29
30
31
logger = init_logger(__name__)


32
class GritLMMeanPool(SequencePoolingMethod):
33
    """As `MeanPool`, but only includes non-instruction tokens."""
34

35
    def __init__(self, model_config: ModelConfig):
36
37
        super().__init__()

38
        self.model_config = model_config
39

40
        tokenizer = cached_tokenizer_from_config(self.model_config)
41
42
43
44
45
46
47
48
49
50

        # Collect the tokens needed for pattern matching.
        # "▁<" is different from "_<". The former uses "▁" to indicate that
        # the next token is the start of a word.
        # "<0x0A>" is the newline token (i.e. "\n")."
        self.token_ids = {
            tok: tokenizer.convert_tokens_to_ids([tok])[0]
            for tok in ["<s>", "▁<", "<", "|", "embed", ">", "<0x0A>", "user"]
        }

51
52
        def tokens_to_ids(tokens: list[str]) -> np.ndarray:
            return np.array([self.token_ids[token] for token in tokens])
53

54
        self.user_pattern_ids = tokens_to_ids(["▁<", "|", "user", "|", ">", "<0x0A>"])
55
        self.embed_newline_pattern_ids = tokens_to_ids(
56
57
58
            ["<0x0A>", "<", "|", "embed", "|", ">", "<0x0A>"]
        )
        self.embed_pattern_ids = tokens_to_ids(["▁<", "|", "embed", "|", ">", "<0x0A>"])
59

60
61
62
63
64
    def _find_array(
        self,
        arr: np.ndarray,
        target: np.ndarray,
        start_idx: int = 0,
65
        end_idx: int | None = None,
66
    ) -> int:
67
        """
68
69
        Find the first occurrence of `target` in `arr` starting from
        `start_idx`.
70
71

        Args:
72
73
74
75
            arr: The array to search within.
            target: The consecutive subsequence to find.
            start_idx: The starting index to search from (inclusive).
            end_idx: The ending index to search from (exclusive).
76
77

        Returns:
78
            The index of the first occurrence of `target` in `arr`.
79
80
        """
        if start_idx < 0:
81
82
83
            raise ValueError("`start_idx` must be non-negative")
        if len(arr) == 0 or len(target) == 0:
            raise ValueError("Empty `arr` or `target` not allowed")
84

85
        arr_len = len(arr)
86
        target_len = len(target)
87
88
89
90
91

        if end_idx is None:
            end_idx = arr_len

        for i in range(start_idx, min(end_idx, arr_len - target_len + 1)):
92
            if (arr[i : i + target_len] == target).all():
93
                return i
94

95
96
        return -1

97
    def _get_instruction_len(self, prompt_token_ids: np.ndarray) -> int:
98
99
100
101
102
103
104
105
106
107
108
109
110
        """
        Get the length of the instruction in the prompt.

        We do a pattern matching to find the instruction in the prompt,
        and then return the length of the instruction.

        The pattern matching is done using integers instead of strings
        because the prompt is given as a list of token IDs.
        """
        instruction_len = 0

        # Return no instruction in case of missing BOS token.
        if prompt_token_ids[0] != self.token_ids["<s>"]:
111
112
113
114
115
            logger.warning(
                "BOS token not found in prompt, "
                "thus using empty string for instruction. "
                "GritLM requires BOS token in prompt."
            )
116
117
118
119
120
            return instruction_len

        # If user pattern is found in the prompt, that means there should be
        # a newline token before the embed pattern.
        embed_pattern_ids = self.embed_pattern_ids
121
122
123
124
125
126
        if (
            self._find_array(
                prompt_token_ids, self.user_pattern_ids, start_idx=1, end_idx=2
            )
            == 1
        ):
127
128
129
            embed_pattern_ids = self.embed_newline_pattern_ids

        # Find the embed pattern in the prompt.
130
131
132
        found_embed_pattern_idx = self._find_array(
            prompt_token_ids, embed_pattern_ids, start_idx=1
        )
133
134
135
136

        if found_embed_pattern_idx != -1:
            instruction_len = found_embed_pattern_idx + len(embed_pattern_ids)
        else:
137
138
139
140
141
            logger.warning(
                "Query instruction not found in prompt, "
                "thus using BOS token as instruction instead. "
                "GritLM requires query instruction in prompt."
            )
142
143
144
145
            instruction_len = 1

        return instruction_len

146
    def get_supported_tasks(self) -> Set[PoolingTask]:
147
        return {"embed"}
148

149
150
    def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
        return PoolingParamsUpdate(requires_token_ids=True)
151
152
153

    def forward(
        self,
154
        hidden_states: torch.Tensor,
155
        pooling_metadata: PoolingMetadata,
156
    ) -> SequencePoolingMethodOutput:
157
        prompt_lens = pooling_metadata.prompt_lens
158
        instr_lens = torch.tensor(
159
            [
160
                self._get_instruction_len(token_ids.cpu().numpy())
161
                for token_ids in pooling_metadata.get_prompt_token_ids()
162
            ],
163
            device="cpu",
164
165
        )

166
167
168
169
170
171
        offset = 0
        pooled_data = list[torch.Tensor]()
        for prompt_len, instr_len in zip(prompt_lens, instr_lens):
            pooled_data.append(
                hidden_states[offset + instr_len : offset + prompt_len].mean(
                    dim=0, dtype=torch.float32
172
                )
173
174
            )
            offset += prompt_len
175

176
        return pooled_data
177
178


179
class GritLMPooler(SequencePooler):
180
    def __init__(self, model_config: ModelConfig):
181
182
183
184
        super().__init__(
            pooling=GritLMMeanPool(model_config),
            head=self.head,
        )
185

186
        self.activation = PoolerNormalize()
187

188
189
    def head(
        self,
190
        pooled_data: SequencePoolingMethodOutput,
191
        pooling_metadata: PoolingMetadata,
192
    ) -> SequencePoolerHeadOutput:
193
194
        return self.activation(pooled_data)

195

196
@default_pooling_type(seq_pooling_type="MEAN")
197
class GritLM(LlamaForCausalLM):
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
    """This class implements the embedding model for parasail-ai/GritLM-7B-vllm.

    The class inherits from LlamaForCausalLM and provides a custom pooling
    layer.

    The main difference between the pooling layer in GritLM and the one in
    LlamaForCausalLM is that GritLM ignores the query instruction in the prompt
    when pooling the hidden states.

    Embedding prompts should be in the following format:
    - With instruction: "<|user|>\nINSTRUCTION\n<|embed|>\nPROMPT".
    - Without instruction: "<|embed|>\nPROMPT".

    Generation prompts should be in the following format:
    - "<|user|>\nPROMPT\n<|assistant|>\n"
    """

215
216
    is_pooling_model = True

217
218
219
220
221
222
    def __init__(
        self,
        vllm_config: VllmConfig,
        prefix: str = "",
        **kwargs,
    ) -> None:
223
224
225
        if vllm_config.model_config.runner_type == "pooling":
            hf_config = vllm_config.model_config.hf_config
            hf_config.is_causal = False
226

227
            vllm_config.cache_config.sliding_window = None
228

229
            hf_config.sliding_window = None
230

231
        super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
232

233
234
        pooler_config = vllm_config.model_config.pooler_config
        if pooler_config is not None:
235
236
            self.pooler = DispatchPooler(
                {
237
                    "token_embed": pooler_for_token_embed(pooler_config),
238
                    "embed": GritLMPooler(vllm_config.model_config),
239
240
                }
            )