U ,-eq@sldZddlZddlZddlmZmZmZddlZddlZddlm Z ddl m Z m Z m Z mZddl mZddlmZmZmZdd lmZmZmZmZmZdd lmZdd lmZd d lm Z e!e"Z#dZ$dZ%dddddddgZ&ej'ej(e)ej*dddZ+ej,e)ej*dddZ-ej,e)ej.ej,dddZ/ej,ej,e0e1ej,dd d!Z2ej,ej,d"d#d$Z3ej,ej,ej,d%d&d'Z4Gd(d)d)ej5j6Z7Gd*d+d+e j8Z9Gd,d-d-e j8Z:Gd.d/d/e j8Z;Gd0d1d1e j8Zd5Z?ed6e>Gd7d8d8e=Z@ed9e>Gd:d;d;e=ZAedGd=d>d>e=ZBed?e>Gd@dAdAe=ZCedBe>GdCdDdDe=ZDdS)EzPyTorch BLOOM model.N)OptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLoss LayerNormMSELoss) functional)add_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forward))BaseModelOutputWithPastAndCrossAttentions!CausalLMOutputWithCrossAttentionsQuestionAnsweringModelOutput SequenceClassifierOutputWithPastTokenClassifierOutput)PreTrainedModel)logging) BloomConfigzbigscience/bloom-560mrz#bigscience/bigscience-small-testingzbigscience/bloom-1b1zbigscience/bloom-1b7zbigscience/bloom-3bzbigscience/bloom-7b1zbigscience/bloom)input_ids_shapedevicepast_key_values_lengthreturncCs|\}}tj|||ftj|d}tj||d}|dddf|dddfk|dd|df<|dkr|d|ddd|f<|ddddddf|d|||}|S)z3 Make causal mask used for self-attention. )dtyperrNrFr)torchemptyboolarangeexpand)rrr batch_sizeZ target_lengthmaskZseq_ids expanded_maskr&i/var/www/html/Darija-Ai-Train/env/lib/python3.8/site-packages/transformers/models/bloom/modeling_bloom.py_make_causal_mask8s0(r()r$ tgt_lengthrcCsL|j\}}|dk r|n|}|ddddddftj}||d||S)zn Expands attention_mask from `[batch_size, src_length]` to `[batch_size, 1, tgt_length, src_length]`. Nr)shapetorr r")r$r)r# src_lengthr%r&r&r' _expand_maskKs "r-attention_mask num_headsrrcCs6|j\}}dtt|}tjddt|d  |jtjd}tjdd||jtj d}t ||}||krtjddtd|d  |jtjd} t |||} tjddd| d|jtj d} tj |t | | gdd}|j ddd|dddddf} |d | } | ||d||S) a Link to paper: https://arxiv.org/abs/2108.12409 Alibi tensor is not causal as the original paper mentions, it relies on a translation invariance of softmax for quick implementation: with l being a tensor, and a fixed value `softmax(l+a) = softmax(l)`. Based on https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742 TODO @thomasw21 this doesn't work as nicely due to the masking strategy, and so masking varies slightly. Args: Returns tensor shaped (batch_size * num_heads, 1, max_seq_len) attention_mask (`torch.Tensor`): Token-wise attention mask, this should be of shape (batch_size, max_seq_len). num_heads (`int`, *required*): number of heads dtype (`torch.dtype`, *optional*, default=`torch.bfloat16`): dtype of the output tensor r )rrrrdimN).N)r*mathfloorlog2rZtensorrfloat32r!Zint32powmincatZcumsumreshaper+)r/r0rr# seq_lengthZclosest_power_of_2baseZpowersZslopesZ extra_baseZnum_remaining_headsZ extra_powersZ arange_tensoralibir&r&r'build_alibi_tensorVs*   & r@)xresidualprobtrainingrcCstj|||d}||}|S)a: Dropout add function Args: x (`torch.tensor`, *required*): input tensor residual (`torch.tensor`, *required*): residual tensor prob (`float`, *required*): dropout probability training (`bool`, *required*): training mode )prD)Fdropout)rArBrCrDoutr&r&r' dropout_addsrIrArcCs*|ddtd|dd||S)z Custom bias GELU function. Adapted from Megatron-DeepSpeed code. Here we use a simple implementation (inference) to make the model jitable. Args: x (`torch.tensor`, *required*): input hidden states ??e3E?rHm?rtanh)rAr&r&r'bloom_gelu_forwards rQ)grArcCs^|d}td|dd||}d|d||dd||dd|}||S)a8 gradient of tanh approximation of gelu gradient of actual gelu is: 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x) Args: g (`torch.tensor`, *required*): gradient output tensor x (`torch.tensor`, *required*): input tensor rrMrrNrKg6vf?rO)rRrAZtanh_outffr&r&r'bloom_gelu_backs 0rTc@s<eZdZeejejdddZeejejdddZdS) GeLUFunction)inputrcCs||t|SN)Zsave_for_backwardrQ)ctxrVr&r&r'forwards zGeLUFunction.forward) grad_outputrcCs|j}t||}|SrW)Z saved_tensorsrT)rXrZrVtmpr&r&r'backwards zGeLUFunction.backwardN)__name__ __module__ __qualname__ staticmethodrTensorrYr\r&r&r&r'rUsrUcs4eZdZdZfddZejejdddZZS) BloomGelua BloomBiasGelu wrapper function that make use of the simple function on inference mode to make the model torchscriptable and use the autograd function in training mode to get the accurate results of the gradients Partly copied from Megatron-DeepSpeed code and adapted for our needs See here why autograd functions are not torchscriptable: https://github.com/pytorch/pytorch/issues/22329 cstdSrWsuper__init__self __class__r&r'reszBloomGelu.__init__rJcCs|jrt|St|SdSrW)rDrUapplyrQ)rgrAr&r&r'rYs zBloomGelu.forward) r]r^r___doc__rerrarY __classcell__r&r&rhr'rbs rbc seZdZedfdd ZejeejejejfdddZejejddd Z dejejejeje eejejfe eje e d d dZ Z S)BloomAttentionconfigcst|j|_|j|_|j|_|j|_|j|j|_|j|_|j |_ |j|j|jkrtt d|jd|jddt |j|_ d|_tj|jd|jdd|_t|j|j|_t|j|_dS)NzA`hidden_size` must be divisible by num_heads (got `hidden_size`: z and `num_heads`: z).rLr Tbias)rdrepretraining_tpslow_but_exact hidden_sizen_headr0head_dimZ split_sizehidden_dropout ValueErrorr5sqrtinv_norm_factorbetarLinearquery_key_valuedenseDropoutattention_dropoutrgrorhr&r'res" zBloomAttention.__init__) fused_qkvrcCsV|j\}}}||||jd|j}|ddddf|ddddf|ddddffS)a Split the last dimension into (num_heads, head_dim) without making any copies, results share same memory storage as `fused_qkv` Args: fused_qkv (`torch.tensor`, *required*): [batch_size, seq_length, num_heads * 3 * head_dim] Returns: query: [batch_size, seq_length, num_heads, head_dim] key: [batch_size, seq_length, num_heads, head_dim] value: [batch_size, seq_length, num_heads, head_dim] r .rNrr1)r*viewr0rv)rgrr#r=Zthree_times_hidden_sizer&r&r' _split_headss zBloomAttention._split_headsrJcCsP|j\}}}||j}|||j||j}|dddd}||||j|jS)a Merge heads together over the last dimension Args: x (`torch.tensor`, *required*): [batch_size * num_heads, seq_length, head_dim] Returns: torch.tensor: [batch_size, seq_length, num_heads * head_dim] rr1rr )r*r0rrvpermuter<)rgrAZbatch_size_and_num_headsr=_r#r&r&r' _merge_headss  zBloomAttention._merge_headsNF) hidden_statesrBr?r/ layer_past head_mask use_cacheoutput_attentionsc  Csp||} || \} } } | j\} }}}| dd| |j||j} | dddd| |j|j|} | dd| |j||j} |dk r|\}}tj || fdd} tj || fdd} | j\}}}|dkr| | f}nd}|j | | |j |j d}| | |j||}|j}|tjkr"|tj}t||t|jj}tj|dtjd |}||}|dk rn||}| | |j||}t|| }||}|jdkr8|jr8|j|j}t|}t|jD]h}|t |ddddt!||t!|d|f|j"j#ddt!||t!|d|f}qn |"|}t$|||j%|j&}||f}|rl||f7}|S) Nrr1rr r2T)Zbatch1Zbatch2r{alphar4)r3r)'r}rr*Z transposer<r0rvrrr;Zbaddbmmr{rzrrZfloat16r+floatZ masked_fillZfinfor:rFZsoftmaxr8rZbmmrrrrsrt zeros_likerangelinearintr~weightrIrwrD)rgrrBr?r/rrrrrZ query_layerZ key_layerZ value_layerr#Zq_lengthrZpast_keyZ past_valueZ kv_lengthZpresentZ matmul_resultZattention_scoresZ input_dtypeZ attn_weightsZattention_probsZattention_probs_reshapedZ context_layerslicesZ output_tensorioutputsr&r&r'rYs\ "          ,*   zBloomAttention.forward)NNFF)r]r^r_rrerrarrrrr rYrlr&r&rhr'rms""rmcs:eZdZedfdd ZejejejdddZZS)BloomMLPrncsXt|j}|j|_|j|_t|d||_t|_ td|||_ |j |_ dS)N) rdrertrrrsrr| dense_h_to_4hrb gelu_impl dense_4h_to_hrwrgrortrhr&r'rets zBloomMLP.__init__)rrBrc Cs|||}|jdkr|jrt|}|jjjd|j}t |jD]f}|t |ddddt ||t |d|f|jjddt ||t |d|f}qHn ||}t |||j|j}|S)Nrr4)rrrrrsrrrrr*rrFrrrIrwrD)rgrrBZintermediate_outputrroutputr&r&r'rYs ,*  zBloomMLP.forward) r]r^r_rrerrarYrlr&r&rhr'rss rc sZeZdZedfdd Zd ejejejeeejejfeeje e dddZ Z S) BloomBlockrncs`t|j}t||jd|_|j|_t||_ t||jd|_ t ||_ |j |_ |j|_dS)NZeps)rdrertrlayer_norm_epsiloninput_layernormrur0rmself_attentionpost_attention_layernormrmlp(apply_residual_connection_post_layernormrwrrhr&r'res   zBloomBlock.__init__NF)rr?r/rrrrc Cs||}|jr|} n|} |j|| ||||||d} | d} | dd} || }|jr^|} n| } ||| } |r~| f| } n| f| dd} | S)N)rr/r?rrrrr)rrrrr)rgrr?r/rrrrZlayernorm_outputrBZ attn_outputsZattention_outputrrr&r&r'rYs2      zBloomBlock.forward)NNFF) r]r^r_rrerrarrr rYrlr&r&rhr'rsrcseZdZeZdZdZdgZdZfddZ e j ddd Z de j e d d d ZeeeejejfeeeejejfdddZeeeejejfeeejejfdddZZS)BloomPreTrainedModel transformerTrpast_key_valuescstj||dSrWrc)rginputskwargsrhr&r'reszBloomPreTrainedModel.__init__modulecCst|tjr:|jjjd|jjd|jdk r|jj ndt|tj rz|jjjd|jjd|j dk r|jj|j  n$t|t r|jj |jj ddS)zInitialize the weights.g)ZmeanZstdNrL) isinstancerr|rdataZnormal_roZinitializer_rangerqZzero_ EmbeddingZ padding_idxrZfill_)rgrr&r&r' _init_weightss      z"BloomPreTrainedModel._init_weightsF)rvaluecCst|tr||_dSrW)r BloomModelgradient_checkpointing)rgrrr&r&r'_set_gradient_checkpointings z0BloomPreTrainedModel._set_gradient_checkpointing)past_key_valuer#rcs8|ddj\}|tfdd|DS)z Standardizes the format of the cache so as to match most implementations, i.e. to tuple(tuple([batch_size, num_heads, ...])) rc3s6|].}|d|dfVqdSrrNr.0rr#rvr0r=r&r' szBBloomPreTrainedModel._convert_to_standard_cache..r*tuple)rr#batch_size_times_num_headsr&rr'_convert_to_standard_caches z/BloomPreTrainedModel._convert_to_standard_cache)rrcs8|ddj\}}||tfdd|DS)zy Converts the cache to the format expected by Bloom, i.e. to tuple(tuple([batch_size * num_heads, ...])) rc3s2|]*}|d|dfVqdSrrrrrvr=r&r'rsz?BloomPreTrainedModel._convert_to_bloom_cache..r)rr#r0r&rr'_convert_to_bloom_cache s z,BloomPreTrainedModel._convert_to_bloom_cache)F)r]r^r_r config_classZbase_model_prefixZsupports_gradient_checkpointingZ_no_split_modulesZ_skip_keys_device_placementrerModulerr rr`rrrarrrrlr&r&rhr'rs" ra, This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings etc.) This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`BloomConfig`]): Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. a Args: input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`): `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values[0][0].shape[2]` (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary. If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as `input_ids`. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) past_key_values (`Tuple[Tuple[torch.Tensor]]` of length `config.n_layers`): Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have their past given to this model should not be passed as `input_ids` as they have already been computed. Each element of `past_key_values` is a tuple (past_key, past_value): - past_key: [batch_size * num_heads, head_dim, kv_length] - past_value: [batch_size * num_heads, kv_length, head_dim] attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see `past_key_values`). use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. z_The bare Bloom Model transformer outputting raw hidden-states without any specific head on top.cseZdZedfdd ZejeejejdddZ ddZ eje eefeej d d d Z ejd d dZeeeeeeddeejee e ejejfdfeejeejeejeeeeeeeeee ejdfefd ddZZS)rrncstj|_j|_tj|j|_ t |jj d|_ t fddtjD|_t |jj d|_d|_|dS)Nrcsg|] }tqSr&)r)rrrnr&r' usz'BloomModel.__init__..F)rdrertZ embed_dimrur0rr vocab_sizeword_embeddingsrrword_embeddings_layernormZ ModuleListrZnum_hidden_layershln_fr post_initrrhrnr'rejs  zBloomModel.__init__r.cCs t|||SrW)r@)rgr/r0rr&r&r'r@szBloomModel.build_alibi_tensorcCs|jSrWrrfr&r&r'get_input_embeddingsszBloomModel.get_input_embeddings)r/ input_shaperrc CsLd}|j}|\}}|dkr(t|||d}t||d}|dkr@|n||B}|S)Nr)rr)r))rr(r-) rgr/rrZcombined_attention_maskrrr,Zexpanded_attn_maskr&r&r'_prepare_attn_masks zBloomModel._prepare_attn_masknew_embeddingscCs ||_dSrWrrgrr&r&r'set_input_embeddingsszBloomModel.set_input_embeddings checkpoint output_typerN.) input_idsrr/r inputs_embedsrroutput_hidden_states return_dictrc  s| dddk rtdtt| dkr6td| dk rBn|jj|dk rV|n|jj}dk rjn|jj | dk r~| n|jj } |dk r|dk rtdn2|dk r|j \} } n|dk r|j \} } } ntd|dkrt dgt|j }|||jj}|dkr||}||}r&dnd}r4dnd}|rBdnd}|jrj|jrjrjtd d| }d}|ddk r|ddj d }||}|dkrtj| |f|jd }n ||j}|j||j|jd }|j|| | f|d }tt|j |D]\}\}}|r||f}|jrZ|jrZfdd}tj j!!||||||||}n|||||||d}|d}dkr||df}r||rd ndf}q|"|}|r||f}| st dd||||fDSt#||||dS)N position_idsFz`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore passing `position_ids`.rGot unexpected arguments: zDYou cannot specify both input_ids and inputs_embeds at the same timez5You have to specify either input_ids or inputs_embedsr&zZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...r1r)r)rrcsfdd}|S)Ncs|dS)N)rrr&)r)rrrr&r'custom_forwardszIBloomModel.forward..create_custom_forward..custom_forwardr&)rrrrrr'create_custom_forwardsz1BloomModel.forward..create_custom_forward)rr/rrrr?Trcss|]}|dk r|VqdSrWr&)rvr&r&r'r%sz%BloomModel.forward..)Zlast_hidden_staterr attentions)$popwarningswarn FutureWarninglenrxrorrruse_return_dictr*rrZ get_head_maskZn_layerrrrrDloggerZ warning_oncerZonesrr+r@r0rr enumerateziputilsrrr)rgrrr/rrrrrrdeprecated_argumentsr#r=rrZpresentsZall_self_attentionsZall_hidden_statesZseq_length_with_pastrr?Z causal_maskrblockrrrr&rr'rYs              zBloomModel.forward) NNNNNNNNN)r]r^r_rrerrarrr@rr BoolTensorrrrBLOOM_INPUTS_DOCSTRINGr _CHECKPOINT_FOR_DOCr_CONFIG_FOR_DOCr LongTensorr rrYrlr&r&rhr'resH  rz The Bloom Model transformer with a language modeling head on top (linear layer with weights tied to the input embeddings). cs6eZdZdgZedfdd ZddZejddd Z dej e eje eje eje d d d Z eeeeeedde ej e eeejejfdfe eje eje eje eje ee ee ee eeeejefd ddZeeejejfdfej eeejejfdfdddZZS)BloomForCausalLMzlm_head.weightrncs8t|t||_tj|j|jdd|_| dSNFrp) rdrerrrr|rtrlm_headrrrhr&r're9s  zBloomForCausalLM.__init__cCs|jSrWrrfr&r&r'get_output_embeddingsAsz&BloomForCausalLM.get_output_embeddingsrcCs ||_dSrWrrr&r&r'set_output_embeddingsDsz&BloomForCausalLM.set_output_embeddingsN)rrr/rrcKs~|r@|dddfd}|ddjd|jdkr@||}|dk rZ|dkrZd|i}nd|i}|||d|d|S)Nr4rrrr)rrr/)Z unsqueezer*rupdateget)rgrrr/rrZ model_inputsr&r&r'prepare_inputs_for_generationGs   z.BloomForCausalLM.prepare_inputs_for_generationr. rrr/rrlabelsrrrrrc  Ks8| dddk rtdtt| dkr6td| | dk rB| n|jj} |j|||||||| | d } | d} | | }d}|dk r| |j }|ddd ddf }|dd df }|j \}}}t}||||||||}| s |f| d d}|dk r|f|S|St||| j| j| jd S) a labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]` rFrrrNrr/rrrrrr.r4rlosslogitsrrr)rrrrrrxrorrrr+r contiguousr*rrrrrr)rgrrr/rrrrrrrrtransformer_outputsrZ lm_logitsrZ shift_logitsZ shift_labelsr#r=rloss_fctrr&r&r'rYfsT      zBloomForCausalLM.forward)pastbeam_idxrcsD|j|td}fdd|Dtfdd|D}||S)aL This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct beam_idx at every generation step. Output shares the same memory storage as `past`. )r#cs&i|]}|D]}|j|jq qSr&)rr+)rrZ past_state)rr&r' s z3BloomForCausalLM._reorder_cache..c3sB|]:}|dd|dj|dd|djfVqdSr)Z index_selectrr)device_to_beam_idxr&r'rsz2BloomForCausalLM._reorder_cache..)rrrr)rgrrZstandardized_pastZreordered_pastr&)rrr'_reorder_caches   zBloomForCausalLM._reorder_cache)NNN) NNNNNNNNNN)r]r^r_Z_tied_weights_keysrrerrrarrrdictrrrr rrrrr rrYr rlr&r&rhr'r/s^ Jra The Bloom Model transformer with a sequence classification head on top (linear layer). [`BloomForSequenceClassification`] uses the last token in order to do the classification, as other causal models (e.g. GPT-1) do. Since it does classification on the last token, it requires to know the position of the last token. If a `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in each row of the batch). cseZdZedfdd Zeeeee e dd e e j e eee je jfdfe e je e je e je e je ee ee ee eeee je fd dd ZZS) BloomForSequenceClassificationrncs@t||j|_t||_tj|j|jdd|_| dSr) rdre num_labelsrrrr|rtscorerrrhr&r'res   z'BloomForSequenceClassification.__init__rN.rc  KsR| dddk rtdtt| dkr6td| | dk rB| n|jj} |j|||||||| | d } | d} | | }|dk r|j d}n |j d}|jj dkr|dkrtd |jj dkrd }nD|dk rt ||jj d d|j}nd }t|jjd |t j||jd |f}d}|dk r |jjdkr|jdkrRd |j_n:|jdkr|jt jksz|jt jkrd|j_nd|j_|jjd krt}|jdkr|||}n |||}n>|jjdkrt}|||}n|jjdkr t}|||}| s:|f| dd}|dk r6|f|S|St||| j | j!| j"dS) labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). rFrrrNrrz=Cannot handle batch sizes > 1 if no padding token is defined.r4z will not detect padding tokens in `inputs_embeds`. Results may be unexpected if using padding tokens in conjunction with `inputs_embeds.`rZ regressionZsingle_label_classificationZmulti_label_classificationr)#rrrrrrxrorrr r*Z pad_token_idrnesumr+rrwarningrir]r!Z problem_typer rlongrr squeezerrrrrr)rgrrr/rrrrrrrrrrrr#Zsequence_lengthsZ pooled_logitsrrrr&r&r'rYs      $    (     z&BloomForSequenceClassification.forward) NNNNNNNNNN)r]r^r_rrerrr rrrrrrrrar rrYrlr&r&rhr'r s< r z Bloom Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. cseZdZedfdd Zeeeee e dd e e j e eee je jfdfe e je e je e je e je ee ee ee eeee je fd dd ZZS) BloomForTokenClassificationrncst||j|_t||_t|dr:|jdk r:|j}n t|drV|jdk rV|j}nd}t ||_ t |j |j|_ |dS)Nclassifier_dropoutrwg?)rdrer rrhasattrrrwrrrGr|rt classifierr)rgrorrhr&r're]s   z$BloomForTokenClassification.__init__rN.rc  Ks | dddk rtdtt| dkr6td| | dk rB| n|jj} |j|||||||| | d } | d} | | } | | }d}|dk r| |j }|j \}}t}|||||j|||}| s|f| dd}|dk r|f|S|St||| j| jd S) rrFrrrNrr1)rrrr)rrrrrrxrorrrGrr+rr*rrr rrr)rgrrr/rrrrrrrrrrrrr#r=rrr&r&r'rYnsP       z#BloomForTokenClassification.forward) NNNNNNNNNN)r]r^r_rrerrr rrrrrrrrar rrYrlr&r&rhr'rUs<rz The BLOOM Model transformer with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`). cseZdZfddZeeddeej eej eej eej eej eej eej ee ee ee e e efd ddZZS) BloomForQuestionAnsweringcs2t|t||_t|jd|_|dS)Nr1) rdrerrrr|rt qa_outputsrrrhr&r'res  z"BloomForQuestionAnswering.__init__zbatch_size, sequence_lengthN) rr/rrrstart_positions end_positionsrrrrc  CsN| dk r | n|jj} |j||||||| | d} | d} || } | jddd\}}|d}|d}d}|dk r|dk rt|dkr|d}t|dkr|d}|d}| d|}| d|}t |d}|||}|||}||d}| s8||f| dd}|dk r4|f|S|St |||| j | j d S) a start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for position (index) of the start of the labelled span for computing the token classification loss. Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence are not taken into account for computing the loss. end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for position (index) of the end of the labelled span for computing the token classification loss. Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence are not taken into account for computing the loss. N)r/rrrrrrrrr4r2)Z ignore_indexr1)r start_logits end_logitsrr)rorrrsplitrrrsizeclamprrrr)rgrr/rrrrrrrrrZsequence_outputrrrZ total_lossZ ignored_indexrZ start_lossZend_lossrr&r&r'rYsN           z!BloomForQuestionAnswering.forward) NNNNNNNNNN)r]r^r_rerrformatrrrZ FloatTensorr rrrrYrlr&r&rhr'rs2   r)Erkr5rtypingrrrrZtorch.utils.checkpointrZtorch.nnrrrr r rFZ file_utilsr r rZmodeling_outputsrrrrrZmodeling_utilsrrrZconfiguration_bloomrZ get_loggerr]rrrZ#BLOOM_PRETRAINED_MODEL_ARCHIVE_LISTSizerrrr(rar-rr@rr rIrQrTZautogradFunctionrUrrbrmrrrZBLOOM_START_DOCSTRINGrrrr rrr&r&r&r's         ,  !GF7G w^