o Zhg@s$dZddlmZmZmZmZmZddlZddlZddlm Z ddl m Z m Z m Z ddlmZddlmZmZdd lmZdd lmZdd lmZmZdd lmZmZmZmZdd lm Z m!Z!ddl"m#Z#ddl$m%Z%m&Z&m'Z'm(Z(m)Z)ddl*m+Z+e(rddl,m-Z-ddl.m/Z/er e)0e1Z2Gddde j3Z4 d5de j5dej6dej6dej6deej6de7de7fddZ8Gd d!d!e j5Z9Gd"d#d#e j5Z:e&Gd$d%d%e!Z;Gd&d'd'e;ZGd,d-d-e;eZ?e&d.d/Gd0d1d1e;Z@e&Gd2d3d3e;ZAgd4ZBdS)6zPyTorch OPT model.)CallableListOptionalTupleUnionN)nn)BCEWithLogitsLossCrossEntropyLossMSELoss)ACT2FN)Cache DynamicCache)GenerationMixin)AttentionMaskConverter)FlashAttentionKwargsis_flash_attn_available)BaseModelOutputWithPastCausalLMOutputWithPastQuestionAnsweringModelOutput SequenceClassifierOutputWithPast)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack) LossKwargsauto_docstringcan_return_tupleis_torch_flex_attn_availablelogging) OPTConfig) BlockMask)make_flex_block_causal_maskcsReZdZdZdedeffdd Z  d dejd ed eejffd d Z Z S)OPTLearnedPositionalEmbeddingzN This module learns positional embeddings up to a fixed maximum size. num_embeddings embedding_dimcsd|_t||j|dSN)offsetsuper__init__)selfr$r% __class__S/var/www/auris/lib/python3.10/site-packages/transformers/models/opt/modeling_opt.pyr*;sz&OPTLearnedPositionalEmbedding.__init__rNattention_maskpast_key_values_length position_idscsL|durtj|dd}||d}|dd|df}t||jS)z3`input_ids_shape` is expected to be [bsz x seqlen].Nrdim)torchcumsumlongr)forwardr()r+r0r1r2r,r.r/r8As z%OPTLearnedPositionalEmbedding.forwardrN) __name__ __module__ __qualname____doc__intr*r5 LongTensorrr8 __classcell__r.r.r,r/r#6s r#modulequerykeyvaluer0scalingdropoutc Ks|t||dd|}|dur||}tjj|dtjd|j}tjj |||j d}t||} | dd } | |fS)N)r4dtypeptrainingrr') r5matmul transposer functionalZsoftmaxZfloat32torJrGrM contiguous) rBrCrDrEr0rFrGkwargs attn_weights attn_outputr.r.r/eager_attention_forwardSs  rVcseZdZdZ ddedeeffdd Z     ddej d ee ej d eej d eej d e d eej de ej eej ee ffddZ ZS) OPTAttentionz=Multi-headed attention from 'Attention Is All You Need' paperNconfig layer_idxc st||_|j|_|j|_|j|_|j |_ ||_ |dur*t d|j jd|j|j|_d|_|j|j|jkrJtd|jd|jd|jd|_tj|j|j|j d|_tj|j|j|j d|_tj|j|j|j d|_tj|j|j|j d|_dS) NzInstantiating z without passing a `layer_idx` is not recommended and will lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` when creating this class.Tz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).g࿩bias)r)r*rX hidden_size embed_dimZnum_attention_heads num_headsZattention_dropoutrG enable_biasrYlogger warning_oncer-r:head_dimZ is_causal ValueErrorrFrLineark_projv_projq_projout_proj)r+rXrYrSr,r.r/r*ms0   zOPTAttention.__init__F hidden_statespast_key_valuer0layer_head_maskoutput_attentionscache_positionreturncKs>|\}} } |||j} | |d|j|jdd} ||} ||} | |d|j|jdd} | |d|j|jdd} |durT| | | |j d|i\} } t }|j j dkrp|j j dkrj|rjtdnt|j j }||| | | |f|js|d n|j|jd |\}}||| d}||}|sd}|||fS) z#Input shape: Batch x Time x ChannelrHrr'Nrmeagersdpaz`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.rA)rGrF)sizergrFviewr^rbrOrerfupdaterYrVrX_attn_implementationr`rarrMrGreshaperRrh)r+rirjr0rkrlrmrSZbszZtgt_len_Z query_statesZ key_statesZ value_statesZattention_interfacerUrTr.r.r/r8sF        zOPTAttention.forwardN)NNNFN)r:r;r<r=r rr>r*r5Tensorrboolr r8r@r.r.r,r/rWjs8&  rWcseZdZddedeeffdd Z       ddejdeejd eejd ee ejd ee d ee d eej deejde e de ejee ejejfffddZZS)OPTDecoderLayerNrXrYcst|j|_t||d|_|j|_|j|_t|j |_ t j |j|j d|_t j|j|j|jd|_t j|j|j|jd|_t j |j|j d|_dS)N)rXrYZelementwise_affinerZ)r)r*r\r]rW self_attndo_layer_norm_beforerGr Zactivation_function activation_fnr LayerNormlayer_norm_elementwise_affineself_attn_layer_normrdZffn_dimr_fc1fc2final_layer_norm)r+rXrYr,r.r/r*s  zOPTDecoderLayer.__init__Frir0rkrjrl use_cacher2rmrSrnc  Ks |} |jr ||}|jd|||||||d| \}} } tjj||j|jd}| |}|js4||}|j} |d| d}|} |jrJ| |}| |}| |}| |}tjj||j|jd}| || }|jss| |}|f}|r}|| f7}|r|| f7}|S)a Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`torch.FloatTensor`, *optional*): mask for attention heads in a given layer of size `(encoder_attention_heads,)`. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence.. )rirjr2r0rkrlrmrKrHNr.)r}rr|rrPrGrMshaperurqrrr~rrr)r+rir0rkrjrlrr2rmrSZresidualZself_attn_weightsZpresent_key_valueZhidden_states_shapeoutputsr.r.r/r8sH           zOPTDecoderLayer.forwardrw)NNNFFNN)r:r;r<r rr>r*r5rxrryr?rr FloatTensorr8r@r.r.r,r/rzs<    rzc@sBeZdZeZdZdZdgZdZdZ dZ dZ dZ dZ dZddZdS)OPTPreTrainedModelmodelTrzcCs|jj}t|tjr"|jjjd|d|jdur |jj dSdSt|tj rC|jjjd|d|j durA|jj|j  dSdSt|tj rX|jj d|jj dSdS)NrA)meanstdg?)rXZinit_std isinstancerrdweightdataZnormal_r[Zzero_ Embedding padding_idxrZfill_)r+rBrr.r.r/ _init_weightsAs     z OPTPreTrainedModel._init_weightsN)r:r;r<r Z config_classZbase_model_prefixZsupports_gradient_checkpointingZ_no_split_modulesZ_supports_attention_backendZ_supports_flash_attn_2Z_supports_sdpaZ_supports_flex_attnZ_supports_cache_classZ_supports_quantized_cacheZ_supports_static_cacherr.r.r.r/r3s rcs6eZdZdZdeffdd ZddZddZ d$d ee j d fd e j d e j de de f ddZ ed e j dedede jd e j def ddZe           d%dee jd ee j dee j deee jdee jdee dee dee dee dee jd ee j d eed!eeeffd"d#ZZS)& OPTDecoderz Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`OPTDecoderLayer`] Args: config: OPTConfig rXcstj|_j|_j|_j|_j|_t jj |j|_ t jj|_j jkr?t jjj dd|_nd|_j jkrTt jj jdd|_nd|_jrhjsht jjjd|_nd|_t fddtjD|_d|_|dS)NFrZr{csg|]}t|dqS))rY)rz).0irXr.r/ wsz'OPTDecoder.__init__..)r)r*rG layerdrop pad_token_idrZmax_position_embeddingsZmax_target_positions vocab_sizerrword_embed_proj_dim embed_tokensr#r\embed_positionsrd project_out project_inr}Z_remove_final_layer_normrrrZ ModuleListrangeZnum_hidden_layerslayersgradient_checkpointing post_initr+rXr,rr/r*Xs,       zOPTDecoder.__init__cC|jSrwrr+r.r.r/get_input_embeddings}zOPTDecoder.get_input_embeddingscC ||_dSrwrr+rEr.r.r/set_input_embeddings zOPTDecoder.set_input_embeddingsFr0r! input_tensorrmpast_key_valuesrlc Cs:|jjdkr|dur|dkr|SdS|jjdkr&t|tjr$t|}|S|dur.|nd}|dur7|jnd}|jjdkrO|sO|sOt j ||||j drOdS|j }|j d} |r^|} nt|tjri|j d n|| d} |j|| | |||j dd } |jjdkr|dur|jjd vr|st|j} t | | } | S) NZflash_attention_2rAZflex_attentionrFrp) inputs_embedsr1Z is_trainingrrH)sequence_length target_lengthrJrm batch_size)cudaZxpuZnpu)rXrtanyrr5rxr"get_seq_lengthZis_compileablerZ_ignore_causal_mask_sdparMrJrZget_max_cache_shape5_prepare_4d_causal_attention_mask_with_cache_positiondevicetypefinfominZ_unmask_unattended) r+r0rrmrrlpast_seen_tokensZusing_compilable_cacherJrr causal_mask min_dtyper.r.r/_update_causal_masksT           zOPTDecoder._update_causal_maskrrrJrc KsD|dur|dkr|}|St|j}tj||f|||jd}|dkr+tj|dd}|tj||jd|ddk9}|ddddddf |ddd}|dur| }|j d} |ddddddd| f|ddddddf |j} | dk} |ddddddd| f | ||ddddddd| f<|S) aM Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. Args: attention_mask (`torch.Tensor`): A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. sequence_length (`int`): The sequence length being processed. target_length (`int`): The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. dtype (`torch.dtype`): The dtype to use for the 4D attention mask. cache_position (`torch.Tensor`): Indices depicting the position of the input sequence tokens in the sequence. batch_size (`torch.Tensor`): Batch size. N)Z fill_valuerJrr)ZdiagonalrrHr)r4r5rrfullrZtriuarangeruexpandclonerrQZ masked_fill) r0rrrJrmrrSrrZ mask_lengthZ padding_maskr.r.r/rs,  $ 6  z@OPTDecoder._prepare_4d_causal_attention_mask_with_cache_positionN input_ids head_maskrroutput_hidden_states return_dictr2rSrnc  Ks|dur|n|jj}|dur|n|jj}|dur|n|jj}| dur$| n|jj} |du|duAr4td|jrC|jrC|rCt dd}|durP| d|j d}|durY| |}d} |rrt |tsrd} t|}|durrt d|durz|nd}| durtj|||j d |jd } |dur||j d }tj|j d||jd }|||| ||}| durtj|d d } | |d } | dd|df} |j||| d }|jdur||}|||j}|rd nd}|rd nd}d}t|gdgD]*\}}|dur$|dt|jkr$td|dt|jd|ddqt |jD]q\}}|r6||f7}|jrGt!g}||j"krGq*|jrg|jrg|#|j$|||dur^||ndd||| | }n||f|| |duru||nd|||| d| }|d}|r||rdnd }|r||d f7}q*|j%dur|%|}|j&dur|&|}|r||f7}|r|nd}| r|'}t(||||dS)a Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *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) head_mask (`torch.Tensor` of shape `(num_hidden_layers, num_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `decoder_input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, 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. position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, config.n_positions - 1]`. for padding use -1. [What are position IDs?](../glossary#position-ids) cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*): Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`, this tensor is not affected by padding. It is used to update the cache in the correct position and to infer the complete sequence length. Nz:You must specify exactly one of input_ids or inputs_embedszX`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.FrHTzPassing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.53.0. You should pass an instance of `DynamicCache` instead, e.g. `past_key_values=DynamicCache.from_legacy_cache(past_key_values)`.rrrr3)r2r.rzThe `z` should be specified for z layers, but it is for .)r0r2rkrjrlrrmr'last_hidden_staterri attentions))rXrlrruse_return_dictrcrrMr`rarrrrrr rZfrom_legacy_cacherr5rrZonesrr6r7rrrQziprqlenr enumerateZrandrZ_gradient_checkpointing_func__call__rrZto_legacy_cacher)r+rr0rrrrrlrrr2rmrSZreturn_legacy_cacherZ seq_lengthrZ pos_embedsriZall_hidden_statesZall_self_attnsZnext_decoder_cacheZ attn_maskZ mask_nameidxZ decoder_layerZdropout_probabilityZ layer_outputsZ next_cacher.r.r/r8sH                    zOPTDecoder.forward)F NNNNNNNNNNN)r:r;r<r=r r*rrrr5rxr ryr staticmethodr>rJrrrr?rrrrrrr8r@r.r.r,r/rPs%   D 6       rcseZdZdeffdd ZddZddZdd Zee dd e e j d e e j d e e j de eee jefde e jde ede ede ede ede e j de e j deedeeeffddZZS)OPTModelrXcs"t|t||_|dSrw)r)r*rdecoderrrr,r.r/r*s   zOPTModel.__init__cC|jjSrwrrrr.r.r/rzOPTModel.get_input_embeddingscC ||j_dSrwrrr.r.r/r zOPTModel.set_input_embeddingscCrrw)rrr.r.r/ get_decoderrzOPTModel.get_decoderNrr0rrrrrlrrr2rmrSrnc Ks|dur|n|jj}|dur|n|jj}|dur|n|jj}| dur$| n|jj} |jd||| ||||||d| d | } t| j| j| j | j dS)NT rr0r2rrrrrlrrrmrr.) rXrlrrrrrrrrir)r+rr0rrrrrlrrr2rmrSZdecoder_outputsr.r.r/r8s4 zOPTModel.forwardr)r:r;r<r r*rrrrrrr5r?rxrrrr ryrrrrr8r@r.r.r,r/rsZ      rc@s eZdZdS)KwargsForCausalLMN)r:r;r<r.r.r.r/rs rc!seZdZdgZfddZddZddZdd Zd d Zd d Z ddZ e e            d#de ejde ejde ejde eeejefde ejde ejde ede ede ede ede ejde ejdeedeeeffdd Zed!d"ZZS)$OPTForCausalLMzlm_head.weightcs8t|t||_tj|j|jdd|_| dSNFrZ) r)r*rrrrdrrlm_headrrr,r.r/r*s   zOPTForCausalLM.__init__cC |jjjSrwrrrrr.r.r/r'rz#OPTForCausalLM.get_input_embeddingscC||jj_dSrwrrr.r.r/r*z#OPTForCausalLM.set_input_embeddingscCrrwrrr.r.r/get_output_embeddings-rz$OPTForCausalLM.get_output_embeddingscCrrwr)r+Znew_embeddingsr.r.r/set_output_embeddings0rz$OPTForCausalLM.set_output_embeddingscCrrwrr)r+rr.r.r/ set_decoder3rzOPTForCausalLM.set_decodercCrrwrrr.r.r/r6rzOPTForCausalLM.get_decoderNrr0rrrlabelsrrlrrr2rmrSrnc Ks|dur|n|jj}| dur| n|jj} | dur| n|jj} |jjd||| |||||| d| d | }||d}d}|durV||j }|j ||fd|jj i| }t |||j |j|jdS)an labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (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]`. Example: ```python >>> from transformers import AutoTokenizer, OPTForCausalLM >>> model = OPTForCausalLM.from_pretrained("facebook/opt-350m") >>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") >>> prompt = "Hey, are you conscious? Can you talk to me?" >>> inputs = tokenizer(prompt, return_tensors="pt") >>> # Generate >>> generate_ids = model.generate(inputs.input_ids, max_length=30) >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0] "Hey, are you conscious? Can you talk to me?\nI'm not conscious. I'm just a little bit of a weirdo." ```NTrrrlosslogitsrrirr.)rXrlrrrrrrRrQrZ loss_functionrrrrir)r+rr0rrrrrrlrrr2rmrSrrrr.r.r/r89sL)  zOPTForCausalLM.forwardcs.d}|D]}|tfdd|Df7}q|S)Nr.c3s$|] }|d|jVqdSr9)Z index_selectrQr)rZ past_statebeam_idxr.r/ s"z0OPTForCausalLM._reorder_cache..)tuple)rrZreordered_pastZ layer_pastr.rr/_reorder_caches zOPTForCausalLM._reorder_cache NNNNNNNNNNNN)r:r;r<Z_tied_weights_keysr*rrrrrrrrrr5r?rxrrrr ryrrrrr8rrr@r.r.r,r/rsl        Rra The OPT Model transformer with a sequence classification head on top (linear layer). [`OPTForSequenceClassification`] uses the last token in order to do the classification, as other causal models (e.g. GPT-2) do. Since it does classification on the last token, it requires to know the position of the last token. If a `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in each row of the batch). )Z custom_introcseZdZdeffdd Ze           ddeejdeej deej dee e ej e fd eej d eejd ee d ee d ee dee deejde eeffddZddZddZZS)OPTForSequenceClassificationrXcs@t||j|_t||_tj|j|jdd|_| dSr) r)r* num_labelsrrrrdrscorerrr,r.r/r*s   z%OPTForSequenceClassification.__init__Nrr0rrrrrrlrrr2rnc  CsL| dur| n|jj} |j|||| ||||| | d } | d} || }|dur0|jdd\}}n |jdd\}}|jjdurG|dkrGtd|jjdurPd}n1|duru||jjk|jt j }t j |jd|jt j d}|| d}n d}t |jjd |t j ||jd |f}d}|dur|jjdur|jdkrd |j_n|jdkr|jt jks|jt jkrd |j_nd |j_|jjd krt}|jdkr|||}n,|||}n&|jjd krt}||d|j|d}n|jjd krt}|||}| s|f| dd}|dur|f|S|St||| j| j| jdS)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). N rr0r2rrrrlrrrr'rz=Cannot handle batch sizes > 1 if no padding token is defined.rH)rrJz will not detect padding tokens in `inputs_embeds`. Results may be unexpected if using padding tokens in conjunction with `inputs_embeds.`rZ regressionZsingle_label_classificationZmulti_label_classificationr)rXrrrrrrcrQrr5Zint32rZargmaxr`rar-r:Z problem_typerrJr7r>r squeezer rrrrrrir)r+rr0rrrrrrlrrr2transformer_outputsrirrrZlast_non_pad_tokenZ non_pad_maskZ token_indicesZ pooled_logitsrloss_fctoutputr.r.r/r8sx       "      z$OPTForSequenceClassification.forwardcCrrwrrr.r.r/rrz1OPTForSequenceClassification.get_input_embeddingscCrrwrrr.r.r/rrz1OPTForSequenceClassification.set_input_embeddingsr)r:r;r<r r*rrr5r?rrrr ryrrr8rrr@r.r.r,r/rsR      ^rcseZdZdeffdd Ze            ddeejdeej deej dee e ej e fd eej d eejd eejd ee d ee dee dee deejde eeffddZddZddZZS)OPTForQuestionAnsweringrXcs2t|t||_t|jd|_|dSr&) r)r*rrrrdr qa_outputsrrr,r.r/r*s   z OPTForQuestionAnswering.__init__Nrr0rrrstart_positions end_positionsrrlrrr2rnc  CsZ| dur| n|jj} |j|||| |||| | | d } | d}||}|jddd\}}|d}|d}d}|dur|durt|dkrP|d}t|dkr]|d}|d}| d| |j }| d| |j }t |d}|||}|||}||d}| s||f| dd}|dur|f|S|St |||| j| jd S) a Example: ```python >>> from transformers import AutoTokenizer, OPTForQuestionAnswering >>> import torch >>> torch.manual_seed(4) # doctest: +IGNORE_RESULT >>> tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") >>> # note: we are loading a OPTForQuestionAnswering from the hub here, >>> # so the head will be randomly initialized, hence the predictions will be random >>> model = OPTForQuestionAnswering.from_pretrained("facebook/opt-350m") >>> question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet" >>> inputs = tokenizer(question, text, return_tensors="pt") >>> with torch.no_grad(): ... outputs = model(**inputs) >>> answer_start_index = outputs.start_logits.argmax() >>> answer_end_index = outputs.end_logits.argmax() >>> answer_offset = len(tokenizer(question)[0]) >>> predict_answer_tokens = inputs.input_ids[ ... 0, answer_offset + answer_start_index : answer_offset + answer_end_index + 1 ... ] >>> predicted = tokenizer.decode(predict_answer_tokens) >>> predicted ' a nice puppet' ```NrrrrHr3)Z ignore_indexr')r start_logits end_logitsrir)rXrrrsplitrrRrrqclamprQrr rrir)r+rr0rrrrrrrlrrr2rrirrrZ total_lossZ ignored_indexrZ start_lossZend_lossrr.r.r/r8sR0        zOPTForQuestionAnswering.forwardcCrrwrrr.r.r/rrz,OPTForQuestionAnswering.get_input_embeddingscCrrwrrr.r.r/rrz,OPTForQuestionAnswering.set_input_embeddingsr)r:r;r<r r*rrr5r?rrrr ryrrr8rrr@r.r.r,r/rsX       ar)rrrrr)rA)Cr=typingrrrrrr5Ztorch.utils.checkpointrZtorch.nnrr r Z activationsr Z cache_utilsr rZ generationrZmodeling_attn_mask_utilsrZmodeling_flash_attention_utilsrrZmodeling_outputsrrrrZmodeling_utilsrrZprocessing_utilsrutilsrrrrrZconfiguration_optr Z!torch.nn.attention.flex_attentionr!Zintegrations.flex_attentionr"Z get_loggerr:r`rr#ModulerxfloatrVrWrzrrrrrrr__all__r.r.r.r/st         $ ah @}p r