o - i@sZddlmZddlZddlmZddlmZdgZGdddejZ GdddejZ dS)) annotationsNTransformerBlockDecoderOnlyTransformercs,eZdZdZd fdd Zdd d ZZS)AbsolutePositionalEmbeddingzAbsolute positional embedding. Args: max_seq_len: Maximum sequence length. embedding_dim: Dimensionality of the embedding. max_seq_lenint embedding_dimreturnNonecs(t||_||_t|||_dSN)super__init__rr nn Embedding embedding)selfrr  __class__a/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/transformer.pyrs z$AbsolutePositionalEmbedding.__init__x torch.TensorcCs2|\}}tj||jd|d}||}|S)N)device)sizetorcharangerrepeatr)rr batch_sizeseq_len positionsrrrrforward$s  z#AbsolutePositionalEmbedding.forward)rrr rr r )rrr r)__name__ __module__ __qualname____doc__rr" __classcell__rrrrrsrcsFeZdZdZ     d!d"fdd Zd#d$ddZd%d&dd ZZS)'raDecoder-only (Autoregressive) Transformer model. Args: num_tokens: Number of tokens in the vocabulary. max_seq_len: Maximum sequence length. attn_layers_dim: Dimensionality of the attention layers. attn_layers_depth: Number of attention layers. attn_layers_heads: Number of attention heads. with_cross_attention: Whether to use cross attention for conditioning. embedding_dropout_rate: Dropout rate for the embedding. include_fc: whether to include the final linear layer. Default to True. use_combined_linear: whether to use a single linear layer for qkv projection, default to True. use_flash_attention: if True, use Pytorch's inbuilt flash attention for a memory efficient attention mechanism (see https://pytorch.org/docs/2.2/generated/torch.nn.functional.scaled_dot_product_attention.html). FT num_tokensrrattn_layers_dimattn_layers_depthattn_layers_headswith_cross_attentionboolembedding_dropout_ratefloat include_fcuse_combined_linearuse_flash_attentionr r c st||_|_|_||_|_|_t ||_ t d|_ t ||_tfddt|D|_t||_dS)N)rr cs.g|]}tddddd qS)r(FT) hidden_sizemlp_dim num_heads dropout_rateqkv_biascausalsequence_lengthr-r1r2r3r).0_r*r,r1rr2r3r-rr Vs z3DecoderOnlyTransformer.__init__..)r rr)rr*r+r,r-rrtoken_embeddingsrposition_embeddingsDropoutembedding_dropout ModuleListrangeblocksLinear to_logits) rr)rr*r+r,r-r/r1r2r3rr>rr<s  zDecoderOnlyTransformer.__init__Nrrcontexttorch.Tensor | NonecCsH||}||}|||}|jD]}|||d}q||}|S)N)rI)r@rArCrFrH)rrrIZtok_embpos_embblocklogitsrrrr"js    zDecoderOnlyTransformer.forwardold_state_dictdictcs|tfdd|Drtd||dS|rBD]}||vr,td|dqtd|D]}|vrAtd|dq3D] }||vrQ|||<qDt|D] }d |vri|||d d <d |vrx|||d d <qX|rtd ||dS) z Load a state dict from a DecoderOnlyTransformer trained with [MONAI Generative](https://github.com/Project-MONAI/GenerativeModels). Args: old_state_dict: state dict from the old DecoderOnlyTransformer model. c3s|]}|vVqdSr r)r<knew_state_dictrr sz=DecoderOnlyTransformer.load_old_state_dict..z#All keys match, loading state dict.Nzkey z not found in old state dictz.----------------------------------------------z not found in new state dictnorm2norm_cross_attnnorm3z!remaining keys in old_state_dict:) state_dictallprintload_state_dictpoplistkeysreplace)rrNverboserPrrQrload_old_state_dictts8  z*DecoderOnlyTransformer.load_old_state_dict)Fr(TFF)r)rrrr*rr+rr,rr-r.r/r0r1r.r2r.r3r.r r r )rrrIrJr r)F)rNrOr r )r#r$r%r&rr"r`r'rrrrr+s . ) __future__rrtorch.nnrmonai.networks.blocksr__all__Modulerrrrrrs