codingutf8 2022 the baai teams s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch altclip model import math from dataclasses import dataclass from typing import any list optional tuple union import torch import torch nn as nn import torch utils checkpoint from activations import act2fn from modelingoutputs import basemodeloutput basemodeloutputwithpastandcrossattentions basemodeloutputwithpooling basemodeloutputwithpoolingandcrossattentions basemodeloutputwithpoolingandprojection from modelingutils import pretrainedmodel from pytorchutils import applychunkingtoforward findpruneableheadsandindices prunelinearlayer from utils import modeloutput addstartdocstringstomodelforward logging replacereturndocstrings from configurationaltclip import altclipconfig altcliptextconfig altclipvisionconfig logger logging getloggername checkpointfordoc baaialtclip configfordoc altclipconfig altclippretrainedmodelarchivelist baaialtclip see all altclip models at https huggingface comodels filteraltclip altclipstartdocstring r this model inherits from pretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a pytorch torch nn modulehttps pytorch orgdocsstablenn htmltorch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config clipconfig 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 frompretrained method to load the model weights altclipvisioninputsdocstring r args pixelvalues torch floattensor of shape batchsize numchannels height width pixel values padding will be ignored by default should you provide it pixel values can be obtained using autoimageprocessor see clipimageprocessor call for details outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple contrastive loss function adapted from https sachinruk github ioblogpytorchpytorch20lightningloss20functiongpu20210307clip html def contrastivelosslogits torch tensor torch tensor return nn functional crossentropylogits torch arangelenlogits devicelogits device def cliplosssimilarity torch tensor torch tensor captionloss contrastivelosssimilarity imageloss contrastivelosssimilarity t return captionloss imageloss 2 0 dataclass copied from transformers models clip modelingclip clipoutput with clipaltclip class altclipoutputmodeloutput loss optionaltorch floattensor none logitsperimage torch floattensor none logitspertext torch floattensor none textembeds torch floattensor none imageembeds torch floattensor none textmodeloutput basemodeloutputwithpooling none visionmodeloutput basemodeloutputwithpooling none def totupleself tupleany return tuple selfk if k not in textmodeloutput visionmodeloutput else getattrself k totuple for k in self keys copied from transformers models roberta modelingroberta robertaembeddings with robertaaltroberta class altrobertaembeddingsnn module copied from transformers models bert modelingbert bertembeddings init def initself config super init self wordembeddings nn embeddingconfig vocabsize config hiddensize paddingidxconfig padtokenid self positionembeddings nn embeddingconfig maxpositionembeddings config hiddensize self tokentypeembeddings nn embeddingconfig typevocabsize config hiddensize self layernorm is not snakecased to stick with tensorflow model variable name and be able to load any tensorflow checkpoint file self layernorm nn layernormconfig hiddensize epsconfig layernormeps self dropout nn dropoutconfig hiddendropoutprob positionids 1 len position emb is contiguous in memory and exported when serialized self positionembeddingtype getattrconfig positionembeddingtype absolute self registerbuffer positionids torch arangeconfig maxpositionembeddings expand1 1 persistentfalse self registerbuffer tokentypeids torch zerosself positionids size dtypetorch long persistentfalse end copy self paddingidx config padtokenid self positionembeddings nn embedding config maxpositionembeddings config hiddensize paddingidxself paddingidx def forward self inputidsnone tokentypeidsnone positionidsnone inputsembedsnone pastkeyvalueslength0 if positionids is none if inputids is not none create the position ids from the input token ids any padded tokens remain padded positionids createpositionidsfrominputidsinputids self paddingidx pastkeyvalueslength else positionids self createpositionidsfrominputsembedsinputsembeds if inputids is not none inputshape inputids size else inputshape inputsembeds size 1 seqlength inputshape1 setting the tokentypeids to the registered buffer in constructor where it is all zeros which usually occurs when its autogenerated registered buffer helps users when tracing the model without passing tokentypeids solves issue 5664 if tokentypeids is none if hasattrself tokentypeids bufferedtokentypeids self tokentypeids seqlength bufferedtokentypeidsexpanded bufferedtokentypeids expandinputshape0 seqlength tokentypeids bufferedtokentypeidsexpanded else tokentypeids torch zerosinputshape dtypetorch long deviceself positionids device if inputsembeds is none inputsembeds self wordembeddingsinputids tokentypeembeddings self tokentypeembeddingstokentypeids embeddings inputsembeds tokentypeembeddings if self positionembeddingtype absolute positionembeddings self positionembeddingspositionids embeddings positionembeddings embeddings self layernormembeddings embeddings self dropoutembeddings return embeddings def createpositionidsfrominputsembedsself inputsembeds inputshape inputsembeds size 1 sequencelength inputshape1 positionids torch arange self paddingidx 1 sequencelength self paddingidx 1 dtypetorch long deviceinputsembeds device return positionids unsqueeze0 expandinputshape copied from transformers models roberta modelingroberta robertaselfattention with robertaaltroberta class altrobertaselfattentionnn module def initself config positionembeddingtypenone super init if config hiddensize config numattentionheads 0 and not hasattrconfig embeddingsize raise valueerror fthe hidden size config hiddensize is not a multiple of the number of attention fheads config numattentionheads self numattentionheads config numattentionheads self attentionheadsize intconfig hiddensize config numattentionheads self allheadsize self numattentionheads self attentionheadsize self query nn linearconfig hiddensize self allheadsize self key nn linearconfig hiddensize self allheadsize self value nn linearconfig hiddensize self allheadsize self dropout nn dropoutconfig attentionprobsdropoutprob self positionembeddingtype positionembeddingtype or getattr config positionembeddingtype absolute if self positionembeddingtype relativekey or self positionembeddingtype relativekeyquery self maxpositionembeddings config maxpositionembeddings self distanceembedding nn embedding2 config maxpositionembeddings 1 self attentionheadsize self isdecoder config isdecoder def transposeforscoresself x torch tensor torch tensor newxshape x size 1 self numattentionheads self attentionheadsize x x viewnewxshape return x permute0 2 1 3 def forward self hiddenstates torch tensor attentionmask optionaltorch floattensor none headmask optionaltorch floattensor none encoderhiddenstates optionaltorch floattensor none encoderattentionmask optionaltorch floattensor none pastkeyvalue optionaltupletupletorch floattensor none outputattentions optionalbool false tupletorch tensor mixedquerylayer self queryhiddenstates if this is instantiated as a crossattention module the keys and values come from an encoder the attention mask needs to be such that the encoder s padding tokens are not attended to iscrossattention encoderhiddenstates is not none if iscrossattention and pastkeyvalue is not none reuse k v crossattentions keylayer pastkeyvalue0 valuelayer pastkeyvalue1 attentionmask encoderattentionmask elif iscrossattention keylayer self transposeforscoresself keyencoderhiddenstates valuelayer self transposeforscoresself valueencoderhiddenstates attentionmask encoderattentionmask elif pastkeyvalue is not none keylayer self transposeforscoresself keyhiddenstates valuelayer self transposeforscoresself valuehiddenstates keylayer torch catpastkeyvalue0 keylayer dim2 valuelayer torch catpastkeyvalue1 valuelayer dim2 else keylayer self transposeforscoresself keyhiddenstates valuelayer self transposeforscoresself valuehiddenstates querylayer self transposeforscoresmixedquerylayer usecache pastkeyvalue is not none if self isdecoder if crossattention save tupletorch tensor torch tensor of all cross attention keyvaluestates further calls to crossattention layer can then reuse all crossattention keyvaluestates first if case if unidirectional selfattention decoder save tupletorch tensor torch tensor of all previous decoder keyvaluestates further calls to unidirectional selfattention can concat previous decoder keyvaluestates to current projected keyvaluestates third elif case if encoder bidirectional selfattention pastkeyvalue is always none pastkeyvalue keylayer valuelayer take the dot product between query and key to get the raw attention scores attentionscores torch matmulquerylayer keylayer transpose1 2 if self positionembeddingtype relativekey or self positionembeddingtype relativekeyquery querylength keylength querylayer shape2 keylayer shape2 if usecache positionidsl torch tensorkeylength 1 dtypetorch long devicehiddenstates device view 1 1 else positionidsl torch arangequerylength dtypetorch long devicehiddenstates device view1 1 positionidsr torch arangekeylength dtypetorch long devicehiddenstates device view1 1 distance positionidsl positionidsr positionalembedding self distanceembeddingdistance self maxpositionembeddings 1 positionalembedding positionalembedding todtypequerylayer dtype fp16 compatibility if self positionembeddingtype relativekey relativepositionscores torch einsumbhld lrdbhlr querylayer positionalembedding attentionscores attentionscores relativepositionscores elif self positionembeddingtype relativekeyquery relativepositionscoresquery torch einsumbhld lrdbhlr querylayer positionalembedding relativepositionscoreskey torch einsumbhrd lrdbhlr keylayer positionalembedding attentionscores attentionscores relativepositionscoresquery relativepositionscoreskey attentionscores attentionscores math sqrtself attentionheadsize if attentionmask is not none apply the attention mask is precomputed for all layers in altrobertamodel forward function attentionscores attentionscores attentionmask normalize the attention scores to probabilities attentionprobs nn functional softmaxattentionscores dim1 this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper attentionprobs self dropoutattentionprobs mask heads if we want to if headmask is not none attentionprobs attentionprobs headmask contextlayer torch matmulattentionprobs valuelayer contextlayer contextlayer permute0 2 1 3 contiguous newcontextlayershape contextlayer size 2 self allheadsize contextlayer contextlayer viewnewcontextlayershape outputs contextlayer attentionprobs if outputattentions else contextlayer if self isdecoder outputs outputs pastkeyvalue return outputs copied from transformers models roberta modelingroberta robertaselfoutput class altrobertaselfoutputnn module def initself config super init self dense nn linearconfig hiddensize config hiddensize self layernorm nn layernormconfig hiddensize epsconfig layernormeps self dropout nn dropoutconfig hiddendropoutprob def forwardself hiddenstates torch tensor inputtensor torch tensor torch tensor hiddenstates self densehiddenstates hiddenstates self dropouthiddenstates hiddenstates self layernormhiddenstates inputtensor return hiddenstates copied from transformers models roberta modelingroberta robertaattention with robertaaltroberta class altrobertaattentionnn module def initself config positionembeddingtypenone super init self self altrobertaselfattentionconfig positionembeddingtypepositionembeddingtype self output altrobertaselfoutputconfig self prunedheads set def pruneheadsself heads if lenheads 0 return heads index findpruneableheadsandindices heads self self numattentionheads self self attentionheadsize self prunedheads prune linear layers self self query prunelinearlayerself self query index self self key prunelinearlayerself self key index self self value prunelinearlayerself self value index self output dense prunelinearlayerself output dense index dim1 update hyper params and store pruned heads self self numattentionheads self self numattentionheads lenheads self self allheadsize self self attentionheadsize self self numattentionheads self prunedheads self prunedheads unionheads def forward self hiddenstates torch tensor attentionmask optionaltorch floattensor none headmask optionaltorch floattensor none encoderhiddenstates optionaltorch floattensor none encoderattentionmask optionaltorch floattensor none pastkeyvalue optionaltupletupletorch floattensor none outputattentions optionalbool false tupletorch tensor selfoutputs self self hiddenstates attentionmask headmask encoderhiddenstates encoderattentionmask pastkeyvalue outputattentions attentionoutput self outputselfoutputs0 hiddenstates outputs attentionoutput selfoutputs1 add attentions if we output them return outputs copied from transformers models roberta modelingroberta robertaintermediate with robertaaltroberta class altrobertaintermediatenn module def initself config super init self dense nn linearconfig hiddensize config intermediatesize if isinstanceconfig hiddenact str self intermediateactfn act2fnconfig hiddenact else self intermediateactfn config hiddenact def forwardself hiddenstates torch tensor torch tensor hiddenstates self densehiddenstates hiddenstates self intermediateactfnhiddenstates return hiddenstates copied from transformers models roberta modelingroberta robertaoutput class altrobertaoutputnn module def initself config super init self dense nn linearconfig intermediatesize config hiddensize self layernorm nn layernormconfig hiddensize epsconfig layernormeps self dropout nn dropoutconfig hiddendropoutprob def forwardself hiddenstates torch tensor inputtensor torch tensor torch tensor hiddenstates self densehiddenstates hiddenstates self dropouthiddenstates hiddenstates self layernormhiddenstates inputtensor return hiddenstates copied from transformers models roberta modelingroberta robertalayer with robertaaltroberta class altrobertalayernn module def initself config super init self chunksizefeedforward config chunksizefeedforward self seqlendim 1 self attention altrobertaattentionconfig self isdecoder config isdecoder self addcrossattention config addcrossattention if self addcrossattention if not self isdecoder raise valueerrorfself should be used as a decoder model if cross attention is added self crossattention altrobertaattentionconfig positionembeddingtypeabsolute self intermediate altrobertaintermediateconfig self output altrobertaoutputconfig def forward self hiddenstates torch tensor attentionmask optionaltorch floattensor none headmask optionaltorch floattensor none encoderhiddenstates optionaltorch floattensor none encoderattentionmask optionaltorch floattensor none pastkeyvalue optionaltupletupletorch floattensor none outputattentions optionalbool false tupletorch tensor decoder unidirectional selfattention cached keyvalues tuple is at positions 1 2 selfattnpastkeyvalue pastkeyvalue 2 if pastkeyvalue is not none else none selfattentionoutputs self attention hiddenstates attentionmask headmask outputattentionsoutputattentions pastkeyvalueselfattnpastkeyvalue attentionoutput selfattentionoutputs0 if decoder the last output is tuple of selfattn cache if self isdecoder outputs selfattentionoutputs1 1 presentkeyvalue selfattentionoutputs1 else outputs selfattentionoutputs1 add self attentions if we output attention weights crossattnpresentkeyvalue none if self isdecoder and encoderhiddenstates is not none if not hasattrself crossattention raise valueerror fif encoderhiddenstates are passed self has to be instantiated with crossattention layers by setting config addcrossattentiontrue crossattn cached keyvalues tuple is at positions 3 4 of pastkeyvalue tuple crossattnpastkeyvalue pastkeyvalue2 if pastkeyvalue is not none else none crossattentionoutputs self crossattention attentionoutput attentionmask headmask encoderhiddenstates encoderattentionmask crossattnpastkeyvalue outputattentions attentionoutput crossattentionoutputs0 outputs outputs crossattentionoutputs1 1 add cross attentions if we output attention weights add crossattn cache to positions 3 4 of presentkeyvalue tuple crossattnpresentkeyvalue crossattentionoutputs1 presentkeyvalue presentkeyvalue crossattnpresentkeyvalue layeroutput applychunkingtoforward self feedforwardchunk self chunksizefeedforward self seqlendim attentionoutput outputs layeroutput outputs if decoder return the attn keyvalues as the last output if self isdecoder outputs outputs presentkeyvalue return outputs def feedforwardchunkself attentionoutput intermediateoutput self intermediateattentionoutput layeroutput self outputintermediateoutput attentionoutput return layeroutput copied from transformers models roberta modelingroberta robertaencoder with robertaaltroberta class altrobertaencodernn module def initself config super init self config config self layer nn modulelistaltrobertalayerconfig for in rangeconfig numhiddenlayers self gradientcheckpointing false def forward self hiddenstates torch tensor attentionmask optionaltorch floattensor none headmask optionaltorch floattensor none encoderhiddenstates optionaltorch floattensor none encoderattentionmask optionaltorch floattensor none pastkeyvalues optionaltupletupletorch floattensor none usecache optionalbool none outputattentions optionalbool false outputhiddenstates optionalbool false returndict optionalbool true uniontupletorch tensor basemodeloutputwithpastandcrossattentions allhiddenstates if outputhiddenstates else none allselfattentions if outputattentions else none allcrossattentions if outputattentions and self config addcrossattention else none if self gradientcheckpointing and self training if usecache logger warningonce usecachetrue is incompatible with gradient checkpointing setting usecachefalse usecache false nextdecodercache if usecache else none for i layermodule in enumerateself layer if outputhiddenstates allhiddenstates allhiddenstates hiddenstates layerheadmask headmaski if headmask is not none else none pastkeyvalue pastkeyvaluesi if pastkeyvalues is not none else none if self gradientcheckpointing and self training layeroutputs self gradientcheckpointingfunc layermodule call hiddenstates attentionmask layerheadmask encoderhiddenstates encoderattentionmask pastkeyvalue outputattentions else layeroutputs layermodule hiddenstates attentionmask layerheadmask encoderhiddenstates encoderattentionmask pastkeyvalue outputattentions hiddenstates layeroutputs0 if usecache nextdecodercache layeroutputs1 if outputattentions allselfattentions allselfattentions layeroutputs1 if self config addcrossattention allcrossattentions allcrossattentions layeroutputs2 if outputhiddenstates allhiddenstates allhiddenstates hiddenstates if not returndict return tuple v for v in hiddenstates nextdecodercache allhiddenstates allselfattentions allcrossattentions if v is not none return basemodeloutputwithpastandcrossattentions lasthiddenstatehiddenstates pastkeyvaluesnextdecodercache hiddenstatesallhiddenstates attentionsallselfattentions crossattentionsallcrossattentions copied from transformers models roberta modelingroberta robertapooler class altrobertapoolernn module def initself config super init self dense nn linearconfig hiddensize config hiddensize self activation nn tanh def forwardself hiddenstates torch tensor torch tensor we pool the model by simply taking the hidden state corresponding to the first token firsttokentensor hiddenstates 0 pooledoutput self densefirsttokentensor pooledoutput self activationpooledoutput return pooledoutput copied from transformers models clip modelingclip clipattention with clipaltclip class altclipattentionnn module input shape batch x time x channel bsz tgtlen embeddim hiddenstates size get query proj querystates self qprojhiddenstates self scale keystates self shapeself kprojhiddenstates 1 bsz valuestates self shapeself vprojhiddenstates 1 bsz projshape bsz self numheads 1 self headdim querystates self shapequerystates tgtlen bsz viewprojshape keystates keystates viewprojshape valuestates valuestates viewprojshape srclen keystates size1 attnweights torch bmmquerystates keystates transpose1 2 if attnweights size bsz self numheads tgtlen srclen raise valueerror fattention weights should be of size bsz self numheads tgtlen srclen but is f attnweights size apply the causalattentionmask first if causalattentionmask is not none if causalattentionmask size bsz 1 tgtlen srclen raise valueerror fattention mask should be of size bsz 1 tgtlen srclen but is f causalattentionmask size attnweights attnweights viewbsz self numheads tgtlen srclen causalattentionmask attnweights attnweights viewbsz self numheads tgtlen srclen if attentionmask is not none if attentionmask size bsz 1 tgtlen srclen raise valueerror fattention mask should be of size bsz 1 tgtlen srclen but is attentionmask size attnweights attnweights viewbsz self numheads tgtlen srclen attentionmask attnweights attnweights viewbsz self numheads tgtlen srclen attnweights nn functional softmaxattnweights dim1 if outputattentions this operation is a bit akward but it s required to make sure that attnweights keeps its gradient in order to do so attnweights have to reshaped twice and have to be reused in the following attnweightsreshaped attnweights viewbsz self numheads tgtlen srclen attnweights attnweightsreshaped viewbsz self numheads tgtlen srclen else attnweightsreshaped none attnprobs nn functional dropoutattnweights pself dropout trainingself training attnoutput torch bmmattnprobs valuestates if attnoutput size bsz self numheads tgtlen self headdim raise valueerror fattnoutput should be of size bsz self numheads tgtlen self headdim but is f attnoutput size attnoutput attnoutput viewbsz self numheads tgtlen self headdim attnoutput attnoutput transpose1 2 attnoutput attnoutput reshapebsz tgtlen embeddim attnoutput self outprojattnoutput return attnoutput attnweightsreshaped copied from transformers models clip modelingclip clipmlp with clipaltclip class altclipmlpnn module def initself config super init self config config self activationfn act2fnconfig hiddenact self fc1 nn linearconfig hiddensize config intermediatesize self fc2 nn linearconfig intermediatesize config hiddensize def forwardself hiddenstates torch tensor torch tensor hiddenstates self fc1hiddenstates hiddenstates self activationfnhiddenstates hiddenstates self fc2hiddenstates return hiddenstates copied from transformers models clip modelingclip clipencoderlayer with clipaltclip class altclipencoderlayernn module def initself config altclipconfig super init self embeddim config hiddensize self selfattn altclipattentionconfig self layernorm1 nn layernormself embeddim epsconfig layernormeps self mlp altclipmlpconfig self layernorm2 nn layernormself embeddim epsconfig layernormeps def forward self hiddenstates torch tensor attentionmask torch tensor causalattentionmask torch tensor outputattentions optionalbool false tupletorch floattensor residual hiddenstates hiddenstates self layernorm1hiddenstates hiddenstates attnweights self selfattn hiddenstateshiddenstates attentionmaskattentionmask causalattentionmaskcausalattentionmask outputattentionsoutputattentions hiddenstates residual hiddenstates residual hiddenstates hiddenstates self layernorm2hiddenstates hiddenstates self mlphiddenstates hiddenstates residual hiddenstates outputs hiddenstates if outputattentions outputs attnweights return outputs copied from transformers models clip modelingclip clipencoder with clipaltclip class altclipencodernn module def initself config altclipconfig super init self config config self layers nn modulelistaltclipencoderlayerconfig for in rangeconfig numhiddenlayers self gradientcheckpointing false def forward self inputsembeds attentionmask optionaltorch tensor none causalattentionmask optionaltorch tensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutput r args inputsembeds torch floattensor of shape batchsize sequencelength hiddensize optionally instead of passing inputids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert inputids indices into associated vectors than the model s internal embedding lookup matrix attentionmask torch tensor of shape batchsize sequencelength 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 glossaryattentionmask causalattentionmask torch tensor of shape batchsize sequencelength optional causal mask for the text model mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked what are attention masks glossaryattentionmask outputattentions bool optional whether or not to return the attentions tensors of all attention layers see attentions under returned tensors for more detail outputhiddenstates bool optional whether or not to return the hidden states of all layers see hiddenstates under returned tensors for more detail returndict bool optional whether or not to return a utils modeloutput instead of a plain tuple copied from transformers models clip modelingclip clipvisionembeddings with clipaltclip an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights factor self config initializerfactor if isinstancemodule altclipvisionembeddings factor self config initializerfactor nn init normalmodule classembedding mean0 0 stdmodule embeddim0 5 factor nn init normalmodule patchembedding weight stdmodule config initializerrange factor nn init normalmodule positionembedding weight stdmodule config initializerrange factor elif isinstancemodule altclipattention factor self config initializerfactor inprojstd module embeddim0 5 2 module config numhiddenlayers 0 5 factor outprojstd module embeddim0 5 factor nn init normalmodule qproj weight stdinprojstd nn init normalmodule kproj weight stdinprojstd nn init normalmodule vproj weight stdinprojstd nn init normalmodule outproj weight stdoutprojstd elif isinstancemodule altclipmlp factor self config initializerfactor inprojstd module config hiddensize0 5 2 module config numhiddenlayers 0 5 factor fcstd 2 module config hiddensize 0 5 factor nn init normalmodule fc1 weight stdfcstd nn init normalmodule fc2 weight stdinprojstd elif isinstancemodule altclipmodel nn init normal module textprojection weight stdmodule textembeddim0 5 self config initializerfactor module textprojection ishfinitialized true nn init normal module visualprojection weight stdmodule visionembeddim0 5 self config initializerfactor module visualprojection ishfinitialized true elif isinstancemodule nn layernorm module bias data zero module weight data fill1 0 elif isinstancemodule nn linear module weight data normalmean0 0 stdself config initializerfactor if module bias is not none module bias data zero elif isinstancemodule nn embedding module weight data normalmean0 0 stdself config initializerfactor if module paddingidx is not none module weight datamodule paddingidx zero copied from transformers models clip modelingclip clipvisiontransformer with clipvisiontransformeraltclipvisiontransformer clipvisionconfigaltclipvisionconfig clipvisionembeddingsaltclipvisionembeddings clipencoderaltclipencoder clipvisioninputsdocstringaltclipvisioninputsdocstring class altclipvisiontransformernn module def initself config altclipvisionconfig super init self config config embeddim config hiddensize self embeddings altclipvisionembeddingsconfig self prelayrnorm nn layernormembeddim epsconfig layernormeps self encoder altclipencoderconfig self postlayernorm nn layernormembeddim epsconfig layernormeps addstartdocstringstomodelforwardaltclipvisioninputsdocstring replacereturndocstringsoutputtypebasemodeloutputwithpooling configclassaltclipvisionconfig def forward self pixelvalues optionaltorch floattensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple basemodeloutputwithpooling r returns initialize weights and apply final processing the model can behave as an encoder with only selfattention as well as a decoder in which case a layer of crossattention is added between the selfattention layers following the architecture described in attention is all you need by ashish vaswani noam shazeer niki parmar jakob uszkoreit llion jones aidan n gomez lukasz kaiser and illia polosukhin to behave as an decoder the model needs to be initialized with the isdecoder argument of the configuration set to true to be used in a seq2seq model the model needs to initialized with both isdecoder argument and addcrossattention set to true an encoderhiddenstates is then expected as an input to the forward pass attention is all you need https arxiv orgabs1706 03762 copied from transformers models bert modelingbert bertmodel init with bertaltroberta initialize weights and apply final processing prunes heads of the model headstoprune dict of layernum list of heads to prune in this layer see base class pretrainedmodel copied from transformers models bert modelingbert bertmodel forward outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict if self config isdecoder usecache usecache if usecache is not none else self config usecache else usecache false if inputids is not none and inputsembeds is not none raise valueerroryou cannot specify both inputids and inputsembeds at the same time elif inputids is not none self warnifpaddingandnoattentionmaskinputids attentionmask inputshape inputids size elif inputsembeds is not none inputshape inputsembeds size 1 else raise valueerroryou have to specify either inputids or inputsembeds batchsize seqlength inputshape device inputids device if inputids is not none else inputsembeds device pastkeyvalueslength pastkeyvalueslength pastkeyvalues00 shape2 if pastkeyvalues is not none else 0 if attentionmask is none attentionmask torch onesbatchsize seqlength pastkeyvalueslength devicedevice if tokentypeids is none if hasattrself embeddings tokentypeids bufferedtokentypeids self embeddings tokentypeids seqlength bufferedtokentypeidsexpanded bufferedtokentypeids expandbatchsize seqlength tokentypeids bufferedtokentypeidsexpanded else tokentypeids torch zerosinputshape dtypetorch long devicedevice we can provide a selfattention mask of dimensions batchsize fromseqlength toseqlength ourselves in which case we just need to make it broadcastable to all heads extendedattentionmask torch tensor self getextendedattentionmaskattentionmask inputshape if a 2d or 3d attention mask is provided for the crossattention we need to make broadcastable to batchsize numheads seqlength seqlength if self config isdecoder and encoderhiddenstates is not none encoderbatchsize encodersequencelength encoderhiddenstates size encoderhiddenshape encoderbatchsize encodersequencelength if encoderattentionmask is none encoderattentionmask torch onesencoderhiddenshape devicedevice encoderextendedattentionmask self invertattentionmaskencoderattentionmask else encoderextendedattentionmask none prepare head mask if needed 1 0 in headmask indicate we keep the head attentionprobs has shape bsz x nheads x n x n input headmask has shape numheads or numhiddenlayers x numheads and headmask is converted to shape numhiddenlayers x batch x numheads x seqlength x seqlength headmask self getheadmaskheadmask self config numhiddenlayers embeddingoutput self embeddings inputidsinputids positionidspositionids tokentypeidstokentypeids inputsembedsinputsembeds pastkeyvalueslengthpastkeyvalueslength encoderoutputs self encoder embeddingoutput attentionmaskextendedattentionmask headmaskheadmask encoderhiddenstatesencoderhiddenstates encoderattentionmaskencoderextendedattentionmask pastkeyvaluespastkeyvalues usecacheusecache outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict sequenceoutput encoderoutputs0 pooledoutput self poolersequenceoutput if self pooler is not none else none if not returndict return sequenceoutput pooledoutput encoderoutputs1 return basemodeloutputwithpoolingandcrossattentions lasthiddenstatesequenceoutput pooleroutputpooledoutput pastkeyvaluesencoderoutputs pastkeyvalues hiddenstatesencoderoutputs hiddenstates attentionsencoderoutputs attentions crossattentionsencoderoutputs crossattentions class altcliptextmodelaltclippretrainedmodel configclass altcliptextconfig def initself config super initconfig self roberta altrobertamodelconfig addpoolinglayerfalse self transformation nn linearconfig hiddensize config projectdim self preln nn layernormconfig hiddensize epsconfig layernormeps self postinit def getinputembeddingsself nn module return self roberta embeddings wordembeddings def setinputembeddingsself value nn embedding none self roberta embeddings wordembeddings value def resizetokenembeddingsself newnumtokens optionalint none nn embedding return super resizetokenembeddingsnewnumtokens addstartdocstringstomodelforwardaltcliptextinputsdocstring replacereturndocstringsoutputtypebasemodeloutputwithpoolingandprojection configclassaltcliptextconfig def forward self inputids optionaltorch tensor none attentionmask optionaltorch tensor none tokentypeids optionaltorch tensor none positionids optionaltorch tensor none headmask optionaltorch tensor none inputsembeds optionaltorch tensor none encoderhiddenstates optionaltorch tensor none encoderattentionmask optionaltorch tensor none outputattentions optionalbool none returndict optionalbool none outputhiddenstates optionalbool none uniontuple basemodeloutputwithpoolingandprojection r returns examples python from transformers import autoprocessor altcliptextmodel model altcliptextmodel frompretrainedbaaialtclip processor autoprocessor frompretrainedbaaialtclip texts it s a cat it s a dog inputs processortexttexts paddingtrue returntensorspt outputs modelinputs lasthiddenstate outputs lasthiddenstate pooledoutput outputs pooleroutput pooled cls states returndict returndict if returndict is not none else self config usereturndict outputs self roberta inputidsinputids attentionmaskattentionmask tokentypeidstokentypeids positionidspositionids headmaskheadmask inputsembedsinputsembeds encoderhiddenstatesencoderhiddenstates encoderattentionmaskencoderattentionmask outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict last module outputs sequenceoutput outputs0 project every module sequenceoutput self prelnsequenceoutput pooler projectionstate self transformationsequenceoutput pooleroutput projectionstate 0 if not returndict return projectionstate pooleroutput outputs2 4 return basemodeloutputwithpoolingandprojection lasthiddenstateprojectionstate pooleroutputpooleroutput hiddenstatesoutputs hiddenstates attentionsoutputs attentions class altclipmodelaltclippretrainedmodel configclass altclipconfig def initself config altclipconfig super initconfig if not isinstanceconfig visionconfig altclipvisionconfig raise valueerror config visionconfig is expected to be of type altclipvisionconfig but is of type f typeconfig visionconfig if not isinstanceconfig textconfig altcliptextconfig raise valueerror config textconfig is expected to be of type altcliptextconfig but is of type f typeconfig textconfig textconfig config textconfig visionconfig config visionconfig self projectiondim config projectiondim self textembeddim textconfig projectdim self visionembeddim visionconfig hiddensize self textmodel altcliptextmodeltextconfig self visionmodel altclipvisiontransformervisionconfig self visualprojection nn linearself visionembeddim self projectiondim biasfalse self textprojection nn linearself textembeddim self projectiondim biasfalse self logitscale nn parametertorch tensorself config logitscaleinitvalue initialize weights and apply final processing self postinit addstartdocstringstomodelforwardaltcliptextinputsdocstring def gettextfeatures self inputids optionaltorch tensor none attentionmask optionaltorch tensor none positionids optionaltorch tensor none tokentypeidsnone outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none torch floattensor r returns textfeatures torch floattensor of shape batchsize outputdim the text embeddings obtained by applying the projection layer to the pooled output of altcliptextmodel examples python from transformers import autoprocessor altclipmodel model altclipmodel frompretrainedbaaialtclip processor autoprocessor frompretrainedbaaialtclip inputs processortexta photo of a cat a photo of a dog paddingtrue returntensorspt textfeatures model gettextfeaturesinputs use altclip model s config for some fields if specified instead of those of vision text components outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict textoutputs self textmodel inputidsinputids attentionmaskattentionmask positionidspositionids tokentypeidstokentypeids outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict pooledoutput textoutputs1 textfeatures self textprojectionpooledoutput return textfeatures addstartdocstringstomodelforwardaltclipvisioninputsdocstring def getimagefeatures self pixelvalues optionaltorch floattensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none torch floattensor r returns imagefeatures torch floattensor of shape batchsize outputdim the image embeddings obtained by applying the projection layer to the pooled output of altclipvisionmodel examples python from pil import image import requests from transformers import autoprocessor altclipmodel model altclipmodel frompretrainedbaaialtclip processor autoprocessor frompretrainedbaaialtclip url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw inputs processorimagesimage returntensorspt imagefeatures model getimagefeaturesinputs use altclip model s config for some fields if specified instead of those of vision text components outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict visionoutputs self visionmodel pixelvaluespixelvalues outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict pooledoutput visionoutputs1 pooledoutput imagefeatures self visualprojectionpooledoutput return imagefeatures addstartdocstringstomodelforwardaltclipinputsdocstring replacereturndocstringsoutputtypealtclipoutput configclassaltclipconfig def forward self inputids optionaltorch longtensor none pixelvalues optionaltorch floattensor none attentionmask optionaltorch tensor none positionids optionaltorch longtensor none tokentypeids optionaltorch tensor none returnloss optionalbool none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple altclipoutput r returns examples python from pil import image import requests from transformers import autoprocessor altclipmodel model altclipmodel frompretrainedbaaialtclip processor autoprocessor frompretrainedbaaialtclip url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw inputs processor texta photo of a cat a photo of a dog imagesimage returntensorspt paddingtrue outputs modelinputs logitsperimage outputs logitsperimage this is the imagetext similarity score probs logitsperimage softmaxdim1 we can take the softmax to get the label probabilities use altclip model s config for some fields if specified instead of those of vision text components outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates returndict returndict if returndict is not none else self config usereturndict textoutputs self textmodel inputidsinputids attentionmaskattentionmask tokentypeidstokentypeids positionidspositionids outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict visionoutputs self visionmodel pixelvaluespixelvalues outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates returndictreturndict imageembeds visionoutputs1 imageembeds self visualprojectionimageembeds textembeds textoutputs1 textembeds self textprojectiontextembeds normalized features imageembeds imageembeds imageembeds normp2 dim1 keepdimtrue textembeds textembeds textembeds normp2 dim1 keepdimtrue cosine similarity as logits logitscale self logitscale exp logitspertext torch matmultextembeds imageembeds t logitscale logitsperimage logitspertext t loss none if returnloss loss cliplosslogitspertext if not returndict output logitsperimage logitspertext textembeds imageembeds textoutputs visionoutputs return loss output if loss is not none else output return altclipoutput lossloss logitsperimagelogitsperimage logitspertextlogitspertext textembedstextembeds imageembedsimageembeds textmodeloutputtextoutputs visionmodeloutputvisionoutputs copied from transformers models roberta modelingroberta createpositionidsfrominputids def createpositionidsfrominputidsinputids paddingidx pastkeyvalueslength0 the series of casts and typeconversions here are carefully balanced to both work with onnx export and xla mask inputids nepaddingidx int incrementalindices torch cumsummask dim1 typeasmask pastkeyvalueslength mask return incrementalindices long paddingidx coding utf 8 2022 the baai teams s and the huggingface inc team licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license pytorch altclip model see all altclip models at https huggingface co models filter altclip this model inherits from pretrainedmodel check the superclass documentation for the generic methods the library implements for all its model such as downloading or saving resizing the input embeddings pruning heads etc this model is also a pytorch torch nn module https pytorch org docs stable nn html torch nn module subclass use it as a regular pytorch module and refer to the pytorch documentation for all matter related to general usage and behavior parameters config clipconfig 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 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 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 max_position_embeddings 1 what are position ids glossary position ids 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 args pixel_values torch floattensor of shape batch_size num_channels height width pixel values padding will be ignored by default should you provide it pixel values can be obtained using autoimageprocessor see clipimageprocessor __call__ for details 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 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 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 max_position_embeddings 1 what are position ids glossary position ids pixel_values torch floattensor of shape batch_size num_channels height width pixel values padding will be ignored by default should you provide it pixel values can be obtained using autoimageprocessor see clipimageprocessor __call__ for details return_loss bool optional whether or not to return the contrastive loss 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 contrastive loss function adapted from https sachinruk github io blog pytorch pytorch 20lightning loss 20function gpu 2021 03 07 clip html copied from transformers models clip modeling_clip clipoutput with clip altclip args loss torch floattensor of shape 1 optional returned when return_loss is true contrastive loss for image text similarity logits_per_image torch floattensor of shape image_batch_size text_batch_size the scaled dot product scores between image_embeds and text_embeds this represents the image text similarity scores logits_per_text torch floattensor of shape text_batch_size image_batch_size the scaled dot product scores between text_embeds and image_embeds this represents the text image similarity scores text_embeds torch floattensor of shape batch_size output_dim the text embeddings obtained by applying the projection layer to the pooled output of altcliptextmodel image_embeds torch floattensor of shape batch_size output_dim the image embeddings obtained by applying the projection layer to the pooled output of altclipvisionmodel text_model_output basemodeloutputwithpooling the output of the altcliptextmodel vision_model_output basemodeloutputwithpooling the output of the altclipvisionmodel copied from transformers models roberta modeling_roberta robertaembeddings with roberta altroberta same as bertembeddings with a tiny tweak for positional embeddings indexing copied from transformers models bert modeling_bert bertembeddings __init__ self layernorm is not snake cased to stick with tensorflow model variable name and be able to load any tensorflow checkpoint file position_ids 1 len position emb is contiguous in memory and exported when serialized end copy create the position ids from the input token ids any padded tokens remain padded setting the token_type_ids to the registered buffer in constructor where it is all zeros which usually occurs when its auto generated registered buffer helps users when tracing the model without passing token_type_ids solves issue 5664 we are provided embeddings directly we cannot infer which are padded so just generate sequential position ids args inputs_embeds torch tensor returns torch tensor copied from transformers models roberta modeling_roberta robertaselfattention with roberta altroberta if this is instantiated as a cross attention module the keys and values come from an encoder the attention mask needs to be such that the encoder s padding tokens are not attended to reuse k v cross_attentions if cross_attention save tuple torch tensor torch tensor of all cross attention key value_states further calls to cross_attention layer can then reuse all cross attention key value_states first if case if uni directional self attention decoder save tuple torch tensor torch tensor of all previous decoder key value_states further calls to uni directional self attention can concat previous decoder key value_states to current projected key value_states third elif case if encoder bi directional self attention past_key_value is always none take the dot product between query and key to get the raw attention scores fp16 compatibility apply the attention mask is precomputed for all layers in altrobertamodel forward function normalize the attention scores to probabilities this is actually dropping out entire tokens to attend to which might seem a bit unusual but is taken from the original transformer paper mask heads if we want to copied from transformers models roberta modeling_roberta robertaselfoutput copied from transformers models roberta modeling_roberta robertaattention with roberta altroberta prune linear layers update hyper params and store pruned heads add attentions if we output them copied from transformers models roberta modeling_roberta robertaintermediate with roberta altroberta copied from transformers models roberta modeling_roberta robertaoutput copied from transformers models roberta modeling_roberta robertalayer with roberta altroberta decoder uni directional self attention cached key values tuple is at positions 1 2 if decoder the last output is tuple of self attn cache add self attentions if we output attention weights cross_attn cached key values tuple is at positions 3 4 of past_key_value tuple add cross attentions if we output attention weights add cross attn cache to positions 3 4 of present_key_value tuple if decoder return the attn key values as the last output copied from transformers models roberta modeling_roberta robertaencoder with roberta altroberta copied from transformers models roberta modeling_roberta robertapooler we pool the model by simply taking the hidden state corresponding to the first token copied from transformers models clip modeling_clip clipattention with clip altclip multi headed attention from attention is all you need paper input shape batch x time x channel get query proj apply the causal_attention_mask first this operation is a bit akward but it s required to make sure that attn_weights keeps its gradient in order to do so attn_weights have to reshaped twice and have to be reused in the following copied from transformers models clip modeling_clip clipmlp with clip altclip copied from transformers models clip modeling_clip clipencoderlayer with clip altclip args hidden_states torch floattensor input to the layer of shape batch seq_len embed_dim attention_mask torch floattensor attention mask of size batch 1 tgt_len src_len where padding elements are indicated by very large negative values config 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 copied from transformers models clip modeling_clip clipencoder with clip altclip transformer encoder consisting of config num_hidden_layers self attention layers each layer is a altclipencoderlayer args config altclipconfig args inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size 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 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 causal_attention_mask torch tensor of shape batch_size sequence_length optional causal mask for the text model 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 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 copied from transformers models clip modeling_clip clipvisionembeddings with clip altclip shape width grid grid an abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models initialize the weights copied from transformers models clip modeling_clip clipvisiontransformer with clipvisiontransformer altclipvisiontransformer clipvisionconfig altclipvisionconfig clipvisionembeddings altclipvisionembeddings clipencoder altclipencoder clip_vision_inputs_docstring altclip_vision_inputs_docstring returns initialize weights and apply final processing returns examples python from pil import image import requests from transformers import autoprocessor altclipvisionmodel model altclipvisionmodel from_pretrained baai altclip processor autoprocessor from_pretrained baai altclip url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw inputs processor images image return_tensors pt outputs model inputs last_hidden_state outputs last_hidden_state pooled_output outputs pooler_output pooled cls states the model can behave as an encoder with only self attention as well as a decoder in which case a layer of cross attention is added between the self attention layers following the architecture described in attention is all you need _ by ashish vaswani noam shazeer niki parmar jakob uszkoreit llion jones aidan n gomez lukasz kaiser and illia polosukhin to behave as an decoder the model needs to be initialized with the is_decoder argument of the configuration set to true to be used in a seq2seq model the model needs to initialized with both is_decoder argument and add_cross_attention set to true an encoder_hidden_states is then expected as an input to the forward pass _ attention is all you need https arxiv org abs 1706 03762 copied from transformers models bert modeling_bert bertmodel __init__ with bert altroberta initialize weights and apply final processing prunes heads of the model heads_to_prune dict of layer_num list of heads to prune in this layer see base class pretrainedmodel copied from transformers models bert modeling_bert bertmodel forward encoder_hidden_states torch floattensor of shape batch_size sequence_length hidden_size optional sequence of hidden states at the output of the last layer of the encoder used in the cross attention if the model is configured as a decoder encoder_attention_mask torch floattensor of shape batch_size sequence_length optional mask to avoid performing attention on the padding token indices of the encoder input this mask is used in the cross attention if the model is configured as a decoder mask values selected in 0 1 1 for tokens that are not masked 0 for tokens that are masked past_key_values tuple tuple torch floattensor of length config n_layers with each tuple having 4 tensors of shape batch_size num_heads sequence_length 1 embed_size_per_head contains precomputed key and value hidden states of the attention blocks can be used to speed up 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 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_values_length we can provide a self attention mask of dimensions batch_size from_seq_length to_seq_length ourselves in which case we just need to make it broadcastable to all heads if a 2d or 3d attention mask is provided for the cross attention we need to make broadcastable to batch_size num_heads seq_length seq_length prepare head mask if needed 1 0 in head_mask indicate we keep the head attention_probs has shape bsz x n_heads x n x n input head_mask has shape num_heads or num_hidden_layers x num_heads and head_mask is converted to shape num_hidden_layers x batch x num_heads x seq_length x seq_length returns examples python from transformers import autoprocessor altcliptextmodel model altcliptextmodel from_pretrained baai altclip processor autoprocessor from_pretrained baai altclip texts it s a cat it s a dog inputs processor text texts padding true return_tensors pt outputs model inputs last_hidden_state outputs last_hidden_state pooled_output outputs pooler_output pooled cls states last module outputs project every module pooler initialize weights and apply final processing returns text_features torch floattensor of shape batch_size output_dim the text embeddings obtained by applying the projection layer to the pooled output of altcliptextmodel examples python from transformers import autoprocessor altclipmodel model altclipmodel from_pretrained baai altclip processor autoprocessor from_pretrained baai altclip inputs processor text a photo of a cat a photo of a dog padding true return_tensors pt text_features model get_text_features inputs use altclip model s config for some fields if specified instead of those of vision text components returns image_features torch floattensor of shape batch_size output_dim the image embeddings obtained by applying the projection layer to the pooled output of altclipvisionmodel examples python from pil import image import requests from transformers import autoprocessor altclipmodel model altclipmodel from_pretrained baai altclip processor autoprocessor from_pretrained baai altclip url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw inputs processor images image return_tensors pt image_features model get_image_features inputs use altclip model s config for some fields if specified instead of those of vision text components pooled_output returns examples python from pil import image import requests from transformers import autoprocessor altclipmodel model altclipmodel from_pretrained baai altclip processor autoprocessor from_pretrained baai altclip url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw inputs processor text a photo of a cat a photo of a dog images image return_tensors pt padding true outputs model inputs logits_per_image outputs logits_per_image this is the image text similarity score probs logits_per_image softmax dim 1 we can take the softmax to get the label probabilities use altclip model s config for some fields if specified instead of those of vision text components normalized features cosine similarity as logits copied from transformers models roberta modeling_roberta create_position_ids_from_input_ids replace non padding symbols with their position numbers position numbers begin at padding_idx 1 padding symbols are ignored this is modified from fairseq s utils make_positions args x torch tensor x returns torch tensor the series of casts and type conversions here are carefully balanced to both work with onnx export and xla
import math
from dataclasses import dataclass
from typing import Any, List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.utils.checkpoint

from ...activations import ACT2FN
from ...modeling_outputs import (
    BaseModelOutput,
    BaseModelOutputWithPastAndCrossAttentions,
    BaseModelOutputWithPooling,
    BaseModelOutputWithPoolingAndCrossAttentions,
    BaseModelOutputWithPoolingAndProjection,
)
from ...modeling_utils import PreTrainedModel
from ...pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
from ...utils import ModelOutput, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
from .configuration_altclip import AltCLIPConfig, AltCLIPTextConfig, AltCLIPVisionConfig


logger = logging.get_logger(__name__)

_CHECKPOINT_FOR_DOC = "BAAI/AltCLIP"
_CONFIG_FOR_DOC = "AltCLIPConfig"

ALTCLIP_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "BAAI/AltCLIP",
    
]


ALTCLIP_START_DOCSTRING = r

ALTCLIP_TEXT_INPUTS_DOCSTRING = r

ALTCLIP_VISION_INPUTS_DOCSTRING = r

ALTCLIP_INPUTS_DOCSTRING = r




def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:
    return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))


def clip_loss(similarity: torch.Tensor) -> torch.Tensor:
    caption_loss = contrastive_loss(similarity)
    image_loss = contrastive_loss(similarity.t())
    return (caption_loss + image_loss) / 2.0


@dataclass

class AltCLIPOutput(ModelOutput):
    

    loss: Optional[torch.FloatTensor] = None
    logits_per_image: torch.FloatTensor = None
    logits_per_text: torch.FloatTensor = None
    text_embeds: torch.FloatTensor = None
    image_embeds: torch.FloatTensor = None
    text_model_output: BaseModelOutputWithPooling = None
    vision_model_output: BaseModelOutputWithPooling = None

    def to_tuple(self) -> Tuple[Any]:
        return tuple(
            self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()
            for k in self.keys()
        )



class AltRobertaEmbeddings(nn.Module):
    

    
    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

        
        
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        
        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
        self.register_buffer(
            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
        )
        self.register_buffer(
            "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
        )

        
        self.padding_idx = config.pad_token_id
        self.position_embeddings = nn.Embedding(
            config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
        )

    def forward(
        self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
    ):
        if position_ids is None:
            if input_ids is not None:
                
                position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length)
            else:
                position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)

        if input_ids is not None:
            input_shape = input_ids.size()
        else:
            input_shape = inputs_embeds.size()[:-1]

        seq_length = input_shape[1]

        
        
        
        if token_type_ids is None:
            if hasattr(self, "token_type_ids"):
                buffered_token_type_ids = self.token_type_ids[:, :seq_length]
                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
                token_type_ids = buffered_token_type_ids_expanded
            else:
                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)

        if inputs_embeds is None:
            inputs_embeds = self.word_embeddings(input_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = inputs_embeds + token_type_embeddings
        if self.position_embedding_type == "absolute":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings += position_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings

    def create_position_ids_from_inputs_embeds(self, inputs_embeds):
        
        input_shape = inputs_embeds.size()[:-1]
        sequence_length = input_shape[1]

        position_ids = torch.arange(
            self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
        )
        return position_ids.unsqueeze(0).expand(input_shape)



class AltRobertaSelfAttention(nn.Module):
    def __init__(self, config, position_embedding_type=None):
        super().__init__()
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
            raise ValueError(
                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
                f"heads ({config.num_attention_heads})"
            )

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
        self.position_embedding_type = position_embedding_type or getattr(
            config, "position_embedding_type", "absolute"
        )
        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            self.max_position_embeddings = config.max_position_embeddings
            self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)

        self.is_decoder = config.is_decoder

    def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        mixed_query_layer = self.query(hidden_states)

        
        
        
        is_cross_attention = encoder_hidden_states is not None

        if is_cross_attention and past_key_value is not None:
            
            key_layer = past_key_value[0]
            value_layer = past_key_value[1]
            attention_mask = encoder_attention_mask
        elif is_cross_attention:
            key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
            value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
            attention_mask = encoder_attention_mask
        elif past_key_value is not None:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))
            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
        else:
            key_layer = self.transpose_for_scores(self.key(hidden_states))
            value_layer = self.transpose_for_scores(self.value(hidden_states))

        query_layer = self.transpose_for_scores(mixed_query_layer)

        use_cache = past_key_value is not None
        if self.is_decoder:
            
            
            
            
            
            
            
            past_key_value = (key_layer, value_layer)

        
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
            query_length, key_length = query_layer.shape[2], key_layer.shape[2]
            if use_cache:
                position_ids_l = torch.tensor(key_length - 1, dtype=torch.long, device=hidden_states.device).view(
                    -1, 1
                )
            else:
                position_ids_l = torch.arange(query_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
            position_ids_r = torch.arange(key_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
            distance = position_ids_l - position_ids_r

            positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
            positional_embedding = positional_embedding.to(dtype=query_layer.dtype)  

            if self.position_embedding_type == "relative_key":
                relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores
            elif self.position_embedding_type == "relative_key_query":
                relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
                relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
                attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            
            attention_scores = attention_scores + attention_mask

        
        attention_probs = nn.functional.softmax(attention_scores, dim=-1)

        
        
        attention_probs = self.dropout(attention_probs)

        
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(new_context_layer_shape)

        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)

        if self.is_decoder:
            outputs = outputs + (past_key_value,)
        return outputs



class AltRobertaSelfOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states



class AltRobertaAttention(nn.Module):
    def __init__(self, config, position_embedding_type=None):
        super().__init__()
        self.self = AltRobertaSelfAttention(config, position_embedding_type=position_embedding_type)
        self.output = AltRobertaSelfOutput(config)
        self.pruned_heads = set()

    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        heads, index = find_pruneable_heads_and_indices(
            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
        )

        
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)

        
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        self_outputs = self.self(
            hidden_states,
            attention_mask,
            head_mask,
            encoder_hidden_states,
            encoder_attention_mask,
            past_key_value,
            output_attentions,
        )
        attention_output = self.output(self_outputs[0], hidden_states)
        outputs = (attention_output,) + self_outputs[1:]  
        return outputs



class AltRobertaIntermediate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
        if isinstance(config.hidden_act, str):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states



class AltRobertaOutput(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states



class AltRobertaLayer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.chunk_size_feed_forward = config.chunk_size_feed_forward
        self.seq_len_dim = 1
        self.attention = AltRobertaAttention(config)
        self.is_decoder = config.is_decoder
        self.add_cross_attention = config.add_cross_attention
        if self.add_cross_attention:
            if not self.is_decoder:
                raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
            self.crossattention = AltRobertaAttention(config, position_embedding_type="absolute")
        self.intermediate = AltRobertaIntermediate(config)
        self.output = AltRobertaOutput(config)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor]:
        
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        self_attention_outputs = self.attention(
            hidden_states,
            attention_mask,
            head_mask,
            output_attentions=output_attentions,
            past_key_value=self_attn_past_key_value,
        )
        attention_output = self_attention_outputs[0]

        
        if self.is_decoder:
            outputs = self_attention_outputs[1:-1]
            present_key_value = self_attention_outputs[-1]
        else:
            outputs = self_attention_outputs[1:]  

        cross_attn_present_key_value = None
        if self.is_decoder and encoder_hidden_states is not None:
            if not hasattr(self, "crossattention"):
                raise ValueError(
                    f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
                    " by setting `config.add_cross_attention=True`"
                )

            
            cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
            cross_attention_outputs = self.crossattention(
                attention_output,
                attention_mask,
                head_mask,
                encoder_hidden_states,
                encoder_attention_mask,
                cross_attn_past_key_value,
                output_attentions,
            )
            attention_output = cross_attention_outputs[0]
            outputs = outputs + cross_attention_outputs[1:-1]  

            
            cross_attn_present_key_value = cross_attention_outputs[-1]
            present_key_value = present_key_value + cross_attn_present_key_value

        layer_output = apply_chunking_to_forward(
            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
        )
        outputs = (layer_output,) + outputs

        
        if self.is_decoder:
            outputs = outputs + (present_key_value,)

        return outputs

    def feed_forward_chunk(self, attention_output):
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        return layer_output



class AltRobertaEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.layer = nn.ModuleList([AltRobertaLayer(config) for _ in range(config.num_hidden_layers)])
        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = False,
        output_hidden_states: Optional[bool] = False,
        return_dict: Optional[bool] = True,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None

        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                )
                use_cache = False

        next_decoder_cache = () if use_cache else None
        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None
            past_key_value = past_key_values[i] if past_key_values is not None else None

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    layer_module.__call__,
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_mask,
                    layer_head_mask,
                    encoder_hidden_states,
                    encoder_attention_mask,
                    past_key_value,
                    output_attentions,
                )

            hidden_states = layer_outputs[0]
            if use_cache:
                next_decoder_cache += (layer_outputs[-1],)
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (layer_outputs[2],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    next_decoder_cache,
                    all_hidden_states,
                    all_self_attentions,
                    all_cross_attentions,
                ]
                if v is not None
            )
        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_decoder_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )



class AltRobertaPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        
        
        first_token_tensor = hidden_states[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output



class AltCLIPAttention(nn.Module):
    

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.embed_dim // self.num_heads
        if self.head_dim * self.num_heads != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
                f" {self.num_heads})."
            )
        self.scale = self.head_dim**-0.5
        self.dropout = config.attention_dropout

        self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
        self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)

    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        causal_attention_mask: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        

        bsz, tgt_len, embed_dim = hidden_states.size()

        
        query_states = self.q_proj(hidden_states) * self.scale
        key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
        value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

        proj_shape = (bsz * self.num_heads, -1, self.head_dim)
        query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
        key_states = key_states.view(*proj_shape)
        value_states = value_states.view(*proj_shape)

        src_len = key_states.size(1)
        attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))

        if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
            raise ValueError(
                f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
                f" {attn_weights.size()}"
            )

        
        if causal_attention_mask is not None:
            if causal_attention_mask.size() != (bsz, 1, tgt_len, src_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is"
                    f" {causal_attention_mask.size()}"
                )
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + causal_attention_mask
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        if attention_mask is not None:
            if attention_mask.size() != (bsz, 1, tgt_len, src_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
                )
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)

        attn_weights = nn.functional.softmax(attn_weights, dim=-1)

        if output_attentions:
            
            
            
            
            attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
            attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
        else:
            attn_weights_reshaped = None

        attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)

        attn_output = torch.bmm(attn_probs, value_states)

        if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
            raise ValueError(
                f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
                f" {attn_output.size()}"
            )

        attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
        attn_output = attn_output.transpose(1, 2)
        attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)

        attn_output = self.out_proj(attn_output)

        return attn_output, attn_weights_reshaped



class AltCLIPMLP(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.activation_fn = ACT2FN[config.hidden_act]
        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.fc1(hidden_states)
        hidden_states = self.activation_fn(hidden_states)
        hidden_states = self.fc2(hidden_states)
        return hidden_states



class AltCLIPEncoderLayer(nn.Module):
    def __init__(self, config: AltCLIPConfig):
        super().__init__()
        self.embed_dim = config.hidden_size
        self.self_attn = AltCLIPAttention(config)
        self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
        self.mlp = AltCLIPMLP(config)
        self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        causal_attention_mask: torch.Tensor,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[torch.FloatTensor]:
        
        residual = hidden_states

        hidden_states = self.layer_norm1(hidden_states)
        hidden_states, attn_weights = self.self_attn(
            hidden_states=hidden_states,
            attention_mask=attention_mask,
            causal_attention_mask=causal_attention_mask,
            output_attentions=output_attentions,
        )
        hidden_states = residual + hidden_states

        residual = hidden_states
        hidden_states = self.layer_norm2(hidden_states)
        hidden_states = self.mlp(hidden_states)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)

        if output_attentions:
            outputs += (attn_weights,)

        return outputs



class AltCLIPEncoder(nn.Module):
    

    def __init__(self, config: AltCLIPConfig):
        super().__init__()
        self.config = config
        self.layers = nn.ModuleList([AltCLIPEncoderLayer(config) for _ in range(config.num_hidden_layers)])
        self.gradient_checkpointing = False

    def forward(
        self,
        inputs_embeds,
        attention_mask: Optional[torch.Tensor] = None,
        causal_attention_mask: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutput]:
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        encoder_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        hidden_states = inputs_embeds
        for idx, encoder_layer in enumerate(self.layers):
            if output_hidden_states:
                encoder_states = encoder_states + (hidden_states,)
            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    encoder_layer.__call__,
                    hidden_states,
                    attention_mask,
                    causal_attention_mask,
                    output_attentions,
                )
            else:
                layer_outputs = encoder_layer(
                    hidden_states,
                    attention_mask,
                    causal_attention_mask,
                    output_attentions=output_attentions,
                )

            hidden_states = layer_outputs[0]

            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)

        if not return_dict:
            return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
        return BaseModelOutput(
            last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
        )



class AltCLIPVisionEmbeddings(nn.Module):
    def __init__(self, config: AltCLIPVisionConfig):
        super().__init__()
        self.config = config
        self.embed_dim = config.hidden_size
        self.image_size = config.image_size
        self.patch_size = config.patch_size

        self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))

        self.patch_embedding = nn.Conv2d(
            in_channels=config.num_channels,
            out_channels=self.embed_dim,
            kernel_size=self.patch_size,
            stride=self.patch_size,
            bias=False,
        )

        self.num_patches = (self.image_size // self.patch_size) ** 2
        self.num_positions = self.num_patches + 1
        self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
        self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)

    def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
        batch_size = pixel_values.shape[0]
        target_dtype = self.patch_embedding.weight.dtype
        patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype))  
        patch_embeds = patch_embeds.flatten(2).transpose(1, 2)

        class_embeds = self.class_embedding.expand(batch_size, 1, -1)
        embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
        embeddings = embeddings + self.position_embedding(self.position_ids)
        return embeddings


class AltCLIPPreTrainedModel(PreTrainedModel):
    

    config_class = AltCLIPConfig
    base_model_prefix = "altclip"
    supports_gradient_checkpointing = True

    def _init_weights(self, module):
        
        factor = self.config.initializer_factor
        if isinstance(module, AltCLIPVisionEmbeddings):
            factor = self.config.initializer_factor
            nn.init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor)
            nn.init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor)
            nn.init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor)
        elif isinstance(module, AltCLIPAttention):
            factor = self.config.initializer_factor
            in_proj_std = (module.embed_dim**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
            out_proj_std = (module.embed_dim**-0.5) * factor
            nn.init.normal_(module.q_proj.weight, std=in_proj_std)
            nn.init.normal_(module.k_proj.weight, std=in_proj_std)
            nn.init.normal_(module.v_proj.weight, std=in_proj_std)
            nn.init.normal_(module.out_proj.weight, std=out_proj_std)
        elif isinstance(module, AltCLIPMLP):
            factor = self.config.initializer_factor
            in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
            fc_std = (2 * module.config.hidden_size) ** -0.5 * factor
            nn.init.normal_(module.fc1.weight, std=fc_std)
            nn.init.normal_(module.fc2.weight, std=in_proj_std)
        elif isinstance(module, AltCLIPModel):
            nn.init.normal_(
                module.text_projection.weight,
                std=module.text_embed_dim**-0.5 * self.config.initializer_factor,
            )
            module.text_projection._is_hf_initialized = True
            nn.init.normal_(
                module.visual_projection.weight,
                std=module.vision_embed_dim**-0.5 * self.config.initializer_factor,
            )
            module.visual_projection._is_hf_initialized = True
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
        elif isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_factor)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_factor)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()



class AltCLIPVisionTransformer(nn.Module):
    def __init__(self, config: AltCLIPVisionConfig):
        super().__init__()
        self.config = config
        embed_dim = config.hidden_size

        self.embeddings = AltCLIPVisionEmbeddings(config)
        self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
        self.encoder = AltCLIPEncoder(config)
        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)

    @add_start_docstrings_to_model_forward(ALTCLIP_VISION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=AltCLIPVisionConfig)
    def forward(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if pixel_values is None:
            raise ValueError("You have to specify pixel_values")

        hidden_states = self.embeddings(pixel_values)
        hidden_states = self.pre_layrnorm(hidden_states)

        encoder_outputs = self.encoder(
            inputs_embeds=hidden_states,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        last_hidden_state = encoder_outputs[0]
        pooled_output = last_hidden_state[:, 0, :]
        pooled_output = self.post_layernorm(pooled_output)

        if not return_dict:
            return (last_hidden_state, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPooling(
            last_hidden_state=last_hidden_state,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )


class AltCLIPVisionModel(AltCLIPPreTrainedModel):
    config_class = AltCLIPVisionConfig
    main_input_name = "pixel_values"

    def __init__(self, config: AltCLIPVisionConfig):
        super().__init__(config)
        self.vision_model = AltCLIPVisionTransformer(config)
        
        self.post_init()

    def get_input_embeddings(self) -> nn.Module:
        return self.vision_model.embeddings.patch_embedding

    @add_start_docstrings_to_model_forward(ALTCLIP_VISION_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPooling, config_class=AltCLIPVisionConfig)
    def forward(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPooling]:
        r
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        return self.vision_model(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )


class AltRobertaModel(AltCLIPPreTrainedModel):
    

    config_class = AltCLIPTextConfig

    
    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

        self.embeddings = AltRobertaEmbeddings(config)
        self.encoder = AltRobertaEncoder(config)

        self.pooler = AltRobertaPooler(config) if add_pooling_layer else None

        
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def _prune_heads(self, heads_to_prune):
        
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

    
    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
        r
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if self.config.is_decoder:
            use_cache = use_cache if use_cache is not None else self.config.use_cache
        else:
            use_cache = False

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
            input_shape = input_ids.size()
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        batch_size, seq_length = input_shape
        device = input_ids.device if input_ids is not None else inputs_embeds.device

        
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if attention_mask is None:
            attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)

        if token_type_ids is None:
            if hasattr(self.embeddings, "token_type_ids"):
                buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
                buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
                token_type_ids = buffered_token_type_ids_expanded
            else:
                token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)

        
        
        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)

        
        
        if self.config.is_decoder and encoder_hidden_states is not None:
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
        else:
            encoder_extended_attention_mask = None

        
        
        
        
        
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        embedding_output = self.embeddings(
            input_ids=input_ids,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            inputs_embeds=inputs_embeds,
            past_key_values_length=past_key_values_length,
        )
        encoder_outputs = self.encoder(
            embedding_output,
            attention_mask=extended_attention_mask,
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPoolingAndCrossAttentions(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            past_key_values=encoder_outputs.past_key_values,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
            cross_attentions=encoder_outputs.cross_attentions,
        )


class AltCLIPTextModel(AltCLIPPreTrainedModel):
    config_class = AltCLIPTextConfig

    def __init__(self, config):
        super().__init__(config)
        self.roberta = AltRobertaModel(config, add_pooling_layer=False)
        self.transformation = nn.Linear(config.hidden_size, config.project_dim)
        self.pre_LN = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.post_init()

    def get_input_embeddings(self) -> nn.Module:
        return self.roberta.embeddings.word_embeddings

    def set_input_embeddings(self, value: nn.Embedding) -> None:
        self.roberta.embeddings.word_embeddings = value

    def resize_token_embeddings(self, new_num_tokens: Optional[int] = None) -> nn.Embedding:
        return super().resize_token_embeddings(new_num_tokens)

    @add_start_docstrings_to_model_forward(ALTCLIP_TEXT_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=BaseModelOutputWithPoolingAndProjection, config_class=AltCLIPTextConfig)
    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPoolingAndProjection]:
        r

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.roberta(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        
        sequence_output = outputs[0]

        
        sequence_output = self.pre_LN(sequence_output)

        
        projection_state = self.transformation(sequence_output)
        pooler_output = projection_state[:, 0]

        if not return_dict:
            return (projection_state, pooler_output) + outputs[2:4]

        return BaseModelOutputWithPoolingAndProjection(
            last_hidden_state=projection_state,
            pooler_output=pooler_output,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class AltCLIPModel(AltCLIPPreTrainedModel):
    config_class = AltCLIPConfig

    def __init__(self, config: AltCLIPConfig):
        super().__init__(config)

        if not isinstance(config.vision_config, AltCLIPVisionConfig):
            raise ValueError(
                "config.vision_config is expected to be of type AltCLIPVisionConfig but is of type"
                f" {type(config.vision_config)}."
            )
        if not isinstance(config.text_config, AltCLIPTextConfig):
            raise ValueError(
                "config.text_config is expected to be of type AltCLIPTextConfig but is of type"
                f" {type(config.text_config)}."
            )

        text_config = config.text_config
        vision_config = config.vision_config

        self.projection_dim = config.projection_dim
        self.text_embed_dim = text_config.project_dim
        self.vision_embed_dim = vision_config.hidden_size

        self.text_model = AltCLIPTextModel(text_config)
        self.vision_model = AltCLIPVisionTransformer(vision_config)

        self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False)
        self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim, bias=False)
        self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))

        
        self.post_init()

    @add_start_docstrings_to_model_forward(ALTCLIP_TEXT_INPUTS_DOCSTRING)
    def get_text_features(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        token_type_ids=None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> torch.FloatTensor:
        r
        
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        text_outputs = self.text_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            token_type_ids=token_type_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        pooled_output = text_outputs[1]
        text_features = self.text_projection(pooled_output)

        return text_features

    @add_start_docstrings_to_model_forward(ALTCLIP_VISION_INPUTS_DOCSTRING)
    def get_image_features(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> torch.FloatTensor:
        r
        
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        vision_outputs = self.vision_model(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        pooled_output = vision_outputs[1]  
        image_features = self.visual_projection(pooled_output)

        return image_features

    @add_start_docstrings_to_model_forward(ALTCLIP_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=AltCLIPOutput, config_class=AltCLIPConfig)
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.FloatTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        return_loss: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, AltCLIPOutput]:
        r
        
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        text_outputs = self.text_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        vision_outputs = self.vision_model(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        image_embeds = vision_outputs[1]
        image_embeds = self.visual_projection(image_embeds)

        text_embeds = text_outputs[1]
        text_embeds = self.text_projection(text_embeds)

        
        image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
        text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)

        
        logit_scale = self.logit_scale.exp()
        logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
        logits_per_image = logits_per_text.T

        loss = None
        if return_loss:
            loss = clip_loss(logits_per_text)

        if not return_dict:
            output = (logits_per_image, logits_per_text, text_embeds, image_embeds, text_outputs, vision_outputs)
            return ((loss,) + output) if loss is not None else output

        return AltCLIPOutput(
            loss=loss,
            logits_per_image=logits_per_image,
            logits_per_text=logits_per_text,
            text_embeds=text_embeds,
            image_embeds=image_embeds,
            text_model_output=text_outputs,
            vision_model_output=vision_outputs,
        )



def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
    
    
    mask = input_ids.ne(padding_idx).int()
    incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
    return incremental_indices.long() + padding_idx