U &ºcÉ)ã@s@dZddlZddlmZddlmmZGdd„dejƒZdS)zt Utilities for PyTorch Transformer XL model. Directly adapted from https://github.com/kimiyoung/transformer-xl. éNcs8eZdZd ‡fdd„ Zdd„Zd dd „Zd d „Z‡ZS)ÚProjectedAdaptiveLogSoftmaxéFc sŒtƒ ¡||_||_||_||g|_dg|j|_||_|jd|_t |jƒd|_ |j|j |_ |j dkrœt   t |j |j¡¡|_t   t |j ¡¡|_t  ¡|_t  ¡|_|dkrtt |jƒƒD]4}||krð|j t   t ||¡¡¡qÈ|j d¡qÈ|j t  ||¡¡nntt |jƒƒD]^}|j||j|d}} |||} |j t   t || ¡¡¡|j t  | | |¡¡q"||_dS)Nrr)ÚsuperÚ__init__Ún_tokenÚd_embedÚd_projÚcutoffsÚ cutoff_endsÚdiv_valZshortlist_sizeÚlenÚ n_clustersZ head_sizeÚnnÚ ParameterÚtorchÚzerosÚcluster_weightÚ cluster_biasZ ModuleListÚ out_layersZ ParameterListÚ out_projsÚrangeÚappendZ FloatTensorZLinearÚ keep_order) Úselfrrrr r rÚiÚl_idxÚr_idxZd_emb_i©Ú __class__©úN/tmp/pip-unpacked-wheel-ymerj3tt/transformers/modeling_transfo_xl_utilities.pyrs6        z$ProjectedAdaptiveLogSoftmax.__init__cCsB|dkrtj|||d}n$t || ¡ ¡¡}tj|||d}|S)N)Úbias)ÚFZlinearÚtÚ contiguous)rÚhiddenÚweightr!ZprojÚlogitZproj_hidrrr Ú_compute_logitHs z*ProjectedAdaptiveLogSoftmax._compute_logitNcCsV|dk rp|ddd…dd…f ¡}|ddd…f ¡}| d| d¡¡}| d¡}| d¡| d¡kr‚tdƒ‚n| d| d¡¡}|jdkrò| ||jdj|jdj|j d¡}|dk ràt j |dd  d|  d¡¡ d¡ }nt j |dd}n`gg}}tt|jƒƒD]´}|jdkr^|j||j|d} } |jdj| | …} |jdj| | …} n|j|j} |j|j} |dkr¨tj| |jgdd} tj| |jgdd} | | ¡| | ¡q |d|d|j d} }}| || ||¡}t j |dd}|dkr| | d¡|jf¡}ntj||j|jd}d}dg|j}tt|ƒdƒD]ü}||||d} } |dk rÐ|| k|| k@}| ¡ ¡}| ¡dkr¦qR|  d|¡| }|  d|¡}|  d|¡}n|}|dkr:|dk r|  d|dd…df¡ d¡}n0|dd…d|jd…f|dd…d|jd…f<n²|||||j |} } }| || | |¡}t j |dd}|jd|d}|dk rÂ|dd…|f|  d|dd…df¡ d¡}n*|dd…|df|}||dd…| | …f<|dk rRt!|d ƒr |j"s|r"| #d|| ¡n|||| d¡… $| ¡|| d¡7}qR|S) a€ Params: hidden :: [len*bsz x d_proj] labels :: [len*bsz] Return: if labels is None: out :: [len*bsz x n_tokens] log probabilities of tokens over the vocabulary else: out :: [(len-1)*bsz] Negative log likelihood We could replace this implementation by the native PyTorch one if their's had an option to set bias on all clusters in the native one. here: https://github.com/pytorch/pytorch/blob/dbe6a7a9ff1a364a8706bf5df58a1ca96d2fd9da/torch/nn/modules/adaptive.py#L138 N.éÿÿÿÿrrzBInput and labels should have the same size in the batch dimension.©Zdim)ÚdtypeÚdevicer)%r$ÚviewÚsizeÚ RuntimeErrorr r(rr&r!rr"Ú log_softmaxZgatherZ unsqueezeZsqueezerr r r r rÚcatrrrÚ new_emptyrZ zeros_liker+r,ZnonzeroZnumelZ index_selectÚhasattrrZ index_copy_Zcopy_)rr%Úlabelsrr'ÚoutÚweightsÚbiasesrrrÚweight_iÚbias_iÚ head_weightÚ head_biasÚ head_projÚ head_logitÚ head_logprobÚoffsetÚ cutoff_valuesZmask_iZ indices_iZtarget_iZhead_logprob_iZhidden_iÚ logprob_iÚproj_iÚ tail_logit_iÚtail_logprob_iZcluster_prob_idxrrr ÚforwardVsˆ   &&              2 ÿþ z#ProjectedAdaptiveLogSoftmax.forwardcCs:|jdkr>| ||jdj|jdj|jd¡}tj|ddSgg}}tt |j ƒƒD]®}|j dkr¨|j ||j |d}}|jdj||…}|jdj||…} n|j|j}|j|j} |dkrðt j||jgdd}t j| |jgdd} | |¡| | ¡qV|d|d|jd} } } | || | | ¡} | |  d¡|jf¡}tj| dd}dg|j }tt |ƒdƒD]¼}||||d}}|dkrÊ|dd…d|j d…f|dd…d|j d…f<nd|||||j|}} }| ||| |¡}tj|dd}|dd…| f|}||dd…||f<qt|SdS)a0 Computes log probabilities for all :math:`n\_classes` From: https://github.com/pytorch/pytorch/blob/master/torch/nn/modules/adaptive.py Args: hidden (Tensor): a minibatch of examples Returns: log-probabilities of for each class :math:`c` in range :math:`0 <= c <= n\_classes`, where :math:`n\_classes` is a parameter passed to ``AdaptiveLogSoftmaxWithLoss`` constructor. Shape: - Input: :math:`(N, in\_features)` - Output: :math:`(N, n\_classes)` rr)r*rN)r r(rr&r!rr"r0rr r r r rr1rrrr2r.r)rr%r'r6r7rrrr8r9r:r;r<r=r5r>r@Z start_idxZstop_idxrBrCrDrArrr Úlog_probÁs> &        2z$ProjectedAdaptiveLogSoftmax.log_prob)rF)NF)Ú__name__Ú __module__Ú __qualname__rr(rErFÚ __classcell__rrrr rs) kr) Ú__doc__rZtorch.nnrZtorch.nn.functionalZ functionalr"ÚModulerrrrr Ús