Sentence Similarity
Transformers
Safetensors
English
llama
feature-extraction
text-embedding
embeddings
information-retrieval
beir
text-classification
language-model
text-clustering
text-semantic-similarity
text-evaluation
text-reranking
Sentence Similarity
natural_questions
ms_marco
fever
hotpot_qa
mteb
custom_code
text-generation-inference
| from typing import List, Optional, Tuple, Union | |
| import torch | |
| from transformers import LlamaModel, LlamaPreTrainedModel | |
| from transformers.models.llama.modeling_llama import LlamaDecoderLayer, LlamaRMSNorm, LlamaConfig, LlamaMLP, LlamaAttention, LlamaFlashAttention2, LlamaSdpaAttention | |
| from transformers.utils import logging | |
| from torch import nn | |
| import torch.nn.functional as F | |
| from transformers.modeling_outputs import BaseModelOutputWithPast | |
| from transformers.cache_utils import Cache, DynamicCache | |
| from .attn_mask_utils import _prepare_4d_attention_mask_for_sdpa, _prepare_4d_attention_mask | |
| logger = logging.get_logger(__name__) | |
| class ModifiedLlamaAttention(LlamaAttention): | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| self.is_causal = False | |
| class ModifiedLlamaFlashAttention2(LlamaFlashAttention2): | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| self.is_causal = False | |
| class ModifiedLlamaSdpaAttention(LlamaSdpaAttention): | |
| def __init__(self, *args, **kwargs): | |
| super().__init__(*args, **kwargs) | |
| self.is_causal = False | |
| LLAMA_ATTENTION_CLASSES = { | |
| "eager": ModifiedLlamaAttention, | |
| "flash_attention_2": ModifiedLlamaFlashAttention2, | |
| "sdpa": ModifiedLlamaSdpaAttention, | |
| } | |
| class ModifiedLlamaDecoderLayer(LlamaDecoderLayer): | |
| def __init__(self, config: LlamaConfig, layer_idx: int): | |
| nn.Module.__init__(self) | |
| self.hidden_size = config.hidden_size | |
| self.self_attn = LLAMA_ATTENTION_CLASSES[config._attn_implementation](config=config, layer_idx=layer_idx) | |
| self.mlp = LlamaMLP(config) | |
| self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) | |
| self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) | |
| class LlamaEncoderModel(LlamaModel): | |
| def __init__(self, config): | |
| LlamaPreTrainedModel.__init__(self, 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( | |
| [ModifiedLlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] | |
| ) | |
| self._use_sdpa = config._attn_implementation == "sdpa" | |
| self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" | |
| self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) | |
| self.gradient_checkpointing = False | |
| # Initialize weights and apply final processing | |
| self.post_init() | |
| def forward( | |
| self, | |
| input_ids: torch.LongTensor = None, | |
| 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, | |
| ) -> Union[Tuple, BaseModelOutputWithPast]: | |
| 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 input_ids and inputs_embeds at the same time") | |
| elif input_ids is not None: | |
| batch_size, seq_length = input_ids.shape[:2] | |
| elif inputs_embeds is not None: | |
| batch_size, seq_length = inputs_embeds.shape[:2] | |
| else: | |
| raise ValueError("You have to specify either input_ids or 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) | |
| if inputs_embeds is None: | |
| inputs_embeds = self.embed_tokens(input_ids) | |
| if self._use_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._use_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, | |
| (batch_size, seq_length), | |
| inputs_embeds, | |
| past_key_values_length, | |
| ) | |
| else: | |
| # 4d mask is passed through the layers | |
| attention_mask = _prepare_4d_attention_mask( | |
| attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length | |
| ) | |
| # embed positions | |
| 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, | |
| ) | |