U ,-e @sdZddlZddlZddlZddlmZddlmZmZm Z ddl Z ddl Z ddl m Z ddl mZmZmZddlmZdd lmZmZmZmZmZmZdd lmZdd lmZmZmZdd l m!Z!m"Z"m#Z#m$Z$m%Z%d dl&m'Z'e$(e)Z*dZ+dZ,ddgZ-ddddddddddddddd d!gZ.eGd"d#d#eZ/d$d%Z0Gd&d'd'e j1Z2Gd(d)d)e j1Z3Gd*d+d+e j1Z4Gd,d-d-e j1Z5Gd.d/d/e j1Z6Gd0d1d1e j1Z7Gd2d3d3e j1Z8Gd4d5d5e j1Z9Gd6d7d7e j1Z:Gd8d9d9e j1Z;Gd:d;d;e j1Zd?d?e j1Z>Gd@dAdAe j1Z?GdBdCdCeZ@dDZAdEZBe"dFeAGdGdHdHe@ZCe"dIeAGdJdKdKe@ZDe"dLeAGdMdNdNe@ZEe"dOeAGdPdQdQe@ZFe"dReAGdSdTdTe@ZGdS)Uz PyTorch CANINE model.N) dataclass)OptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)BaseModelOutput ModelOutputMultipleChoiceModelOutputQuestionAnsweringModelOutputSequenceClassifierOutputTokenClassifierOutput)PreTrainedModel)apply_chunking_to_forward find_pruneable_heads_and_indicesprune_linear_layer)add_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forwardloggingreplace_return_docstrings) CanineConfigzgoogle/canine-srzgoogle/canine-r+;=Iagqc@sZeZdZUdZdZejed<dZejed<dZ e e ejed<dZ e e ejed<dS)CanineModelOutputWithPoolinga Output type of [`CanineModel`]. Based on [`~modeling_outputs.BaseModelOutputWithPooling`], but with slightly different `hidden_states` and `attentions`, as these also include the hidden states and attentions of the shallow Transformer encoders. Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model (i.e. the output of the final shallow Transformer encoder). pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`): Hidden-state of the first token of the sequence (classification token) at the last layer of the deep Transformer encoder, further processed by a Linear layer and a Tanh activation function. The Linear layer weights are trained from the next sentence prediction (classification) objective during pretraining. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the input to each encoder + one for the output of each layer of each encoder) of shape `(batch_size, sequence_length, hidden_size)` and `(batch_size, sequence_length // config.downsampling_rate, hidden_size)`. Hidden-states of the model at the output of each layer plus the initial input to each Transformer encoder. The hidden states of the shallow encoders have length `sequence_length`, but the hidden states of the deep encoder have length `sequence_length` // `config.downsampling_rate`. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of the 3 Transformer encoders of shape `(batch_size, num_heads, sequence_length, sequence_length)` and `(batch_size, num_heads, sequence_length // config.downsampling_rate, sequence_length // config.downsampling_rate)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. Nlast_hidden_state pooler_output hidden_states attentions) __name__ __module__ __qualname____doc__r.torch FloatTensor__annotations__r/r0rrr1r9r9k/var/www/html/Darija-Ai-Train/env/lib/python3.8/site-packages/transformers/models/canine/modeling_canine.pyr-As r-c CsRzddl}ddl}ddl}Wn tk r<tdYnXtj|}t d||j |}g}g} |D]@\} } t d| d| |j || } | | | | qrt|| D]\} } | d} tdd | Drt d d| q| dd krd | d<n| d dkr2| | d nh| d dkrJd| d <nP| d dkrldg| dd} n.| ddkr| d dkrdg| d d} |} | D]}|d|rd|kr|d|}n|g}|ddks|ddkrt| d} n|ddks|ddkr$t| d} n`|dd kr>t| d} nFzt| |d} Wn2tk rt d d| YqYnXt|d!krt|d }| |} q|d"dd#krt| d} n@|d$dd%d&td'Dkrt| d} n|dkr|| } | j| jkr0td(| jd)| jd*t d+| t| | _q|S),z'Load tf checkpoints in a pytorch model.rNzLoading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see https://www.tensorflow.org/install/ for installation instructions.z&Converting TensorFlow checkpoint from zLoading TF weight z with shape /css|]}|dkVqdS))Zadam_vZadam_mZAdamWeightDecayOptimizerZAdamWeightDecayOptimizer_1Z global_stepclsZautoregressive_decoderZchar_output_weightsNr9).0nr9r9r: s z,load_tf_weights_in_canine..z Skipping Zbertencoderr embeddingsZsegment_embeddingstoken_type_embeddingsinitial_char_encoderchars_to_moleculesfinal_char_encoder) LayerNormconv projectionz [A-Za-z]+_\d+ZEmbedderz_(\d+)ZkernelgammaweightZ output_biasbetabiasZoutput_weightsiZ _embeddingsicSsg|]}d|qS)Z Embedder_r9)r=ir9r9r: sz-load_tf_weights_in_canine..zPointer shape z and array shape z mismatchedzInitialize PyTorch weight )renumpyZ tensorflow ImportErrorloggererrorospathabspathinfotrainZlist_variablesZ load_variableappendzipsplitanyjoinremove fullmatchgetattrAttributeErrorlenintrange transposeshape ValueErrorr6Z from_numpydata)modelconfigZtf_checkpoint_pathrRnptfZtf_pathZ init_varsnamesZarraysnameriarrayZpointerZm_nameZ scope_namesnumr9r9r:load_tf_weights_in_canineds                      rtcsteZdZdZfddZeedddZeeeddd Zdee j ee j ee j ee j e j d d d Z Z S)CanineEmbeddingsztj|ddS) zDConverts IDs (e.g. codepoints) into embeddings via multiple hashing.rzExpected `embedding_size` (z) % `num_hashes` (z) == 0rrvrydim)rjr enumeratercr\r6cat) rrrrrZhash_bucket_tensorsZembedding_shardsrOZhash_bucket_idsrqZshard_embeddingsr9r9r:_embed_hash_bucketss   z$CanineEmbeddings._embed_hash_bucketsN)rtoken_type_idsrx inputs_embedsreturnc Cs|dk r|}n|dd}|d}|dkrH|jddd|f}|dkrftj|tj|jjd}|dkr|||jj|jj |jj }| |}||}|j dkr| |} || 7}||}||}|S)Nryrdtypedevicer|)sizerxr6zeroslongrrrmrrrrBr{rrGr) rrrrxr input_shape seq_lengthrBrAZposition_embeddingsr9r9r:forwards.      zCanineEmbeddings.forward)NNNN)r2r3r4r5r~rfrrrr6 LongTensorr7r __classcell__r9r9rr:rus rucs4eZdZdZfddZejejdddZZS)CharactersToMoleculeszeConvert character sequence to initial molecule sequence (i.e. downsample) using strided convolutions.csJttj|j|j|j|jd|_t|j|_ tj |j|j d|_ dS)NZ in_channelsZ out_channels kernel_sizestriderw) r}r~rConv1drdownsampling_raterHr hidden_act activationrGrrrmrr9r:r~-s  zCharactersToMolecules.__init__) char_encodingrcCs|ddddddf}t|dd}||}t|dd}||}|ddddddf}tj||gdd}||}|S)NrrrNryr)r6rhrHrrrG)rrZ cls_encodingZ downsampledZdownsampled_truncatedresultr9r9r:r<s   zCharactersToMolecules.forward) r2r3r4r5r~r6Tensorrrr9r9rr:r*s rcs>eZdZdZfddZdejeejejdddZZ S) ConvProjectionz Project representations from hidden_size*2 back to hidden_size across a window of w = config.upsampling_kernel_size characters. cs`t||_tj|jd|j|jdd|_t|j |_ tj |j|j d|_ t |j|_dS)NrNrrrw)r}r~rmrrrupsampling_kernel_sizerHr rrrGrrrrrrr9r:r~^s  zConvProjection.__init__N)inputsfinal_seq_char_positionsrc Cst|dd}|jjd}|d}||}t||fd}|||}t|dd}||}||}| |}|}|dk rt dn|} | S)NrrNrz,CanineForMaskedLM is currently not supported) r6rhrmrrZ ConstantPad1drHrrGrNotImplementedError) rrrZ pad_totalZpad_begZpad_endpadrZfinal_char_seqZ query_seqr9r9r:rms     zConvProjection.forward)N) r2r3r4r5r~r6rrrrr9r9rr:rXs rc sbeZdZfddZddZd ejejeejeejee e ejeejfddd Z Z S) CanineSelfAttentioncst|j|jdkr>t|ds>td|jd|jd|j|_t|j|j|_|j|j|_t |j|j|_ t |j|j|_ t |j|j|_ t |j|_t|dd|_|jdks|jd kr|j|_t d |jd |j|_dS) NrrzThe hidden size (z6) is not a multiple of the number of attention heads ()r{r| relative_keyrelative_key_queryrNr)r}r~rnum_attention_headshasattrrjrfattention_head_size all_head_sizerLinearquerykeyvaluerZattention_probs_dropout_probrrcr{rrdistance_embeddingrrr9r:r~s  zCanineSelfAttention.__init__cCs6|dd|j|jf}|j|}|ddddS)NryrrNrr )rrrviewpermute)rxZ new_x_shaper9r9r:transpose_for_scoress z(CanineSelfAttention.transpose_for_scoresNF) from_tensor to_tensorattention_mask head_maskoutput_attentionsrcCs||}|||}|||}||} t| |dd} |jdks^|jdkr"|d} tj | tj |j d dd} tj | tj |j d dd} | | }| ||jd}|j| jd}|jdkrtd| |}| |} n4|jdkr"td| |}td ||}| ||} | t|j} |dk rx|jd krptj|dd }d |t| jj}| |} tjj| dd }||}|dk r||}t||}|d ddd }|dd|j f}|j |}|r||fn|f}|S)NryrErrrr)rzbhld,lrd->bhlrzbhrd,lrd->bhlrr r?rrN)!rrrrr6matmulrhr{rrrrrrrtorZeinsummathsqrtrndimZ unsqueezefloatZfinfominrZ functionalZsoftmaxrr contiguousr)rrrrrrZmixed_query_layerZ key_layerZ value_layerZ query_layerZattention_scoresrZposition_ids_lZposition_ids_rZdistanceZpositional_embeddingZrelative_position_scoresZrelative_position_scores_queryZrelative_position_scores_keyZattention_probsZ context_layerZnew_context_layer_shapeoutputsr9r9r:rsF             zCanineSelfAttention.forward)NNF) r2r3r4r~rr6rrr7boolrrrr9r9rr:rs  rcsBeZdZfddZeejejeejejfdddZZS)CanineSelfOutputcsBtt|j|j|_tj|j|jd|_t|j |_ dSNrw) r}r~rrrdenserGrrrrrrr9r:r~s zCanineSelfOutput.__init__r0 input_tensorrcCs&||}||}|||}|SNrrrGrr0rr9r9r:rs  zCanineSelfOutput.forward r2r3r4r~rr6r7rrr9r9rr:rs  rc sxeZdZdZd eeeeeedfdd ZddZdee j e e j e e j e eee j e e j fd d d Z Z S)CanineAttentionav Additional arguments related to local attention: - **local** (`bool`, *optional*, defaults to `False`) -- Whether to apply local attention. - **always_attend_to_first_position** (`bool`, *optional*, defaults to `False`) -- Should all blocks be able to attend to the `to_tensor`'s first position (e.g. a [CLS] position)? - **first_position_attends_to_all** (`bool`, *optional*, defaults to `False`) -- Should the *from_tensor*'s first position be able to attend to all positions within the *from_tensor*? - **attend_from_chunk_width** (`int`, *optional*, defaults to 128) -- The width of each block-wise chunk in `from_tensor`. - **attend_from_chunk_stride** (`int`, *optional*, defaults to 128) -- The number of elements to skip when moving to the next block in `from_tensor`. - **attend_to_chunk_width** (`int`, *optional*, defaults to 128) -- The width of each block-wise chunk in *to_tensor*. - **attend_to_chunk_stride** (`int`, *optional*, defaults to 128) -- The number of elements to skip when moving to the next block in `to_tensor`. F)always_attend_to_first_positionfirst_position_attends_to_allattend_from_chunk_widthattend_from_chunk_strideattend_to_chunk_widthattend_to_chunk_stridec sttt||_t||_t|_||_||kr.F) r}r~rmrZ ModuleListrgnum_hidden_layerslayergradient_checkpointingrrrr:r~s  zCanineEncoder.__init__NT)r0rrroutput_hidden_states return_dictrcs|rdnd}rdnd}t|jD]\} } |r8||f}|dk rH|| nd} |jr~|jr~fdd} tjj| | ||| } n| ||| } | d}r"|| df}q"|r||f}|stdd|||fDSt|||dS) Nr9csfdd}|S)Ncs|fSrr9)r)modulerr9r:custom_forward szLCanineEncoder.forward..create_custom_forward..custom_forwardr9)rrr )rr:create_custom_forwardsz4CanineEncoder.forward..create_custom_forwardrrcss|]}|dk r|VqdSrr9r=vr9r9r:r?6sz(CanineEncoder.forward..)r.r0r1) rrrZtrainingr6utils checkpointrr )rr0rrrrrall_hidden_statesall_self_attentionsrOZ layer_moduleZlayer_head_maskrZ layer_outputsr9r r:r s6      zCanineEncoder.forward)FFFrrrr)NNFFT) r2r3r4r~rr6r7rrrr rrr9r9rr:rs,! rcs4eZdZfddZeejejdddZZS) CaninePoolercs*tt|j|j|_t|_dSr)r}r~rrrrZTanhrrrr9r:r~?s zCaninePooler.__init__rcCs(|dddf}||}||}|S)Nr)rr)rr0Zfirst_token_tensor pooled_outputr9r9r:rDs  zCaninePooler.forwardrr9r9rr:r>s rcs4eZdZfddZeejejdddZZS)CaninePredictionHeadTransformcsVtt|j|j|_t|jtr6t |j|_ n|j|_ tj |j|j d|_ dSr) r}r~rrrrrrrr transform_act_fnrGrrrr9r:r~Ns   z&CaninePredictionHeadTransform.__init__rcCs"||}||}||}|Sr)rr"rGrr9r9r:rWs   z%CaninePredictionHeadTransform.forwardrr9r9rr:r!Ms r!cs4eZdZfddZeejejdddZZS)CanineLMPredictionHeadcsLtt||_tj|j|jdd|_t t |j|_ |j |j_ dS)NF)rM) r}r~r! transformrrrZ vocab_sizedecoder Parameterr6rrMrrr9r:r~_s   zCanineLMPredictionHead.__init__rcCs||}||}|Sr)r$r%rr9r9r:rls  zCanineLMPredictionHead.forwardrr9r9rr:r#^s r#cs8eZdZfddZeejeejdddZZS)CanineOnlyMLMHeadcstt||_dSr)r}r~r# predictionsrrr9r:r~ss zCanineOnlyMLMHead.__init__)sequence_outputrcCs||}|Sr)r()rr)Zprediction_scoresr9r9r:rws zCanineOnlyMLMHead.forward) r2r3r4r~rr6rrrr9r9rr:r'rs r'c@s2eZdZdZeZeZdZdZ ddZ d ddZ d S) CaninePreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. canineTcCst|tjtjfr@|jjjd|jjd|j dk r|j j nft|tj r|jjjd|jjd|j dk r|jj|j  n&t|tj r|j j |jjddS)zInitialize the weightsg)ZmeanZstdNr)rrrrrKrkZnormal_rmZinitializer_rangerMZzero_rZ padding_idxrGZfill_)rrr9r9r: _init_weightss     z#CaninePreTrainedModel._init_weightsFcCst|tr||_dSr)rrr)rrrr9r9r:_set_gradient_checkpointings z1CaninePreTrainedModel._set_gradient_checkpointingN)F) r2r3r4r5r config_classrtZload_tf_weightsZbase_model_prefixZsupports_gradient_checkpointingr,r-r9r9r9r:r*sr*aI This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`CanineConfig`]): 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. a5 Args: input_ids (`torch.LongTensor` of shape `({0})`): Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.FloatTensor` of shape `({0})`, *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) token_type_ids (`torch.LongTensor` of shape `({0})`, *optional*): Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0, 1]`: - 0 corresponds to a *sentence A* token, - 1 corresponds to a *sentence B* token. [What are token type IDs?](../glossary#token-type-ids) position_ids (`torch.LongTensor` of shape `({0})`, *optional*): Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.max_position_embeddings - 1]`. [What are position IDs?](../glossary#position-ids) 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 `({0}, 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. 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 [`~utils.ModelOutput`] instead of a plain tuple. z`The bare CANINE Model transformer outputting raw hidden-states without any specific head on top.cseZdZdfdd ZddZddZejedd d Z ejejejd d d Z e e deeeeddeejeejeejeejeejeejeeeeeeeeefd ddZZS) CanineModelTc st|||_t|}d|_t||_t|ddd|j |j |j |j d|_ t ||_ t||_ t||_t||_|rt|nd|_|dS)NrTF)rrrrrrr)r}r~rmcopydeepcopyrruchar_embeddingsrZlocal_transformer_striderCrrDr@rrIrFrpooler post_init)rrmZadd_pooling_layerZshallow_configrr9r:r~s*       zCanineModel.__init__cCs*|D]\}}|jj|j|qdS)z Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base class PreTrainedModel N)itemsr@rr r)rZheads_to_prunerrr9r9r: _prune_headsszCanineModel._prune_headscCs\|jd|jd}}|jd}t||d|f}tj||dftj|jd}||}|S)aP Create 3D attention mask from a 2D tensor mask. Args: from_tensor: 2D or 3D Tensor of shape [batch_size, from_seq_length, ...]. to_mask: int32 Tensor of shape [batch_size, to_seq_length]. Returns: float Tensor of shape [batch_size, from_seq_length, to_seq_length]. rr)rrr)rir6reshaperonesZfloat32r)rrZto_mask batch_sizerrZbroadcast_onesmaskr9r9r:)_create_3d_attention_mask_from_input_masks  z5CanineModel._create_3d_attention_mask_from_input_mask)char_attention_maskrcCsF|j\}}t||d|f}tjj||d|}tj|dd}|S)z[Downsample 2D character attention mask to 2D molecule attention mask using MaxPool1d layer.r)rrryr)rir6r7rZ MaxPool1drsqueeze)rr<rr9Z char_seq_lenZpoolable_char_maskZpooled_molecule_maskmolecule_attention_maskr9r9r:_downsample_attention_mask"s z&CanineModel._downsample_attention_mask) moleculeschar_seq_lengthrc Cs|jj}|ddddddf}tj||dd}|ddddddf}tt|t|}tj|||dd}tj||gddS)zDRepeats molecules to make them the same length as the char sequence.NrrE)Zrepeatsrryr)rmrr6Zrepeat_interleavefmodZtensoritemr) rr@rAZrateZmolecules_without_extra_clsZrepeatedZ last_moleculeZremainder_lengthZremainder_repeatedr9r9r:_repeat_molecules3szCanineModel._repeat_moleculesbatch_size, sequence_lengthr output_typer.N) rrrrxrrrrrrc "Cs|dk r |n|jj}|dk r |n|jj}|r0dnd} |r.)r.r/r0r1)!rmrruse_return_dictrjZ%warn_if_padding_and_no_attention_maskrrr6r8rrZget_extended_attention_maskr?rriZ get_head_maskrr2r;rCr.rDr@r3rDrrIrFr0r1rr-)"rrrrrxrrrrrrrrr9rrZextended_attention_maskr>Z extended_molecule_attention_maskZinput_char_embeddingsr<Zinit_chars_encoder_outputsZinput_char_encodingZinit_molecule_encodingZencoder_outputsZmolecule_sequence_outputr Zrepeated_moleculesconcatr)Zfinal_chars_encoder_outputsZdeep_encoder_hidden_statesZdeep_encoder_self_attentionsrr9r9r:rLs           zCanineModel.forward)T) NNNNNNNNN)r2r3r4r~r6r;r6rrfr?rDrCANINE_INPUTS_DOCSTRINGformatr_CHECKPOINT_FOR_DOCr-_CONFIG_FOR_DOCrrr7rrrrrr9r9rr:r/s@  r/z CANINE Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. cseZdZfddZeedeee e dd e e j e e je e j e e j e e je e je e j e ee ee eeee fd ddZZS) CanineForSequenceClassificationcsJt||j|_t||_t|j|_t |j |j|_ | dSr r}r~ num_labelsr/r+rrrrrr classifierr4rrr9r:r~s   z(CanineForSequenceClassification.__init__rErFN rrrrxrrlabelsrrrrc  Cs|| dk r | n|jj} |j|||||||| | d } | d} || } || } d}|dk r8|jjdkr|jdkrzd|j_n4|jdkr|jtj ks|jtj krd|j_nd|j_|jjdkrt }|jdkr|| | }n || |}nN|jjdkrt }|| d|j|d}n|jjdkr8t}|| |}| sh| f| dd}|dk rd|f|S|St|| | j| jd S) a 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). NrrrxrrrrrrZ regressionZsingle_label_classificationZmulti_label_classificationryrNlosslogitsr0r1)rmrHr+rrQZ problem_typerPrr6rrfr r=rrrrr0r1)rrrrrxrrrSrrrrr rWrVloss_fctrr9r9r:rsV       "     z'CanineForSequenceClassification.forward) NNNNNNNNNN)r2r3r4r~rrJrKrrLrrMrr6rr7rrrrrr9r9rr:rNs<  rNz CANINE Model with a multiple choice classification head on top (a linear layer on top of the pooled output and a softmax) e.g. for RocStories/SWAG tasks. cseZdZfddZeedeee e dd e e j e e je e j e e j e e je e je e j e ee ee eeee fd ddZZS) CanineForMultipleChoicecs@t|t||_t|j|_t|j d|_ | dSr) r}r~r/r+rrrrrrrQr4rrr9r:r~Zs   z CanineForMultipleChoice.__init__z(batch_size, num_choices, sequence_lengthrFNrRc  Cst| dk r | n|jj} |dk r&|jdn|jd} |dk rJ|d|dnd}|dk rh|d|dnd}|dk r|d|dnd}|dk r|d|dnd}|dk r|d|d|dnd}|j|||||||| | d } | d} || } || }|d| }d}|dk r0t}|||}| s`|f| dd}|dk r\|f|S|St ||| j | j dS)aJ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See `input_ids` above) NrryrErTrNrU) rmrHrirrr+rrQrrr0r1)rrrrrxrrrSrrrZ num_choicesrr rWZreshaped_logitsrVrXrr9r9r:rdsL      zCanineForMultipleChoice.forward) NNNNNNNNNN)r2r3r4r~rrJrKrrLrrMrr6rr7rrrrrr9r9rr:rYRs<  rYz CANINE 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. cseZdZfddZeedeee dd e e j e e j e e j e e j e e j e e j e e j e ee ee eeeefd ddZZS) CanineForTokenClassificationcsJt||j|_t||_t|j|_t |j |j|_ | dSrrOrrr9r:r~s   z%CanineForTokenClassification.__init__rE)rGr.NrRc  Cs| dk r | n|jj} |j|||||||| | d } | d} || } || } d}|dk rxt}|| d|j|d}| s| f| dd}|dk r|f|S|St|| | j | j dS)a labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. Returns: Example: ```python >>> from transformers import AutoTokenizer, CanineForTokenClassification >>> import torch >>> tokenizer = AutoTokenizer.from_pretrained("google/canine-s") >>> model = CanineForTokenClassification.from_pretrained("google/canine-s") >>> inputs = tokenizer( ... "HuggingFace is a company based in Paris and New York", add_special_tokens=False, return_tensors="pt" ... ) >>> with torch.no_grad(): ... logits = model(**inputs).logits >>> predicted_token_class_ids = logits.argmax(-1) >>> # Note that tokens are classified rather then input words which means that >>> # there might be more predicted token classes than words. >>> # Multiple token classes might account for the same word >>> predicted_tokens_classes = [model.config.id2label[t.item()] for t in predicted_token_class_ids[0]] >>> predicted_tokens_classes # doctest: +SKIP ``` ```python >>> labels = predicted_token_class_ids >>> loss = model(**inputs, labels=labels).loss >>> round(loss.item(), 2) # doctest: +SKIP ```NrTrryrNrU) rmrHr+rrQrrrPrr0r1)rrrrrxrrrSrrrrr)rWrVrXrr9r9r:rs83   z$CanineForTokenClassification.forward) NNNNNNNNNN)r2r3r4r~rrJrKrrrMrr6rr7rrrrrr9r9rr:rZs4   rZz CANINE Model 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`). cseZdZfddZeededee dddd e e j e e j e e j e e j e e j e e j e e j e e j e ee ee eeeefd d d ZZS) CanineForQuestionAnsweringcs<t||j|_t||_t|j|j|_| dSr) r}r~rPr/r+rrr qa_outputsr4rrr9r:r~s   z#CanineForQuestionAnswering.__init__rEzSplend1dchan/canine-c-squadz 'nice puppet'gQ!@)rrGr.Zexpected_outputZ expected_lossN) rrrrxrrstart_positions end_positionsrrrrc  CsD| 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}| s.||f| dd}|dk r*|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. NrTrrryr)Z ignore_indexrN)rV start_logits end_logitsr0r1) rmrHr+r\r^r=rerZclamp_rrr0r1)rrrrrxrrr]r^rrrrr)rWr_r`Z total_lossZ ignored_indexrXZ start_lossZend_lossrr9r9r:r)sP              z"CanineForQuestionAnswering.forward) NNNNNNNNNNN)r2r3r4r~rrJrKrrrMrr6rr7rrrrrr9r9rr:r[sD   r[)Hr5r0rrW dataclassesrtypingrrrr6Ztorch.utils.checkpointrZtorch.nnrrr Z activationsr Zmodeling_outputsr r rrrrZmodeling_utilsrZ pytorch_utilsrrrrrrrrrZconfiguration_caninerZ get_loggerr2rUrLrMZ$CANINE_PRETRAINED_MODEL_ARCHIVE_LISTrr-rtModulerurrrrrrrrrrr!r#r'r*ZCANINE_START_DOCSTRINGrJr/rNrYrZr[r9r9r9r:s       $"ae.:d:R   2ZTc