File size: 8,165 Bytes
4679932 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 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 43 44 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 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 |
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from array import array
from typing import Optional
import torch
import torch.nn as nn
from vllm.config import ModelConfig, VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.layers.pooler import PoolerHead
from vllm.model_executor.models.llama import LlamaForCausalLM
from vllm.model_executor.pooling_metadata import (PoolingMetadata,
PoolingTensors)
from vllm.sequence import PoolerOutput, PoolingSequenceGroupOutput
from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
from .interfaces import SupportsV0Only
logger = init_logger(__name__)
class GritLMPooler(nn.Module):
def __init__(self, model_config: ModelConfig):
super().__init__()
self.model_config = model_config
tokenizer = cached_tokenizer_from_config(self.model_config)
# 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"]
}
def tokens_to_ids(tokens: list[str]) -> array:
return array("i", [self.token_ids[token] for token in tokens])
self.user_pattern_ids = tokens_to_ids(
["▁<", "|", "user", "|", ">", "<0x0A>"])
self.embed_newline_pattern_ids = tokens_to_ids(
["<0x0A>", "<", "|", "embed", "|", ">", "<0x0A>"])
self.embed_pattern_ids = tokens_to_ids(
["▁<", "|", "embed", "|", ">", "<0x0A>"])
self.head = PoolerHead(normalize=True, softmax=False)
def _find_array(self, arr: array, target: array, start_idx: int) -> int:
"""
Find the first occurrence of target in arr starting from start_idx.
Args:
arr: The array to search within
target: The consecutive subsequence to find
start_idx: The starting index to search from
Returns:
int: The index of the first occurrence of target in arr.
"""
if start_idx < 0:
raise ValueError("start_idx must be non-negative")
if not target or not arr:
raise ValueError("Empty arr or target not allowed")
target_len = len(target)
for i in range(start_idx, len(arr) - target_len + 1):
if arr[i:i + target_len] == target:
return i
return -1
def _get_instruction_len(self, prompt_token_ids: array) -> int:
"""
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>"]:
logger.warning("BOS token not found in prompt, "
"thus using empty string for instruction. "
"GritLM requires BOS token in prompt.")
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
if self._find_array(prompt_token_ids,
self.user_pattern_ids,
start_idx=1) == 1:
embed_pattern_ids = self.embed_newline_pattern_ids
# Find the embed pattern in the prompt.
found_embed_pattern_idx = self._find_array(prompt_token_ids,
embed_pattern_ids,
start_idx=1)
if found_embed_pattern_idx != -1:
instruction_len = found_embed_pattern_idx + len(embed_pattern_ids)
else:
logger.warning("Query instruction not found in prompt, "
"thus using BOS token as instruction instead. "
"GritLM requires query instruction in prompt.")
instruction_len = 1
return instruction_len
def forward(
self,
hidden_states: torch.Tensor,
pooling_metadata: PoolingMetadata,
) -> PoolerOutput:
"""
Pool the hidden states by summing the embeddings of
non-instruction tokens.
"""
prompts_token_ids = [
token_ids.prompt_token_ids_array
for _, token_ids in pooling_metadata.seq_data.items()
]
instruction_lens = torch.tensor(
[
self._get_instruction_len(prompt_token_ids)
for prompt_token_ids in prompts_token_ids
],
device=hidden_states.device,
)
prompt_lens = PoolingTensors.from_pooling_metadata(
pooling_metadata, hidden_states.device).prompt_lens
mask = torch.zeros_like(hidden_states, dtype=torch.bool)
start_idx = 0
for prompt_len, instruction_len in zip(prompt_lens, instruction_lens):
end_idx = start_idx + prompt_len
mask[start_idx + instruction_len:end_idx] = True
start_idx = end_idx
masked_hidden_states = hidden_states.masked_fill(~mask, 0.0)
sum_embeddings = torch.zeros(len(prompt_lens),
hidden_states.size(1),
device=hidden_states.device)
start_idx = 0
for i, prompt_len in enumerate(prompt_lens):
end_idx = start_idx + prompt_len
sum_embeddings[i] = masked_hidden_states[start_idx:end_idx].sum(
dim=0)
start_idx = end_idx
num_non_instruction_tokens = prompt_lens - instruction_lens
mean_embeddings = sum_embeddings / num_non_instruction_tokens.unsqueeze(
1)
pooled_data = self.head(mean_embeddings,
pooling_metadata=pooling_metadata)
pooled_outputs = [
PoolingSequenceGroupOutput(data) for data in pooled_data
]
return PoolerOutput(outputs=pooled_outputs)
class GritLM(LlamaForCausalLM, SupportsV0Only):
"""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"
"""
def __init__(
self,
vllm_config: VllmConfig,
prefix: str = "",
**kwargs,
) -> None:
# Use full attention for pooling
if vllm_config.model_config.runner_type == "pooling":
hf_config = vllm_config.model_config.hf_config
hf_config.is_causal = False
vllm_config.cache_config.sliding_window = None
for attr in ("sliding_window", "interleaved_sliding_window"):
if hasattr(hf_config, attr):
delattr(hf_config, attr)
super().__init__(vllm_config=vllm_config, prefix=prefix, **kwargs)
self._pooler = GritLMPooler(vllm_config.model_config)
def pooler(
self,
hidden_states: torch.Tensor,
pooling_metadata: PoolingMetadata,
) -> Optional[PoolerOutput]:
return self._pooler(hidden_states, pooling_metadata)
|