from typing import Union, Mapping, Optional, Tuple, TypedDict, Dict, List
from functools import partial
import torch
import numpy as np
from transformers import AutoModel, PreTrainedTokenizerFast, BatchEncoding, DataCollatorWithPadding
from transformers.models.auto import AutoTokenizer
from transformers.models.llama.modeling_llama import LLAMA_INPUTS_DOCSTRING
from transformers.modeling_outputs import BaseModelOutputWithPast
from transformers.modeling_attn_mask_utils import _prepare_4d_attention_mask, _prepare_4d_attention_mask_for_sdpa
from transformers import LlamaModel
from transformers.cache_utils import Cache, DynamicCache
from transformers.utils import (
add_start_docstrings_to_model_forward,
logging,
)
from tqdm.auto import tqdm
from datasets import Dataset
from torch.utils.data import DataLoader
from .configuration_conan import ConanEmbedConfig
logger = logging.get_logger(__name__)
class ConanEmbedFeatures(TypedDict):
input_dict: torch.Tensor
attention_mask: torch.Tensor
pool_mask: torch.Tensor
def _move_to_device(maybe_tensor, device: torch.device):
if torch.is_tensor(maybe_tensor):
return maybe_tensor.to(device, non_blocking=device.type == "cuda")
elif isinstance(maybe_tensor, dict):
return {key: _move_to_device(value, device) for key, value in maybe_tensor.items()}
elif isinstance(maybe_tensor, list):
return [_move_to_device(x, device) for x in maybe_tensor]
elif isinstance(maybe_tensor, tuple):
return tuple([_move_to_device(x, device) for x in maybe_tensor])
elif isinstance(maybe_tensor, Mapping):
return type(maybe_tensor)({k: _move_to_device(v, device) for k, v in maybe_tensor.items()})
else:
return maybe_tensor
def move_to_device(sample, device: torch.device):
if device.type == "cpu":
return sample
if len(sample) == 0:
return {}
return _move_to_device(sample, device)
def input_transform_func(
tokenizer: PreTrainedTokenizerFast,
examples: Dict[str, List],
always_add_eos: bool,
max_length: int,
instruction: str,
) -> BatchEncoding:
if always_add_eos:
examples["input_texts"] = [
instruction + input_example + tokenizer.eos_token for input_example in examples["input_texts"]
]
print(examples["input_texts"])
batch_dict = tokenizer(
examples["input_texts"],
max_length=max_length,
padding=True,
return_token_type_ids=False,
return_tensors="pt",
truncation=True,
)
print(examples["input_texts"])
return batch_dict
class ConanEmbedModel(LlamaModel):
config_class = ConanEmbedConfig
def __init__(self, config: ConanEmbedConfig) -> None:
"""
Initialize the model with a given configuration.
Args:
config (ConanEmbedConfig): The configuration for the model.
"""
super().__init__(config)
for layer in self.layers:
layer.self_attn.is_causal = not config.do_dir
self._attn_implementation = "eager"
self.tokenizer = AutoTokenizer.from_pretrained(config._name_or_path)
self.padding_side = config.padding_side
self.is_mask_instruction = config.is_mask_instruction
self.add_eos = config.add_eos
self.mask_type = config.mask_type
self.sentence_pooling_method = config.sentence_pooling_method
if config.add_pad_token and self.tokenizer is not None:
self.add_pad_token()
def add_pad_token(self):
self.tokenizer.pad_token = self.tokenizer.eos_token
self.tokenizer.padding_side = self.padding_side
def _sentence_embedding(self, last_hidden_state, attention_mask=None):
"""Use the pooling method to get the sentence embedding.
Args:
last_hidden_state (torch.Tensor): The model output's last hidden state.
attention_mask (torch.Tensor): Mask out padding tokens during pooling.
Raises:
NotImplementedError: Specified pooling method not implemented.
Returns:
torch.Tensor: The sentence embeddings.
"""
if self.sentence_pooling_method == "cls":
return last_hidden_state[:, 0]
elif self.sentence_pooling_method == "mean":
s = torch.sum(last_hidden_state, dim=1)
# d = attention_mask.sum(dim=1, keepdim=True).float()
return s
elif self.sentence_pooling_method == "last_token":
left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
if left_padding:
return last_hidden_state[:, -1]
else:
sequence_lengths = attention_mask.sum(dim=1) - 1
batch_size = last_hidden_state.shape[0]
return last_hidden_state[
torch.arange(batch_size, device=last_hidden_state.device),
sequence_lengths,
]
else:
raise NotImplementedError(f"pooling method {self.sentence_pooling_method} not implemented")
def prepare_kwargs_from_batch(
self,
batch_dict: Dict[str, torch.Tensor],
instruction_lens: int,
device: torch.device,
) -> ConanEmbedFeatures:
"""
Prepare the batch dictionary for encoding.
Args:
batch_dict: A dictionary containing the input_ids and attention_mask.
instruction_lens: The length of the instruction.
device: The device to move the data to.
Returns:
A ConanEmbedFeatures object with the prepared input_ids and attention_mask.
"""
batch_dict = move_to_device(batch_dict, device)
attention_mask = batch_dict["attention_mask"].clone() if "attention_mask" in batch_dict else None
if (
attention_mask is not None
and self.padding_side == "right"
and self.is_mask_instruction
and instruction_lens > 0
):
# Mask out the instruction tokens for mean-pooling
attention_mask[:, :instruction_lens] = 0
features: ConanEmbedFeatures = {
"input_ids": torch.tensor(batch_dict.get("input_ids").to(batch_dict.get("input_ids")).long()),
"attention_mask": batch_dict["attention_mask"],
}
return features
@torch.no_grad()
def _do_encode(
self,
prompts: List[str],
batch_size: int = 1,
instruction: str = "",
max_length: int = 4096,
num_workers: int = 32,
return_numpy: bool = False,
) -> Union[torch.FloatTensor, np.ndarray]:
"""
Encode a list of prompts using the model.
Args:
prompts: A list of prompts to encode.
batch_size: The batch size to use for encoding. Defaults to 1.
instruction: An instruction to prepend to the prompts. Defaults to "".
max_length: The maximum length of the input_ids. Defaults to 4096.
num_workers: The number of workers to use for encoding. Defaults to 32.
return_numpy: Whether to return the output as a numpy array or a torch tensor. Defaults to False.
Returns:
A tensor or numpy array of shape (len(prompts), hidden_size) containing the encoded prompts.
"""
dataset: Dataset = Dataset.from_dict({"input_texts": prompts})
dataset.set_transform(
partial(
input_transform_func,
self.tokenizer,
always_add_eos=True,
max_length=max_length,
instruction=instruction,
)
)
data_collator = DataCollatorWithPadding(self.tokenizer)
data_loader = DataLoader(
dataset,
batch_size=batch_size,
shuffle=False,
drop_last=False,
num_workers=num_workers,
collate_fn=data_collator,
pin_memory=True,
)
if self.padding_side == "right" and self.is_mask_instruction and len(instruction) > 0:
instruction_lens = len(self.tokenizer.tokenize(instruction))
else:
instruction_lens = 0
encoded_embeds: List[torch.Tensor] = []
device = next(self.parameters()).device
for batch_dict in tqdm(data_loader, desc="encoding", mininterval=10):
features = self.prepare_kwargs_from_batch(batch_dict, instruction_lens, device=device)
embeds = self(**features)["sentence_embeddings"].squeeze(1)
encoded_embeds.append(embeds)
encoded_embeds = torch.cat(encoded_embeds, axis=0)
if return_numpy:
encoded_embeds = encoded_embeds.cpu().detach().numpy()
return encoded_embeds
@add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
def forward(
self,
input_ids: torch.LongTensor,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_values: Optional[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,
token_type_ids: Optional[torch.LongTensor] = None,
) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPast]:
"""
Args:
input_ids: a tensor of shape (batch_size, sequence_length)
attention_mask: a tensor of shape (batch_size, sequence_length)
position_ids: a tensor of shape (batch_size, sequence_length)
past_key_values: a list of tensors of shape (batch_size, key_length, hidden_size)
inputs_embeds: a tensor of shape (batch_size, sequence_length, hidden_size)
use_cache: a boolean indicating whether to use the cache
output_attentions: a boolean indicating whether to output the attention weights
output_hidden_states: a boolean indicating whether to output the hidden states
return_dict: a boolean indicating whether to return a dictionary
Returns:
a tuple of length 4 containing the last hidden state, the cache, the hidden states,
and the attention weights
or a BaseModelOutputWithPast object
"""
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
)
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
# retrieve input_ids and inputs_embeds
if input_ids is not None and inputs_embeds is not None:
raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
elif input_ids is not None:
batch_size, seq_length = input_ids.shape
elif inputs_embeds is not None:
batch_size, seq_length, _ = inputs_embeds.shape
else:
raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
if self.gradient_checkpointing and self.training:
if use_cache:
logger.warning_once(
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
)
use_cache = False
past_key_values_length = 0
if use_cache:
use_legacy_cache = not isinstance(past_key_values, Cache)
if use_legacy_cache:
past_key_values = DynamicCache.from_legacy_cache(past_key_values)
past_key_values_length = past_key_values.get_usable_length(seq_length)
if position_ids is None:
device = input_ids.device if input_ids is not None else inputs_embeds.device
position_ids = torch.arange(
past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
)
position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
else:
position_ids = position_ids.view(-1, seq_length).long()
if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)
if attention_mask is not None and self._attn_implementation == "flash_attention_2" and use_cache:
is_padding_right = attention_mask[:, -1].sum().item() != batch_size
if is_padding_right:
raise ValueError(
"You are attempting to perform batched generation with padding_side='right'"
" this may lead to unexpected behaviour for Flash Attention version of Mistral. Make sure to "
" call `tokenizer.padding_side = 'left'` before tokenizing the input. "
)
if self._attn_implementation == "flash_attention_2":
# 2d mask is passed through the layers
attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
elif self._attn_implementation == "sdpa" and not output_attentions:
# output_attentions=True can not be supported when using SDPA, and we fall back on
# the manual implementation that requires a 4D causal mask in all cases.
attention_mask = _prepare_4d_attention_mask_for_sdpa(attention_mask, inputs_embeds.dtype)
else:
# 4d mask is passed through the layers
attention_mask = _prepare_4d_attention_mask(
attention_mask,
inputs_embeds.dtype,
)
hidden_states = inputs_embeds
# decoder layers
all_hidden_states = () if output_hidden_states else None
all_self_attns = () if output_attentions else None
next_decoder_cache = None
for decoder_layer in self.layers:
if output_hidden_states:
all_hidden_states += (hidden_states,)
if self.gradient_checkpointing and self.training:
layer_outputs = self._gradient_checkpointing_func(
decoder_layer.__call__,
hidden_states,
attention_mask,
position_ids,
past_key_values,
output_attentions,
use_cache,
)
else:
layer_outputs = decoder_layer(
hidden_states,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_value=past_key_values,
output_attentions=output_attentions,
use_cache=use_cache,
)
hidden_states = layer_outputs[0]
if use_cache:
next_decoder_cache = layer_outputs[2 if output_attentions else 1]
if output_attentions:
all_self_attns += (layer_outputs[1],)
hidden_states = self.norm(hidden_states)
# add hidden states from the last decoder layer
if output_hidden_states:
all_hidden_states += (hidden_states,)
next_cache = None
if use_cache:
next_cache = next_decoder_cache.to_legacy_cache() if use_legacy_cache else next_decoder_cache
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 BaseModelOutputWithPast(
last_hidden_state=hidden_states,
past_key_values=next_cache,
hidden_states=all_hidden_states,
attentions=all_self_attns,
)
@torch.no_grad()
def encode(
self,
prompts: List[str],
instruction: str = "",
max_length: int = 4096,
) -> Dict[str, torch.Tensor]:
"""
Encode a list of prompts and an instruction using the model.
Args:
prompts: A list of prompts to encode.
instruction: An instruction to prepend to the prompts. Defaults to "".
max_length: The maximum length of the input_ids. Defaults to 4096.
Returns:
A dictionary containing the sentence embeddings with key "sentence_embeddings".
"""
if self.padding_side == "right" and self.is_mask_instruction and len(instruction) > 0:
instruction_lens = len(self.tokenizer.tokenize(instruction))
else:
instruction_lens = 0
device = next(self.parameters()).device
batch_dict = input_transform_func(
self.tokenizer,
{"input_texts": [prompt for prompt in prompts]},
always_add_eos=False,
max_length=max_length,
instruction=instruction,
)
features: ConanEmbedFeatures = self.prepare_kwargs_from_batch(batch_dict, instruction_lens, device=device)
outputs = self(**features)
embeds = self._sentence_embedding(outputs.last_hidden_state)
return {"sentence_embeddings": embeds}
# AutoModel Register
AutoModel.register(ConanEmbedConfig, ConanEmbedModel)
# Register for auto class
ConanEmbedModel.register_for_auto_class("AutoModel")