o Zh*@sdZddlZddlZddlZddlmZmZmZmZddl Z ddl Z ddl m Z ddl m Z ddlmZmZmZddlmZdd lmZdd lmZmZmZdd lmZdd lmZmZmZdd l m!Z!m"Z"m#Z#ddl$m%Z%e"&e'Z(ddZ)Gddde j*Z+Gddde j*Z,Gddde j*Z-Gddde j*Z.e!GdddeZ/e!Gddde/Z0e!ddGd d!d!e/eZ1e!d"dGd#d$d$e/Z2gd%Z3dS)&zPyTorch OpenAI ImageGPT model.N)AnyOptionalTupleUnion)nn)autocast)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)GenerationMixin))BaseModelOutputWithPastAndCrossAttentions!CausalLMOutputWithCrossAttentions SequenceClassifierOutputWithPast)PreTrainedModel)Conv1D find_pruneable_heads_and_indicesprune_conv1d_layer)auto_docstringlogging torch_float)ImageGPTConfigc Csz ddl}ddl}Wn tytdwtj|}td ||j |}g}g}|D] \} } td | | |j || } | | | | q2t||D]\} } | dd} | d} tdd | Dsw| d d vrtd d| qX|} | d d vrt| d} | D]} |d| r|d| }n| g}|ddks|ddkrt| d} n|ddkrt| d} nw|ddks|ddkrt| |d} t| d} n^|ddvrt| d} t| d} nMt| dkr| ddkr|ddkrt| |d} t| d} n+|ddkr t| d} t| d} n|dd kr2t| d} t| d} nt| |d} t|d!krJt|d}| |} qt| dkrY| ddksn| d dksn| d d ksn| d dkron%z | j| jksyJWnty}z |j| j| jf7_d}~wwtd" | | d d#krt| |j|jj| jddd|jf<qX| d d$krt| |j|jj| jdd|jd!|jf<qX| d d%krt| |j|jj| jddd!|jdf<qXt| dkr| ddkr| d!dkrt| |j|j| _qX| d dkr+t| | _qX| d dkrDt| | jd|j dddf<qX| d d krTt| | jd <qXt| | _qX|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 {}z"Loading TF weight {} with shape {}/css|]}|dvVqdS))Zadam_vZadam_mZAdamWeightDecayOptimizerZAdamWeightDecayOptimizer_1Z global_stepN).0nrr]/var/www/auris/lib/python3.10/site-packages/transformers/models/imagegpt/modeling_imagegpt.py Qs  z.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 ImportErrorloggererrorospathabspathinfoformattrainZlist_variablesZ load_variableappendsqueezezipsplitanyjoingetattr fullmatchlenintshapeAssertionErrorargstorchZ from_numpyreshapen_embdTdata vocab_size)modelconfigZimagegpt_checkpoint_pathr4tfZtf_pathZ init_varsnamesZarraysnamerHarrayZpointerZm_nameZ scope_namesnumerrrload_tf_weights_in_imagegpt0s                   *      F.2*$rYcsBeZdZd deedeffdd Zdejdejfdd Z Z S) ImageGPTLayerNormh㈵> hidden_sizeepscs&t||_tt||_dSN)super__init__r]r ParameterrKTensorr&)selfr\r] __class__rrr`s zImageGPTLayerNorm.__init__tensorreturncCs4|ttjt|ddd|j}||j}|S)Nr!T)ZaxisZkeepdim)rKsqrtmeanZsquarer]r&)rcrfrrrforwards& zImageGPTLayerNorm.forward)r[) __name__ __module__ __qualname__rrGfloatr`rKrbrj __classcell__rrrdrrZsrZcseZdZddeedeeffdd ZddZdd d Zdd d Z d dZ ddZ       dde j deedee j dee j dee j dee j deedeedefddZZS)ImageGPTAttentionFNis_cross_attention layer_idxcsFt|j}|jdttj||ftjddd||dd|jdt ddd|j |_ |j |_ |j |j |_|j |_|j|j |j krUtd|j d |j d |j|_||_|j|_||_|j|_|jr}td |j |j |_t|j |j |_n td |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).r3r )"r_r`max_position_embeddingsZregister_bufferrKZtrilonesboolviewrfr\ embed_dimZnum_attention_heads num_headshead_dim split_size ValueErrorscale_attn_weightsrqscale_attn_by_inverse_layer_idxrrreorder_and_upcast_attnrr.q_attnr0rDropoutZ attn_pdrop attn_dropout resid_pdrop resid_dropoutset pruned_heads)rcrRrqrrZ max_positionsrdrrr`sB   zImageGPTAttention.__init__cCst|dkrdSt||j|j|j\}}t|||j|d|jg}t|j |dd|_ t|j |dd|_ |j|j|jt||_|jt||_|j ||_dS)Nrr3rdim) rFrr{r|rrKcatr}rr.r0union)rcheadsindexZ index_attnrrr prune_headss  zImageGPTAttention.prune_headsc Cst||dd}|jr|t|dd}|jr$|t|jd}|j s]|d|d}}|j dddd|||d|f} t |j j } tj| |j |jd} t| || }|dure||}tjdd|}||j }||}|dur||}t||} | |fS)Nr!?rrtdevicer)rKmatmul transposerrsizerrnrrrqr(finfortminrfrwhererSoftmaxtyper) rcquerykeyvalueattention_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} |jr:| t|jd} t dd1| d|| | dd d| | }}tj | ||d | d } | |||| } Wdn1sswY|j s|d|d}}|jdddd|||d|f}t| jj}tj|| j| jd}t|| |} |dur| |} tjdd | } | jtjkrtd | |j} || } |dur| |} t| |}|| fS) Nr?r!rrF)enabledrr)betaalpharzDError with upcasting, attn_weights does not have dtype torch.float32)rrKemptyZfloat32rrrnrrrrrLrZbaddbmmrqr(rrtrrfrrr RuntimeErrorrrr)rcrrrrrZbszr{Z 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!rr3rr )rrypermutercrfr{Zattn_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 rr3rr Nr)r contiguousrryrrrr _merge_headsFs zImageGPTAttention._merge_heads hidden_states layer_pastrrencoder_hidden_statesencoder_attention_mask use_cacheoutput_attentionsrgc CsT|dur"t|ds td||} ||j|jdd\} } |}n||j|jdd\} } } || |j|j} || |j|j} || |j|j} |dure|\} } t j | | fdd} t j | | fdd} |durn| | f}nd}|j r| | | | ||\}}n | | | | ||\}}|||j|j}||}||}||f}|r||f7}|S)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)`.r3rrT)hasattrr~rr.rAr}rr{r|rKrrrrrr0r)rcrrrrrrrrrrrZpast_keyZ past_valueZpresentrroutputsrrrrjNs:      zImageGPTAttention.forward)FN)NNNNNNNFF)rkrlrmrrxrGr`rrrrrrKrbtuplerjrorrrdrrpsB +  &4   rpcs2eZdZfddZdejdejfddZZS) ImageGPTMLPcsFt|j}t|||_t|||_t|j|_t |j |_ dSr^) r_r`r\rc_fcr0r Zactivation_functionactrrrdropout)rcZintermediate_sizerRrzrdrrr`s    zImageGPTMLP.__init__rrgcCs,||}||}||}||}|Sr^)rrr0r)rcrrrrrjs    zImageGPTMLP.forward)rkrlrmr`rKrbrjrorrrdrrs rcseZdZdfdd Z       ddejdeedeejdeejd eejd eejd eed eed efddZ Z S) ImageGPTBlockNcst|j}|jdur|jnd|}t||jd|_t||d|_t||jd|_ |j r>t|d|d|_ t||jd|_ t |||_dS)Nr]rrT)rqrr)r_r`r\Zn_innerrZlayer_norm_epsilonln_1rpr/ln_2add_cross_attentioncrossattention ln_cross_attnrmlp)rcrRrrr\Z inner_dimrdrrr`s zImageGPTBlock.__init__Frrrrrrrrrgc Cs|} ||}|j||||||d} | d} | dd} | | }|durSt|ds1td|d|} ||}|j||||||d} | d} | | }| | dd} |} ||}||}| |}|f|rl| } | S| dd} | S) N)rrrrrrrrz'If `encoder_hidden_states` are passed, z` has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`)rrrrrr3)rr/rr~rrrr)rcrrrrrrrrZresidualZ attn_outputsrrZcross_attn_outputsZfeed_forward_hidden_statesrrrrjsN       zImageGPTBlock.forwardr^r) rkrlrmr`rKrbrrxrrjrorrrdrrs8  rcs>eZdZeZeZdZdZdZ dgZ fddZ ddZ Z S) ImageGPTPreTrainedModelr# input_idsTrcstj|i|dSr^)r_r`)rcZinputskwargsrdrrr`sz ImageGPTPreTrainedModel.__init__cCst|tjtfr|jjjd|jjd|j dur|j j n,t|tj r?|jjjd|jjd|j dur>|jj|j  n t|t rK|jjd|D]\}}d|vrnd|vrn|jjd|jjtd|jjdqOdS)zInitialize the weights.g)riZstdNrr0r&r3) isinstancerLinearrr&rOZnormal_rRZinitializer_ranger(Zzero_ EmbeddingZ padding_idxrZZfill_Znamed_parametersmathrhn_layer)rcmodulerUprrr _init_weightss"     &z%ImageGPTPreTrainedModel._init_weights)rkrlrmrZ config_classrYZload_tf_weightsZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingZ_no_split_modulesr`rrorrrdrrs rc"seZdZdeffdd ZddZddZdd Ze dd e e j d e e e e j d e e j de e j de e j de e j de e j de e j de e j de e de e de e de e dedee effddZZS) ImageGPTModelrRcstj|_tj|j|_tj|j|_ t j |_ t fddtjD|_t|jjd|_d|_d|_d|_|dS)Ncsg|]}t|dqS)r)r)rirRrr sz*ImageGPTModel.__init__..rF)r_r`r\rzrrrPr*rvr)rZ embd_pdropdropZ ModuleListrangeZnum_hidden_layershrZrln_fmodel_parallel device_mapgradient_checkpointing post_initrcrRrdrrr` s   zImageGPTModel.__init__cC|jSr^r*rcrrrget_input_embeddingsz"ImageGPTModel.get_input_embeddingscC ||_dSr^rrcZnew_embeddingsrrrset_input_embeddings z"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)itemsrr/r)rcZheads_to_prunelayerrrrr _prune_heads"szImageGPTModel._prune_headsNrpast_key_valuesrtoken_type_ids position_idsr inputs_embedsrrrroutput_hidden_states return_dictrrgc$ sd|vrtdt|durtd|d}| dur| n|jj} | dur'| n|jj} | dur1| n|jj} | dur;| n|jj } |durK|durKtd|durg| ||| }| d|d}|j d}n|dury| dd}|j d}ntd|dur|jn|j}|dur| d|d}|durd}tdgt|j}n |dd d }|durtj||d|tj|d }|d}|dur|dkrtd | |d}|ddddddf}|j|jd }d |t|jj}|jjr|dur| \}}}||f}| durtj||d} || } nd} |||jj}|dur/||}| |}|||j|durJ||}||!| df}|j"ri|j#ri| rit$%dd} | rndnd}| rudnd}| r|jjrdnd}| rdnd}t&t'|j|D]\}\}} |j(rtj)*j| durtfdd| D} |dur|j}t+|tj,r|j}| r|f}|j"r|j#r|-|j.d||||| | | }!n|| ||||| | | d}!|!d| dur ||!df}| r*||!| rdndf}|jjr*||!| r%dndf}|j(rU|j/0D]!\}"}#||#dkrSdt1|"|j2krSdt1|"dq3q|3j || ri|f}| sztdd||||fDSt4||||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[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 [`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 ``` pixel_values`The `pixel_values` argument is deprecated and will be removed in v4.47, use `input_ids` instead.N_You cannot pass both `pixel_values` and `input_ids`. Please make sure to only pass `input_ids`.zDYou cannot specify both input_ids and inputs_embeds at the same timer!rz5You have to specify either input_ids or inputs_embedsrrz$batch_size has to be defined and > 0rsr)rzZ`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...Frc3s|] }|jVqdSr^)torrZ past_staterrrr sz(ImageGPTModel.forward..)rrrrrrrTrr3r zcuda:css|] }|dur|VqdSr^r)rvrrrr s)Zlast_hidden_staterr attentionscross_attentions)5warningswarn FutureWarningr~poprRrrruse_return_dictZ%warn_if_padding_and_no_attention_maskrryrHrrrFrrKZarangelongZ unsqueezerrtrrrrwZinvert_attention_maskZ get_head_maskrr*r)rrZtrainingr6Z warning_once enumerater@rcudaZ set_devicerrbZ_gradient_checkpointing_func__call__rrstrZ last_devicerr)$rcrrrrrrrrrrrrrrZ input_shapeZ batch_sizerZ past_lengthZencoder_batch_sizeZencoder_sequence_lengthrZencoder_hidden_shapeZposition_embedsZtoken_type_embedsZ output_shapeZpresentsZall_self_attentionsZall_cross_attentionsZall_hidden_statesrblockrrrrrrrrj)s/                      "    zImageGPTModel.forward) NNNNNNNNNNNNN)rkrlrmrr`rrrrrrKrbrrxrrrrjrorrrdrrsd      rz The ImageGPT Model transformer with a language modeling head on top (linear layer with weights tied to the input embeddings). )Z custom_introc$s0eZdZdgZdeffdd ZddZddZe dd e e j d e e e e j d e e j d e e j de e j de e j de e j de e j de e j de e j de e de e de e de e dedee eff ddZed e e e j de j de e e j fddZZS) ImageGPTForCausalImageModelingzlm_head.weightrRcsHt|t||_tj|j|jddd|_d|_ d|_ | dS)NrFr() r_r`rr#rrrMrPr1rrrrrdrrr`s   z'ImageGPTForCausalImageModeling.__init__cCrr^r1rrrrget_output_embeddings%rz4ImageGPTForCausalImageModeling.get_output_embeddingscCrr^rrrrrset_output_embeddings(rz4ImageGPTForCausalImageModeling.set_output_embeddingsNrrrrrrrrrlabelsrrrrrrgcKsd|vrtdt|durtd|d}|dur|n|jj}|j||||||||| | | | |d }|d}||}d}| durk|dddddf }| dd df }t }|| d| d| d}|s|f|d d}|dur|f|S|St |||j|j|j|jd S) 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[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 [`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) ```rrNr) rrrrrrrrrrrrr.r!r)losslogitsrrrr)rrrr~rrRrr#r1rr ryrrrrrr)rcrrrrrrrrrrrrrrrtransformer_outputsrZ lm_logitsrZ shift_logitsZ shift_labelsloss_fctoutputrrrrj+sZF  z&ImageGPTForCausalImageModeling.forwardbeam_idxcstfdd|DS)a 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. c3s&|]}tfdd|DVqdS)c3s$|] }|d|jVqdS)rN)Z index_selectrrrrrrr s"zJImageGPTForCausalImageModeling._reorder_cache...Nr)rrrrrr s  z@ImageGPTForCausalImageModeling._reorder_cache..r)rrrrr_reorder_caches z-ImageGPTForCausalImageModeling._reorder_cache)NNNNNNNNNNNNNN)rkrlrmZ_tied_weights_keysrr`rrrrrKrbrrxrrrrj staticmethodrrorrrdrrsz        }  rz 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. c seZdZdeffdd Ze            ddeejdee e ejdeejdeejd eejd eejd eejd eejd ee dee dee dee de de e e ffddZZS)ImageGPTForImageClassificationrRcs@t||j|_t||_tj|j|jdd|_| dS)NFr) r_r` num_labelsrr#rrrMscorerrrdrrr`s   z'ImageGPTForImageClassification.__init__Nrrrrrrrrrrrrrrgc  Ksd| vrtdt|durtd| d}| dur| n|jj} |j|||||||| | | | d }|d}|jdd}| |}d}|dur|jj durp|j dkrVd |j_ n|j dkrl|j t jksg|j t jkrld |j_ nd |j_ |jj d krt}|j dkr|||}n+|||}n%|jj d krt}||d |j |d }n|jj d krt}|||}| s|f|dd}|dur|f|S|St|||j|j|jd S)ax 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 [`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 ```rrNr) rrrrrrrrrrrrrZ regressionZsingle_label_classificationZmulti_label_classificationr!)rrrrr)rrrr~rrRrr#rir Z problem_typerrtrKrrGr r?r ryrrrrr)rcrrrrrrrrrrrrrrrZpooled_hidden_statesrrrrrrrrjsp2      "       z&ImageGPTForImageClassification.forward) NNNNNNNNNNNN)rkrlrmrr`rrrKrbrrxrrrrjrorrrdrrsX       r)rrrrrY)4__doc__rr8rtypingrrrrrKZtorch.utils.checkpointrZtorch.cuda.amprZtorch.nnrr r Z activationsr Z generationr Zmodeling_outputsrrrZmodeling_utilsrZ pytorch_utilsrrrutilsrrrZconfiguration_imagegptrZ get_loggerrkr6rYModulerZrprrrrrr__all__rrrrsR       l \K& "