# we don't want to support mypy for this file for now # type: ignore import numpy as np from typing import List, Optional, Tuple, Union, Dict from tqdm import tqdm from einops import rearrange, repeat import torch from torch import nn from transformers.activations import ACT2FN from transformers.cache_utils import Cache, DynamicCache, StaticCache from transformers.configuration_utils import PretrainedConfig from transformers.modeling_attn_mask_utils import AttentionMaskConverter from transformers.modeling_outputs import ( BaseModelOutputWithPast, ) from transformers.modeling_utils import PreTrainedModel try: from flash_attn.flash_attn_interface import flash_attn_func except Exception as e: print( f"Could not import flash attention. " ) flash_attn_func = None class RotaryConfig(): def __init__( self, dimensions: int = 0, base: int = 10000, max_seq_length: int = 2048 ): self.dimensions = dimensions self.base = base self.max_seq_length = max_seq_length class PhariaAdapterConfig: def __init__( self, hidden_size: int, intermediate_size: int, mlp_bias: bool, hidden_act: str ): self.hidden_size = hidden_size self.intermediate_size = intermediate_size self.mlp_bias = mlp_bias self.hidden_act = hidden_act def to_dict(self): return { "hidden_size": self.hidden_size, "intermediate_size": self.intermediate_size, "mlp_bias": self.mlp_bias, "hidden_act": self.hidden_act, } @classmethod def from_dict(cls, config_dict): return cls(**config_dict) class PhariaConfig(PretrainedConfig): def __init__( self, pad_token_id=None, bos_token_id=1, eos_token_id=2, hidden_act="gelu", hidden_size=512, bias_name=None, initializer_range=0.02, intermediate_size=2048, max_position_embeddings=8192, model_type="pharia-v2", num_attention_heads=4, num_hidden_layers=4, num_key_value_heads=2, torch_dtype="bfloat16", transformers_version="4.31.0.dev0", use_cache=True, vocab_size=128000, mlp_bias=True, attention_bias=True, tie_word_embeddings=False, attention_dropout=0.0, causal_attention=True, rope_theta=1000000, # rotary_embeddingbase, rope_scaling=None, mlp_adapter_config=None, attn_adapter_config=None, _attn_implementation='eager', embedding_head_out=1024, lora_config=None, pooling_method=None, layer_norm_epsilon=1e-05, **kwargs, ): 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, ) self.pad_token_id = pad_token_id self.bos_token_id = bos_token_id self.eos_token_id = eos_token_id self.hidden_act = hidden_act self.hidden_size = hidden_size self.initializer_range = initializer_range self.intermediate_size = intermediate_size self.max_position_embeddings = max_position_embeddings self.model_type = model_type self.num_attention_heads = num_attention_heads self.num_hidden_layers = num_hidden_layers self.num_key_value_heads = num_key_value_heads self.torch_dtype = torch_dtype self.causal_attention = causal_attention self.attn_adapter_config = attn_adapter_config self.mlp_adapter_config = mlp_adapter_config self.bias_name = bias_name self.transformers_version = transformers_version self.use_cache = use_cache self.vocab_size = vocab_size self.mlp_bias = mlp_bias self.attention_bias = attention_bias self.tie_word_embeddings = tie_word_embeddings self.attention_dropout = attention_dropout self.rope_theta = rope_theta self.rope_scaling = rope_scaling self.embedding_head_out = embedding_head_out self.pooling_method = pooling_method self.lora_config = lora_config self._attn_implementation = _attn_implementation self.layer_norm_epsilon = layer_norm_epsilon def to_dict(self): output = super(PhariaConfig, self).to_dict() if self.mlp_adapter_config is not None: output["mlp_adapter_config"] = self.mlp_adapter_config.to_dict() if self.attn_adapter_config is not None: output["attn_adapter_config"] = self.attn_adapter_config.to_dict() return output @classmethod def from_dict(cls, config_dict, **kwargs): if 'use_cache' in config_dict: del config_dict['use_cache'] if 'mlp_adapter_config' in config_dict and config_dict["mlp_adapter_config"] is not None: config_dict["mlp_adapter_config"] = PhariaAdapterConfig.from_dict(config_dict["mlp_adapter_config"]) if 'attn_adapter_config' in config_dict and config_dict["attn_adapter_config"] is not None: config_dict["attn_adapter_config"] = PhariaAdapterConfig.from_dict(config_dict["attn_adapter_config"]) return cls(**config_dict, **kwargs) def reshape_complex_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor) -> torch.Tensor: ndim = x.ndim assert 0 <= 1 < ndim assert freqs_cis.shape[0] == x.shape[1] assert freqs_cis.shape[1] == x.shape[-1] shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] return freqs_cis.view(*shape) def precompute_freqs_cis( dim: int, end: int, theta: float, device: torch.device, ) -> torch.Tensor: theta = float(theta) freqs = 1.0 / (theta ** (torch.arange(0, dim, 2, device=device)[: (dim // 2)].float() / dim)).to(device) t = torch.arange(end, device=device) # type: ignore freqs = torch.outer(t, freqs).float() # type: ignore freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 return freqs_cis.to(device) def apply_complex_rotary_emb( xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor, query_position_ids: Optional[torch.Tensor], key_position_ids: Optional[torch.Tensor], ) -> tuple[torch.Tensor, torch.Tensor]: xq_complex = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) xk_complex = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) if query_position_ids is None: freqs_cis_q = reshape_complex_for_broadcast(freqs_cis, xq_complex) else: freqs_cis_q = vector_gather_complex(freqs_cis, query_position_ids) if key_position_ids is None: freqs_cis_k = reshape_complex_for_broadcast(freqs_cis, xq_complex) else: freqs_cis_k = vector_gather_complex(freqs_cis, key_position_ids) xq_out = torch.view_as_real(xq_complex * freqs_cis_q).flatten(3) xk_out = torch.view_as_real(xk_complex * freqs_cis_k).flatten(3) return xq_out.type_as(xq), xk_out.type_as(xk) class RotaryEmbeddingComplex(torch.nn.Module): """ Relative rotary position embedding based on * RoFormer: Enhanced Transformer with Rotary Position Embedding (https://arxiv.org/abs/2104.09864) * Rotary Embeddings: A Relative Revolution (https://blog.eleuther.ai/rotary-embeddings/) """ def __init__( self, config: RotaryConfig, device: torch.device, ) -> None: super().__init__() assert config.dimensions > 1, "RotaryEmbedding cannot use `dim` == 1, this results in weird reshape errors" freqs_cis = precompute_freqs_cis( dim=config.dimensions, end=config.max_seq_length, theta=config.base, device=device, ) # Store real and imaginary in separate buffers for correct type casting. self.freqs_cis_real = freqs_cis.real self.freqs_cis_imag = freqs_cis.imag def forward( self, query: torch.Tensor, key: torch.Tensor, query_position_ids: Optional[torch.Tensor] = None, key_position_ids: Optional[torch.Tensor] = None, ) -> tuple[torch.Tensor, torch.Tensor]: query, key = apply_complex_rotary_emb( xq=rearrange(query, "sq b nh hh -> b sq nh hh"), xk=rearrange(key, "sq b nh hh -> b sq nh hh"), freqs_cis=torch.complex(self.freqs_cis_real.float(), self.freqs_cis_imag.float()), query_position_ids=query_position_ids, key_position_ids=key_position_ids, ) return rearrange(query, "b sq nh hh -> sq b nh hh"), rearrange(key, "b sq nh hh -> sq b nh hh") def vector_gather(vectors: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: """ Gathers (batched) vectors according to indices. """ vectors = repeat(vectors, "sq b nh d -> sq b B nh d", B=indices.shape[1]).squeeze(1) indices = repeat( indices, "sq b -> sq b nh d", nh=vectors.shape[-2], d=vectors.shape[-1], ) out = torch.gather(vectors, dim=0, index=indices) return out def vector_gather_complex(vectors: torch.Tensor, indices: torch.Tensor) -> torch.Tensor: """ Gathers (batched) vectors according to indices. """ vectors = repeat(vectors, "sq d -> sq B nh d", B=indices.shape[1], nh=1) indices = repeat( indices, "sq b -> sq b nh d", nh=1, d=vectors.shape[-1], ) out = torch.gather(vectors, dim=0, index=indices) out = rearrange(out, "sq b nh hh -> b sq nh hh") return out def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor: """torch.repeat_interleave(x, dim=2, repeats=n_rep)""" bs, slen, n_kv_heads, head_dim = x.shape if n_rep == 1: return x return ( x[:, :, :, None, :] .expand(bs, slen, n_kv_heads, n_rep, head_dim) .reshape(bs, slen, n_kv_heads * n_rep, head_dim) ) class PhariaAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" def __init__(self, config: PhariaConfig, layer_idx: Optional[int] = None): super().__init__() self.config = config self.layer_idx = layer_idx self.attention_dropout = config.attention_dropout self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.hidden_size // self.num_heads self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta self.is_causal = config.causal_attention self.query_key_scaling_factor = 1 / (self.head_dim ** 0.5) if (self.head_dim * self.num_heads) != self.hidden_size: raise ValueError( f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" f" and `num_heads`: {self.num_heads})." ) self.q_proj = nn.Linear( self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias ) self.k_proj = nn.Linear( self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias, ) self.v_proj = nn.Linear( self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias, ) self.o_proj = nn.Linear( self.hidden_size, self.hidden_size, bias=config.attention_bias ) self._init_rope() def _init_rope(self): self.rotary_emb = RotaryEmbeddingComplex( config=RotaryConfig( dimensions=self.head_dim, max_seq_length=self.max_position_embeddings, base=self.rope_theta ), device='cuda:0' ) def prepare_query_key_value( self, hidden_states: torch.Tensor, position_ids: torch.Tensor, past_key_value: Optional[Cache] = None, cache_position: Optional[torch.LongTensor] = None, ): query_states = rearrange(self.q_proj(hidden_states), "b sq (np hn) -> sq b np hn", np=self.num_heads) key_states = rearrange(self.k_proj(hidden_states), "b sq (np hn) -> sq b np hn", np=self.num_key_value_heads) value_states = rearrange(self.v_proj(hidden_states), "b sq (np hn) -> sq b np hn", np=self.num_key_value_heads) # cos, sin = self.rotary_emb(value_states, position_ids) position_ids = rearrange(position_ids, 'b sq -> sq b') query_states, key_states = self.rotary_emb( query_states, key_states, query_position_ids=position_ids, key_position_ids=position_ids ) if past_key_value is not None: # cache_position needed for the static cache cache_kwargs = {"cache_position": cache_position} key_states, value_states = past_key_value.update( key_states, value_states, self.layer_idx, cache_kwargs ) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) return query_states, key_states, value_states def forward ( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, softmax_in_fp32: Optional[bool] = False ): bsz, _, _ = hidden_states.size() query, key, value = self.prepare_query_key_value( hidden_states, position_ids=position_ids, past_key_value=past_key_value, cache_position=cache_position ) seq_length, batch_size, _, head_dim = query.shape query = rearrange(query, "sq bs nh hd -> sq (bs nh) hd") key = rearrange(key, "sq bs nh hd -> sq (bs nh) hd") value = rearrange(value, "sq bs nh hd -> sq (bs nh) hd") matmul_result = torch.empty( query.size(1), query.size(0), key.size(0), dtype=query.dtype, device=query.device, ) # Raw attention scores. [b * np, s_q, s_k] matmul_result = torch.baddbmm( matmul_result, query.transpose(0, 1), # [b * np, s_q, hn] key.transpose(0, 1).transpose(1, 2), # [b * np, hn, s_k] beta=0.0, alpha=self.query_key_scaling_factor, ) attention_scores = rearrange(matmul_result, "(b n) s_q s_k -> b n s_q s_k", b=batch_size) if softmax_in_fp32 and attention_scores.dtype != torch.float32: input_dtype = attention_scores.dtype attention_scores = attention_scores.float() else: input_dtype = None causal_mask = torch.triu( torch.ones(seq_length, seq_length, device=query.device), diagonal=1 ).bool() attention_scores.masked_fill_(causal_mask.to(attention_scores.device), -10000.0) probs = torch.nn.functional.softmax(attention_scores, dim=-1) if softmax_in_fp32 and input_dtype is not None: probs = probs.to(input_dtype) probs = rearrange(probs, "b n s_q s_k -> (b n) s_q s_k") hidden_state = torch.bmm(probs.to(dtype=value.dtype), value.transpose(0, 1)) attn_output = rearrange(hidden_state, "(b np) sq hn -> b sq (np hn)", b=bsz) attn_output = nn.functional.linear(attn_output, self.o_proj.weight, None) + self.o_proj.bias return attn_output, _, past_key_value class PhariaFlashAttention2(PhariaAttention): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @staticmethod def get_max_seq_length(cumulative_seq_lengths: torch.Tensor) -> int: return int((cumulative_seq_lengths[1:] - cumulative_seq_lengths[:-1]).max().item()) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, softmax_in_fp32: Optional[bool] = False ): assert flash_attn_func is not None, "Please install Flash Attention via optimization requirements" query, key, value = self.prepare_query_key_value(hidden_states, position_ids=position_ids) batch_size = query.shape[1] # reshape into format expected by flash attention [sq, b, np, hn] => [b, sq, np, hn] query = rearrange(query, "s_q b n h -> b s_q n h") key = rearrange(key, "s_k b n h -> b s_k n h") value = rearrange(value, "s_k b n h -> b s_k n h") attention_output = flash_attn_func( q=query, k=key, v=value, causal=self.is_causal, softmax_scale=self.query_key_scaling_factor ) attention_output = rearrange(attention_output, "b sq np hn -> b sq (np hn)", b=batch_size) attention_output = nn.functional.linear(attention_output, self.o_proj.weight, None) + self.o_proj.bias if not output_attentions: attn_weights = None return attention_output, attn_weights, past_key_value ATTN_IMPLEMENTATION = { 'flash_attention_2': PhariaFlashAttention2, 'sdpa': PhariaAttention, 'eager': PhariaAttention } class PhariaMLP(nn.Module): def __init__(self, config, layer_idx: int): super().__init__() self.layer_idx = layer_idx self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size self.up_proj = nn.Linear( self.hidden_size, self.intermediate_size, bias=config.mlp_bias ) self.down_proj = nn.Linear( self.intermediate_size, self.hidden_size, bias=config.mlp_bias ) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): x = self.up_proj(x) x = self.act_fn(x) if not self.down_proj.bias is None: # Scaling implements this with bias being seperately added. To match numerics we change this also o = nn.functional.linear(x, self.down_proj.weight, None) + self.down_proj.bias else: o = self.down_proj(x) return o class PhariaDecoderLayer(nn.Module): def __init__(self, config: PhariaConfig, layer_idx: int): super().__init__() self.hidden_size = config.hidden_size self.self_attn = ATTN_IMPLEMENTATION[config._attn_implementation](config=config, layer_idx=layer_idx) self.post_mlp_adapter = None if config.mlp_adapter_config: self.post_mlp_adapter = PhariaMLP(config.mlp_adapter_config, layer_idx=layer_idx) self.post_attn_adapter = None if config.attn_adapter_config: self.post_attn_adapter = PhariaMLP(config.attn_adapter_config, layer_idx=layer_idx) self.mlp = PhariaMLP(config, layer_idx=layer_idx) self.input_layernorm = nn.LayerNorm(config.hidden_size) self.post_attention_layernorm = nn.LayerNorm(config.hidden_size) self.layer_idx = layer_idx def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: Optional[bool] = False, use_cache: Optional[bool] = False, cache_position: Optional[torch.LongTensor] = None, ) -> Tuple[ torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] ]: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) hidden_states, self_attn_weights, present_key_value = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, ) hidden_states = residual + hidden_states if self.post_attn_adapter: hidden_states = self.post_attn_adapter(hidden_states) + hidden_states residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states if self.post_mlp_adapter: hidden_states = self.post_mlp_adapter(hidden_states) + hidden_states outputs = (hidden_states,) if output_attentions: outputs += (self_attn_weights,) if use_cache: outputs += (present_key_value,) return outputs class PhariaPreTrainedModel(PreTrainedModel): config_class = PhariaConfig base_model_prefix = "model" supports_gradient_checkpointing = False _no_split_modules = ["PhariaDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn_2 = True _supports_sdpa = True _supports_cache_class = True _supports_static_cache = 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_() class PhariaModel(PhariaPreTrainedModel): config_class = PhariaConfig def __init__(self, config: PhariaConfig): 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( [ PhariaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers) ] ) self.norm = nn.LayerNorm(config.hidden_size) 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, ) -> 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 ) 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 inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) return_legacy_cache = False if use_cache and not isinstance( past_key_values, Cache ): # kept for BC (non `Cache` `past_key_values` inputs) return_legacy_cache = True past_key_values = DynamicCache.from_legacy_cache(past_key_values) if cache_position is None: past_seen_tokens = ( past_key_values.get_seq_length() if past_key_values is not None else 0 ) cache_position = torch.arange( past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device, ) if position_ids is None: position_ids = cache_position.unsqueeze(0) if self.config.causal_attention: mask = self._update_causal_mask( attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions, ) else: mask = self._create_bidirectional_attention_mask( attention_mask, inputs_embeds.dtype ) # 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,) layer_outputs = decoder_layer( hidden_states, attention_mask=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 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 = next_decoder_cache if use_cache else None if return_legacy_cache: next_cache = next_cache.to_legacy_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, ) def _create_bidirectional_attention_mask(self, attention_mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: bidirectional_mask = attention_mask.unsqueeze(1) * attention_mask.unsqueeze(2).to(dtype) bidirectional_mask = 1 - bidirectional_mask # flip dtype_min_value = torch.finfo(dtype).min attention_mask = bidirectional_mask.masked_fill(bidirectional_mask == 1, dtype_min_value) return attention_mask def _update_causal_mask( self, attention_mask: torch.Tensor, input_tensor: torch.Tensor, cache_position: torch.Tensor, past_key_values: Cache, output_attentions: bool, ): # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length even when the static # KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at each decode steps due to the dynamic shapes. # (`recording cudagraph tree for symint key 13`, etc.), which is VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using # `fullgraph=True`. See more context in https://github.com/huggingface/transformers/pull/29114 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 # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail # to infer the attention mask. past_seen_tokens = ( past_key_values.get_seq_length() if past_key_values is not None else 0 ) using_static_cache = isinstance(past_key_values, StaticCache) # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward if ( self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions ): if AttentionMaskConverter._ignore_causal_mask_sdpa( attention_mask, inputs_embeds=input_tensor, past_key_values_length=past_seen_tokens, is_training=self.training, ): return None dtype, device = input_tensor.dtype, input_tensor.device min_dtype = torch.finfo(dtype).min sequence_length = input_tensor.shape[1] if using_static_cache: target_length = past_key_values.get_max_length() else: target_length = ( attention_mask.shape[-1] if isinstance(attention_mask, torch.Tensor) else past_seen_tokens + sequence_length + 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) if ( self.config._attn_implementation == "sdpa" and attention_mask is not None and attention_mask.device.type == "cuda" and not output_attentions ): # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path. # Details: https://github.com/pytorch/pytorch/issues/110213 causal_mask = AttentionMaskConverter._unmask_unattended( causal_mask, min_dtype ) return causal_mask class Embeddinghead(torch.nn.Module): def __init__( self, pooling_method: str ): super().__init__() self.pooling_method = pooling_method def forward(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor: """ Args: hidden_state: [b, n, d] attention_mask: [b, n] """ hidden_state = hidden_state.to(attention_mask.device) if self.pooling_method == 'cls': embedding = hidden_state[:, 0] elif self.pooling_method == 'lasttoken': b, n, d = hidden_state.size() reversed_mask = torch.flip(attention_mask, dims=(1,)) argmax_reverse = torch.argmax(reversed_mask, dim=1, keepdim=False) gather_indices = attention_mask.size(1) - argmax_reverse - 1 gather_indices = torch.clamp(gather_indices, min=0) gather_indices = gather_indices.unsqueeze(-1).repeat(1, d) gather_indices = gather_indices.unsqueeze(1) assert gather_indices.shape == (b, 1, d) input_mask_expanded = attention_mask.unsqueeze(-1).expand((b, n, d)).float() embedding = torch.gather(hidden_state * input_mask_expanded, 1, gather_indices).squeeze(dim=1) elif self.pooling_method in ['mean', 'weighted_mean']: if self.pooling_method == 'weighted_mean': attention_mask *= attention_mask.cumsum(dim=1) s = torch.sum(hidden_state * attention_mask.unsqueeze(-1).float(), dim=1) d = attention_mask.sum(dim=1, keepdim=True).float() embedding = s / d else: raise NotImplementedError(f"Unknown pooling method: {self.pooling_method}") return embedding class PhariaForEmbedding(PhariaPreTrainedModel): def __init__(self, config, tokenizer): super().__init__(config) self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2" self._use_sdpa = config._attn_implementation == "sdpa" self.model = PhariaModel(config) self.tokenizer = tokenizer self.tokenizer.pad_token_id = 1 self.embedding_head = Embeddinghead(pooling_method=self.config.pooling_method) def encode_queries(self, queries: Union[List[str], str], **kwargs) -> np.ndarray: """Used for encoding the queries of retrieval or reranking tasks""" return self.encode(queries, **kwargs) def encode_corpus(self, corpus: Union[List[str], str, List[Dict[str, str]]], **kwargs) -> np.ndarray: """Used for encoding the corpus of retrieval tasks""" if isinstance(corpus, dict): corpus = [corpus] if isinstance(corpus, list) and isinstance(corpus[0], dict): corpus = [ doc["text"] for doc in corpus ] return self.encode(corpus, **kwargs) @torch.no_grad() def encode( self, sentences: Union[List[str], str], batch_size: int = 256, max_length: int = 512, instruction: str = "", user_token: str = "<|start_header_id|>user<|end_header_id|>", embed_instruction: bool = False, embed_eos_token: str = "\n<|embed|>\n", convert_to_tensor: bool = False, add_special_tokens: bool = True, **kwargs, ) -> np.ndarray: input_was_string = False if isinstance(sentences, str): sentences = [sentences] input_was_string = True all_embeddings, all_kv_caches = [], [] for start_index in tqdm(range(0, len(sentences), batch_size), desc="Batches", disable=len(sentences)<256): sentences_batch = [ user_token + instruction + embed_eos_token + s for s in sentences[start_index:start_index + batch_size] ] # This will prepend the bos token if the tokenizer has `add_bos_token=True` inputs = self.tokenizer( sentences_batch, padding=True, truncation=True, return_tensors='pt', max_length=max_length, add_special_tokens=add_special_tokens, ).to(self.device) last_hidden_state = self.model(inputs['input_ids'])['last_hidden_state'] if ("mean" in self.embedding_head.pooling_method) and not embed_instruction: instruct_with_special_tokens = user_token + instruction + embed_eos_token # Remove instruction tokens from the embeddings by masking them instruction_tokens = self.tokenizer( instruct_with_special_tokens, padding=False, truncation=True, max_length=max_length, add_special_tokens=add_special_tokens, )["input_ids"] inputs['attention_mask'][:, :len(instruction_tokens)] = 0 embeddings = self.embedding_head(last_hidden_state, inputs['attention_mask']) if convert_to_tensor: all_embeddings.append(embeddings) else: # NumPy does not support bfloat16 all_embeddings.append(embeddings.cpu().to(torch.float32).numpy()) all_embeddings = ( torch.cat(all_embeddings, dim=0) if convert_to_tensor else np.concatenate(all_embeddings, axis=0) ) if input_was_string: all_embeddings = all_embeddings[0] return all_embeddings