Pi ddlZddlmZddlZddlmZmZddlmZGddejZ GddejZ dS) N)Optional)nnTensor) _MaskTypeceZdZdZdddddddejd ejd ed ed ed edeffdZdddede e defdZ xZ S)CLIPTextEncodera Text encoder for CLIP. CLIP is a model that encodes text and images into a shared vector space. Blog post: https://openai.com/index/clip/ Paper: https://arxiv.org/abs/2103.00020 Args: layer (nn.Module): A single encoder layer. final_norm (nn.Module): Callable that applies normalization to the output of the encoder vocab_size (int): size of the vocabulary, default 49408 max_seq_len (int): context size, default 77 embed_dim (int): embedding/model dimension size, default 768 num_layers (int): number of transformer layers, default 12 eot_token (int): the id of the end-of-text token (for selecting the final output) iMi i) vocab_size max_seq_len embed_dim num_layers eot_tokenlayer final_normr r r rrcdttjfdt |D|_||_||_||_tj |||_ tj tj |||_dS)Nc8g|]}tjS)copydeepcopy).0irs w/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/torchtune/models/clip/_text_encoder.py z,CLIPTextEncoder.__init__...s#$U$U$UaT]5%9%9$U$U$U)super__init__r ModuleListrangelayersrr r Embeddingtoken_embedding Parametertorchemptyposition_embedding) selfrrr r r rr __class__s ` rrzCLIPTextEncoder.__init__"s m$U$U$U$U5CTCT$U$U$UVV $&"!|J BB"$,u{; /R/R"S"SrNmasktokensr*returnc|j\}}||jkrtd|d|jd|||jz}|jD]}|||}||}||jk d}| | dddd d}|S)a! Args: tokens (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. Default is None. Returns: Tensor: output tensor with shape [b x d] Raises: ValueError: if seq_len of tokens is bigger than max_seq_len Shape notation: - b: batch size - s: token sequence length - d: token embed dim z seq_len (z6) of input tensor should be smaller than max_seq_len ()r))dim) shaper ValueErrorr"r&r rrintargmaxtake_along_dimviewsqueeze)r'r+r*bszseq_lenxreos_token_positionss rforwardzCLIPTextEncoder.forward6s$2| W T% % %9G99%)%5999    ( (4+B B[  EAA OOA   &7<<>>EE"EMM  055b!Q??Q  G G O OTU O V Vr) __name__ __module__ __qualname____doc__rModuler4rrrrr= __classcell__)r(s@rrrs, TTTyTI T  T  TTTTTTTTT0%) 0000y! 0  00000000rrc6eZdZdZdejdejfdZdS) QuickGELUz% Fast approximation of GELU. r;r,c6|tjd|zzS)NgZd;?)r$sigmoid)r'r;s rr=zQuickGELU.forwardns5=++++rN)r>r?r@rAr$rr=rrrrErEisD,,%,,,,,,,rrE) rtypingrr$rr!torchtune.modules.attention_utilsrrBrrErrrrJs  777777VVVVVbiVVVr,,,,, ,,,,,r