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")