o ha@sdZddlZddlZddlmZddlmZmZmZddl Z ddl Z ddl m Z ddl m Z mZmZddlmZmZdd lmZmZmZmZmZmZdd lmZmZdd lmZmZm Z dd l!m"Z"m#Z#m$Z$m%Z%d dl&m'Z'e%(e)Z*dZ+dZ,gdZ-ddZ.Gddde j/Z0GdddeZ1Gddde j/Z2Gddde j/Z3Gddde j/Z4Gddde j/Z5Gd d!d!e j/Z6Gd"d#d#e j/Z7Gd$d%d%e j/Z8Gd&d'd'e j/Z9Gd(d)d)e j/Z:Gd*d+d+e j/Z;d,ZGd1d2d2e j/Z?e#d3e<Gd4d5d5e1Z@Gd6d7d7e j/ZAe#d8e<Gd9d:d:e1ZBe#d;e<Gde<Gd?d@d@e1ZDe#dAe<GdBdCdCe1ZEdS)Dz PyTorch ConvBERT model.N) attrgetter)OptionalTupleUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FNget_activation)"BaseModelOutputWithCrossAttentionsMaskedLMOutputMultipleChoiceModelOutputQuestionAnsweringModelOutputSequenceClassifierOutputTokenClassifierOutput)PreTrainedModelSequenceSummary)apply_chunking_to_forward find_pruneable_heads_and_indicesprune_linear_layer)add_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forwardlogging)ConvBertConfigYituTech/conv-bert-baser)rzYituTech/conv-bert-medium-smallzYituTech/conv-bert-smallcCszddl}Wn tytdwtj|}td||j |}i}|D]\}}td|d||j ||} | ||<q,ddd d d d d d} |j dkrYd} nd} t |j D]"} d| d| d| d<d| d| d| d<d| d| d| d<d| d| d| d<d| d| d| d<d| d| d| d<d| d | d| d!<d| d"| d| d#<d| d$| d| d%<d| d&| d| d'<d| d(| d| d)<d| d*| d| d+<d| d,| d| d-<d| d.| d| d/<d| d0| d| d1<d| d2| d| d3<d| d4| d| d5<d| d6| d7| d| d8<d| d6| d9| d| d:<d| d;| d7| d| d<<d| d;| d9| d| d=<d| d>| d| d?<d| d@| d| dA<q`|D]c} | d}t|}||}| |}t||}tdB|dC|dD|d7r|dEs|dFs|j}|dGr|ddHd}|dIr|dHdd}|dJr|dK}||_q|S)Lz'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 z"electra/embeddings/word_embeddingsz&electra/embeddings/position_embeddingsz(electra/embeddings/token_type_embeddingsz"electra/embeddings/LayerNorm/gammaz!electra/embeddings/LayerNorm/betaz!electra/embeddings_project/kernelzelectra/embeddings_project/bias)z!embeddings.word_embeddings.weightz%embeddings.position_embeddings.weightz'embeddings.token_type_embeddings.weightzembeddings.LayerNorm.weightzembeddings.LayerNorm.biaszembeddings_project.weightzembeddings_project.biasrg_densedensezelectra/encoder/layer_z/attention/self/query/kernelzencoder.layer.z.attention.self.query.weightz/attention/self/query/biasz.attention.self.query.biasz/attention/self/key/kernelz.attention.self.key.weightz/attention/self/key/biasz.attention.self.key.biasz/attention/self/value/kernelz.attention.self.value.weightz/attention/self/value/biasz.attention.self.value.biasz./attention/self/conv_attn_key/depthwise_kernelz4.attention.self.key_conv_attn_layer.depthwise.weightz./attention/self/conv_attn_key/pointwise_kernelz4.attention.self.key_conv_attn_layer.pointwise.weightz"/attention/self/conv_attn_key/biasz(.attention.self.key_conv_attn_layer.biasz'/attention/self/conv_attn_kernel/kernelz(.attention.self.conv_kernel_layer.weightz%/attention/self/conv_attn_kernel/biasz&.attention.self.conv_kernel_layer.biasz&/attention/self/conv_attn_point/kernelz%.attention.self.conv_out_layer.weightz$/attention/self/conv_attn_point/biasz#.attention.self.conv_out_layer.biasz/attention/output/dense/kernelz.attention.output.dense.weightz!/attention/output/LayerNorm/gammaz".attention.output.LayerNorm.weightz/attention/output/dense/biasz.attention.output.dense.biasz /attention/output/LayerNorm/betaz .attention.output.LayerNorm.biasz/intermediate/z/kernelz.intermediate.dense.weightz/biasz.intermediate.dense.biasz/output/z.output.dense.weightz.output.dense.biasz/output/LayerNorm/gammaz.output.LayerNorm.weightz/output/LayerNorm/betaz.output.LayerNorm.biaszTF: z, PT:  z/intermediate/g_dense/kernelz/output/g_dense/kernelz/depthwise_kernelz/pointwise_kernelz/conv_attn_key/bias) tensorflow ImportErrorloggererrorospathabspathinfotrainlist_variables load_variable num_groupsrangenum_hidden_layersnamed_parametersrtorch from_numpyendswithTpermute unsqueezedata)modelconfigtf_checkpoint_pathtftf_path init_varstf_datanameshapearray param_mappinggroup_dense_namejparam param_name retrieverresulttf_namevaluerMd/var/www/html/ai/venv/lib/python3.10/site-packages/transformers/models/convbert/modeling_convbert.pyload_tf_weights_in_convbert8s                                                        rOc sbeZdZdZfddZ    d deejdeejdeejdeejd ejf d d Z Z S) ConvBertEmbeddingszGConstruct the embeddings from word, position and token_type embeddings.csttj|j|j|jd|_t|j|j|_ t|j |j|_ tj |j|j d|_ t|j|_|jdt|jddd|jdtj|jtjddddS) N) padding_idxeps position_ids)rr#F) persistenttoken_type_ids)dtype)super__init__r Embedding vocab_sizeembedding_size pad_token_idword_embeddingsmax_position_embeddingsposition_embeddingstype_vocab_sizetoken_type_embeddings LayerNormlayer_norm_epsDropouthidden_dropout_probdropoutregister_bufferr3arangeexpandzerosrTsizelongselfr; __class__rMrNrYs  zConvBertEmbeddings.__init__N input_idsrVrT inputs_embedsreturnc Cs|dur |}n|dd}|d}|dur$|jddd|f}|durNt|drC|jddd|f}||d|}|}n tj|tj|jjd}|durW| |}| |} | |} || | } | | } | | } | S)Nr#rrVrrWdevice)rlrThasattrrVrjr3rkrmrvr^r`rbrcrg) rorrrVrTrs input_shape seq_lengthbuffered_token_type_ids buffered_token_type_ids_expandedr`rb embeddingsrMrMrNforwards(        zConvBertEmbeddings.forward)NNNN) __name__ __module__ __qualname____doc__rYrr3 LongTensor FloatTensorr} __classcell__rMrMrprNrPs$ rPc@s(eZdZdZeZeZdZdZ ddZ dS)ConvBertPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. convbertTcCst|tjr |jjjd|jjd|jdur|jj dSdSt|tj rC|jjjd|jjd|j durA|jj|j  dSdSt|tj rX|jj |jj ddSdS)zInitialize the weightsmeanstdNg?) isinstancerLinearweightr9normal_r;initializer_rangebiaszero_rZrQrcfill_)romodulerMrMrN _init_weightss      z%ConvBertPreTrainedModel._init_weightsN) r~rrrr config_classrOload_tf_weightsbase_model_prefixsupports_gradient_checkpointingrrMrMrMrNrs rc6eZdZdZfddZdejdejfddZZS)SeparableConv1DzSThis class implements separable convolution, i.e. a depthwise and a pointwise layerc s~ttj|||||ddd|_tj||ddd|_tt|d|_ |jj j j d|j d|jj j j d|j ddS)Nr"F) kernel_sizegroupspaddingrr)rrrr)rXrYrConv1d depthwise pointwise Parameterr3rkrrr9rr)ror; input_filtersoutput_filtersrkwargsrprMrNrYs zSeparableConv1D.__init__ hidden_statesrtcCs"||}||}||j7}|SN)rrr)rorxrMrMrNr}   zSeparableConv1D.forward r~rrrrYr3Tensorr}rrMrMrprNr s rcsxeZdZfddZddZ    ddejdeejd eejd eejd ee d e ejeejff d dZ Z S)ConvBertSelfAttentioncs`t|j|jdkrt|dstd|jd|jd|j|j}|dkr1|j|_d|_n||_|j|_|j|_|j|jdkrHtd|j|jd|_|j|j|_ t |j|j |_ t |j|j |_ t |j|j |_t||j|j |j|_t |j |j|j|_t |j|j |_t j|jdgt|jdddgd |_t |j|_dS) Nrr\zThe hidden size (z6) is not a multiple of the number of attention heads ()rz6hidden_size should be divisible by num_attention_headsr")rr)rXrY hidden_sizenum_attention_headsrw ValueError head_ratioconv_kernel_sizeattention_head_size all_head_sizerrquerykeyrLrkey_conv_attn_layerconv_kernel_layerconv_out_layerUnfoldintunfoldreattention_probs_dropout_probrg)ror;new_num_attention_headsrprMrNrY's<   zConvBertSelfAttention.__init__cCs6|dd|j|jf}|j|}|ddddS)Nr#rr"rr )rlrrviewr7)ror new_x_shaperMrMrNtranspose_for_scoresNs z*ConvBertSelfAttention.transpose_for_scoresNFrattention_mask head_maskencoder_hidden_statesoutput_attentionsrtcCsV||}|d}|dur||}||} n ||}||} ||dd} | dd} ||} ||} || } t| |}| |}t |d|j dg}tj |dd}| |}t ||d|jg}|ddd}tjj||j dgd|j dddgdd}|dd |d|j|j }t |d|j|j g}t||}t |d|jg}t| | dd}|t|j}|dur||}tjj |dd}||}|dur||}t|| }|dddd}t ||d|j|jg}t||gd}|dd|j|jdf}|j|}|r&||f}|S|f}|S) Nrrr"r#dim)rdilationrstrider )rrlrrLr transposerr3multiplyrreshapersoftmaxrr contiguousr8r functionalrrmatmulmathsqrtrgr7rcatr)rorrrrrmixed_query_layer batch_sizemixed_key_layermixed_value_layermixed_key_conv_attn_layer query_layer key_layer value_layerconv_attn_layerrrattention_scoresattention_probs context_layerconv_outnew_context_layer_shapeoutputsrMrMrNr}Ssh                   zConvBertSelfAttention.forwardNNNF) r~rrrYrr3rrrboolrr}rrMrMrprNr&s( 'rc8eZdZfddZdejdejdejfddZZS)ConvBertSelfOutputcsBtt|j|j|_tj|j|jd|_t|j |_ dSNrR) rXrYrrrr rcrdrerfrgrnrprMrNrYs zConvBertSelfOutput.__init__r input_tensorrtcC&||}||}|||}|Srr rgrcrorrrMrMrNr}  zConvBertSelfOutput.forwardr~rrrYr3rr}rrMrMrprNrs $rcsxeZdZfddZddZ    ddejdeejd eejd eejd ee d e ejeejff d dZ Z S)ConvBertAttentioncs*tt||_t||_t|_dSr)rXrYrroroutputset pruned_headsrnrprMrNrYs    zConvBertAttention.__init__cCst|dkrdSt||jj|jj|j\}}t|jj||j_t|jj||j_t|jj ||j_ t|j j |dd|j _ |jjt||j_|jj|jj|j_ |j ||_dS)Nrrr)lenrrorrrrrrrLrr runion)roheadsindexrMrMrN prune_headss zConvBertAttention.prune_headsNFrrrrrrtc Cs8||||||}||d|}|f|dd}|S)Nrr)ror) rorrrrr self_outputsattention_outputrrMrMrNr}szConvBertAttention.forwardr) r~rrrYrr3rrrrrr}rrMrMrprNrs( rc2eZdZfddZdejdejfddZZS)GroupedLinearLayercsjt||_||_||_|j|j|_|j|j|_tt |j|j|j|_ tt ||_ dSr) rXrY input_size output_sizer/ group_in_dim group_out_dimrrr3emptyrr)rorrr/rprMrNrYs zGroupedLinearLayer.__init__rrtcCsrt|d}t|d|j|jg}|ddd}t||j}|ddd}t||d|j g}||j }|S)Nrr#rr") listrlr3rr/rr7rrrr)rorrrrMrMrNr}s zGroupedLinearLayer.forwardrrMrMrprNrs  rcr)ConvBertIntermediatecsft|jdkrt|j|j|_n t|j|j|jd|_t |j t r-t |j |_ dS|j |_ dS)Nrrrr/)rXrYr/rrrintermediate_sizer rr hidden_actstrr intermediate_act_fnrnrprMrNrYs     zConvBertIntermediate.__init__rrtcCs||}||}|Sr)r rrorrMrMrNr}s  zConvBertIntermediate.forwardrrMrMrprNrs  rcr)ConvBertOutputcsdt|jdkrt|j|j|_n t|j|j|jd|_tj |j|j d|_ t |j |_ dS)NrrrR)rXrYr/rrrrr rrcrdrerfrgrnrprMrNrY s   zConvBertOutput.__init__rrrtcCrrrrrMrMrNr}rzConvBertOutput.forwardrrMrMrprNr s $ rcseZdZfddZ     ddejdeejdeejdeejd eejd eed e ejeejffd d Z ddZ Z S) ConvBertLayercsnt|j|_d|_t||_|j|_|j|_|jr+|js&t|dt||_ t ||_ t ||_ dS)Nrz> should be used as a decoder model if cross attention is added)rXrYchunk_size_feed_forward seq_len_dimr attention is_decoderadd_cross_attention TypeErrorcrossattentionr intermediaterrrnrprMrNrYs    zConvBertLayer.__init__NFrrrrencoder_attention_maskrrtc Cs|j||||d}|d}|dd} |jr<|dur_sz,ConvBertEncoder.__init__..F) rXrYr;r ModuleListr0r1layergradient_checkpointingrnrprrNrY\s   zConvBertEncoder.__init__NFTrrrrrroutput_hidden_states return_dictrtc  Cs|rdnd} |r dnd} |r|jjrdnd} t|jD]I\} } |r&| |f} |dur.|| nd}|jrC|jrC|| j||||||}n | ||||||}|d}|rd| |df} |jjrd| |df} q|rl| |f} |s{tdd|| | | fDSt || | | dS)NrMrrr"css|] }|dur|VqdSrrM)rvrMrMrN sz*ConvBertEncoder.forward..)last_hidden_stater attentionscross_attentions) r;r enumeraterr training_gradient_checkpointing_func__call__tupler )rorrrrrrr!r"all_hidden_statesall_self_attentionsall_cross_attentionsi layer_modulelayer_head_mask layer_outputsrMrMrNr}bsV      zConvBertEncoder.forward)NNNNFFT)r~rrrYr3rrrrrrr r}rrMrMrprNr[s8    rcr)ConvBertPredictionHeadTransformcsVtt|j|j|_t|jtrt |j|_ n|j|_ tj |j|j d|_ dSr) rXrYrrrr rrrr transform_act_fnrcrdrnrprMrNrYs  z(ConvBertPredictionHeadTransform.__init__rrtcCs"||}||}||}|Sr)r r5rcrrMrMrNr}rz'ConvBertPredictionHeadTransform.forwardrrMrMrprNr4s  r4aK 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 ([`ConvBertConfig`]): 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. a8 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. zbThe bare ConvBERT Model transformer outputting raw hidden-states without any specific head on top.cseZdZfddZddZddZddZee d e e e e d  dd eejd eejdeejdeejdeejdeejdeedeedeedeee ffddZZS) ConvBertModelcsPt|t||_|j|jkrt|j|j|_t ||_ ||_ | dSr) rXrYrPr|r\rrrembeddings_projectrencoderr; post_initrnrprMrNrYs     zConvBertModel.__init__cCs|jjSrr|r^rorMrMrNget_input_embeddingssz"ConvBertModel.get_input_embeddingscCs ||j_dSrr:)rorLrMrMrNset_input_embeddingss z"ConvBertModel.set_input_embeddingscCs*|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)itemsr8rr r)roheads_to_prunerrrMrMrN _prune_headsszConvBertModel._prune_headsbatch_size, sequence_length checkpoint output_typerNrrrrVrTrrsrr!r"rtc Csr|dur|n|jj}|dur|n|jj}| dur| n|jj} |dur*|dur*td|dur9||||} n|durF|dd} ntd| \} } |durU|jn|j} |durctj | | d}|durt |j dr|j j ddd| f}| | | }|}n tj| tj| d}||| }|||jj}|j ||||d}t |dr||}|j|||||| d }|S) NzDYou cannot specify both input_ids and inputs_embeds at the same timer#z5You have to specify either input_ids or inputs_embeds)rvrVru)rrrTrVrsr7)rrrr!r")r;rr!use_return_dictr%warn_if_padding_and_no_attention_maskrlrvr3onesrwr|rVrjrkrmget_extended_attention_mask get_head_maskr1r7r8)rorrrrVrTrrsrr!r"rxrryrvrzr{extended_attention_maskrrMrMrNr}sL        zConvBertModel.forward) NNNNNNNNN)r~rrrYr<r=r@rCONVBERT_INPUTS_DOCSTRINGformatr_CHECKPOINT_FOR_DOCr _CONFIG_FOR_DOCrr3rrrrrr}rrMrMrprNr6sR      r6cr)ConvBertGeneratorPredictionszAPrediction module for the generator, made up of two dense layers.cs4ttj|j|jd|_t|j|j|_dSr) rXrYrrcr\rdrrr rnrprMrNrYXs z%ConvBertGeneratorPredictions.__init__generator_hidden_statesrtcCs$||}td|}||}|S)Ngelu)r r rc)rorPrrMrMrNr}^s   z$ConvBertGeneratorPredictions.forward) r~rrrrYr3rr}rrMrMrprNrOUs rOz6ConvBERT Model with a `language modeling` head on top.cseZdZdgZfddZddZddZee de e e e d  dd eejd eejd eejdeejdeejdeejdeejdeedeedeedeee ffddZZS)ConvBertForMaskedLMzgenerator.lm_head.weightcs>t|t||_t||_t|j|j |_ | dSr) rXrYr6rrOgenerator_predictionsrrr\r[generator_lm_headr9rnrprMrNrYjs    zConvBertForMaskedLM.__init__cCs|jSrrTr;rMrMrNget_output_embeddingstsz)ConvBertForMaskedLM.get_output_embeddingscCs ||_dSrrU)ror^rMrMrNset_output_embeddingsws z)ConvBertForMaskedLM.set_output_embeddingsrArBNrrrrVrTrrslabelsrr!r"rtc  Cs| dur| n|jj} ||||||||| | } | d} || } || } d}|dur=t}|| d|jj|d}| sS| f| dd}|durQ|f|S|St || | j | j dS)a labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]` Nrr#rlosslogitsrr&) r;rErrSrTrrrr[rrr&)rorrrrVrTrrsrXrr!r"rPgenerator_sequence_outputprediction_scoresrZloss_fctrrMrMrNr}zs8  zConvBertForMaskedLM.forward NNNNNNNNNN)r~rr_tied_weights_keysrYrVrWrrKrLrrMrrNrr3rrrrrr}rrMrMrprNrRfsX       rRcr)ConvBertClassificationHeadz-Head for sentence-level classification tasks.csZtt|j|j|_|jdur|jn|j}t||_ t|j|j |_ ||_ dSr) rXrYrrrr classifier_dropoutrfrerg num_labelsout_projr;ror;rbrprMrNrYs   z#ConvBertClassificationHead.__init__rrtcKsR|dddddf}||}||}t|jj|}||}||}|S)Nr)rgr r r;rrd)rorrrrMrMrNr}s    z"ConvBertClassificationHead.forwardrrMrMrprNras  raz ConvBERT Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled output) e.g. for GLUE tasks. ceZdZfddZeedeee e d          dde e j de e jde e j d e e j d e e jd e e jd e e j d e ede ede edeee ffddZZS)!ConvBertForSequenceClassificationcs:t||j|_||_t||_t||_|dSr) rXrYrcr;r6rra classifierr9rnrprMrNrYs    z*ConvBertForSequenceClassification.__init__rArBNrrrrVrTrrsrXrr!r"rtc  Csh| dur| n|jj} |j|||||||| | d } | d} || } d}|dur|jjdurQ|jdkr7d|j_n|jdkrM|jtjksH|jtj krMd|j_nd|j_|jjdkrot }|jdkri|| | }n+|| |}n%|jjdkrt }|| d|j| d}n|jjdkrt}|| |}| s| f| dd}|dur|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). NrrVrTrrsrr!r"rr regressionsingle_label_classificationmulti_label_classificationr#rY)r;rErrh problem_typercrWr3rmrr squeezerrrrrr&rorrrrVrTrrsrXrr!r"rsequence_outputr[rZr^rrMrMrNr}sT    "       z)ConvBertForSequenceClassification.forwardr_)r~rrrYrrKrLrrMrrNrr3rrrrrr}rrMrMrprNrgsR      rgz ConvBERT 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. crf)ConvBertForMultipleChoicecs<t|t||_t||_t|jd|_ | dS)Nr) rXrYr6rrsequence_summaryrrrrhr9rnrprMrNrY6s    z"ConvBertForMultipleChoice.__init__z(batch_size, num_choices, sequence_lengthrBNrrrrVrTrrsrXrr!r"rtc  Csn| dur| n|jj} |dur|jdn|jd} |dur%|d|dnd}|dur4|d|dnd}|durC|d|dnd}|durR|d|dnd}|dure|d|d|dnd}|j|||||||| | d } | d} || }||}|d| }d}|durt}|||}| s|f| dd}|dur|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) Nrr#rrirrY) r;rErBrrlrrrrhrrrr&)rorrrrVrTrrsrXrr!r" num_choicesrrp pooled_outputr[reshaped_logitsrZr^rrMrMrNr}@sL    z!ConvBertForMultipleChoice.forwardr_)r~rrrYrrKrLrrMrrNrr3rrrrrr}rrMrMrprNrq.sV      rqz ConvBERT 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. crf)ConvBertForTokenClassificationcs^t||j|_t||_|jdur|jn|j}t||_ t |j |j|_ | dSr)rXrYrcr6rrbrfrrergrrrhr9rerprMrNrYs    z'ConvBertForTokenClassification.__init__rArBNrrrrVrTrrsrXrr!r"rtc  Cs| dur| n|jj} |j|||||||| | d } | d} || } || } d}|durs     |<-=E 5`PXVJ