PiynddlZddlmZmZmZmZmZddlZddlmZddl m Z ddl m Z Gddej ZGdd ej Zd ej d ed ejfd ZGddej ZdS)N)CallableDictListOptionalUnion)nn)MultiHeadAttention) _MaskTypecLeZdZdZddddddedejdeejdeejdeejd eejd dffd Zd e d e j de de d df dZ d e fdZd e fdZdZdddde jdeedee jded e jf dZxZS)TransformerSelfAttentionLayeraF Transformer layer derived from the Llama2 model. Normalization is applied before the attention **and** FF layer. Args: attn (MultiHeadAttention): Attention module. mlp (nn.Module): Feed-forward module. sa_norm (Optional[nn.Module]): Normalization to be applied before self-attention. mlp_norm (Optional[nn.Module]): Normalization to be applied before the feed-forward layer. sa_scale (Optional[nn.Module]): Module to scale self-attention output. mlp_scale (Optional[nn.Module]): Module to scale the feed-forward output. N)sa_normmlp_normsa_scale mlp_scaleattnmlpr rrrreturnc4t||_||_|pt j|_|pt j|_|pt j|_|pt j|_ dSN) super__init__rrrIdentityr rrr)selfrrr rrr __class__s q/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/torchtune/modules/transformer.pyrz&TransformerSelfAttentionLayer.__init__su  /"+--  1BKMM  1BKMM "3bkmm batch_sizedtypeencoder_max_seq_lendecoder_max_seq_lenc@|j|||dS)aOSetup key value caches for attention calculation. Args: batch_size (int): batch size for the caches. dtype (torch.dtype): dtype for the caches. encoder_max_seq_len (int): this parameter is ignored in this layer. decoder_max_seq_len (int): maximum cache sequence length. ) max_seq_lenNr setup_cacherrrrr s r setup_cachesz*TransformerSelfAttentionLayer.setup_caches.s( j%=PQQQQQrc|jjduSz Check if the key value caches are setup on ``self.attn``. See :func:~torchtune.modules.TransformerDecoder.caches_are_setup`. Nrkv_cachers rcaches_are_setupz.TransformerSelfAttentionLayer.caches_are_setup@ y!--rc|jjSz Checks if the key value caches on ``self.attn`` are enabled. See :func:~torchtune.modules.TransformerDecoder.caches_are_enabled`. r cache_enabledr+s rcaches_are_enabledz0TransformerSelfAttentionLayer.caches_are_enabledG y&&rc8|jdSzReset the key value caches.Nr reset_cacher+s rr7z)TransformerSelfAttentionLayer.reset_cacheN rmask input_posxr:r;kwargsc ||}|||||}|||z}|||}|||z}|S)a Args: x (torch.Tensor): input tensor with shape [batch_size x seq_length x embed_dim] mask (Optional[_MaskType]): Used to mask the scores after the query-key multiplication and before the softmax. Either: A boolean tensor with shape ``[b x s x s]``, ``[b x s x self.encoder_max_cache_seq_len]``, or ``[b x s x self.encoder_max_cache_seq_len]`` if using KV-cacheing with encoder/decoder layers. A value of True in row ``i`` and column ``j`` means token ``i`` attends to token ``j``. A value of False means token ``i`` does not attend to token ``j``. If no mask is specified, a causal mask is used by default. A :class:`~torch.nn.attention.flex_attention.BlockMask` for document masking in a packed sequence created via `create_block_mask `_. We use :func:`~torch.nn.attention.flex_attention.flex_attention` when computing attention with block masks. Default is None. input_pos (Optional[torch.Tensor]): Optional tensor which contains the position ids of each token. During training, this is used to indicate the positions of each token relative to its sample when packed, shape [b x s]. During inference, this indicates the position of the current token. If none, assume the index of the token is its position id. Default is None. **kwargs (Dict): transformer layer inputs not relevant to self attention. Returns: torch.Tensor: output tensor with same shape as input [batch_size x seq_length x embed_dim] r9)r rrrrr) rr<r:r;r=hattn_outmlp_outouts rforwardz%TransformerSelfAttentionLayer.forwardRsN LLOO99Q 9BB MM( # #a '((4==++,,$..))) r)__name__ __module__ __qualname____doc__r rModulerrinttorchrr&boolr,r2r7Tensorr rrC __classcell__rs@rr r s  "(,(,(,)-444 4Y4 ")$ 4 29% 429%4BI&4 444444$RR{R ! R ! R RRRR$.$....'D''''   %),0 222 <2y! 2 EL) 2  2 22222222rr ceZdZdZddddddedejdeejdeejdeejd eejd dffd Zd e d e j de de d df dZ d e fdZd e fdZdZdee jd ee jfdZdddde jdee jdee jded e jf dZxZS)TransformerCrossAttentionLayera Cross attention Transformer layer following the same conventions as the TransformerSelfAttentionLayer. Normalization is applied before the attention **and** FF layer. Args: attn (MultiHeadAttention): Attention module. mlp (nn.Module): Feed-forward module. ca_norm (Optional[nn.Module]): Normalization to be applied before cross-attention. mlp_norm (Optional[nn.Module]): Normalization to be applied before the feed-forward layer. ca_scale (Optional[nn.Module]): Module to scale cross-attention output. mlp_scale (Optional[nn.Module]): Module to scale the feed-forward output. Raises: AssertionError: if attn.pos_embeddings is set. N)ca_normrca_scalerrrrQrrRrrc`t|jtd||_||_|pt j|_|pt j|_ |pt j|_ |pt j|_ dS)NzsDoesn't support positional embeddings for cross attention, because q and k are different sequences.) rrpos_embeddingsAssertionErrorrrrrrQrrRr)rrrrQrrRrrs rrz'TransformerCrossAttentionLayer.__init__s    * :  /"+--  1BKMM  1BKMM "3bkmmrrrrr c>|j|||dS)aOSetup key value caches for attention calculation. Args: batch_size (int): batch size for the caches. dtype (torch.dtype): dtype for the caches. encoder_max_seq_len (int): maximum cache sequence length. decoder_max_seq_len (int): this parameter is ignored in this layer. Nr#r%s rr&z+TransformerCrossAttentionLayer.setup_cachess% j%1DEEEEErc|jjduSr(r)r+s rr,z/TransformerCrossAttentionLayer.caches_are_setupr-rc|jjSr/r0r+s rr2z1TransformerCrossAttentionLayer.caches_are_enabledr3rc8|jdSr5r6r+s rr7z*TransformerCrossAttentionLayer.reset_cacher8rr:c|dS|jtjkr|}ntj|}tj|dd}|S)a0Some tokens in x may not attend to any encoder inputs due to the cross attention mask (encoder_mask). This results in a full row of the attention matrix being masked out. In the example below, the word "the" is masked from every embedding. The False value means a token can't attend to an embedding. .. code-block:: text |emb||emb||emb| |The| F F F |red| T F T |car| F T T This results in no inputs into the softmax layer which causes a NaN. The skip mask is used to mask the outputs of attention and mlp resulting in the token being skipped. The above example would result in a skip mask of: [[True], [False], [False]] which specifies which tokens to fully mask out. NT)dimkeepdim)rrJrKisneginfall)rr:s r _skip_maskz)TransformerCrossAttentionLayer._skip_masksP0 <4 : # #5DD>$''Dy2t444 r) encoder_input encoder_maskr<rarbr=c &| p|jjjdk}||r|S||}|||d}|||||}|||d}|||z}|| |} || |d} || | z} | S)aZ Args: x (torch.Tensor): input tensor with shape [batch_size x seq_length x embed_dim] encoder_input (Optional[torch.Tensor]): Optional input embeds from the encoder. Shape [batch_size x token_sequence x embed_dim] encoder_mask (Optional[torch.Tensor]): Boolean tensor defining a relational matrix between tokens and encoder embeddings. A True value at position i,j means token i can attend to embedding j in the decoder. Mask has shape [batch_size x token_sequence x embed_sequence]. Default is None. **kwargs (Dict): transformer layer inputs not relevant to self attention. Returns: torch.Tensor: output tensor with same shape as input [batch_size x seq_length x embed_dim] rNT)r:) r2rr*sizer` masked_fillrQrRrrr) rr<rarbr= empty_cache skip_maskr@r?rArBs rrCz&TransformerCrossAttentionLayer.forwards411333Sty7I7NRS7S   [ HOOL11  #(33ItDDL 99T\\!__m,9OO  ++Iq99H MM( # #a '((4==++,,  )))Q77G$..))) r)rDrErFrGr rrHrrrIrJrr&rKr,r2r7rLr`rrCrMrNs@rrPrPs*(,(,(,)-444 4Y4 ")$ 4 29% 429%4BI&4 444444.FF{F ! F ! F FFFF$.$....'D''''   !x 5!(5<:P!!!!N15/3 ::: <:  - : u|, :  : ::::::::rrPmodulenrc^tjfdt|DS)z Return a list of ``n`` identical layers. Args: module (nn.Module): module to be cloned n (int): number of clones Returns: nn.ModuleList: list of ``n`` identical layers c8g|]}tjS)copydeepcopy).0irhs r z_get_clones..?s#BBBA$-//BBBr)r ModuleListrange)rhris` r _get_clonesrt3s0 =BBBBqBBB C CCrceZdZdZddddejdeejeejej fde de de d ejd eej e fd e e d e ee d dffdZde d dfdZdddde dejde e de e fdZd efdZd efdZdZejjdejd eejfdZ d&de de ejde ejde ejd e ejf d!Zddddd"d#ejde ede ejde ejd e ejd eejeejff d$Zd%ZxZS)'TransformerDecodera Transformer Decoder derived from the Llama2 architecture. Args: tok_embeddings (nn.Embedding): PyTorch embedding layer, to be used to move tokens to an embedding space. layers (Union[nn.Module, List[nn.Module], nn.ModuleList]): A single transformer Decoder layer, an nn.ModuleList of layers or a list of layers. It is recommended to use an nn.ModuleList. max_seq_len (int): maximum sequence length the model will be run with, as used by :func:`~torchtune.modules.KVCache` num_heads (int): number of query heads. For MHA this is also the number of heads for key and value. This is used to setup the :func:`~torchtune.modules.KVCache` head_dim (int): embedding dimension for each head in self-attention. This is used to setup the :func:`~torchtune.modules.KVCache` norm (nn.Module): Callable that applies normalization to the output of the decoder, before final MLP. output (Union[nn.Linear, Callable]): Callable that applies a linear transformation to the output of the decoder. num_layers (Optional[int]): Number of Transformer Decoder layers, only define when layers is not a list. output_hidden_states (Optional[List[int]]): List of layers (indices) to include in the output Raises: AssertionError: If ``num_layers`` is set and layer is a list, **or** ``num_layers`` is not set and layer is an ``nn.Module``. Note: Arg values are checked for correctness (eg: ``attn_dropout`` belongs to [0,1]) in the module where they are used. This helps reduces the number of raise statements in code and improves readability. N) num_layersoutput_hidden_statestok_embeddingslayersr" num_headshead_dimnormoutputrwrxrc tt|tjrntt|t rtj|}nJt|tjstd|tdt||}||_ ||_ ||_ ||_ | pg|_ ||_||_||_d|_d|_d|_d|_dS)Nz.num_layers is defined, layers must be a modulez0num_layers is not defined, layers must be a listr)rr isinstancerrrlistrHrUrtryrzr}r~rxr"r{r| causal_masknum_output_chunksencoder_max_cache_seq_lendecoder_max_cache_seq_len) rryrzr"r{r|r}r~rwrxrs rrzTransformerDecoder.__init__es  fbm , , 5   % % 5]6**FFfbi00 W$%UVVV!$%WXXX 44F,   $8$>B!&"  !"*.&)-&&&rrc||_dS)zUsed to save memory in combination with :class:`~torchtune.modules.loss.CEWithChunkedOutputLoss`. This should be called before the first forward pass, in the recipe.N)r)rrs rset_num_output_chunksz(TransformerDecoder.set_num_output_chunkss"3rrr rrrr cltd|D}td|D}|r|||_n |j|_|r|||_n |j|_|jD]%}||||j|j&dS)a Sets up key-value attention caches for inference. For each layer in ``self.layers``: - :class:`~torchtune.modules.TransformerSelfAttentionLayer` will use ``decoder_max_seq_len``. - :class:`~torchtune.modules.TransformerCrossAttentionLayer` will use ``encoder_max_seq_len``. - :class:`~torchtune.modules.model_fusion.FusionLayer` will use ``decoder_max_seq_len`` and ``encoder_max_seq_len``. Args: batch_size (int): batch size for the caches. dtype (torch.dtype): dtype for the caches. encoder_max_seq_len (Optional[int]): maximum encoder cache sequence length. decoder_max_seq_len (Optional[int]): maximum decoder cache sequence length. c3@K|]}t|tVdSr)rrProms r z2TransformerDecoder.setup_caches..s>! ! >?Jq8 9 9! ! ! ! ! ! rc3@K|]}t|tVdSr)rr rs rrz2TransformerDecoder.setup_caches..s>! ! =>Jq7 8 8! ! ! ! ! ! rNr)anymodulesrr"rrzr&)rrrrr has_encoder_layershas_decoder_layerslayers rr&zTransformerDecoder.setup_cachess (!! ! CG<<>>! ! !   !! ! BF,,..! ! !     B".1D..151A.  B".1D..151A.[  E   $($B$($B       rc@|jdS)z Check if the key value caches are setup. This means ``setup_caches`` has been called, and the relevant attention modules in the model have created their ``KVCache``. r)rzr,r+s rr,z#TransformerDecoder.caches_are_setups {1~..000rc@|jdS)a Checks if the key value caches are enabled. Once KV-caches have been setup, the relevant attention modules will be "enabled" and all forward passes will update the caches. This behaviour can be disabled without altering the state of the KV-caches by "disabling" the KV-caches using :func:`torchtune.modules.common_utils.disable_kv_cache`, upon which ``caches_are_enabled`` would return False. r)rzr2r+s rr2z%TransformerDecoder.caches_are_enableds{1~00222rc|std|jD]}|dS)aX Resets KV-cache buffers on relevant attention modules to zero, and reset cache positions to zero, without deleting or reallocating cache tensors. Raises: RuntimeError: if KV-caches are not setup. Use :func:`~torchtune.modules.TransformerDecoder.setup_caches` to setup caches first. z>Key value caches are not setup. Call model.setup_caches first.N)r2 RuntimeErrorrzr7)rrs r reset_cacheszTransformerDecoder.reset_cachess_&&(( P [  E         rlast_hidden_statecTfd|jdDS)a Apply output projection in chunks. This should be applied in conjunction with :class:`~torchtune.modules.loss.CEWithChunkedOutputLoss` as upcasting to fp32 is done there. To use this method, you should first call :func:`~torchtune.modules.TransformerDecoder.set_num_output_chunks`. Args: last_hidden_state (torch.Tensor): last hidden state of the decoder, having shape [b, seq_len, embed_dim]. Returns: List[torch.Tensor]: List of num_chunks output tensors, each with shape [b, seq_len/num_chunks, out_dim], where out_dim is usually the vocab size. c:g|]}|Srl)r~)rochunkrs rrqz5TransformerDecoder.chunked_output..s5    KK     r)r\)rr)rrs` rchunked_outputz!TransformerDecoder.chunked_outputsD"    *001GQ0OO    rseq_lenr:rarbr;c||jkrtd|d|jd|r5|td||td|tddSdS)a Validates inputs for ``forward``. Args: seq_len (int): Input tensor sequence length. mask (Optional[torch.Tensor]): Attention mask used for inference and for sequence packing. encoder_input (Optional[torch.Tensor]): Encoder input for cross-attention. encoder_mask (Optional[torch.Tensor]): Encoder attention mask for cross-embedding attention. input_pos (Optional[torch.Tensor]): Input tensor position IDs. Raises: ValueError: If seq_len of x is bigger than max_seq_len, **or** if the model has caches which have been setup with self-attention layers and ``mask`` is not provided, **or** if the model has caches which have been setup with encoder layers and ``encoder_mask`` is not provided, **or** if the model has caches which have been setup ``input_pos`` is not provided. z seq_len (z6) of input tensor should be smaller than max_seq_len ()NzKV-caches for self-attention layers are setup for inference mode, causal masks must be provided! Use the `mask` arg to provide a causal mask.zKV-caches for cross-attention/fusion layers are setup for inference mode and you seem to be using encoder_input, causal masks must be provided! Use the `encoder_mask` arg to provide a causal mask.zIKV-caches are setup for inference mode, input positions must be provided!)r" ValueErrorr2)rrr:rarbr;s r_validate_inputsz#TransformerDecoder._validate_inputss2 T% % %9G99%)%5999   " " $ $ | D (\-A z   _  ! rr:rarbr;tokenscP|jd}||||||||}g}t|jD]3\} } | |jvr||| |||||}4||} |s| ng|| } | S)a Args: tokens (torch.Tensor): input tensor with shape ``[b x s]`` mask (Optional[_MaskType]): Used to mask the scores after the query-key multiplication and before the softmax. This parameter is required during inference if caches have been setup. Either: A boolean tensor with shape ``[b x s x s]``, ``[b x s x self.encoder_max_cache_seq_len]``, or ``[b x s x self.encoder_max_cache_seq_len]`` if using KV-cacheing with encoder/decoder layers. A value of True in row ``i`` and column ``j`` means token ``i`` attends to token ``j``. A value of False means token ``i`` does not attend to token ``j``. If no mask is specified, a causal mask is used by default. A :class:`~torch.nn.attention.flex_attention.BlockMask` for document masking in a packed sequence created via `create_block_mask `_. We use :func:`~torch.nn.attention.flex_attention.flex_attention` when computing attention with block masks. Default is None. encoder_input (Optional[torch.Tensor]): Optional input embeds from the encoder. Shape ``[b x s_e x d_e]`` encoder_mask (Optional[torch.Tensor]): Boolean tensor defining a relational matrix between tokens and encoder embeddings. A True value at position ``i,j`` means token ``i`` can attend to embedding ``j`` in the decoder. Mask has shape ``[b x s x s_e]``. Default is None, but this is required during inference if the model has been setup with any layers which use encoder embeddings and caches have been setup. input_pos (Optional[torch.Tensor]): Optional tensor which contains the position ids of each token. During training, this is used to indicate the positions of each token relative to its sample when packed, shape ``[b x s]``. During inference, this indicates the position of the current token. This parameter is required during inference if caches have been setup. Default is None. Returns: Union[torch.Tensor, List[torch.Tensor]]: output tensor with shape ``[b x s x v]`` or a list of layer output tensors defined by ``output_hidden_states`` with the final output tensor appended to the list. Note: At the very first step of inference, when the model is provided with a prompt, ``input_pos`` should contain the positions of all of the tokens in the prompt. For a single-batch prompt, or a batch of prompts with identical lengths, this will be ``torch.arange(prompt_length)``. For a batch of varying-length prompts, shorter prompts are left-padded and position ids are correspondingly right-shifted, thus positional ids should be of shape ``[b, padded_prompt_length]``. This is because we will need to retrieve the positional embeddings for each input id. In the subsequent steps, if the model has been setup with KV-caches, ``input_pos`` will contain the position(s) of the current token(s) ``torch.tensor([padded_prompt_length])``. Otherwise, ``input_pos`` will contain all the position ids up to the current token. Shape notation: - b: batch size - s: token sequence length - s_e: encoder sequence length - v: vocab size - d: token embed dim - d_e: encoder embed dim - m_s: max seq len rr)shaperry enumeraterzrxappendunembed) rrr:rarbr;rr?hiddenrprr~s rrCzTransformerDecoder.forward(sB,q/  '%        ' '!$+..  HAuD--- a   +)# AAa &<+rsW  88888888888888 000000777777uuuuuBIuuupiiiiiRYiiiX D Dc Dbm D D D DSSSSSSSSSSr