wehub/flagopen--flagembedding
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:39:21 +08:00
commit bc34f6df14
1149 changed files with 328099 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
FlagEmbedding/inference/reranker/decoder_only/__init__.py
+9
View File
@@ -0,0 +1,9 @@
from .base import BaseLLMReranker as FlagLLMReranker
from .layerwise import LayerWiseLLMReranker as LayerWiseFlagLLMReranker
from .lightweight import LightweightLLMReranker as LightWeightFlagLLMReranker
__all__ = [
"FlagLLMReranker",
"LayerWiseFlagLLMReranker",
"LightWeightFlagLLMReranker"
]
FlagEmbedding/inference/reranker/decoder_only/base.py
+506
View File
@@ -0,0 +1,506 @@
import torch
import warnings
import numpy as np
from tqdm import tqdm, trange
from typing import Any, List, Union, Tuple, Optional
from peft import PeftModel
from torch import Tensor
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.utils.data import Dataset, DataLoader
from FlagEmbedding.abc.inference import AbsReranker
from FlagEmbedding.inference.reranker.encoder_only.base import sigmoid
def last_logit_pool(logits: Tensor,
attention_mask: Tensor) -> Tensor:
"""Pool the last logit.
Args:
logits (torch.Tensor): The output logits of the model.
attention_mask (torch.Tensor): Attention mask.
Returns:
torch.Tensor: The tensor after pooling.
"""
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
if left_padding:
return logits[:, -1, :]
else:
sequence_lengths = attention_mask.sum(dim=1) - 1
batch_size = logits.shape[0]
return torch.stack([logits[i, sequence_lengths[i], :] for i in range(batch_size)], dim=0)
class DatasetForReranker(Dataset):
"""Prepare the dataset for dataloader.
Args:
all_queries_inputs (_type_): All the input queries.
all_passages_inputs (_type_): All the input passages.
tokenizer_path (str): Path to the tokenizer to use.
max_len (int, optional): Maximum length of tokens. Defaults to :data:`512`.
cache_dir (Optional[str], optional): Cache directory for the tokenzier. Defaults to :data:`None`.
prompt (Optional[str], optional): Prompt for the specific task, will use the default if not provided.
Defaults to `None`.
"""
def __init__(
self,
all_queries_inputs,
all_passages_inputs,
tokenizer_path: str,
max_len: int = 512,
cache_dir: Optional[str] = None,
prompt: Optional[str] = None,
**kwargs: Any,
):
self.tokenizer = AutoTokenizer.from_pretrained(
tokenizer_path,
trust_remote_code=True,
cache_dir=cache_dir
)
self.all_queries_inputs = all_queries_inputs
self.all_passages_inputs = all_passages_inputs
self.max_len = max_len
self.total_len = len(self.all_queries_inputs)
self.kwargs = kwargs
if prompt is None:
prompt = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'."
self.prompt_inputs = self.tokenizer(
prompt,
return_tensors=None,
add_special_tokens=False
)['input_ids']
sep = "\n"
self.sep_inputs = self.tokenizer(
sep,
return_tensors=None,
add_special_tokens=False
)['input_ids']
self.encode_max_length = self.max_len + len(self.sep_inputs) + len(self.prompt_inputs)
def __len__(self):
return self.total_len
def __getitem__(self, item):
query_inputs = self.all_queries_inputs[item]
passage_inputs = self.all_passages_inputs[item]
if self.tokenizer.bos_token_id is not None and self.tokenizer.bos_token_id != self.tokenizer.pad_token_id:
item = self.tokenizer.prepare_for_model(
[self.tokenizer.bos_token_id] + query_inputs['input_ids'],
self.sep_inputs + passage_inputs['input_ids'],
truncation='only_second',
max_length=self.encode_max_length,
padding=False,
return_attention_mask=False,
return_token_type_ids=False,
add_special_tokens=False
)
else:
item = self.tokenizer.prepare_for_model(
query_inputs['input_ids'],
self.sep_inputs + passage_inputs['input_ids'],
truncation='only_second',
max_length=self.encode_max_length,
padding=False,
return_attention_mask=False,
return_token_type_ids=False,
add_special_tokens=False
)
item['input_ids'] = item['input_ids'] + self.sep_inputs + self.prompt_inputs
item['attention_mask'] = [1] * len(item['input_ids'])
item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
if 'position_ids' in item.keys():
item['position_ids'] = list(range(len(item['input_ids'])))
return item
class Collater:
"""
Collator of the reranker.
Args:
tokenizer (transformers.AutoTokenizer): The tokenizer for reranker.
max_len (int): Maximum length of tokens.
"""
def __init__(self, tokenizer, max_len):
self.tokenizer = tokenizer
self.max_len = max_len
self.pad_to_multiple_of = 8
self.label_pad_token_id = -100
warnings.filterwarnings("ignore",
message="`max_length` is ignored when `padding`=`True` and there is no truncation strategy.")
def __call__(self, data):
labels = [feature["labels"] for feature in data] if "labels" in data[0].keys() else None
# We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the
# same length to return tensors.
if labels is not None:
max_label_length = max(len(l) for l in labels)
if self.pad_to_multiple_of is not None:
max_label_length = (
(max_label_length + self.pad_to_multiple_of - 1)
// self.pad_to_multiple_of
* self.pad_to_multiple_of
)
padding_side = self.tokenizer.padding_side
for feature in data:
remainder = [self.label_pad_token_id] * (max_label_length - len(feature["labels"]))
if isinstance(feature["labels"], list):
feature["labels"] = (
feature["labels"] + remainder if padding_side == "right" else remainder + feature["labels"]
)
elif padding_side == "right":
feature["labels"] = np.concatenate([feature["labels"], remainder]).astype(np.int64)
else:
feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64)
return self.tokenizer.pad(
data,
padding=True,
pad_to_multiple_of=8,
return_tensors='pt',
)
class BaseLLMReranker(AbsReranker):
"""Base reranker class for LLM like decoder only models.
Args:
model_name_or_path (str): If it's a path to a local model, it loads the model from the path. Otherwise tries to download and
load a model from HuggingFace Hub with the name.
peft_path (Optional[str], optional): Path to the PEFT config. Defaults to :data:`None`.
use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance
degradation. Defaults to :data:`False`. Defaults to :data:`False`.
use_bf16 (bool, optional): Another type of half-precision floating-point, you can use bf16 if the hardware supports.
Defaults to :data:False.
query_instruction_for_rerank (str, optional): Query instruction for retrieval tasks, which will be used with
with :attr:`query_instruction_format`. Defaults to :data:`"A: "`.
query_instruction_format (str, optional): The template for :attr:`query_instruction_for_rerank`. Defaults to :data:`"{}{}"`.
passage_instruction_for_rerank (str, optional): Passage instruction for retrieval tasks, which will be used with
with :attr:`passage_instruction_format`. Defaults to :data:`"B: "`.
passage_instruction_format (str, optional): The template for passage. Defaults to "{}{}".
cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
trust_remote_code (bool, optional): trust_remote_code. Defaults to :data:`False`.
devices (Union[str, List[str], List[int]], optional): Devices to use for model inference, such as ["cuda:0"] or ["0"].
Defaults to :data:`None`.
prompt (Optional[str], optional): Prompt for the specific task. Defaults to :data:`None`.
batch_size (int, optional): Batch size for inference. Defaults to :data:`128`.
query_max_length (int, optional): Maximum length for queries. If not specified, will be 3/4 of :attr:`max_length`.
Defaults to :data:`None`.
max_length (int, optional): Maximum length of passages. Defaults to :data`512`.
normalize (bool, optional): If True, use Sigmoid to normalize the results. Defaults to :data:`False`.
"""
def __init__(
self,
model_name_or_path: str,
peft_path: Optional[str] = None,
use_fp16: bool = False,
use_bf16: bool = False,
query_instruction_for_rerank: str = "A: ",
query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_rerank
passage_instruction_for_rerank: str = "B: ",
passage_instruction_format: str = "{}{}", # specify the format of passage_instruction_for_rerank
cache_dir: Optional[str] = None,
trust_remote_code: bool = False,
devices: Union[str, List[str], List[int]] = None, # specify devices, such as ["cuda:0"] or ["0"]
# inference
prompt: Optional[str] = None,
batch_size: int = 128,
query_max_length: int = None,
max_length: int = 512,
normalize: bool = False,
**kwargs: Any,
) -> None:
super().__init__(
model_name_or_path=model_name_or_path,
use_fp16=use_fp16,
query_instruction_for_rerank=query_instruction_for_rerank,
query_instruction_format=query_instruction_format,
passage_instruction_for_rerank=passage_instruction_for_rerank,
passage_instruction_format=passage_instruction_format,
devices=devices,
batch_size=batch_size,
query_max_length=query_max_length,
max_length=max_length,
normalize=normalize,
prompt=prompt,
**kwargs
)
self.prompt = prompt
self.tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path,
cache_dir=cache_dir,
trust_remote_code=trust_remote_code
)
self.model = AutoModelForCausalLM.from_pretrained(
model_name_or_path,
cache_dir=cache_dir,
trust_remote_code=trust_remote_code,
torch_dtype=torch.bfloat16 if use_bf16 else torch.float32
)
if peft_path:
self.model = PeftModel.from_pretrained(self.model, peft_path)
self.model = self.model.merge_and_unload()
self.yes_loc = self.tokenizer('Yes', add_special_tokens=False)['input_ids'][0]
@torch.no_grad()
def compute_score_single_gpu(
self,
sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
batch_size: Optional[int] = None,
query_max_length: Optional[int] = None,
max_length: Optional[int] = None,
prompt: Optional[str] = None,
normalize: Optional[bool] = None,
use_dataloader: bool = False,
num_workers: int = None,
device: Optional[str] = None,
**kwargs: Any
) -> List[float]:
"""Compute the relevance scores using a single GPU.
Args:
sentence_pairs (Union[List[Tuple[str, str]], Tuple[str, str]]): Input sentence pairs to compute scores.
batch_size (Optional[int], optional): Number of inputs for each iter. Defaults to :data:`None`.
query_max_length (Optional[int], optional): Maximum length of tokens of queries. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
prompt (Optional[str], optional): Prompt for the specific task. Defaults to :data:`None`.
normalize (Optional[bool], optional): If True, use Sigmoid to normalize the results. Defaults to :data:`None`.
use_dataloader (bool, optional): If True, will use the dataloader to load the datasets. Defaults to :data:`False`.
num_workers (int, optional): Number of workers for dataloader. Defaults to :data:`None`.
device (Optional[str], optional): Device to use for computation. Defaults to :data:`None`.
Returns:
List[float]: The computed scores.
"""
if prompt is None: prompt = self.prompt
if batch_size is None: batch_size = self.batch_size
if max_length is None: max_length = self.max_length
if query_max_length is None:
if self.query_max_length is not None:
query_max_length = self.query_max_length
else:
query_max_length = max_length * 3 // 4
if normalize is None: normalize = self.normalize
if device is None:
device = self.target_devices[0]
if device == "cpu": self.use_fp16 = False
if self.use_fp16: self.model.half()
self.model.to(device)
self.model.eval()
assert isinstance(sentence_pairs, list)
if isinstance(sentence_pairs[0], str):
sentence_pairs = [sentence_pairs]
# tokenize without padding to get the correct length
all_queries_inputs = []
all_passages_inputs = []
for start_index in trange(0, len(sentence_pairs), batch_size, desc="pre tokenize",
disable=len(sentence_pairs) < batch_size):
sentences_batch = sentence_pairs[start_index:start_index + batch_size]
queries = [s[0] for s in sentences_batch]
passages = [s[1] for s in sentences_batch]
queries_inputs_batch = self.tokenizer(
queries,
return_tensors=None,
add_special_tokens=False,
max_length=query_max_length,
truncation=True,
**kwargs
)
passages_inputs_batch = self.tokenizer(
passages,
return_tensors=None,
add_special_tokens=False,
max_length=max_length,
truncation=True,
**kwargs
)
queries_inputs_batch = [{
k: queries_inputs_batch[k][i] for k in queries_inputs_batch.keys()
} for i in range(len(sentences_batch))]
passages_inputs_batch = [{
k: passages_inputs_batch[k][i] for k in passages_inputs_batch.keys()
} for i in range(len(sentences_batch))]
all_queries_inputs.extend(queries_inputs_batch)
all_passages_inputs.extend(passages_inputs_batch)
# sort by length for less padding
length_sorted_idx = np.argsort([-len(x['input_ids']) - len(y['input_ids']) for (x, y) in zip(all_queries_inputs, all_passages_inputs)])
all_queries_inputs_sorted = [all_queries_inputs[i] for i in length_sorted_idx]
all_passages_inputs_sorted = [all_passages_inputs[i] for i in length_sorted_idx]
# other inputs
if prompt is None:
prompt = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'."
prompt_inputs = self.tokenizer(
prompt,
return_tensors=None,
add_special_tokens=False
)['input_ids']
sep = "\n"
sep_inputs = self.tokenizer(
sep,
return_tensors=None,
add_special_tokens=False
)['input_ids']
encode_max_length = max_length + len(sep_inputs) + len(prompt_inputs)
# adjust batch size
flag = False
while flag is False:
try:
batch_inputs = []
for query_inputs, passage_inputs in zip(
all_queries_inputs_sorted[:min(len(all_queries_inputs_sorted), batch_size)],
all_passages_inputs_sorted[:min(len(all_passages_inputs_sorted), batch_size)]
):
if self.tokenizer.bos_token_id is not None and self.tokenizer.bos_token_id != self.tokenizer.pad_token_id:
item = self.tokenizer.prepare_for_model(
[self.tokenizer.bos_token_id] + query_inputs['input_ids'],
sep_inputs + passage_inputs['input_ids'],
truncation='only_second',
max_length=encode_max_length,
padding=False,
return_attention_mask=False,
return_token_type_ids=False,
add_special_tokens=False
)
else:
item = self.tokenizer.prepare_for_model(
query_inputs['input_ids'],
sep_inputs + passage_inputs['input_ids'],
truncation='only_second',
max_length=encode_max_length,
padding=False,
return_attention_mask=False,
return_token_type_ids=False,
add_special_tokens=False
)
item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
item['attention_mask'] = [1] * len(item['input_ids'])
item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
if 'position_ids' in item.keys():
item['position_ids'] = list(range(len(item['input_ids'])))
batch_inputs.append(item)
collater_instance = Collater(self.tokenizer, encode_max_length)
batch_inputs = collater_instance([{
'input_ids': item['input_ids'],
'attention_mask': item['attention_mask']
} for item in batch_inputs]
)
batch_inputs = {key: val.to(device) for key, val in batch_inputs.items()}
self.model(**batch_inputs, output_hidden_states=True)
flag = True
except RuntimeError as e:
batch_size = batch_size * 3 // 4
except torch.cuda.OutOfMemoryError as e:
batch_size = batch_size * 3 // 4
dataset, dataloader = None, None
if use_dataloader:
if num_workers is None:
num_workers = min(batch_size, 16)
dataset = DatasetForReranker(
all_queries_inputs_sorted,
all_passages_inputs_sorted,
self.model_name_or_path,
max_length,
cache_dir=self.cache_dir,
prompt=prompt,
**kwargs
)
dataloader = DataLoader(
dataset, shuffle=False, batch_size=batch_size, drop_last=False,
num_workers=num_workers,
collate_fn=Collater(self.tokenizer, encode_max_length)
)
all_scores = []
if dataloader is not None:
for inputs in tqdm(dataloader):
inputs = inputs.to(device)
outputs = self.model(**inputs, output_hidden_states=True)
logits = outputs.logits
scores = last_logit_pool(logits, inputs['attention_mask'])
scores = scores[:, self.yes_loc]
all_scores.extend(scores.cpu().float().tolist())
else:
for batch_start in trange(0, len(all_queries_inputs_sorted), batch_size):
queries_inputs = all_queries_inputs_sorted[batch_start:batch_start+batch_size]
passages_inputs = all_passages_inputs_sorted[batch_start:batch_start+batch_size]
batch_inputs = []
for query_inputs, passage_inputs in zip(queries_inputs, passages_inputs):
if self.tokenizer.bos_token_id is not None and self.tokenizer.bos_token_id != self.tokenizer.pad_token_id:
item = self.tokenizer.prepare_for_model(
[self.tokenizer.bos_token_id] + query_inputs['input_ids'],
sep_inputs + passage_inputs['input_ids'],
truncation='only_second',
max_length=encode_max_length,
padding=False,
return_attention_mask=False,
return_token_type_ids=False,
add_special_tokens=False
)
else:
item = self.tokenizer.prepare_for_model(
query_inputs['input_ids'],
sep_inputs + passage_inputs['input_ids'],
truncation='only_second',
max_length=encode_max_length,
padding=False,
return_attention_mask=False,
return_token_type_ids=False,
add_special_tokens=False
)
item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
item['attention_mask'] = [1] * len(item['input_ids'])
item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
if 'position_ids' in item.keys():
item['position_ids'] = list(range(len(item['input_ids'])))
batch_inputs.append(item)
collater_instance = Collater(self.tokenizer, encode_max_length)
batch_inputs = collater_instance([{
'input_ids': item['input_ids'],
'attention_mask': item['attention_mask']
} for item in batch_inputs]
)
batch_inputs = {key: val.to(device) for key, val in batch_inputs.items()}
outputs = self.model(**batch_inputs, output_hidden_states=True)
logits = outputs.logits
scores = last_logit_pool(logits, batch_inputs['attention_mask'])
scores = scores[:, self.yes_loc]
all_scores.extend(scores.cpu().float().tolist())
all_scores = [all_scores[idx] for idx in np.argsort(length_sorted_idx)]
if normalize:
all_scores = [sigmoid(score) for score in all_scores]
# if len(all_scores) == 1:
# return all_scores[0]
return all_scores
FlagEmbedding/inference/reranker/decoder_only/layerwise.py
+380
View File
@@ -0,0 +1,380 @@
import torch
import warnings
import numpy as np
from tqdm import tqdm, trange
from typing import Any, List, Union, Tuple, Optional
from peft import PeftModel
from torch import Tensor
from transformers import AutoModelForCausalLM, AutoTokenizer
from torch.utils.data import DataLoader
from FlagEmbedding.abc.inference import AbsReranker
from FlagEmbedding.inference.reranker.encoder_only.base import sigmoid
from FlagEmbedding.inference.reranker.decoder_only.base import DatasetForReranker, Collater
from .models.modeling_minicpm_reranker import LayerWiseMiniCPMForCausalLM
def last_logit_pool_layerwise(logits: Tensor,
attention_mask: Tensor) -> Tensor:
"""Pool the last logit.
Args:
logits (torch.Tensor): The output logits of the model.
attention_mask (torch.Tensor): Attention mask.
Returns:
torch.Tensor: The tensor after pooling.
"""
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
if left_padding:
return logits[:, -1]
else:
sequence_lengths = attention_mask.sum(dim=1) - 1
batch_size = logits.shape[0]
return logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
class LayerWiseLLMReranker(AbsReranker):
"""Base reranker class for layerwise LLM like decoder only models.
Args:
model_name_or_path (str): If it's a path to a local model, it loads the model from the path. Otherwise tries to download and
load a model from HuggingFace Hub with the name.
peft_path (Optional[str], optional): Path to the PEFT config. Defaults to :data:`None`.
use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance
degradation. Defaults to :data:`False`. Defaults to :data:`False`.
use_bf16 (bool, optional): Another type of half-precision floating-point, you can use bf16 if the hardware supports.
Defaults to :data:False.
query_instruction_for_rerank (str, optional): Query instruction for retrieval tasks, which will be used with
with :attr:`query_instruction_format`. Defaults to :data:`"A: "`.
query_instruction_format (str, optional): The template for :attr:`query_instruction_for_rerank`. Defaults to :data:`"{}{}"`.
passage_instruction_for_rerank (str, optional): Passage instruction for retrieval tasks, which will be used with
with :attr:`passage_instruction_format`. Defaults to :data:`"B: "`.
passage_instruction_format (str, optional): The template for passage. Defaults to "{}{}".
cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
trust_remote_code (bool, optional): trust_remote_code. Defaults to :data:`False`.
devices (Union[str, List[str], List[int]], optional): Devices to use for model inference, such as ["cuda:0"] or ["0"].
Defaults to :data:`None`.
cutoff_layers (Optional[List[int]]): Pick which layers are used for computing the score. Defaults to :data:`None`.
prompt (Optional[str], optional): Prompt for the specific task. Defaults to :data:`None`.
batch_size (int, optional): Batch size for inference. Defaults to :data:`128`.
query_max_length (int, optional): Maximum length for queries. If not specified, will be 3/4 of :attr:`max_length`.
Defaults to :data:`None`.
max_length (int, optional): Maximum length of passages. Defaults to :data`512`.
normalize (bool, optional): If True, use Sigmoid to normalize the results. Defaults to :data:`False`.
"""
def __init__(
self,
model_name_or_path: str,
peft_path: Optional[str] = None,
use_fp16: bool = False,
use_bf16: bool = False,
query_instruction_for_rerank: str = "A: ",
query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_rerank
passage_instruction_for_rerank: str = "B: ",
passage_instruction_format: str = "{}{}", # specify the format of passage_instruction_for_rerank
cache_dir: Optional[str] = None,
trust_remote_code: bool = False,
devices: Optional[Union[str, List[str], List[int]]] = None, # specify devices, such as ["cuda:0"] or ["0"]
# inference
cutoff_layers: Optional[List[int]] = None,
prompt: Optional[str] = None,
batch_size: int = 128,
query_max_length: Optional[int] = None,
max_length: int = 512,
normalize: bool = False,
**kwargs: Any,
) -> None:
super().__init__(
model_name_or_path=model_name_or_path,
use_fp16=use_fp16,
query_instruction_for_rerank=query_instruction_for_rerank,
query_instruction_format=query_instruction_format,
passage_instruction_for_rerank=passage_instruction_for_rerank,
passage_instruction_format=passage_instruction_format,
devices=devices,
batch_size=batch_size,
query_max_length=query_max_length,
max_length=max_length,
normalize=normalize,
**kwargs
)
self.cutoff_layers = cutoff_layers
self.prompt = prompt
self.tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path,
cache_dir=cache_dir,
trust_remote_code=trust_remote_code
)
if use_bf16 is False and use_fp16 is False:
warnings.warn("Due to model constraints, `use_bf16` and `use_fp16` cannot both be `False`. Here, `use_fp16` is set to `True` by default.", UserWarning)
self.use_fp16 = True
try:
self.model = LayerWiseMiniCPMForCausalLM.from_pretrained(
model_name_or_path,
cache_dir=cache_dir,
trust_remote_code=trust_remote_code,
torch_dtype=torch.bfloat16 if use_bf16 else torch.float32
)
except:
self.model = AutoModelForCausalLM.from_pretrained(
model_name_or_path,
cache_dir=cache_dir,
trust_remote_code=trust_remote_code,
torch_dtype=torch.bfloat16 if use_bf16 else torch.float32
)
if peft_path:
self.model = PeftModel.from_pretrained(self.model,peft_path)
self.model = self.model.merge_and_unload()
@torch.no_grad()
def compute_score_single_gpu(
self,
sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
batch_size: Optional[int] = None,
query_max_length: Optional[int] = None,
max_length: Optional[int] = None,
cutoff_layers: Optional[List[int]] = None,
prompt: Optional[str] = None,
normalize: Optional[bool] = None,
use_dataloader: bool = False,
num_workers: Optional[int] = None,
device: Optional[str] = None,
**kwargs: Any
) -> List[float]:
"""Compute the relevance scores using a single GPU.
Args:
sentence_pairs (Union[List[Tuple[str, str]], Tuple[str, str]]): Input sentence pairs to compute scores.
batch_size (Optional[int], optional): Number of inputs for each iter. Defaults to :data:`None`.
query_max_length (Optional[int], optional): Maximum length of tokens of queries. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
cutoff_layers (Optional[List[int]], optional): Pick which layers are used for computing the score. Defaults to :data:`None`.
prompt (Optional[str], optional): Prompt for the specific task. Defaults to :data:`None`.
normalize (Optional[bool], optional): If True, use Sigmoid to normalize the results. Defaults to :data:`None`.
use_dataloader (bool, optional): If True, will use the dataloader to load the datasets. Defaults to :data:`False`.
num_workers (int, optional): Number of workers for dataloader. Defaults to :data:`None`.
device (Optional[str], optional): Device to use for computation. Defaults to :data:`None`.
Returns:
List[float]: The computed scores.
"""
if cutoff_layers is None: cutoff_layers = self.cutoff_layers
if prompt is None: prompt = self.prompt
if batch_size is None: batch_size = self.batch_size
if max_length is None: max_length = self.max_length
if query_max_length is None:
if self.query_max_length is not None:
query_max_length = self.query_max_length
else:
query_max_length = max_length * 3 // 4
if normalize is None: normalize = self.normalize
if device is None:
device = self.target_devices[0]
if device == "cpu": self.use_fp16 = False
if self.use_fp16: self.model.half()
self.model.to(device)
self.model.eval()
assert isinstance(sentence_pairs, list)
if isinstance(sentence_pairs[0], str):
sentence_pairs = [sentence_pairs]
# tokenize without padding to get the correct length
all_queries_inputs = []
all_passages_inputs = []
for start_index in trange(0, len(sentence_pairs), batch_size, desc="pre tokenize",
disable=len(sentence_pairs) < batch_size):
sentences_batch = sentence_pairs[start_index:start_index + batch_size]
queries = [s[0] for s in sentences_batch]
passages = [s[1] for s in sentences_batch]
queries_inputs_batch = self.tokenizer(
queries,
return_tensors=None,
add_special_tokens=False,
max_length=query_max_length,
truncation=True,
**kwargs
)
passages_inputs_batch = self.tokenizer(
passages,
return_tensors=None,
add_special_tokens=False,
max_length=max_length,
truncation=True,
**kwargs
)
queries_inputs_batch = [{
k: queries_inputs_batch[k][i] for k in queries_inputs_batch.keys()
} for i in range(len(sentences_batch))]
passages_inputs_batch = [{
k: passages_inputs_batch[k][i] for k in passages_inputs_batch.keys()
} for i in range(len(sentences_batch))]
all_queries_inputs.extend(queries_inputs_batch)
all_passages_inputs.extend(passages_inputs_batch)
# sort by length for less padding
length_sorted_idx = np.argsort([-len(x['input_ids']) - len(y['input_ids']) for (x, y) in zip(all_queries_inputs, all_passages_inputs)])
all_queries_inputs_sorted = [all_queries_inputs[i] for i in length_sorted_idx]
all_passages_inputs_sorted = [all_passages_inputs[i] for i in length_sorted_idx]
# other inputs
if prompt is None:
prompt = "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'."
prompt_inputs = self.tokenizer(
prompt,
return_tensors=None,
add_special_tokens=False
)['input_ids']
sep = "\n"
sep_inputs = self.tokenizer(
sep,
return_tensors=None,
add_special_tokens=False
)['input_ids']
encode_max_length = max_length + len(sep_inputs) + len(prompt_inputs)
# adjust batch size
flag = False
while flag is False:
try:
batch_inputs = []
for query_inputs, passage_inputs in zip(
all_queries_inputs_sorted[:min(len(all_queries_inputs_sorted), batch_size)],
all_passages_inputs_sorted[:min(len(all_passages_inputs_sorted), batch_size)]
):
item = self.tokenizer.prepare_for_model(
[self.tokenizer.bos_token_id] + query_inputs['input_ids'],
sep_inputs + passage_inputs['input_ids'],
truncation='only_second',
max_length=encode_max_length,
padding=False,
return_attention_mask=False,
return_token_type_ids=False,
add_special_tokens=False
)
item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
item['attention_mask'] = [1] * len(item['input_ids'])
item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
if 'position_ids' in item.keys():
item['position_ids'] = list(range(len(item['input_ids'])))
batch_inputs.append(item)
collater_instance = Collater(self.tokenizer, encode_max_length)
batch_inputs = collater_instance([{
'input_ids': item['input_ids'],
'attention_mask': item['attention_mask']
} for item in batch_inputs]
)
batch_inputs = {key: val.to(device) for key, val in batch_inputs.items()}
self.model(**batch_inputs, output_hidden_states=True, cutoff_layers=cutoff_layers)
flag = True
except RuntimeError as e:
batch_size = batch_size * 3 // 4
except torch.cuda.OutOfMemoryError as e:
batch_size = batch_size * 3 // 4
dataset, dataloader = None, None
if use_dataloader:
if num_workers is None:
num_workers = min(batch_size, 16)
dataset = DatasetForReranker(
all_queries_inputs_sorted,
all_passages_inputs_sorted,
self.model_name_or_path,
max_length,
cache_dir=self.cache_dir,
prompt=prompt,
**kwargs
)
dataloader = DataLoader(
dataset, shuffle=False, batch_size=batch_size, drop_last=False,
num_workers=num_workers,
collate_fn=Collater(self.tokenizer, encode_max_length)
)
all_scores = []
if dataloader is not None:
for inputs in tqdm(dataloader):
inputs = inputs.to(device)
outputs = self.model(**inputs, output_hidden_states=True, cutoff_layers=cutoff_layers)
all_logits = outputs.logits
tmp_all_scores = []
for logits in all_logits:
scores = last_logit_pool_layerwise(logits, inputs['attention_mask'])
tmp_all_scores.append(scores.contiguous())
if len(all_scores) == 0:
for _ in range(len(tmp_all_scores)):
all_scores.append([])
for i in range(len(tmp_all_scores)):
all_scores[i].extend(tmp_all_scores[i].cpu().float().tolist())
else:
for batch_start in trange(0, len(all_queries_inputs_sorted), batch_size):
queries_inputs = all_queries_inputs_sorted[batch_start:batch_start+batch_size]
passages_inputs = all_passages_inputs_sorted[batch_start:batch_start+batch_size]
batch_inputs = []
for query_inputs, passage_inputs in zip(queries_inputs, passages_inputs):
item = self.tokenizer.prepare_for_model(
[self.tokenizer.bos_token_id] + query_inputs['input_ids'],
sep_inputs + passage_inputs['input_ids'],
truncation='only_second',
max_length=encode_max_length,
padding=False,
return_attention_mask=False,
return_token_type_ids=False,
add_special_tokens=False
)
item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
item['attention_mask'] = [1] * len(item['input_ids'])
item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
if 'position_ids' in item.keys():
item['position_ids'] = list(range(len(item['input_ids'])))
batch_inputs.append(item)
collater_instance = Collater(self.tokenizer, encode_max_length)
batch_inputs = collater_instance([{
'input_ids': item['input_ids'],
'attention_mask': item['attention_mask']
} for item in batch_inputs]
)
batch_inputs = {key: val.to(device) for key, val in batch_inputs.items()}
outputs = self.model(**batch_inputs, output_hidden_states=True, cutoff_layers=cutoff_layers)
all_logits = outputs.logits
tmp_all_scores = []
for logits in all_logits:
scores = last_logit_pool_layerwise(logits, batch_inputs['attention_mask'])
tmp_all_scores.append(scores.contiguous())
if len(all_scores) == 0:
for _ in range(len(tmp_all_scores)):
all_scores.append([])
for i in range(len(tmp_all_scores)):
all_scores[i].extend(tmp_all_scores[i].cpu().float().tolist())
for i in range(len(all_scores)):
all_scores[i] = [all_scores[i][idx] for idx in np.argsort(length_sorted_idx)]
if normalize:
all_scores[i] = [sigmoid(score) for score in all_scores[i]]
if len(all_scores) == 1 and isinstance(all_scores[0], list):
all_scores = all_scores[0]
return all_scores
FlagEmbedding/inference/reranker/decoder_only/lightweight.py
+449
View File
@@ -0,0 +1,449 @@
import torch
import sys
import warnings
import numpy as np
from tqdm import trange
from typing import Any, List, Union, Tuple, Optional
from peft import PeftModel
from torch import Tensor
from transformers import AutoModelForCausalLM, AutoTokenizer
from FlagEmbedding.abc.inference import AbsReranker
from FlagEmbedding.inference.reranker.encoder_only.base import sigmoid
def last_logit_pool_lightweight(logits: Tensor,
attention_mask: Tensor) -> Tensor:
"""Pool the last logit.
Args:
logits (torch.Tensor): The output logits of the model.
attention_mask (torch.Tensor): Attention mask.
Returns:
torch.Tensor: The tensor after pooling.
"""
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
if left_padding:
return logits[:, -1]
else:
sequence_lengths = attention_mask.sum(dim=1) - 1
batch_size = logits.shape[0]
return torch.stack([logits[i, sequence_lengths[i]] for i in range(batch_size)], dim=0)
class Collater_for_lightweight:
"""
Collator of the lightweight LLM reranker.
Args:
tokenizer (transformers.AutoTokenizer): The tokenizer for reranker.
max_len (int): Maximum length of tokens.
"""
def __init__(self, tokenizer, max_len):
self.tokenizer = tokenizer
self.max_len = max_len
self.pad_to_multiple_of = 8
self.label_pad_token_id = -100
warnings.filterwarnings("ignore",
message="`max_length` is ignored when `padding`=`True` and there is no truncation strategy.")
def __call__(self, data):
features = data[0]
query_lengths = data[1]
prompt_lengths = data[2]
labels = [feature["labels"] for feature in features] if "labels" in features[0].keys() else None
# We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the
# same length to return tensors.
if labels is not None:
max_label_length = max(len(l) for l in labels)
if self.pad_to_multiple_of is not None:
max_label_length = (
(max_label_length + self.pad_to_multiple_of - 1)
// self.pad_to_multiple_of
* self.pad_to_multiple_of
)
padding_side = self.tokenizer.padding_side
for feature in features:
remainder = [self.label_pad_token_id] * (max_label_length - len(feature["labels"]))
if isinstance(feature["labels"], list):
feature["labels"] = (
feature["labels"] + remainder if padding_side == "right" else remainder + feature["labels"]
)
elif padding_side == "right":
feature["labels"] = np.concatenate([feature["labels"], remainder]).astype(np.int64)
else:
feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64)
collected = self.tokenizer.pad(
features,
padding=True,
pad_to_multiple_of=8,
return_tensors='pt',
)
return collected, query_lengths, prompt_lengths
class LightweightLLMReranker(AbsReranker):
"""Base reranker class for light weight LLM like decoder only models.
Args:
model_name_or_path (str): If it's a path to a local model, it loads the model from the path. Otherwise tries to download and
load a model from HuggingFace Hub with the name.
peft_path (Optional[str], optional): Path to the PEFT config. Defaults to :data:`None`.
use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance
degradation. Defaults to :data:`False`. Defaults to :data:`False`.
use_bf16 (bool, optional): Another type of half-precision floating-point, you can use bf16 if the hardware supports.
Defaults to :data:False.
query_instruction_for_rerank (str, optional): Query instruction for retrieval tasks, which will be used with
with :attr:`query_instruction_format`. Defaults to :data:`"A: "`.
query_instruction_format (str, optional): The template for :attr:`query_instruction_for_rerank`. Defaults to :data:`"{}{}"`.
passage_instruction_for_rerank (str, optional): Passage instruction for retrieval tasks, which will be used with
with :attr:`passage_instruction_format`. Defaults to :data:`"B: "`.
passage_instruction_format (str, optional): The template for passage. Defaults to "{}{}".
cache_dir (Optional[str], optional): Cache directory for the model. Defaults to :data:`None`.
trust_remote_code (bool, optional): trust_remote_code. Defaults to :data:`False`.
devices (Union[str, List[str], List[int]], optional): Devices to use for model inference, such as ["cuda:0"] or ["0"].
Defaults to :data:`None`.
cutoff_layers (Optional[List[int]]): Pick which layers are used for computing the score. Defaults to :data:`None`.
compress_layers (List[int], optional): Choose the layers to compress. Defaults to :data:`[8]`.
compress_ratio (int, optional): Ratio to compress the selected layers, supported ratios: :data:`[1, 2, 4, 8]`.
Defaults to :data:`1`.
prompt (Optional[str], optional): Prompt for the specific task. Defaults to :data:`None`.
batch_size (int, optional): Batch size for inference. Defaults to :data:`128`.
query_max_length (int, optional): Maximum length for queries. If not specified, will be 3/4 of :attr:`max_length`.
Defaults to :data:`None`.
max_length (int, optional): Maximum length of passages. Defaults to :data`512`.
normalize (bool, optional): If True, use Sigmoid to normalize the results. Defaults to :data:`False`.
"""
def __init__(
self,
model_name_or_path: str,
peft_path: Optional[str] = None,
use_fp16: bool = False,
use_bf16: bool = False,
query_instruction_for_rerank: str = "A: ",
query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_rerank
passage_instruction_for_rerank: str = "B: ",
passage_instruction_format: str = "{}{}", # specify the format of passage_instruction_for_rerank
cache_dir: Optional[str] = None,
trust_remote_code: bool = False,
devices: Union[str, List[str], List[int]] = None, # specify devices, such as ["cuda:0"] or ["0"]
# inference
cutoff_layers: Optional[List[int]] = None,
compress_layers: List[int] = [8],
compress_ratio: int = 1,
prompt: Optional[str] = None,
batch_size: int = 128,
query_max_length: Optional[int] = None,
max_length: int = 512,
normalize: bool = False,
**kwargs: Any,
) -> None:
try:
from .models.gemma_model import CostWiseGemmaForCausalLM
except:
print('*') * 20
print('*') * 20
print('error for load lightweight reranker, please install transformers==4.46.0')
print('*') * 20
print('*') * 20
sys.exit()
super().__init__(
model_name_or_path=model_name_or_path,
use_fp16=use_fp16,
query_instruction_for_rerank=query_instruction_for_rerank,
query_instruction_format=query_instruction_format,
passage_instruction_for_rerank=passage_instruction_for_rerank,
passage_instruction_format=passage_instruction_format,
devices=devices,
batch_size=batch_size,
query_max_length=query_max_length,
max_length=max_length,
normalize=normalize,
**kwargs
)
self.cutoff_layers = cutoff_layers
self.compress_layers = compress_layers
self.compress_ratio = compress_ratio
self.prompt = prompt
self.tokenizer = AutoTokenizer.from_pretrained(
model_name_or_path,
cache_dir=cache_dir,
trust_remote_code=trust_remote_code
)
self.tokenizer.padding_side = 'right'
if use_bf16 is False and use_fp16 is False:
warnings.warn("Due to model constraints, `use_bf16` and `use_fp16` cannot both be `False`. Here, `use_fp16` is set to `True` by default.", UserWarning)
use_fp16 = True
try:
self.model = CostWiseGemmaForCausalLM.from_pretrained(
model_name_or_path,
cache_dir=cache_dir,
trust_remote_code=trust_remote_code,
torch_dtype=torch.bfloat16 if use_bf16 else torch.float32
)
except:
self.model = AutoModelForCausalLM.from_pretrained(
model_name_or_path,
cache_dir=cache_dir,
trust_remote_code=trust_remote_code,
torch_dtype=torch.bfloat16 if use_bf16 else torch.float32
)
if peft_path:
self.model = PeftModel.from_pretrained(self.model,peft_path)
self.model = self.model.merge_and_unload()
@torch.no_grad()
def compute_score_single_gpu(
self,
sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
batch_size: Optional[int] = None,
query_max_length: Optional[int] = None,
max_length: Optional[int] = None,
cutoff_layers: Optional[List[int]] = None,
compress_layer: Optional[List[int]] = None,
compress_layers: Optional[List[int]] = None,
compress_ratio: Optional[int] = None,
prompt: Optional[str] = None,
normalize: Optional[bool] = None,
device: Optional[str] = None,
**kwargs: Any
) -> List[float]:
"""Compute the relevance scores using a single GPU.
Args:
sentence_pairs (Union[List[Tuple[str, str]], Tuple[str, str]]): Input sentence pairs to compute scores.
batch_size (Optional[int], optional): Number of inputs for each iter. Defaults to :data:`None`.
query_max_length (Optional[int], optional): Maximum length of tokens of queries. Defaults to :data:`None`.
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
cutoff_layers (Optional[List[int]], optional): Pick which layers are used for computing the score. Defaults to :data:`None`.
compress_layer (Optional[List[int]]): Deprecated, use :attr:`compress_layers` instead. Defaults to :data:`None`.
compress_layers (Optional[List[int]]): Selected layers to compress. Defaults to :data:`None`.
compress_ratio (Optional[int]): Ratio to compress the selected layers, supported ratios: :data:`[1, 2, 4, 8]`.
Defaults to :data:`None`.
prompt (Optional[str], optional): Prompt for the specific task. Defaults to :data:`None`.
normalize (Optional[bool], optional): If True, use Sigmoid to normalize the results. Defaults to :data:`None`.
device (Optional[str], optional): Device to use for computation. Defaults to :data:`None`.
Returns:
List[float]: The computed scores.
"""
if cutoff_layers is None: cutoff_layers = self.cutoff_layers
if compress_layers is None: compress_layers = self.compress_layers
if compress_layer is not None:
print('Try not to use the parameter `compress_layer`; use `compress_layers` instead.')
compress_layers = compress_layer
if compress_ratio is None: compress_ratio = self.compress_ratio
if prompt is None: prompt = self.prompt
if batch_size is None: batch_size = self.batch_size
if max_length is None: max_length = self.max_length
if query_max_length is None:
if self.query_max_length is not None:
query_max_length = self.query_max_length
else:
query_max_length = max_length * 3 // 4
if normalize is None: normalize = self.normalize
if device is None:
device = self.target_devices[0]
if device == "cpu": self.use_fp16 = False
if self.use_fp16: self.model.half()
self.model.to(device)
self.model.eval()
assert isinstance(sentence_pairs, list)
if isinstance(sentence_pairs[0], str):
sentence_pairs = [sentence_pairs]
# tokenize without padding to get the correct length
all_queries_inputs = []
all_passages_inputs = []
for start_index in trange(0, len(sentence_pairs), batch_size, desc="pre tokenize",
disable=len(sentence_pairs) < batch_size):
sentences_batch = sentence_pairs[start_index:start_index + batch_size]
queries = [s[0] for s in sentences_batch]
passages = [s[1] for s in sentences_batch]
queries_inputs_batch = self.tokenizer(
queries,
return_tensors=None,
add_special_tokens=False,
max_length=query_max_length,
truncation=True,
**kwargs
)
passages_inputs_batch = self.tokenizer(
passages,
return_tensors=None,
add_special_tokens=False,
max_length=max_length,
truncation=True,
**kwargs
)
queries_inputs_batch = [{
k: queries_inputs_batch[k][i] for k in queries_inputs_batch.keys()
} for i in range(len(sentences_batch))]
passages_inputs_batch = [{
k: passages_inputs_batch[k][i] for k in passages_inputs_batch.keys()
} for i in range(len(sentences_batch))]
all_queries_inputs.extend(queries_inputs_batch)
all_passages_inputs.extend(passages_inputs_batch)
# sort by length for less padding
length_sorted_idx = np.argsort([-len(x['input_ids']) - len(y['input_ids']) for (x, y) in zip(all_queries_inputs, all_passages_inputs)])
all_queries_inputs_sorted = [all_queries_inputs[i] for i in length_sorted_idx]
all_passages_inputs_sorted = [all_passages_inputs[i] for i in length_sorted_idx]
# other inputs
if prompt is None:
prompt = "Predict whether passage B contains an answer to query A."
prompt_inputs = self.tokenizer(
prompt,
return_tensors=None,
add_special_tokens=False
)['input_ids']
sep = "\n"
sep_inputs = self.tokenizer(
sep,
return_tensors=None,
add_special_tokens=False
)['input_ids']
encode_max_length = max_length + len(sep_inputs) + len(prompt_inputs)
# adjust batch size
flag = False
while flag is False:
try:
batch_inputs = []
query_lengths = []
prompt_lengths = []
for query_inputs, passage_inputs in zip(
all_queries_inputs_sorted[:min(len(all_queries_inputs_sorted), batch_size)],
all_passages_inputs_sorted[:min(len(all_passages_inputs_sorted), batch_size)]
):
item = self.tokenizer.prepare_for_model(
[self.tokenizer.bos_token_id] + query_inputs['input_ids'],
sep_inputs + passage_inputs['input_ids'],
truncation='only_second',
max_length=encode_max_length,
padding=False,
return_attention_mask=False,
return_token_type_ids=False,
add_special_tokens=False
)
item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
item['attention_mask'] = [1] * len(item['input_ids'])
item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
if 'position_ids' in item.keys():
item['position_ids'] = list(range(len(item['input_ids'])))
batch_inputs.append(item)
query_lengths.append(len([self.tokenizer.bos_token_id] + query_inputs['input_ids'] + sep_inputs))
prompt_lengths.append(len(sep_inputs + prompt_inputs))
collater_instance = Collater_for_lightweight(self.tokenizer, max_length)
batch_inputs = collater_instance([
[{
'input_ids': item['input_ids'],
'attention_mask': item['attention_mask']
} for item in batch_inputs],
query_lengths,
prompt_lengths
])[0]
batch_inputs = {key: val.to(device) for key, val in batch_inputs.items()}
self.model(
**batch_inputs,
output_hidden_states=True,
compress_layer=compress_layers,
compress_ratio=compress_ratio,
query_lengths=query_lengths,
prompt_lengths=prompt_lengths,
cutoff_layers=cutoff_layers
)
flag = True
except RuntimeError as e:
batch_size = batch_size * 3 // 4
except torch.cuda.OutOfMemoryError as e:
batch_size = batch_size * 3 // 4
all_scores = []
for batch_start in trange(0, len(all_queries_inputs_sorted), batch_size):
queries_inputs = all_queries_inputs_sorted[batch_start:batch_start+batch_size]
passages_inputs = all_passages_inputs_sorted[batch_start:batch_start+batch_size]
batch_inputs = []
query_lengths = []
prompt_lengths = []
for query_inputs, passage_inputs in zip(queries_inputs, passages_inputs):
item = self.tokenizer.prepare_for_model(
[self.tokenizer.bos_token_id] + query_inputs['input_ids'],
sep_inputs + passage_inputs['input_ids'],
truncation='only_second',
max_length=encode_max_length,
padding=False,
return_attention_mask=False,
return_token_type_ids=False,
add_special_tokens=False
)
item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
item['attention_mask'] = [1] * len(item['input_ids'])
item.pop('token_type_ids') if 'token_type_ids' in item.keys() else None
if 'position_ids' in item.keys():
item['position_ids'] = list(range(len(item['input_ids'])))
batch_inputs.append(item)
query_lengths.append(len([self.tokenizer.bos_token_id] + query_inputs['input_ids'] + sep_inputs))
prompt_lengths.append(len(sep_inputs + prompt_inputs))
collater_instance = Collater_for_lightweight(self.tokenizer, max_length)
batch_inputs = collater_instance([
[{
'input_ids': item['input_ids'],
'attention_mask': item['attention_mask']
} for item in batch_inputs],
query_lengths,
prompt_lengths
])[0]
batch_inputs = {key: val.to(device) for key, val in batch_inputs.items()}
outputs = self.model(
**batch_inputs,
output_hidden_states=True,
compress_layer=compress_layers,
compress_ratio=compress_ratio,
query_lengths=query_lengths,
prompt_lengths=prompt_lengths,
cutoff_layers=cutoff_layers
)
scores = []
for i in range(len(outputs.logits)):
logits = last_logit_pool_lightweight(outputs.logits[i], outputs.attention_masks[i])
scores.append(logits.cpu().float().tolist())
if len(all_scores) == 0:
for i in range(len(scores)):
all_scores.append([])
for i in range(len(scores)):
all_scores[i].extend(scores[i])
for i in range(len(all_scores)):
all_scores[i] = [all_scores[i][idx] for idx in np.argsort(length_sorted_idx)]
if normalize:
all_scores[i] = [sigmoid(score) for score in all_scores[i]]
if len(all_scores) == 1 and isinstance(all_scores[0], list):
all_scores = all_scores[0]
return all_scores
FlagEmbedding/inference/reranker/decoder_only/models/__init__.py
View File
FlagEmbedding/inference/reranker/decoder_only/models/configuration_minicpm_reranker.py
+208
View File
@@ -0,0 +1,208 @@
# coding=utf-8
# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
#
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
# and OPT implementations in this library. It has been modified from its
# original forms to accommodate minor architectural differences compared
# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
""" MiniCPM model configuration"""
from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging
logger = logging.get_logger(__name__)
MINICPM_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
class LayerWiseMiniCPMConfig(PretrainedConfig):
r"""
This is the configuration class to store the configuration of a [`MiniCPMModel`]. It is used to instantiate an MiniCPM
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
defaults will yield a similar configuration to that of the MiniCPM-7B.
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.
Args:
vocab_size (`int`, *optional*, defaults to 32000):
Vocabulary size of the MiniCPM model. Defines the number of different tokens that can be represented by the
`inputs_ids` passed when calling [`MiniCPMModel`]
hidden_size (`int`, *optional*, defaults to 4096):
Dimension of the hidden representations.
intermediate_size (`int`, *optional*, defaults to 11008):
Dimension of the MLP representations.
num_hidden_layers (`int`, *optional*, defaults to 32):
Number of hidden layers in the Transformer decoder.
num_attention_heads (`int`, *optional*, defaults to 32):
Number of attention heads for each attention layer in the Transformer decoder.
num_key_value_heads (`int`, *optional*):
This is the number of key_value heads that should be used to implement Grouped Query Attention. If
`num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
`num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
by meanpooling all the original heads within that group. For more details checkout [this
paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
`num_attention_heads`.
hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
The non-linear activation function (function or string) in the decoder.
max_position_embeddings (`int`, *optional*, defaults to 2048):
The maximum sequence length that this model might ever be used with. MiniCPM 1 supports up to 2048 tokens,
MiniCPM 2 up to 4096, CodeMiniCPM up to 16384.
initializer_range (`float`, *optional*, defaults to 0.02):
The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
rms_norm_eps (`float`, *optional*, defaults to 1e-06):
The epsilon used by the rms normalization layers.
use_cache (`bool`, *optional*, defaults to `True`):
Whether or not the model should return the last key/values attentions (not used by all models). Only
relevant if `config.is_decoder=True`.
pad_token_id (`int`, *optional*):
Padding token id.
bos_token_id (`int`, *optional*, defaults to 1):
Beginning of stream token id.
eos_token_id (`int`, *optional*, defaults to 2):
End of stream token id.
pretraining_tp (`int`, *optional*, defaults to 1):
Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
issue](https://github.com/pytorch/pytorch/issues/76232).
tie_word_embeddings (`bool`, *optional*, defaults to `False`):
Whether to tie weight embeddings
rope_theta (`float`, *optional*, defaults to 10000.0):
The base period of the RoPE embeddings.
rope_scaling (`Dict`, *optional*):
Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
`{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
`max_position_embeddings` to the expected new maximum. See the following thread for more information on how
these scaling strategies behave:
https://www.reddit.com/r/LocalMiniCPM/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
experimental feature, subject to breaking API changes in future versions.
attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
Whether to use a bias in the query, key, value and output projection layers during self-attention.
attention_dropout (`float`, *optional*, defaults to 0.0):
The dropout ratio for the attention probabilities.
```python
>>> from transformers import MiniCPMModel, MiniCPMConfig
>>> # Initializing a MiniCPM minicpm-7b style configuration
>>> configuration = MiniCPMConfig()
>>> # Initializing a model from the minicpm-7b style configuration
>>> model = MiniCPMModel(configuration)
>>> # Accessing the model configuration
>>> configuration = model.config
```"""
model_type = "minicpm"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
vocab_size=32000,
hidden_size=4096,
intermediate_size=11008,
num_hidden_layers=32,
num_attention_heads=32,
num_key_value_heads=None,
hidden_act="silu",
max_position_embeddings=2048,
initializer_range=0.02,
rms_norm_eps=1e-6,
use_cache=True,
pad_token_id=None,
bos_token_id=1,
eos_token_id=2,
pretraining_tp=1,
tie_word_embeddings=True,
rope_theta=10000.0,
rope_scaling=None,
attention_bias=False,
attention_dropout=0.0,
scale_emb=1,
dim_model_base=1,
scale_depth=1,
start_layer=8,
head_multi=True,
head_type="simple",
**kwargs,
):
self.vocab_size = vocab_size
self.max_position_embeddings = max_position_embeddings
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.num_hidden_layers = num_hidden_layers
self.num_attention_heads = num_attention_heads
# for backward compatibility
if num_key_value_heads is None:
num_key_value_heads = num_attention_heads
self.num_key_value_heads = num_key_value_heads
self.hidden_act = hidden_act
self.initializer_range = initializer_range
self.rms_norm_eps = rms_norm_eps
self.pretraining_tp = pretraining_tp
self.use_cache = use_cache
self.rope_theta = rope_theta
self.rope_scaling = rope_scaling
self._rope_scaling_validation()
self.attention_bias = attention_bias
self.attention_dropout = attention_dropout
self.scale_emb = scale_emb
self.dim_model_base = dim_model_base
self.scale_depth = scale_depth
self.start_layer = start_layer
self.head_multi = head_multi
self.head_type = head_type
super().__init__(
pad_token_id=pad_token_id,
bos_token_id=bos_token_id,
eos_token_id=eos_token_id,
tie_word_embeddings=tie_word_embeddings,
**kwargs,
)
try:
import flash_attn
self._attn_implementation = "flash_attention_2"
except:
pass
def _rope_scaling_validation(self):
"""
Validate the `rope_scaling` configuration.
"""
if self.rope_scaling is None:
return
if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
raise ValueError(
"`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
f"got {self.rope_scaling}"
)
rope_scaling_type = self.rope_scaling.get("type", None)
rope_scaling_factor = self.rope_scaling.get("factor", None)
if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
raise ValueError(
f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
)
if rope_scaling_factor is None or not isinstance(rope_scaling_factor, float) or rope_scaling_factor <= 1.0:
raise ValueError(f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}")
FlagEmbedding/inference/reranker/decoder_only/models/gemma_config.py
+67
View File
@@ -0,0 +1,67 @@
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# This file was automatically generated from <path_to_diff_file.py>.
# Do NOT edit this file manually as any edits will be overwritten by the generation of
# the file from the diff. If any change should be done, please apply the change to the
# diff.py file directly.
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# coding=utf-8
# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved.
#
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from transformers.models.gemma2.configuration_gemma2 import Gemma2Config
class CostWiseGemmaConfig(Gemma2Config):
r"""
This is the configuration class to store the configuration of a [`GemmaModel`]. It is used to instantiate an Gemma
model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
defaults will yield a similar configuration to that of the Gemma-7B.
e.g. [google/gemma-7b](https://huggingface.co/google/gemma-7b)
Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
documentation from [`PretrainedConfig`] for more information.
Args:
start_layer (`int`, *optional*, defaults to 28):
The start layer to output score.
layer_sep (`int`, *optional*, defaults to 28):
The sep layer from the start layer to output score.
layer_wise (`bool`, *optional*, defaults to `False`):
Whether or not the model should be layerwise.
```python
>>> from transformers import Gemma2Model, Gemma2Config
>>> # Initializing a Gemma2 gemma2-9b style configuration
>>> configuration = Gemma2Config()
>>> # Initializing a model from the gemma2-9b style configuration
>>> model = Gemma2Model(configuration)
>>> # Accessing the model configuration
>>> configuration = model.config
```"""
model_type = "cost_wise_gemma"
keys_to_ignore_at_inference = ["past_key_values"]
def __init__(
self,
start_layer: int = 28,
layer_sep: int = 28,
layer_wise: bool = False,
**kwargs,
):
self.start_layer = start_layer
self.layer_sep = layer_sep
self.layer_wise = layer_wise
super().__init__(
**kwargs,
)
FlagEmbedding/inference/reranker/decoder_only/models/gemma_model.py
+745
View File
@@ -0,0 +1,745 @@
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# This file was automatically generated from <path_to_diff_file.py>.
# Do NOT edit this file manually as any edits will be overwritten by the generation of
# the file from the diff. If any change should be done, please apply the change to the
# diff.py file directly.
# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
# coding=utf-8
# Copyright 2024 Google Inc. HuggingFace Inc. team. All rights reserved.
#
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from dataclasses import dataclass
import math
from typing import List, Optional, Tuple, Union
import inspect
import torch
import torch.nn.functional as F
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
from transformers.activations import ACT2FN
from transformers.cache_utils import Cache, DynamicCache, StaticCache
from transformers.modeling_attn_mask_utils import AttentionMaskConverter
from transformers.modeling_outputs import (
BaseModelOutputWithPast,
CausalLMOutputWithPast,
SequenceClassifierOutputWithPast,
TokenClassifierOutput,
)
from transformers.modeling_utils import PreTrainedModel
from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
from transformers.utils import (
add_start_docstrings,
add_start_docstrings_to_model_forward,
is_flash_attn_2_available,
is_flash_attn_greater_or_equal_2_10,
logging,
replace_return_docstrings,
ModelOutput,
)
from .gemma_config import CostWiseGemmaConfig
from transformers.models.gemma2.modeling_gemma2 import Gemma2RMSNorm, Gemma2RotaryEmbedding, rotate_half, apply_rotary_pos_emb
from transformers.models.gemma2.modeling_gemma2 import Gemma2MLP, repeat_kv, Gemma2Attention, Gemma2DecoderLayer, GEMMA2_START_DOCSTRING
from transformers.models.gemma2.modeling_gemma2 import GEMMA2_INPUTS_DOCSTRING
if is_flash_attn_2_available():
from flash_attn import flash_attn_func, flash_attn_varlen_func
from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
_flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters)
logger = logging.get_logger(__name__)
def _get_unpad_data(attention_mask):
seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
max_seqlen_in_batch = seqlens_in_batch.max().item()
cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))
return (
indices,
cu_seqlens,
max_seqlen_in_batch,
)
@add_start_docstrings(
"The bare Gemma2 Model outputting raw hidden-states without any specific head on top.",
GEMMA2_START_DOCSTRING,
)
class CostWiseGemma2PreTrainedModel(PreTrainedModel):
config_class = CostWiseGemmaConfig
base_model_prefix = "model"
supports_gradient_checkpointing = True
_no_split_modules = ["Gemma2DecoderLayer"]
_skip_keys_device_placement = ["past_key_values"]
_supports_flash_attn_2 = True
_supports_sdpa = True
_supports_cache_class = False
_supports_quantized_cache = False
_supports_static_cache = True
_is_stateful = True
def _init_weights(self, module):
std = self.config.initializer_range
if isinstance(module, nn.Linear):
module.weight.data.normal_(mean=0.0, std=std)
if module.bias is not None:
module.bias.data.zero_()
elif isinstance(module, nn.Embedding):
module.weight.data.normal_(mean=0.0, std=std)
if module.padding_idx is not None:
module.weight.data[module.padding_idx].zero_()
_CONFIG_FOR_DOC = "CostWiseGemmaConfig"
@dataclass
class CostWiseModelOutputWithPast(ModelOutput):
last_hidden_state: torch.FloatTensor = None
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
hidden_states: Optional[Tuple[torch.FloatTensor]] = None
attentions: Optional[Tuple[torch.FloatTensor]] = None
attention_masks: Optional[Tuple[torch.FloatTensor]] = None
@dataclass
class CostWiseCausalLMOutputWithPast(ModelOutput):
loss: Optional[torch.FloatTensor] = None
logits: torch.FloatTensor = None
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
hidden_states: Optional[Tuple[torch.FloatTensor]] = None
attentions: Optional[Tuple[torch.FloatTensor]] = None
attention_masks: Optional[Tuple[torch.FloatTensor]] = None
def token_compress(compress_ratio,
hidden_states,
attention_mask,
query_lengths,
prompt_lengths):
"""
compress_ratio: int
hidden_states: (b, s, h)
attention_mask: (b, s)
query_lengths: (b)
prompt_lengths: (b)
"""
# get some specific parameters
passage_lengths = torch.sum(attention_mask, dim=1, dtype=torch.int) - query_lengths - prompt_lengths # the raw passage lengths (b)
retain_passage_lengths = (passage_lengths + compress_ratio - 1) // compress_ratio # the passage lengths need to be retained (b)
final_useful_lengths = query_lengths + prompt_lengths + retain_passage_lengths # the final useful length after compress (b)
max_passage_length = torch.max(passage_lengths) # the max passage lengths (1)
max_final_lengths = torch.max(final_useful_lengths) # the max useful lengths after compress (1)
# make new hidden states and new attention masks
new_hidden_states = torch.zeros((hidden_states.shape[0], max_final_lengths,
hidden_states.shape[-1]), dtype=hidden_states.dtype).to(hidden_states.device) # (b, s', h)
new_attention_mask = torch.ones((hidden_states.shape[0], max_final_lengths), dtype=attention_mask.dtype).to(attention_mask.device) # (b, s')
# get new attention mask
mask_attention_index = torch.arange(max_final_lengths, device=hidden_states.device).unsqueeze(0) >= final_useful_lengths[:, None]
new_attention_mask[mask_attention_index] = 0
# get new hidden states
# add query into new hidden states
query_index = torch.arange(max_final_lengths, device=hidden_states.device).unsqueeze(0)
mask_query_index = query_index < query_lengths[:, None]
new_hidden_states[mask_query_index] = hidden_states[:, : max_final_lengths, :][mask_query_index]
# add prompt into new hidden states
# get the index of the prompt in new hidden states
new_prompt_start_length = query_lengths + retain_passage_lengths
new_prompt_end_length = new_prompt_start_length + prompt_lengths
new_prompt_index = torch.arange(max_final_lengths, device=hidden_states.device).unsqueeze(0)
new_mask_prompt_index_start = new_prompt_index >= new_prompt_start_length[:, None]
new_mask_prompt_index_end = new_prompt_index < new_prompt_end_length[:, None]
new_mask_prompt_index = new_mask_prompt_index_start & new_mask_prompt_index_end
# get the index of the prompt in hidden states
raw_prompt_start_length = query_lengths + passage_lengths
raw_prompt_end_length = raw_prompt_start_length + prompt_lengths
raw_prompt_index = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
raw_mask_prompt_index_start = raw_prompt_index >= raw_prompt_start_length[:, None]
raw_mask_prompt_index_end = raw_prompt_index < raw_prompt_end_length[:, None]
raw_mask_prompt_index = raw_mask_prompt_index_start & raw_mask_prompt_index_end
# replace the prompt hidden states
new_hidden_states[new_mask_prompt_index] = hidden_states[raw_mask_prompt_index]
# 以上均没问题
# print(new_hidden_states.view(len(new_hidden_states), -1))
# print(new_attention_mask)
# get the index of the passage in new hidden states
new_passage_start_length = query_lengths
new_passage_end_length = new_passage_start_length + retain_passage_lengths
new_passage_index = torch.arange(max_final_lengths, device=hidden_states.device).unsqueeze(0)
new_mask_passage_index_start = new_passage_index >= new_passage_start_length[:, None]
new_mask_passage_index_end = new_passage_index < new_passage_end_length[:, None]
new_mask_passage_index = new_mask_passage_index_start & new_mask_passage_index_end
# print(query_lengths, prompt_lengths, retain_passage_lengths, final_useful_lengths)
# add passage into new hidden states
# get mask hidden states
psg_start_length = query_lengths
psg_end_length = query_lengths + passage_lengths
psg_index = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
mask_psg_index_start = psg_index >= psg_start_length[:, None]
mask_psg_index_end = psg_index < psg_end_length[:, None]
mask_psg_index = mask_psg_index_start & mask_psg_index_end
hidden_states = hidden_states * mask_psg_index.unsqueeze(-1)
passage_hidden_states = torch.zeros((hidden_states.shape[0],
(max_passage_length + compress_ratio - 1) // compress_ratio * compress_ratio,
hidden_states.shape[-1]), dtype=hidden_states.dtype).to(hidden_states.device)
passage_end_length = passage_lengths
passage_index = torch.arange(passage_hidden_states.shape[1], device=hidden_states.device).unsqueeze(0) # maybe exceed the max passage length
mask_passage_index = passage_index < passage_end_length[:, None]
raw_passage_end_length = query_lengths + passage_lengths
raw_passage_start_length = query_lengths
raw_passage_index = torch.arange(hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
raw_mask_passage_index_start = raw_passage_index >= raw_passage_start_length[:, None]
raw_mask_passage_index_end = raw_passage_index < raw_passage_end_length[:, None]
raw_mask_passage_index = raw_mask_passage_index_start & raw_mask_passage_index_end
passage_hidden_states[mask_passage_index] = hidden_states[raw_mask_passage_index]
passage_weights = torch.zeros((hidden_states.shape[0],
(max_passage_length + compress_ratio - 1) // compress_ratio * compress_ratio)
, dtype=hidden_states.dtype).to(hidden_states.device)
passage_weights[mask_passage_index] = 1
passage_weights = passage_weights.view(passage_weights.shape[0], -1, compress_ratio)
passage_weights = passage_weights / torch.sum(passage_weights, dim=-1
).view(passage_weights.shape[0], -1, 1)
passage_weights = passage_weights.view(passage_weights.shape[0], -1)
# passage_weights = torch.where(passage_weights == torch.nan, 0, passage_weights)
passage_hidden_states = passage_hidden_states * passage_weights.unsqueeze(-1)
passage_hidden_states = passage_hidden_states.view(passage_hidden_states.shape[0], -1, compress_ratio,
passage_hidden_states.shape[-1])
passage_hidden_states = torch.sum(passage_hidden_states, dim=2)
passage_end_length = retain_passage_lengths
passage_index = torch.arange(passage_hidden_states.shape[1], device=hidden_states.device).unsqueeze(0)
mask_passage_index = passage_index < passage_end_length[:, None]
new_hidden_states[new_mask_passage_index] = passage_hidden_states[mask_passage_index]
return new_hidden_states, new_attention_mask
@add_start_docstrings(
"The bare Gemma2 Model outputting raw hidden-states without any specific head on top.",
GEMMA2_START_DOCSTRING,
)
class CostWiseGemmaModel(CostWiseGemma2PreTrainedModel):
"""
Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`GemmaDecoderLayer`]
Args:
config: GemmaConfig
"""
def __init__(self, config: CostWiseGemmaConfig):
super().__init__(config)
self.padding_idx = config.pad_token_id
self.vocab_size = config.vocab_size
self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
self.layers = nn.ModuleList(
[Gemma2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
)
self.norm = Gemma2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
self.gradient_checkpointing = False
# Initialize weights and apply final processing
self.post_init()
def get_input_embeddings(self):
return self.embed_tokens
def set_input_embeddings(self, value):
self.embed_tokens = value
@add_start_docstrings_to_model_forward(GEMMA2_INPUTS_DOCSTRING)
def forward(
self,
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
cache_position: Optional[torch.LongTensor] = None,
compress_layer: Optional[int] = None,
compress_ratio: Optional[int] = None,
cutoff_layers: Optional[List[int]] = None,
query_lengths: Optional[int] = None,
prompt_lengths: Optional[int] = None,
) -> Union[Tuple, CostWiseModelOutputWithPast]:
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
compress_ratio = None if compress_ratio == 1 else compress_ratio
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
if self.config.layer_wise:
output_hidden_states = True
use_cache = use_cache if use_cache is not None else self.config.use_cache
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
if (input_ids is None) ^ (inputs_embeds is not None):
raise ValueError(
"You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
)
if self.gradient_checkpointing and self.training and use_cache:
logger.warning_once(
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
)
use_cache = False
if compress_layer is not None and compress_ratio is not None:
logger.warning_once(
"`use_cache=True` is incompatible with reranker. Setting `use_cache=False`."
)
use_cache = False
if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
if cache_position is None:
cache_position = torch.arange(0, inputs_embeds.shape[1], device=inputs_embeds.device)
if position_ids is None:
position_ids = cache_position.unsqueeze(0)
causal_mask = self._update_causal_mask(
attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
)
# embed positions
hidden_states = inputs_embeds
# normalized
# Gemma downcasts the below to float16, causing sqrt(3072)=55.4256 to become 55.5
# See https://github.com/huggingface/transformers/pull/29402
normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype)
hidden_states = hidden_states * normalizer
# decoder layers
all_hidden_states = () if output_hidden_states else None
all_attention_masks = ()
all_self_attns = () if output_attentions else None
next_decoder_cache = None
is_padding_left = (attention_mask[:, -1].sum() == attention_mask.shape[0]) and (
torch.sum(attention_mask) != attention_mask.shape[0] * attention_mask.shape[1])
query_lengths = [0] * hidden_states.shape[0] if query_lengths is None else query_lengths
prompt_lengths = [0] * hidden_states.shape[0] if prompt_lengths is None else prompt_lengths
if not isinstance(query_lengths, torch.Tensor):
query_lengths = torch.tensor(query_lengths, device=hidden_states.device)
if not isinstance(prompt_lengths, torch.Tensor):
prompt_lengths = torch.tensor(prompt_lengths, device=hidden_states.device)
if cutoff_layers is None:
max_layer = self.config.num_hidden_layers
cutoff_layers = [max_layer]
if isinstance(cutoff_layers, int):
max_layer = cutoff_layers
cutoff_layers = [cutoff_layers]
else:
max_layer = max(cutoff_layers)
for idx, decoder_layer in enumerate(self.layers):
if self.config.layer_wise:
if idx in cutoff_layers and output_hidden_states:
all_hidden_states += (self.norm(hidden_states),)
all_attention_masks += (attention_mask,)
if idx == max_layer:
break
elif output_hidden_states:
all_hidden_states += (hidden_states,)
if compress_layer is not None and compress_ratio is not None and idx in compress_layer and idx != 0:
if is_padding_left:
raise ValueError('You must use right padding...')
hidden_states, attention_mask = token_compress(compress_ratio, hidden_states, attention_mask,
query_lengths, prompt_lengths)
seq_length = hidden_states.shape[1]
cache_position = torch.arange(0, seq_length, device=hidden_states.device)
position_ids = cache_position.unsqueeze(0)
causal_mask = self._update_causal_mask(
attention_mask, hidden_states, cache_position, past_key_values, output_attentions
)
if self.gradient_checkpointing and self.training:
layer_outputs = self._gradient_checkpointing_func(
decoder_layer.__call__,
hidden_states,
causal_mask,
position_ids,
past_key_values,
output_attentions,
use_cache,
cache_position,
)
else:
layer_outputs = decoder_layer(
hidden_states,
attention_mask=causal_mask,
position_ids=position_ids,
past_key_value=past_key_values,
output_attentions=output_attentions,
use_cache=use_cache,
cache_position=cache_position,
)
hidden_states = layer_outputs[0]
if output_attentions:
all_self_attns += (layer_outputs[1],)
hidden_states = self.norm(hidden_states)
# add hidden states from the last decoder layer
if not self.config.layer_wise:
if output_hidden_states:
all_hidden_states += (hidden_states,)
all_attention_masks += (attention_mask,)
else:
if output_hidden_states and self.config.num_hidden_layers == max_layer:
all_hidden_states += (hidden_states,)
all_attention_masks += (attention_mask,)
next_cache = next_decoder_cache if use_cache else None
if not return_dict:
return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
return CostWiseModelOutputWithPast(
last_hidden_state=hidden_states,
past_key_values=next_cache,
hidden_states=all_hidden_states,
attentions=all_self_attns,
attention_masks=all_attention_masks
)
def _update_causal_mask(
self,
attention_mask: torch.Tensor,
input_tensor: torch.Tensor,
cache_position: torch.Tensor,
past_key_values: Cache,
output_attentions: bool,
):
if self.config._attn_implementation == "flash_attention_2":
if attention_mask is not None and 0.0 in attention_mask:
return attention_mask
return None
dtype, device = input_tensor.dtype, input_tensor.device
min_dtype = torch.finfo(dtype).min
sequence_length = input_tensor.shape[1]
if past_key_values is not None:
target_length = past_key_values.get_max_length()
else:
target_length = attention_mask.shape[-1] if attention_mask is not None else input_tensor.shape[1]
if attention_mask is not None and attention_mask.dim() == 4:
# in this case we assume that the mask comes already in inverted form and requires no inversion or slicing
if attention_mask.max() != 0:
raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`")
causal_mask = attention_mask
else:
causal_mask = torch.full(
(sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device
)
if sequence_length != 1:
causal_mask = torch.triu(causal_mask, diagonal=1)
causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
if attention_mask is not None:
causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
mask_length = attention_mask.shape[-1]
padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
padding_mask = padding_mask == 0
causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
padding_mask, min_dtype
)
return causal_mask
class CostWiseHead(nn.Module):
"""Head for sentence-level classification tasks."""
def __init__(self, input_size, output_size):
super().__init__()
self.linear_head = nn.Linear(input_size, output_size, bias=False)
def forward(self, **kwargs):
return self.linear_head(**kwargs)
class CostWiseGemmaForCausalLM(CostWiseGemma2PreTrainedModel):
_tied_weights_keys = ["lm_head.weight"]
def __init__(self, config: CostWiseGemmaConfig):
super().__init__(config)
self.model = CostWiseGemmaModel(config)
self.vocab_size = config.vocab_size
if not config.layer_wise:
self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
else:
self.lm_head = nn.ModuleList(
[CostWiseHead(config.hidden_size, 1) for _ in range(
config.start_layer, config.num_hidden_layers + 1, config.layer_sep
)]
)
# Initialize weights and apply final processing
self.post_init()
def get_input_embeddings(self):
return self.model.embed_tokens
def set_input_embeddings(self, value):
self.model.embed_tokens = value
def get_output_embeddings(self):
return self.lm_head
def set_output_embeddings(self, new_embeddings):
self.lm_head = new_embeddings
def set_decoder(self, decoder):
self.model = decoder
def get_decoder(self):
return self.model
@add_start_docstrings_to_model_forward(GEMMA2_INPUTS_DOCSTRING)
@replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
def forward(
self,
input_ids: torch.LongTensor = None,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
inputs_embeds: Optional[torch.FloatTensor] = None,
labels: Optional[torch.LongTensor] = None,
use_cache: Optional[bool] = None,
output_attentions: Optional[bool] = None,
output_hidden_states: Optional[bool] = None,
return_dict: Optional[bool] = None,
cache_position: Optional[torch.LongTensor] = None,
compress_layer: Optional[int] = None,
compress_ratio: Optional[int] = None,
cutoff_layers: Optional[List[int]] = None,
query_lengths: Optional[int] = None,
prompt_lengths: Optional[int] = None,
) -> Union[Tuple, CostWiseCausalLMOutputWithPast]:
r"""
Args:
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
Labels for computing the masked language modeling loss. Indices should either be in `[0, transformers.,
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
(masked), the loss is only computed for the tokens with labels in `[0, transformers., config.vocab_size]`.
Returns:
Example:
```python
>>> from transformers import AutoTokenizer, GemmaForCausalLM
>>> model = GemmaForCausalLM.from_pretrained("google/gemma-2-9b")
>>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
>>> prompt = "What is your favorite condiment?"
>>> inputs = tokenizer(prompt, return_tensors="pt")
>>> # Generate
>>> generate_ids = model.generate(inputs.input_ids, max_length=30)
>>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
"What is your favorite condiment?"
```"""
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
output_hidden_states = (
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
)
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
if compress_ratio is not None and compress_ratio == 1:
compress_ratio = None
if self.config.layer_wise:
if cutoff_layers is None:
cutoff_layers = [self.config.num_hidden_layers]
elif isinstance(cutoff_layers, int):
cutoff_layers = [cutoff_layers]
can_use_layers = list(range(self.config.start_layer, self.config.num_hidden_layers + 1, self.config.layer_sep))
remove_layers = [i for i in cutoff_layers if i not in can_use_layers]
if len(remove_layers) > 0:
logger.warning_once(
f"layers {remove_layers} are incompatible with the setting. They will be removed..."
)
cutoff_layers = [i for i in cutoff_layers if i not in remove_layers]
if len(cutoff_layers) == 0:
raise ValueError(f"Your cutoff layers must in [{self.config.start_layer}, {self.config.num_hidden_layers}]")
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
outputs = self.model(
input_ids=input_ids,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
inputs_embeds=inputs_embeds,
use_cache=use_cache,
output_attentions=output_attentions,
output_hidden_states=output_hidden_states,
return_dict=return_dict,
cache_position=cache_position,
compress_layer=compress_layer,
compress_ratio=compress_ratio,
query_lengths=query_lengths,
prompt_lengths=prompt_lengths,
cutoff_layers=cutoff_layers,
)
if not self.config.layer_wise:
hidden_states = outputs[0]
logits = self.lm_head(hidden_states)
if self.config.final_logit_softcapping is not None:
logits = logits / self.config.final_logit_softcapping
logits = torch.tanh(logits)
logits = logits * self.config.final_logit_softcapping
logits = logits.float()
loss = None
if labels is not None:
# Shift so that tokens < n predict n
shift_logits = logits[..., :-1, :].contiguous()
shift_labels = labels[..., 1:].contiguous()
# Flatten the tokens
loss_fct = CrossEntropyLoss()
shift_logits = shift_logits.view(-1, self.config.vocab_size)
shift_labels = shift_labels.view(-1)
# Enable model parallelism
shift_labels = shift_labels.to(shift_logits.device)
loss = loss_fct(shift_logits, shift_labels)
else:
hidden_states = outputs.hidden_states
logits = ()
for i in range(len(hidden_states)):
tmp_logits = self.lm_head[i].linear_head(hidden_states[i])
if self.config.final_logit_softcapping is not None:
tmp_logits = tmp_logits / self.config.final_logit_softcapping
tmp_logits = torch.tanh(tmp_logits)
tmp_logits = tmp_logits * self.config.final_logit_softcapping
tmp_logits = tmp_logits.float()
tmp_logits = tmp_logits.reshape(hidden_states[i].shape[0], -1)
logits = logits + (tmp_logits,)
loss = None
if not return_dict:
output = (logits,) + outputs[1:]
return (loss,) + output if loss is not None else output
return CostWiseCausalLMOutputWithPast(
loss=loss,
logits=logits,
past_key_values=outputs.past_key_values,
hidden_states=outputs.hidden_states,
attentions=outputs.attentions,
attention_masks=outputs[-1] if self.model.config.layer_wise else outputs[-1][-1]
)
def prepare_inputs_for_generation(
self,
input_ids,
past_key_values=None,
attention_mask=None,
inputs_embeds=None,
cache_position=None,
use_cache=True,
**kwargs,
):
past_length = 0
if past_key_values is not None:
# Past key values are always initialized with a `Cache` object -> no need for if-else anymore
past_length = cache_position[0] if cache_position is not None else torch.tensor(0, device=input_ids.device)
max_cache_length = (
torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
if past_key_values.get_max_length() is not None
else None
)
cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length)
# Keep only the unprocessed tokens:
# 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
# some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as input)
if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
# 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
# input_ids based on the past_length.
elif past_length < input_ids.shape[1]:
input_ids = input_ids[:, past_length:]
# 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
# If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
if (
max_cache_length is not None
and attention_mask is not None
and cache_length + input_ids.shape[1] > max_cache_length
):
attention_mask = attention_mask[:, -max_cache_length:]
position_ids = kwargs.get("position_ids", None)
if attention_mask is not None and position_ids is None:
# create position_ids on the fly for batch generation
position_ids = attention_mask.long().cumsum(-1) - 1
position_ids.masked_fill_(attention_mask == 0, 1)
if past_key_values:
position_ids = position_ids[:, -input_ids.shape[1] :]
# if `inputs_embeds` are passed, we only want to use them in the 1st generation step
if inputs_embeds is not None and past_length == 0:
model_inputs = {"inputs_embeds": inputs_embeds}
else:
# The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
# recompiles graphs as the stride of the inputs is a guard. Ref: https://github.com/huggingface/transformers/pull/29114
# TODO: use `next_tokens` directly instead.
model_inputs = {"input_ids": input_ids.contiguous()}
input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
if cache_position is None:
cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
elif use_cache:
cache_position = cache_position[-input_length:]
model_inputs.update(
{
"position_ids": position_ids,
"cache_position": cache_position,
"past_key_values": past_key_values,
"use_cache": use_cache,
"attention_mask": attention_mask,
}
)
return model_inputs
@staticmethod
def _reorder_cache(past_key_values, beam_idx):
reordered_past = ()
for layer_past in past_key_values:
reordered_past += (
tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
)
return reordered_past
FlagEmbedding/inference/reranker/decoder_only/models/modeling_minicpm_reranker.py
+1500
View File
File diff suppressed because it is too large Load Diff