# Licensed to the Apache Software Foundation (ASF) under one # or more contributor license agreements. See the NOTICE file # distributed with this work for additional information # regarding copyright ownership. The ASF licenses this file # to you 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 typing import Optional import torch import torch.nn.functional as F from torch import Tensor from torch.nn import ( Module, Parameter, Linear, SiLU, Conv1d, Embedding, ) from transformers.cache_utils import ( Cache, DynamicCache, ) from transformers.modeling_outputs import ( BaseModelOutputWithPast, CausalLMOutputWithPast, ) from transformers.modeling_utils import ( PreTrainedModel, ) from transformers.generation import GenerationMixin from transformers.pytorch_utils import ( ALL_LAYERNORM_LAYERS, ) from .configuration_iii import IIIConfig _CONFIG_FOR_DOC = "IIIConfig" # ============================================================ # RMSNorm # ============================================================ class IIIRMSNorm(Module): def __init__( self, dim, eps=1e-6, ): super().__init__() self.weight = Parameter( torch.ones(dim) ) self.epsilon = eps def forward( self, tensor, ): input_dtype = tensor.dtype tensor = tensor.to( torch.float32 ) variance = tensor.pow(2).mean( -1, keepdim=True, ) tensor = tensor * torch.rsqrt( variance + self.epsilon ) return ( self.weight * tensor.to(input_dtype) ) ALL_LAYERNORM_LAYERS.append( IIIRMSNorm ) # ============================================================ # RoPE # ============================================================ class IIIRoPE(Module): def __init__( self, dim, max_position_embeddings=100000, theta=10000.0, device=None, ): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.theta = theta inv_freq = 1.0 / ( theta ** ( torch.arange( 0, dim, 2, device=device, dtype=torch.float32, ) / dim ) ) self.register_buffer( "inv_freq", inv_freq, persistent=False, ) # Do not build RoPE cache during initialization. # This makes loading model weights independent of the cache. self.register_buffer( "cos_cached", None, persistent=False, ) self.register_buffer( "sin_cached", None, persistent=False, ) self.max_seq_len_cached = 0 def _set_cos_sin_cache( self, seq_len, device, dtype, ): self.max_seq_len_cached = seq_len t = torch.arange( seq_len, device=device, dtype=self.inv_freq.dtype, ) freqs = torch.outer( t, self.inv_freq.to(device=t.device), ) emb = torch.cat( (freqs, freqs), dim=-1, ) self.cos_cached = emb.cos().to(dtype) self.sin_cached = emb.sin().to(dtype) def forward( self, x, seq_len=None, ): if seq_len is None: seq_len = x.shape[-2] # RoPE cache does not exist yet, or is too short. if ( self.cos_cached is None or self.sin_cached is None or seq_len > self.max_seq_len_cached ): self._set_cos_sin_cache( seq_len=seq_len, device=x.device, dtype=x.dtype, ) return ( self.cos_cached[:seq_len].to( device=x.device, dtype=x.dtype, ), self.sin_cached[:seq_len].to( device=x.device, dtype=x.dtype, ), ) def rotate_half(x): x1 = x[ ..., : x.shape[-1] // 2, ] x2 = x[ ..., x.shape[-1] // 2 :, ] return torch.cat( ( -x2, x1, ), dim=-1, ) def apply_rotary_pos_emb( q, k, cos, sin, position_ids, unsqueeze_dim=1, ): cos = cos[ position_ids ].unsqueeze( unsqueeze_dim ) sin = sin[ position_ids ].unsqueeze( unsqueeze_dim ) q_embed = ( q * cos + rotate_half(q) * sin ) k_embed = ( k * cos + rotate_half(k) * sin ) return q_embed, k_embed # ============================================================ # SwiGLU-style MLP # ============================================================ class IIIMLP(Module): def __init__( self, config, dim=None, intermediate_dim=None, ): super().__init__() self.config = config self.dim = ( config.hidden_dim if dim is None else dim ) self.intermediate_dim = ( config.intermediate_dim if intermediate_dim is None else intermediate_dim ) self.gate = Linear( self.dim, self.intermediate_dim, bias=False, ) self.up = Linear( self.dim, self.intermediate_dim, bias=False, ) self.down = Linear( self.intermediate_dim, self.dim, bias=False, ) self.activation = SiLU() def forward(self, x): value = self.up(x) gate = self.activation( self.gate(x) ) return self.down( gate * value ) # ============================================================ # Attention # ============================================================ class IIIAttention(Module): def __init__( self, config, layer_id=None, ): super().__init__() self.config = config self.layer_id = layer_id self.hidden_dim = config.hidden_dim self.lora_dim = config.lora_dim self.num_heads = config.n_heads if self.lora_dim % self.num_heads != 0: raise ValueError( f"lora_dim ({self.lora_dim}) " f"must be divisible by " f"n_heads ({self.num_heads})." ) self.head_dim = self.lora_dim // self.num_heads # ---------------------------------------------------- # QKV # ---------------------------------------------------- self.qkv = Linear( self.hidden_dim, self.lora_dim * 3, bias=config.attention_bias, ) self.o = Linear( self.lora_dim, self.hidden_dim, bias=config.attention_bias, ) # ---------------------------------------------------- # Q/K normalization # ---------------------------------------------------- self.q_norm = IIIRMSNorm( self.head_dim, eps=config.rms_norm_eps, ) self.k_norm = IIIRMSNorm( self.head_dim, eps=config.rms_norm_eps, ) # ---------------------------------------------------- # RoPE # ---------------------------------------------------- self.rotary_emb = IIIRoPE( dim=self.head_dim, max_position_embeddings=config.max_position_embedding, theta=config.rope_theta, ) # ============================================================ # Causal mask # ============================================================ def _make_causal_mask( self, q_len, kv_len, device, ): """ Returns a boolean causal mask of shape: [1, 1, q_len, kv_len] True = attention is allowed False = attention is blocked Works for both: 1. Full sequence q_len = kv_len = S 2. KV-cache decoding q_len < kv_len Example: q_len=1, kv_len=5 [1 1 1 1 1] Example: q_len=2, kv_len=5 [1 1 1 1 0] [1 1 1 1 1] """ q_positions = torch.arange( kv_len - q_len, kv_len, device=device, ) k_positions = torch.arange( kv_len, device=device, ) causal = ( k_positions.unsqueeze(0) <= q_positions.unsqueeze(1) ) return causal.unsqueeze(0).unsqueeze(0) # ============================================================ # Prepare user attention mask # ============================================================ def _prepare_attention_mask( self, attention_mask, device, ): """ Converts attention_mask into an attention-compatible shape. Supported input shapes: [B, KV] [B, Q, KV] [B, 1, Q, KV] Boolean masks: True = allowed False = blocked Integer masks: 1 = allowed 0 = blocked Float masks: treated as additive attention bias. """ if attention_mask is None: return None # -------------------------------------------------------- # [B, KV] # -------------------------------------------------------- if attention_mask.dim() == 2: if attention_mask.dtype == torch.bool: return attention_mask[ :, None, None, : ].to(device) # Integer 0/1 attention masks if not torch.is_floating_point(attention_mask): return ( attention_mask != 0 )[:, None, None, :].to(device) # Float mask. # Hugging Face commonly uses 0 for valid tokens # and a large negative value for masked tokens. return attention_mask[ :, None, None, : ].to(device) # -------------------------------------------------------- # [B, Q, KV] # -------------------------------------------------------- if attention_mask.dim() == 3: return attention_mask[ :, None, :, : ].to(device) # -------------------------------------------------------- # [B, 1, Q, KV] # -------------------------------------------------------- if attention_mask.dim() == 4: return attention_mask.to(device) raise ValueError( "attention_mask must be " "2D, 3D, or 4D." ) # ============================================================ # Combine causal + attention mask # ============================================================ def _combine_masks( self, attention_mask, q_len, kv_len, device, dtype, ): """ Combines the user's attention/padding mask with the causal mask. This is the important part. We NEVER allow an attention_mask to replace causality. Returns: bool mask OR additive float mask suitable for SDPA with is_causal=False. """ causal_mask = self._make_causal_mask( q_len=q_len, kv_len=kv_len, device=device, ) # -------------------------------------------------------- # No user mask # -------------------------------------------------------- if attention_mask is None: return causal_mask mask = self._prepare_attention_mask( attention_mask, device, ) # -------------------------------------------------------- # Boolean mask # -------------------------------------------------------- if mask.dtype == torch.bool: return mask & causal_mask # -------------------------------------------------------- # Integer mask # -------------------------------------------------------- if not torch.is_floating_point(mask): mask = mask != 0 return mask & causal_mask # -------------------------------------------------------- # Additive float mask # # Typical form: # # valid -> 0 # masked -> -inf / large negative # # We impose causality by setting all future positions # to the minimum representable value. # -------------------------------------------------------- min_value = torch.finfo(dtype).min # Convert causal mask to additive form. causal_bias = torch.zeros( (1, 1, q_len, kv_len), dtype=dtype, device=device, ) causal_bias = causal_bias.masked_fill( ~causal_mask, min_value, ) # Make sure the supplied mask has the correct dtype. mask = mask.to(dtype) return mask + causal_bias # ============================================================ # Forward # ============================================================ def forward( self, hidden_states, attention_mask=None, position_ids=None, kv_cache: Optional[Cache] = None, use_cache=False, output_attentions=False, ): batch_size, seq_len, _ = hidden_states.shape # -------------------------------------------------------- # Validate cache usage # -------------------------------------------------------- if ( use_cache and kv_cache is not None and self.layer_id is None ): raise ValueError( "layer_id must be set when " "using KV cache." ) # -------------------------------------------------------- # QKV # -------------------------------------------------------- qkv = self.qkv(hidden_states) q, k, v = qkv.chunk( 3, dim=-1, ) # -------------------------------------------------------- # Heads # -------------------------------------------------------- q = q.view( batch_size, seq_len, self.num_heads, self.head_dim, ).transpose(1, 2) k = k.view( batch_size, seq_len, self.num_heads, self.head_dim, ).transpose(1, 2) v = v.view( batch_size, seq_len, self.num_heads, self.head_dim, ).transpose(1, 2) # -------------------------------------------------------- # Q/K RMSNorm # -------------------------------------------------------- q = self.q_norm(q) k = self.k_norm(k) # -------------------------------------------------------- # Cache length # -------------------------------------------------------- if ( kv_cache is not None and kv_cache.get_seq_length( self.layer_id ) > 0 ): past_len = kv_cache.get_seq_length( self.layer_id ) else: past_len = 0 # -------------------------------------------------------- # Position IDs # -------------------------------------------------------- if position_ids is None: position_ids = torch.arange( past_len, past_len + seq_len, device=hidden_states.device, ) position_ids = ( position_ids .unsqueeze(0) .expand( batch_size, -1, ) ) max_position = ( int(position_ids.max().item()) + 1 ) # -------------------------------------------------------- # RoPE # -------------------------------------------------------- cos, sin = self.rotary_emb( q, seq_len=max_position, ) q, k = apply_rotary_pos_emb( q, k, cos, sin, position_ids, ) # -------------------------------------------------------- # Update KV cache # -------------------------------------------------------- if ( use_cache and kv_cache is not None ): k, v = kv_cache.update( k, v, self.layer_id, ) # -------------------------------------------------------- # Actual attention sizes # -------------------------------------------------------- kv_len = k.shape[-2] q_len = q.shape[-2] # ======================================================== # Build combined mask # ======================================================== mask = self._combine_masks( attention_mask=attention_mask, q_len=q_len, kv_len=kv_len, device=hidden_states.device, dtype=q.dtype, ) # We explicitly constructed the causal mask above, # so don't also ask SDPA to construct another one. # # This is also necessary because SDPA does not allow # simultaneously supplying an attn_mask and is_causal=True. is_causal = False # ======================================================== # Attention # ======================================================== attention_weights = None # -------------------------------------------------------- # Debug / visualization path # -------------------------------------------------------- if output_attentions: scale = self.head_dim ** -0.5 scores = torch.matmul( q, k.transpose(-2, -1), ) * scale # ---------------------------------------------------- # Boolean mask # ---------------------------------------------------- if mask.dtype == torch.bool: scores = scores.masked_fill( ~mask, torch.finfo( scores.dtype ).min, ) # ---------------------------------------------------- # Additive mask # ---------------------------------------------------- else: scores = scores + mask.to( scores.dtype ) attention_weights = torch.softmax( scores, dim=-1, ) attn_output = torch.matmul( attention_weights, v, ) # -------------------------------------------------------- # SDPA path # -------------------------------------------------------- else: attn_output = ( F.scaled_dot_product_attention( q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=is_causal, ) ) # ======================================================== # Merge heads # ======================================================== attn_output = ( attn_output .transpose(1, 2) .contiguous() .view( batch_size, seq_len, self.lora_dim, ) ) # ======================================================== # Output # ======================================================== attn_output = self.o( attn_output ) return ( attn_output, kv_cache, attention_weights, ) # ============================================================ # Gated Causal CNN # ============================================================ class GatedCausalCNN(Module): """ Gated causal dilated CNN. value = W_v(x) score = W_s(x) gated = value * sigmoid(score) output = CausalConv1D(gated) No RoPE is applied here. The convolution itself encodes relative local position through its kernel offsets and dilation. """ def __init__( self, config, layer_id, ): super().__init__() self.config = config self.layer_id = layer_id self.hidden_dim = ( config.hidden_dim ) self.window_size = ( config.layer_window_sizes[ layer_id ] ) self.dilation = ( config.layer_dilation_size[ layer_id ] ) if self.window_size <= 0: raise ValueError( "GatedCausalCNN requires " "window_size > 0." ) if self.dilation <= 0: raise ValueError( "GatedCausalCNN requires " "dilation > 0." ) # Receptive field: # # 1 + (kernel_size - 1) * dilation # self.receptive_field = ( 1 + ( self.window_size - 1 ) * self.dilation ) # ---------------------------------------------------- # Gated projection # # One projection gives us: # # [value | score] # ---------------------------------------------------- self.gate_proj = Linear( self.hidden_dim, self.hidden_dim * 2, bias=False, ) # ---------------------------------------------------- # Depthwise causal convolution # ---------------------------------------------------- self.conv = Conv1d( in_channels=self.hidden_dim, out_channels=self.hidden_dim, kernel_size=self.window_size, dilation=self.dilation, groups=self.hidden_dim, bias=config.attention_bias, ) # ---------------------------------------------------- # Output projection # ---------------------------------------------------- self.out_proj = Linear( self.hidden_dim, self.hidden_dim, bias=config.attention_bias, ) def _gate( self, hidden_states, ): """ hidden_states: [B, S, D] returns: [B, S, D] """ value, score = self.gate_proj( hidden_states ).chunk( 2, dim=-1, ) return ( value * torch.sigmoid(score) ) def _causal_conv( self, x, ): """ x: [B, S, D] Returns: [B, S, D] """ left_padding = ( self.receptive_field - 1 ) # [B, S, D] # ↓ # [B, D, S] x = x.transpose( 1, 2, ) x = F.pad( x, ( left_padding, 0, ), ) x = self.conv(x) # [B, D, S] # ↓ # [B, S, D] return x.transpose( 1, 2, ) def forward( self, hidden_states, kv_cache: Optional[Cache] = None, use_cache=False, ): batch_size, seq_len, _ = ( hidden_states.shape ) # ==================================================== # No cache # ==================================================== if ( not use_cache or kv_cache is None ): gated = self._gate( hidden_states ) out = self._causal_conv( gated ) out = self.out_proj(out) return ( out, kv_cache, ) # ==================================================== # Cache # ==================================================== if self.layer_id is None: raise ValueError( "layer_id must be set when " "using a cache." ) past_len = ( kv_cache.get_seq_length( self.layer_id ) ) # ---------------------------------------------------- # Retrieve cached raw hidden states # # CNN caches store: # # [B, 1, S, D] # # in the key slot. # ---------------------------------------------------- if past_len > 0: layer_cache = ( kv_cache.layers[ self.layer_id ] ) past_states = ( layer_cache.keys[:, 0] ) # We only need the actual receptive # field before the new token(s). history_len = ( self.receptive_field - 1 ) past_states = ( past_states[ :, -history_len:, : ] ) else: past_states = ( hidden_states[:, :0, :] ) # ---------------------------------------------------- # Concatenate history + new states # ---------------------------------------------------- states = torch.cat( ( past_states, hidden_states, ), dim=1, ) # ---------------------------------------------------- # GATE FIRST # ---------------------------------------------------- gated = self._gate( states ) # ---------------------------------------------------- # THEN CONV1D # ---------------------------------------------------- out = self._causal_conv( gated ) # Only return the outputs belonging # to the newly supplied tokens. out = out[ :, -seq_len:, :, ] out = self.out_proj(out) # ---------------------------------------------------- # Cache RAW hidden states # # Do not cache the gated representation. # # The next token needs: # # value(history) * # sigmoid(score(history)) # # before convolution. # ---------------------------------------------------- cache_keys = ( hidden_states.unsqueeze(1) ) cache_values = ( hidden_states[ ..., :1 ].unsqueeze(1) ) kv_cache.update( cache_keys, cache_values, self.layer_id, ) return ( out, kv_cache, ) # ============================================================ # III Decoder Layer # ============================================================ class IIILayer(Module): def __init__( self, config, layer_id, ): super().__init__() self.config = config self.layer_id = layer_id self.window_size = ( config.layer_window_sizes[ layer_id ] ) self.dilation = ( config.layer_dilation_size[ layer_id ] ) # ---------------------------------------------------- # Pre-mixer RMSNorm # ---------------------------------------------------- self.input_layernorm = ( IIIRMSNorm( config.hidden_dim, eps=config.rms_norm_eps, ) ) # ---------------------------------------------------- # Mixer # ---------------------------------------------------- if self.window_size == 0: self.is_attention = True self.mixer = IIIAttention( config, layer_id=layer_id, ) else: self.is_attention = False self.mixer = GatedCausalCNN( config, layer_id=layer_id, ) # ---------------------------------------------------- # Post-mixer RMSNorm # ---------------------------------------------------- self.post_mixer_layernorm = ( IIIRMSNorm( config.hidden_dim, eps=config.rms_norm_eps, ) ) # ---------------------------------------------------- # MLP # ---------------------------------------------------- self.mlp = IIIMLP(config) def forward( self, hidden_states, attention_mask=None, position_ids=None, kv_cache=None, use_cache=False, output_attentions=False, ): # ---------------------------------------------------- # Mixer # ---------------------------------------------------- residual = hidden_states hidden_states = ( self.input_layernorm( hidden_states ) ) if self.is_attention: ( mixer_output, kv_cache, attention_weights, ) = self.mixer( hidden_states, attention_mask=attention_mask, position_ids=position_ids, kv_cache=kv_cache, use_cache=use_cache, output_attentions=( output_attentions ), ) else: ( mixer_output, kv_cache, ) = self.mixer( hidden_states, kv_cache=kv_cache, use_cache=use_cache, ) attention_weights = None hidden_states = ( residual + mixer_output ) # ---------------------------------------------------- # MLP # ---------------------------------------------------- residual = hidden_states hidden_states = ( self.post_mixer_layernorm( hidden_states ) ) hidden_states = self.mlp( hidden_states ) hidden_states = ( residual + hidden_states ) return ( hidden_states, kv_cache, attention_weights, ) # ============================================================ # Pretrained base # ============================================================ class IIIPreTrainedModel( PreTrainedModel ): config_class = IIIConfig base_model_prefix = "model" _supports_cache_class = True def _init_weights( self, module, ): std = getattr( self.config, "initializer_range", 0.02, ) if isinstance( module, Linear, ): torch.nn.init.normal_( module.weight, mean=0.0, std=std, ) if module.bias is not None: torch.nn.init.zeros_( module.bias ) elif isinstance( module, Embedding, ): torch.nn.init.normal_( module.weight, mean=0.0, std=std, ) if ( module.padding_idx is not None ): torch.nn.init.zeros_( module.weight[ module.padding_idx ] ) elif isinstance( module, Conv1d, ): torch.nn.init.normal_( module.weight, mean=0.0, std=std, ) if module.bias is not None: torch.nn.init.zeros_( module.bias ) elif isinstance( module, IIIRMSNorm, ): torch.nn.init.ones_( module.weight ) # ============================================================ # III Model # ============================================================ class IIIModel( IIIPreTrainedModel ): def __init__( self, config, ): super().__init__( config ) # ---------------------------------------------------- # Find an attention layer. # # This is important because CNN cache length can be # local, while attention cache length represents the # complete generated sequence. # ---------------------------------------------------- attention_layers = [ i for i, window in enumerate( config.layer_window_sizes ) if window == 0 ] if not attention_layers: raise ValueError( "IIIModel requires at least " "one attention layer." ) self.first_attention_layer = ( attention_layers[0] ) # ---------------------------------------------------- # Token embeddings # ---------------------------------------------------- self.embed_tokens = Embedding( config.vocab_size, config.hidden_dim, padding_idx=config.pad_token_id, ) # ---------------------------------------------------- # Decoder layers # ---------------------------------------------------- self.layers = ( torch.nn.ModuleList( [ IIILayer( config, layer_id=i, ) for i in range( config.n_layers ) ] ) ) # ---------------------------------------------------- # Final norm # ---------------------------------------------------- self.norm = IIIRMSNorm( config.hidden_dim, eps=config.rms_norm_eps, ) self.gradient_checkpointing = ( False ) self.post_init() def get_input_embeddings( self, ): return self.embed_tokens def set_input_embeddings( self, value, ): self.embed_tokens = value def _get_past_length( self, past_key_values, ): if past_key_values is None: return 0 # IMPORTANT: # # Layer 0 is a CNN and may conceptually only need # a local history. Attention layers contain the # complete sequence, so use the first attention # layer to determine the real sequence position. return past_key_values.get_seq_length( self.first_attention_layer ) def forward( self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, cache_position=None, use_cache=None, output_attentions=False, output_hidden_states=False, return_dict=None, ): if ( input_ids is not None and inputs_embeds is not None ): raise ValueError( "Specify either input_ids " "or inputs_embeds, not both." ) if ( input_ids is None and inputs_embeds is None ): raise ValueError( "Either input_ids or " "inputs_embeds must be provided." ) if use_cache is None: use_cache = ( self.config.use_cache ) if return_dict is None: return_dict = True # ---------------------------------------------------- # Cache # ---------------------------------------------------- if use_cache: if past_key_values is None: past_key_values = ( DynamicCache() ) elif not isinstance( past_key_values, Cache, ): raise TypeError( "past_key_values must be " "a Transformers Cache." ) else: past_key_values = None # ---------------------------------------------------- # Inputs # ---------------------------------------------------- if inputs_embeds is None: hidden_states = ( self.embed_tokens( input_ids ) ) else: hidden_states = ( inputs_embeds ) batch_size, seq_len, _ = ( hidden_states.shape ) # ---------------------------------------------------- # Positions # ---------------------------------------------------- past_length = ( self._get_past_length( past_key_values ) ) if position_ids is None: if cache_position is not None: position_ids = ( cache_position ) if position_ids.dim() == 1: position_ids = ( position_ids .unsqueeze(0) .expand( batch_size, -1, ) ) else: position_ids = ( torch.arange( past_length, past_length + seq_len, device=( hidden_states.device ), ) .unsqueeze(0) .expand( batch_size, -1, ) ) # ---------------------------------------------------- # Outputs # ---------------------------------------------------- all_hidden_states = ( () if output_hidden_states else None ) all_attentions = ( () if output_attentions else None ) # ---------------------------------------------------- # Decoder # ---------------------------------------------------- for layer in self.layers: if output_hidden_states: all_hidden_states += ( hidden_states, ) ( hidden_states, past_key_values, attention_weights, ) = layer( hidden_states, attention_mask=attention_mask, position_ids=position_ids, kv_cache=past_key_values, use_cache=use_cache, output_attentions=( output_attentions ), ) if output_attentions: all_attentions += ( attention_weights, ) # ---------------------------------------------------- # Final norm # ---------------------------------------------------- hidden_states = self.norm( hidden_states ) if output_hidden_states: all_hidden_states += ( hidden_states, ) # ---------------------------------------------------- # Tuple output # ---------------------------------------------------- if not return_dict: return tuple( value for value in ( hidden_states, past_key_values, all_hidden_states, all_attentions, ) if value is not None ) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=( past_key_values if use_cache else None ), hidden_states=( all_hidden_states ), attentions=( all_attentions ), ) # ============================================================ # III Causal LM # ============================================================ class IIIForCausalLM( IIIPreTrainedModel, GenerationMixin, ): _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} def __init__( self, config, ): super().__init__( config ) self.model = IIIModel( config ) self.lm_head = Linear( config.hidden_dim, config.vocab_size, bias=False, ) self.post_init() if config.tie_word_embedding: self.lm_head.weight = ( self.model.embed_tokens.weight ) # -------------------------------------------------------- # Embeddings # -------------------------------------------------------- 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 # -------------------------------------------------------- # Forward # -------------------------------------------------------- def forward( self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, labels=None, cache_position=None, use_cache=None, output_attentions=False, output_hidden_states=False, return_dict=None, ): if return_dict is None: return_dict = True 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, cache_position=cache_position, use_cache=use_cache, output_attentions=( output_attentions ), output_hidden_states=( output_hidden_states ), return_dict=True, ) # ---------------------------------------------------- # LM logits # ---------------------------------------------------- logits = self.lm_head( outputs.last_hidden_state ) # ---------------------------------------------------- # Causal LM loss # ---------------------------------------------------- loss = None if labels is not None: shift_logits = ( logits[ ..., :-1, :, ] .contiguous() ) shift_labels = ( labels[ ..., 1:, ] .contiguous() ) loss = F.cross_entropy( shift_logits.view( -1, shift_logits.size(-1), ), shift_labels.view( -1 ), ) # ---------------------------------------------------- # Tuple output # ---------------------------------------------------- if not return_dict: result = ( logits, outputs.past_key_values, outputs.hidden_states, outputs.attentions, ) if loss is not None: result = ( loss, *result, ) return result return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=( outputs.past_key_values ), hidden_states=( outputs.hidden_states ), attentions=( outputs.attentions ), ) # -------------------------------------------------------- # Generation # -------------------------------------------------------- def prepare_inputs_for_generation( self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, **kwargs, ): """ Prepare only the unprocessed tokens for generation. When a cache exists, the model should not recompute hidden states for tokens already represented in the cache. """ if ( past_key_values is not None ): past_length = ( past_key_values.get_seq_length( self.model.first_attention_layer ) ) if ( input_ids is not None and input_ids.shape[1] > past_length ): input_ids = input_ids[ :, past_length:, ] elif input_ids is not None: # Usually generation reaches this path with # one new token already supplied. input_ids = input_ids[ :, -1:, ] if ( inputs_embeds is not None ): inputs_embeds = ( inputs_embeds[ :, -input_ids.shape[1]:, :, ] ) model_inputs = {} if ( inputs_embeds is not None and past_key_values is None ): model_inputs[ "inputs_embeds" ] = inputs_embeds if input_ids is not None: model_inputs[ "input_ids" ] = input_ids else: model_inputs[ "input_ids" ] = input_ids model_inputs[ "past_key_values" ] = past_key_values model_inputs[ "attention_mask" ] = attention_mask model_inputs[ "cache_position" ] = cache_position model_inputs[ "use_cache" ] = kwargs.get( "use_cache", self.config.use_cache, ) return model_inputs