a h@sdZddlZddlZddlmZmZmZddlZddlZddlm Z ddl m Z m Z m Z ddlmZddlmZmZmZdd lmZdd lmZdd lmZmZmZdd lmZdd lmZm Z m!Z!ddl"m#Z#m$Z$m%Z%ddl&m'Z'e$(e)Z*ddZ+Gddde j,Z-Gddde j,Z.Gddde j,Z/GdddeZ0e#GdddeZ1e#Gddde1Z2e#dd Gd!d"d"e1eZ3e#d#d Gd$d%d%e1Z4gd&Z5dS)'zPyTorch OpenAI ImageGPT model.N)AnyOptionalUnion)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)Cache DynamicCacheEncoderDecoderCache)GenerationMixin)GradientCheckpointingLayer))BaseModelOutputWithPastAndCrossAttentions!CausalLMOutputWithCrossAttentions SequenceClassifierOutputWithPast)PreTrainedModel)Conv1D find_pruneable_heads_and_indicesprune_conv1d_layer)auto_docstringlogging torch_float)ImageGPTConfigc Cszddl}ddl}Wnty2tdYn0tj|}td||j |}g}g}|D]D\} } td| d| |j || } | | | | qht||D]&\} } | dd} | d} td d | Ds| d d vrtd d| q|} | d dvr,t| d} | D]} |d| rR|d| }n| g}|ddkst|ddkrt| d} n|ddkrt| d} n|ddks|ddkrt| |d} t| d} n|ddvrt| d} t| d} nt| dkr:| ddkr:|ddkr:t| |d} t| d} nV|ddkr^t| d } t| d} n2|dd!krt| d} t| d} nt| |d} t|d"kr0t|d}| |} q0t| dkr| ddksT| d dksT| d d!ksT| d dkrnVz| j| jksJWn>tyR}z$|j| j| jf7_WYd}~n d}~00td#| | d d$krt| |j|jj| jddd|jf<q| d d%krt| |j|jj| jdd|jd"|jf<q| d d&kr"t| |j|jj| jddd"|jdf<qt| dkrf| ddkrf| d"dkrft| |j|j| _q| d dkrt| | _q| d dkrt| | jd|j dddf<q| d d!krt| | jd <qt| | _q|S)'z0 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))Zadam_vZadam_mZAdamWeightDecayOptimizerZAdamWeightDecayOptimizer_1Z global_stepN).0nrrj/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/imagegpt/modeling_imagegpt.py Qsz.load_tf_weights_in_imagegpt..)Z_stepz Skipping {})wtet transformerz [A-Za-z]+\d+z(\d+)wgweightbbiaswpewte)q_projk_projv_projc_attnr rattnc_projr$lm_headsoszInitialize PyTorch weight r-r.r/)!reZ tensorflow ImportErrorloggererrorospathabspathinfotrainZlist_variablesZ load_variableappendsqueezezipsplitanyformatjoingetattr fullmatchlenintshapeAssertionErrorargstorchZ from_numpyreshapen_embdTdata vocab_size)modelconfigZimagegpt_checkpoint_pathr6tfZtf_pathZ init_varsnamesZarraysnamerJarrayZpointerZm_nameZ scope_namesnumerrr!load_tf_weights_in_imagegpt0s               *       F."2*$r[cs>eZdZdeeedfdd ZejejdddZ Z S) ImageGPTLayerNormh㈵>) hidden_sizeepscs&t||_tt||_dSN)super__init__r_r ParameterrMTensorr()selfr^r_ __class__rr!rbs zImageGPTLayerNorm.__init__)tensorreturncCs4|ttjt|ddd|j}||j}|S)Nr#T)ZaxisZkeepdim)rMsqrtmeanZsquarer_r()rerhrrr!forwards& zImageGPTLayerNorm.forward)r]) __name__ __module__ __qualname__tuplerIfloatrbrMrdrl __classcell__rrrfr!r\sr\c seZdZdeeeedfdd ZddZddd Zdd d Z d d Z ddZ de j eeee j ee j ee j ee j eeeeee j ed ddZZS)ImageGPTAttentionFNis_cross_attention layer_idxcsFt|j}|jdttj||ftjddd||dd|jdt ddd|j |_ |j |_ |j |j |_|j |_|j|j |j krtd|j d |j d |j|_||_|j|_||_|j|_|jrtd |j |j |_t|j |j |_ntd |j |j |_t|j |j |_t|j|_t|j|_t |_!dS) Nr*dtyperF) persistentZ masked_biasgz=`embed_dim` must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).r5r )"rarbmax_position_embeddingsZregister_bufferrMZtrilonesboolviewrhr^ embed_dimZnum_attention_heads num_headshead_dim split_size ValueErrorscale_attn_weightsruscale_attn_by_inverse_layer_idxrvreorder_and_upcast_attnrr0q_attnr2rDropoutZ attn_pdrop attn_dropout resid_pdrop resid_dropoutset pruned_heads)rerTrurvZ max_positionsrfrr!rbsB  zImageGPTAttention.__init__cCst|dkrdSt||j|j|j\}}t|||j|d|jg}t|j |dd|_ t|j |dd|_ |j|j|jt||_|jt||_|j ||_dS)Nrr5rdim) rHrrrrrMcatrrr0r2union)reheadsindexZ index_attnrrr! prune_headss  zImageGPTAttention.prune_headsc Cst||dd}|jr0|t|dd}|jrH|t|jd}|j s|d|d}}|j dddd|||d|f} t |j j } tj| |j |jd} t| || }|dur||}tjdd|}||j }||}|dur||}t||} | |fS)Nr#?rrxdevicer)rMmatmul transposerrsizerrqrvrur*finforxminrhrwhererSoftmaxtyper) requerykeyvalueattention_mask head_mask attn_weights query_length key_length causal_mask mask_value attn_outputrrr!_attns(&    zImageGPTAttention._attncCs|\}}}} |\} } } } tj|||| tj|jd} d} |jr\| t|dd} |jrt| t|jd} tj |jj ddb| d|| | dd d| | }}tj | ||d | d } | |||| } Wdn1s0Y|jsl|d|d}}|jdddd|||d|f}t| jj}tj|| j| jd}t|| |} |dur~| |} tjdd | } | jtjkrtd | |j} || } |dur| |} t| |}|| fS) Nr?r#rrF)Zenabledrr)betaalpharzDError with upcasting, attn_weights does not have dtype torch.float32)rrMemptyZfloat32rrrqrrvZautocastrrNrZbaddbmmrur*rrxrrhrrr RuntimeErrorrr)rerrrrrbszrZ q_seq_lenZdk_Z k_seq_lenrZ scale_factorqkrrrrrrrr!_upcast_and_reordered_attn s:&.&     z,ImageGPTAttention._upcast_and_reordered_attncCs2|dd||f}|j|}|ddddS)zJ Splits hidden_size dim into attn_head_size and num_heads Nr#rr5rr )rr}permutererhrZattn_head_sizeZ new_shaperrr! _split_heads>s zImageGPTAttention._split_headscCs8|dddd}|dd||f}||S)zS Merges attn_head_size dim and num_attn_heads dim into hidden_size rr5rr Nr)r contiguousrr}rrrr! _merge_headsFszImageGPTAttention._merge_heads hidden_states layer_pastrrencoder_hidden_statesencoder_attention_mask use_cacheoutput_attentionscache_positionric Cs|du} |j\} } } |durLt|trH|j|j}| r@|j}qL|j}n|}| rT|n|}| rt|dspt d|dur|r| |}|j |jj }|j |jj }n\| |}||j|jdd\}}|| d|j|jdd}|| d|j|jdd}nT||j|jdd\}}}|| d|j|jdd}|| d|j|jdd}|dur| sj| nd} ||||jd| i\}}| rd|j|j<|| | |j|jdd}|jr||||||\}}n||||||\}}|||j|j}||}||}||fS) NrzIf class is used as cross attention, the weights `q_attn` have to be defined. Please make sure to instantiate class with `ImageGPTAttention(..., is_cross_attention=True)`.r5rr#rrT)rJ isinstancer is_updatedgetrvZcross_attention_cacheZself_attention_cachehasattrrrZlayerskeysvaluesr0rBrr}rrrupdaterrrrr2r)rerrrrrrrrrrurZseq_lenrrZcurr_past_key_valueZcurrent_statesrrrrrrrr!rlNsN            zImageGPTAttention.forward)FN)NN)NN)NNNNNFFN)rmrnrorr|rIrbrrrrrrMrdr rprlrrrrrfr!rss4+ & 4 rscs0eZdZfddZejejdddZZS) ImageGPTMLPcsFt|j}t|||_t|||_t|j|_t |j |_ dSr`) rarbr^rc_fcr2r Zactivation_functionactrrrdropout)reZintermediate_sizerTr~rfrr!rbs     zImageGPTMLP.__init__)rricCs,||}||}||}||}|Sr`)rrr2r)rerrrr!rls     zImageGPTMLP.forward)rmrnrorbrMrdrlrrrrrfr!rs rc sleZdZdfdd Zd ejeeeejeejeejeejeeeeeeje d ddZ Z S) ImageGPTBlockNcst|j}|jdur |jnd|}t||jd|_t||d|_t||jd|_ |j r|t|d|d|_ t||jd|_ t |||_dS)Nr_rvTrt)rarbr^Zn_innerr\layer_norm_epsilonln_1rsr1ln_2add_cross_attentioncrossattention ln_cross_attnrmlp)rerTrvr^Z inner_dimrfrr!rbs zImageGPTBlock.__init__Frc  Cs|} ||}|j||||||| d} | d} | dd} | | }|durt|dsdtd|d|} ||}|j|||||||| d}|d} | | }| |dd} |} ||}||}| |}|f| S)N)rrrrrrrrrz'If `encoder_hidden_states` are passed, z` has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`)rrrrrrr)rr1rrrrrr)rerrrrrrrrrZresidualZ attn_outputsroutputsZcross_attn_outputsZfeed_forward_hidden_statesrrr!rlsN          zImageGPTBlock.forward)N)NNNNNFFN) rmrnrorbrMrdrr r|rprlrrrrrfr!rs*rcsDeZdZUeed<eZdZdZdZ dgZ fddZ dd Z Z S) ImageGPTPreTrainedModelrTr% input_idsTrcstj|i|dSr`)rarb)reinputskwargsrfrr!rbsz ImageGPTPreTrainedModel.__init__cCst|tjtfr>|jjjd|jjd|j dur|j j nXt|tj r~|jjjd|jjd|j dur|jj|j  nt|t r|jjd|D]>\}}d|vrd|vr|jjd|jjtd|jjdqdS)zInitialize the weights.g)rkZstdNrr2r(r5)rrLinearrr(rQZnormal_rTZinitializer_ranger*Zzero_ EmbeddingZ padding_idxr\Zfill_Znamed_parametersmathrjn_layer)remodulerWprrr! _init_weightss    z%ImageGPTPreTrainedModel._init_weights)rmrnror__annotations__r[Zload_tf_weightsZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingZ_no_split_modulesrbrrrrrrfr!rs  rcseZdZedfdd ZddZddZdd Zede e j e e 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 e e e e e e e e e j eee efd d d ZZS) ImageGPTModelrTcstj|_tj|j|_tj|j|_ t j |_ t fddtjD|_t|jjd|_d|_d|_d|_|dS)Ncsg|]}t|dqS)r)r)rirrr! %z*ImageGPTModel.__init__..rF)rarbr^r~rrrRr,rzr+rZ embd_pdropdropZ ModuleListrangeZnum_hidden_layershr\rln_fmodel_parallel device_mapgradient_checkpointing post_initrerTrfrr!rbs  zImageGPTModel.__init__cCs|jSr`r,)rerrr!get_input_embeddings/sz"ImageGPTModel.get_input_embeddingscCs ||_dSr`r)reZnew_embeddingsrrr!set_input_embeddings2sz"ImageGPTModel.set_input_embeddingscCs(|D]\}}|j|j|qdS)zv Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} N)itemsrr1r)reZheads_to_prunelayerrrrr! _prune_heads5szImageGPTModel._prune_headsN)rpast_key_valuesrtoken_type_ids position_idsr inputs_embedsrrrroutput_hidden_states return_dictrrric$ Ks| dur | n|jj} | dur | n|jj} | dur4| n|jj} | durH| n|jj} |durj|durjtdnd|dur||||}|d|d}|j d}n,|dur|dd}|j d}ntd|dur|j n|j }|j r|j r| rt dd} | r0|dur0tt|jdt|jd}| rVt|trVt d t|}|durh|n|}|dur|d|d}|durtj||d|tj|d }|d}|dur|dkrtd ||d}|ddddddf}|j|jd }d |t|jj}|jjrh|durh|\}}}||f}| dur\tj||d} || } nd} | ||jj!}|dur|"|}|#|}|||j }|dur|"|}||}|$|}||df}| rdnd}| r|jjrdnd}| rdnd}t%|j&D]\}} |j'rntj()|j |durT||j }t|tj*rn||j }| r~||f}| ||||||| | | |d }!|!d}| r||!df}|jjr||!df}|j'r|j+,D]B\}"}#||#dkrdt-|"|j.kr|dt-|"d}qq|/|}|j|}| rN||f}| sptdd|||||fDSt0|||||dS)a 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.get_seq_length()` (`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 [`AutoImageProcessor`]. See [`ImageGPTImageProcessor.__call__`] for details. Examples: ```python >>> from transformers import AutoImageProcessor, ImageGPTModel >>> from PIL import Image >>> import requests >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> image_processor = AutoImageProcessor.from_pretrained("openai/imagegpt-small") >>> model = ImageGPTModel.from_pretrained("openai/imagegpt-small") >>> inputs = image_processor(images=image, return_tensors="pt") >>> outputs = model(**inputs) >>> last_hidden_states = outputs.last_hidden_state ```NzDYou cannot specify both input_ids and inputs_embeds at the same timer#rz5You have to specify either input_ids or inputs_embedszZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...FrzPassing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.58.0. You should pass an instance of `EncoderDecoderCache` instead, e.g. `past_key_values=EncoderDecoderCache.from_legacy_cache(past_key_values)`.rz$batch_size has to be defined and > 0rwr)rr)rrrrrr5zcuda:css|]}|dur|VqdSr`r)rvrrr!r"sz(ImageGPTModel.forward..)Zlast_hidden_staterr attentionscross_attentions)1rTrrruse_return_dictrZ%warn_if_padding_and_no_attention_maskrr}rJrrZtrainingr8Z warning_oncer r rrpZfrom_legacy_cacheZget_seq_lengthrMZarangelongZ unsqueezetorxrrrr{Zinvert_attention_maskZ get_head_maskrr,r+r enumeraterrcudaZ set_devicerdrrstrZ last_devicerr)$rerrrrrrrrrrrrrrrZ input_shapeZ batch_sizerZ past_lengthZencoder_batch_sizeZencoder_sequence_lengthrZencoder_hidden_shapeZposition_embedsrZtoken_type_embedsZ output_shapeZall_self_attentionsZall_cross_attentionsZall_hidden_statesrblockrrrrrr!rl<s0                         "    zImageGPTModel.forward)NNNNNNNNNNNNNN)rmrnrorrbrrrrrrMrdrpr|rrrrlrrrrrfr!rsJ rz The ImageGPT Model transformer with a language modeling head on top (linear layer with weights tied to the input embeddings). )Z custom_introcseZdZdgZedfdd Zed eej ee e ej eej eej eej eej eej eej eej eej ee ee ee ee eej e e e efdddZZS) ImageGPTForCausalImageModelingzlm_head.weightrcsHt|t||_tj|j|jddd|_d|_ d|_ | dS)NrFr*) rarbrr%rrrOrRr3rrrrrfrr!rbs   z'ImageGPTForCausalImageModeling.__init__N)rrrrrrrrrlabelsrrrrrrricKs|dur |n|jj}|j||||||||| | | | ||d}|d}||}d}| dur|dddddf}| dddf}t}||d|d|d}|s|f|dd}|dur|f|S|St|||j |j |j |j dS)a& 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.get_seq_length()` (`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 [`AutoImageProcessor`]. See [`ImageGPTImageProcessor.__call__`] for details. labels (`torch.LongTensor` of shape `(batch_size, input_ids_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]` Examples: ```python >>> from transformers import AutoImageProcessor, ImageGPTForCausalImageModeling >>> import torch >>> import matplotlib.pyplot as plt >>> import numpy as np >>> image_processor = AutoImageProcessor.from_pretrained("openai/imagegpt-small") >>> model = ImageGPTForCausalImageModeling.from_pretrained("openai/imagegpt-small") >>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu") >>> model.to(device) # doctest: +IGNORE_RESULT >>> # unconditional generation of 8 images >>> batch_size = 4 >>> context = torch.full((batch_size, 1), model.config.vocab_size - 1) # initialize with SOS token >>> context = context.to(device) >>> output = model.generate( ... input_ids=context, max_length=model.config.n_positions + 1, temperature=1.0, do_sample=True, top_k=40 ... ) >>> clusters = image_processor.clusters >>> height = image_processor.size["height"] >>> width = image_processor.size["width"] >>> samples = output[:, 1:].detach().cpu().numpy() >>> samples_img = [ ... np.reshape(np.rint(127.5 * (clusters[s] + 1.0)), [height, width, 3]).astype(np.uint8) for s in samples ... ] # convert color cluster tokens back to pixels >>> f, axes = plt.subplots(1, batch_size, dpi=300) >>> for img, ax in zip(samples_img, axes): # doctest: +IGNORE_RESULT ... ax.axis("off") ... ax.imshow(img) ```N) rrrrrrrrrrrrrr.r#r)losslogitsrrrr) rTrr%r3rrr}rrrrrr)rerrrrrrrrrr rrrrrrtransformer_outputsrZ lm_logitsr Z shift_logitsZ shift_labelsloss_fctoutputrrr!rlsHG z&ImageGPTForCausalImageModeling.forward)NNNNNNNNNNNNNNN)rmrnroZ_tied_weights_keysrrbrrrMrdrpr|rrrrlrrrrrfr!r  sJ  r z The ImageGPT Model transformer with an image classification head on top (linear layer). [`ImageGPTForImageClassification`] average-pools the hidden states in order to do the classification. cseZdZedfdd Zedeejee e ejeejeejeejeejeejeejee ee ee ee e e e e fdddZZS) ImageGPTForImageClassificationrcs@t||j|_t||_tj|j|jdd|_| dS)NFr ) rarb num_labelsrr%rrrOscorerrrfrr!rbs   z'ImageGPTForImageClassification.__init__N)rrrrrrrr rrrrrric  Ks| dur | n|jj} |j|||||||| | | | d }|d}|jdd}||}d}|dur>|jjdur|jdkrd|j_n4|jdkr|jtj ks|jtj krd|j_nd|j_|jjdkrt }|jdkr|| | }n |||}nN|jjdkr t }||d |j|d }n|jjdkr>t}|||}| sn|f|dd}|durj|f|S|St|||j|j|jd S) ay 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.get_seq_length()` (`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 [`AutoImageProcessor`]. See [`ImageGPTImageProcessor.__call__`] for details. 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). Examples: ```python >>> from transformers import AutoImageProcessor, ImageGPTForImageClassification >>> from PIL import Image >>> import requests >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" >>> image = Image.open(requests.get(url, stream=True).raw) >>> image_processor = AutoImageProcessor.from_pretrained("openai/imagegpt-small") >>> model = ImageGPTForImageClassification.from_pretrained("openai/imagegpt-small") >>> inputs = image_processor(images=image, return_tensors="pt") >>> outputs = model(**inputs) >>> logits = outputs.logits ```N) rrrrrrrrrrrrrZ regressionZsingle_label_classificationZmulti_label_classificationr#)r rrrr)rTrr%rkrZ problem_typerrxrMrrIrr@rr}rrrrr)rerrrrrrrr rrrrrrrZpooled_hidden_statesrr rrrrr!rls\2       "     z&ImageGPTForImageClassification.forward) NNNNNNNNNNNN)rmrnrorrbrrrMrdrpr|rrrrlrrrrrfr!rs<  r)r rrrr[)6__doc__rr:typingrrrrMZtorch.utils.checkpointrZtorch.nnrrrZ activationsr Z cache_utilsr r r Z generationrZmodeling_layersrZmodeling_outputsrrrZmodeling_utilsrZ pytorch_utilsrrrutilsrrrZconfiguration_imagegptrZ get_loggerrmr8r[Moduler\rsrrrrr r__all__rrrr!sL       l mM&qt