o ۷i @sdZddlZddlZddlZddlmZddlmZmZddl Z ddl m Z ddl m Z m Z mZddlmZdd lmZdd lmZmZmZmZmZmZdd lmZdd lmZmZmZdd l m!Z!m"Z"ddl#m$Z$e"%e&Z'gdZ(ee!ddGdddeZ)ddZ*Gddde j+Z,Gddde j+Z-Gddde j+Z.Gddde j+Z/Gdd d e j+Z0Gd!d"d"e j+Z1Gd#d$d$e j+Z2Gd%d&d&e j+Z3Gd'd(d(eZ4Gd)d*d*e j+Z5Gd+d,d,e j+Z6Gd-d.d.e j+Z7Gd/d0d0e j+Z8Gd1d2d2e j+Z9e!Gd3d4d4eZ:e!Gd5d6d6e:Z;e!d7dGd8d9d9e:Ze!Gd>d?d?e:Z?gd@Z@dS)AzPyTorch CANINE model.N) dataclass)OptionalUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)GradientCheckpointingLayer)BaseModelOutput ModelOutputMultipleChoiceModelOutputQuestionAnsweringModelOutputSequenceClassifierOutputTokenClassifierOutput)PreTrainedModel)apply_chunking_to_forward find_pruneable_heads_and_indicesprune_linear_layer)auto_docstringlogging) CanineConfig)+;=Iagqa 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. ) custom_introc@sbeZdZUdZdZeejed<dZ eejed<dZ ee ejed<dZ ee ejed<dS)CanineModelOutputWithPoolinga 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,rtorch FloatTensor__annotations__r-r.tupler/r8r8`/home/ubuntu/vllm_env/lib/python3.10/site-packages/transformers/models/canine/modeling_canine.pyr+2s  r+c Cs*zddl}ddl}ddl}Wn tytdwtj|}t d||j |}g}g} |D] \} } t d| d| |j || } | | | | q6t|| D]5\} } | d} tdd | Dr{t d d| q\| dd krd | d<n>| d dkr| | d n0| d dkrd| d <n%| d dkrdg| dd} n| ddkr| d dvrdg| d d} |} | D]x}|d|rd|vr|d|}n|g}|ddks|ddkrt| d} n?|ddks|ddkrt| d} n,|dd krt| d} nz t| |d} Wnty.t d d| Yqwt|d!kr@t|d }| |} q|d"dd#krPt| d} n |d$dd%d&td'Dvrft| d} n |dkrp|| } | j| jkrtd(| 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|]}|dvVqdS))adam_vadam_mAdamWeightDecayOptimizerAdamWeightDecayOptimizer_1 global_stepclsautoregressive_decoderchar_output_weightsNr8).0nr8r8r9 us  z,load_tf_weights_in_canine..z Skipping bertencoderr embeddingssegment_embeddingstoken_type_embeddingsinitial_char_encoderchars_to_moleculesfinal_char_encoder) LayerNormconv projectionz [A-Za-z]+_\d+Embedderz_(\d+)kernelgammaweight output_biasbetabiasoutput_weightsi _embeddingsicSsg|]}d|qS) Embedder_r8)rCir8r8r9 sz-load_tf_weights_in_canine..zPointer shape z and array shape z mismatchedzInitialize PyTorch weight )renumpy tensorflow ImportErrorloggererrorospathabspathinfotrainlist_variables load_variableappendzipsplitanyjoinremove fullmatchgetattrAttributeErrorlenintrange transposeshape ValueErrorr4 from_numpydata)modelconfigtf_checkpoint_pathr`nptftf_path init_varsnamesarraysnamerzarraypointerm_name scope_namesnumr8r8r9load_tf_weights_in_canineWs                        rc seZdZdZfddZdedefddZdededefd d Z dd ee j d ee j dee j dee j de j f ddZ Z S)CanineEmbeddingsz