codingutf8 2023 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 fuyu model from typing import list optional tuple union import torch import torch utils checkpoint from torch import nn from modelingoutputs import causallmoutputwithpast from modelingutils import pretrainedmodel from models auto modelingauto import automodelforcausallm from utils import addstartdocstrings addstartdocstringstomodelforward logging replacereturndocstrings from configurationfuyu import fuyuconfig logger logging getloggername configfordoc fuyuconfig fuyustartdocstring 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 fuyuconfig 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 addstartdocstrings fuyu model with a language modeling head on top for causal language model conditioned on image patches and text fuyustartdocstring class fuyuforcausallmfuyupretrainedmodel def initself config fuyuconfig super initconfig self paddingidx config padtokenid self vocabsize config vocabsize self languagemodel automodelforcausallm fromconfigconfig textconfig self visionembedtokens nn linear config patchsize config patchsize config numchannels config hiddensize self gradientcheckpointing false initialize weights and apply final processing self postinit def getinputembeddingsself return self languagemodel getinputembeddings def setinputembeddingsself value self languagemodel setinputembeddingsvalue def gathercontinuousembeddings self wordembeddings torch tensor continuousembeddings listtorch tensor imagepatchinputindices torch tensor torch tensor if not wordembeddings shape0 lencontinuousembeddings raise valueerror fbatch sizes must match got lencontinuousembeddings and wordembeddings shape0 outputembeddings wordembeddings clone for batchidx in rangewordembeddings shape0 first find the positions of all the nonnegative values in imagepatchinputindices those are the positions in wordembeddings that we want to replace with content from continuousembeddings dstindices torch nonzeroimagepatchinputindicesbatchidx 0 astupletrue0 next look up those indices in imagepatchinputindices to find the indices in continuousembeddings that we want to use to replace the values in wordembeddings srcindices imagepatchinputindicesbatchidxdstindices check if we have more indices than embeddings note that we could have fewer indices if images got truncated if srcindices shape0 continuousembeddingsbatchidx shape0 raise valueerror fnumber of continuous embeddings continuousembeddingsbatchidx shape does not match fnumber of continuous token ids srcindices shape in batch element batchidx outputembeddingsbatchidx dstindices continuousembeddingsbatchidxsrcindices return outputembeddings addstartdocstringstomodelforwardfuyuinputsdocstring replacereturndocstringsoutputtypecausallmoutputwithpast configclassconfigfordoc def forward self inputids torch longtensor none imagepatches torch tensor none batchsize numtotalpatches patchsize x patchsize x numchannels imagepatchesindices torch tensor none attentionmask optionaltorch tensor none positionids optionaltorch longtensor none pastkeyvalues optionallisttorch floattensor none inputsembeds optionaltorch floattensor none usecache optionalbool none labels optionaltorch tensor none outputattentions optionalbool none outputhiddenstates optionalbool none returndict optionalbool none uniontuple causallmoutputwithpast r labels torch longtensor of shape batchsize sequencelength optional labels for computing the masked language modeling loss indices should either be in 0 config vocabsize or 100 see inputids docstring tokens with indices set to 100 are ignored masked the loss is only computed for the tokens with labels in 0 config vocabsize returns examples python from transformers import fuyuprocessor fuyuforcausallm from pil import image import requests processor fuyuprocessor frompretrainedadeptfuyu8b model fuyuforcausallm frompretrainedadeptfuyu8b url http images cocodataset orgval2017000000039769 jpg image image openrequests geturl streamtrue raw prompt generate a cocostyle caption n inputs processortexttextprompt imagesimage returntensorspt outputs modelinputs generatedids model generatemodelinputs maxnewtokens7 generationtext processor batchdecodegeneratedids skipspecialtokenstrue printgenerationtext a bus parked on the side of a road outputattentions outputattentions if outputattentions is not none else self config outputattentions outputhiddenstates outputhiddenstates if outputhiddenstates is not none else self config outputhiddenstates usecache usecache if usecache is not none else self config usecache returndict returndict if returndict is not none else self config usereturndict 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 batchsize seqlength inputids shape elif inputsembeds is not none batchsize seqlength inputsembeds shape else raise valueerroryou have to specify either inputis or inputsembeds seqlengthwithpast seqlength pastkeyvalueslength 0 if pastkeyvalues is not none pastkeyvalueslength pastkeyvalues00 shape2 seqlengthwithpast seqlengthwithpast pastkeyvalueslength if positionids is none device inputids device if inputids is not none else inputsembeds device positionids torch arange pastkeyvalueslength seqlength pastkeyvalueslength dtypetorch long devicedevice positionids positionids unsqueeze0 if inputsembeds is none inputsembeds self languagemodel getinputembeddingsinputids if imagepatches is not none and pastkeyvalues is none patchembeddings self visionembedtokenspatch toself visionembedtokens weight dtype squeeze0 for patch in imagepatches inputsembeds self gathercontinuousembeddings wordembeddingsinputsembeds continuousembeddingspatchembeddings imagepatchinputindicesimagepatchesindices outputs self languagemodel inputsembedsinputsembeds attentionmaskattentionmask positionidspositionids pastkeyvaluespastkeyvalues outputattentionsoutputattentions outputhiddenstatesoutputhiddenstates labelslabels usecacheusecache returndictreturndict return outputs def prepareinputsforgeneration self inputids pastkeyvaluesnone attentionmasknone inputsembedsnone imagepatchesnone imagepatchesindicesnone kwargs if pastkeyvalues inputids inputids 1 positionids kwargs getpositionids none if attentionmask is not none and positionids is none create positionids on the fly for batch generation positionids attentionmask long cumsum1 1 positionids maskedfillattentionmask 0 1 if pastkeyvalues positionids positionids 1 unsqueeze1 if inputsembeds are passed we only want to use them in the 1st generation step if inputsembeds is not none and pastkeyvalues is none modelinputs inputsembeds inputsembeds else modelinputs inputids inputids if imagepatchesindices is not none modelinputsimagepatchesindices imagepatchesindices modelinputs update positionids positionids pastkeyvalues pastkeyvalues usecache kwargs getusecache attentionmask attentionmask imagepatchesindices imagepatchesindices if pastkeyvalues is none else none imagepatches imagepatches if pastkeyvalues is none else none return modelinputs coding utf 8 2023 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 fuyu model 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 fuyuconfig 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 indices can be obtained using autotokenizer see pretrainedtokenizer encode and pretrainedtokenizer __call__ for details if past_key_values is used optionally only the last decoder_input_ids have to be input see past_key_values if you want to change padding behavior you should read modeling_opt _prepare_decoder_attention_mask and modify to your needs see diagram 1 in the paper https arxiv org abs 1910 13461 for more information on the default strategy 1 indicates the head is not masked 0 indicates the head is masked image_patches torch floattensor of shape batch_size num_total_patches patch_size_ x patch_size x num_channels optional image patches to be used as continuous embeddings the patches are flattened and then projected to the hidden size of the model image_patches_indices torch longtensor of shape batch_size num_total_patches number_of_newline_tokens number_of_text_tokens patch_size_ x patch_size x num_channels optional indices indicating at which position the image_patches have to be inserted in input_embeds position_ids torch longtensor of shape batch_size sequence_length optional indices of positions of each input sequence tokens in the position embeddings selected in the range 0 config n_positions 1 what are position ids glossary position ids past_key_values tuple tuple torch floattensor optional returned when use_cache true is passed or when config use_cache true tuple of tuple torch floattensor of length config n_layers with each tuple having 2 tensors of shape batch_size num_heads sequence_length embed_size_per_head and 2 additional tensors of shape batch_size num_heads encoder_sequence_length embed_size_per_head contains pre computed hidden states key and values in the self attention blocks and in the cross attention blocks that can be used see past_key_values input to speed up sequential decoding if past_key_values are used the user can optionally input only the last decoder_input_ids those that don t have their past key value states given to this model of shape batch_size 1 instead of all decoder_input_ids of shape batch_size sequence_length inputs_embeds torch floattensor of shape batch_size sequence_length hidden_size optional optionally instead of passing input_ids you can choose to directly pass an embedded representation this is useful if you want more control over how to convert input_ids indices into associated vectors than the model s internal embedding lookup matrix 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 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 initialize weights and apply final processing this function places the continuous_embeddings into the word_embeddings at the locations indicated by image_patch_input_indices different batch elements can have different numbers of continuous embeddings args word_embeddings torch floattensor of shape batch_size sequence_length hidden_size tensor of word embeddings continuous_embeddings torch floattensor of shape batch_size num_patches hidden_size tensor of continuous embeddings the length of the list is the batch size each entry is shape num_image_embeddings hidden and num_image_embeddings needs to match the number of non negative indices in image_patch_input_indices for that batch element image_patch_input_indices torch longtensor of shape batch_size sequence_length tensor of indices of the image patches in the input_ids tensor first find the positions of all the non negative values in image_patch_input_indices those are the positions in word_embeddings that we want to replace with content from continuous_embeddings next look up those indices in image_patch_input_indices to find the indices in continuous_embeddings that we want to use to replace the values in word_embeddings check if we have more indices than embeddings note that we could have fewer indices if images got truncated batch_size num_total_patches patch_size_ x patch_size x num_channels labels torch longtensor of shape batch_size sequence_length optional labels for computing the masked language modeling loss indices should either be in 0 config vocab_size or 100 see input_ids docstring tokens with indices set to 100 are ignored masked the loss is only computed for the tokens with labels in 0 config vocab_size returns examples python from transformers import fuyuprocessor fuyuforcausallm from pil import image import requests processor fuyuprocessor from_pretrained adept fuyu 8b model fuyuforcausallm from_pretrained adept fuyu 8b url http images cocodataset org val2017 000000039769 jpg image image open requests get url stream true raw prompt generate a coco style caption n inputs processor text text_prompt images image return_tensors pt outputs model inputs generated_ids model generate model_inputs max_new_tokens 7 generation_text processor batch_decode generated_ids skip_special_tokens true print generation_text a bus parked on the side of a road create position_ids on the fly for batch generation if inputs_embeds are passed we only want to use them in the 1st generation step
from typing import List, Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn

from ...modeling_outputs import CausalLMOutputWithPast
from ...modeling_utils import PreTrainedModel
from ...models.auto.modeling_auto import AutoModelForCausalLM
from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
from .configuration_fuyu import FuyuConfig


logger = logging.get_logger(__name__)

_CONFIG_FOR_DOC = "FuyuConfig"


FUYU_START_DOCSTRING = r


@add_start_docstrings(
    "The bare Fuyu Model outputting raw hidden-states without any specific head on top.",
    FUYU_START_DOCSTRING,
)
class FuyuPreTrainedModel(PreTrainedModel):
    config_class = FuyuConfig
    base_model_prefix = "fuyu"
    supports_gradient_checkpointing = True
    _no_split_modules = []
    _skip_keys_device_placement = "past_key_values"

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()


FUYU_INPUTS_DOCSTRING = r


@add_start_docstrings(
    "Fuyu Model with a language modeling head on top for causal language model conditioned on image patches and text.",
    FUYU_START_DOCSTRING,
)
class FuyuForCausalLM(FuyuPreTrainedModel):
    def __init__(self, config: FuyuConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size
        self.language_model = AutoModelForCausalLM.from_config(config.text_config)

        self.vision_embed_tokens = nn.Linear(
            config.patch_size * config.patch_size * config.num_channels, config.hidden_size
        )

        self.gradient_checkpointing = False
        
        self.post_init()

    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.language_model.set_input_embeddings(value)

    def gather_continuous_embeddings(
        self,
        word_embeddings: torch.Tensor,
        continuous_embeddings: List[torch.Tensor],
        image_patch_input_indices: torch.Tensor,
    ) -> torch.Tensor:
        
        if not (word_embeddings.shape[0] == len(continuous_embeddings)):
            raise ValueError(
                f"Batch sizes must match! Got {len(continuous_embeddings)=} and {word_embeddings.shape[0]=}"
            )

        output_embeddings = word_embeddings.clone()
        for batch_idx in range(word_embeddings.shape[0]):
            
            
            dst_indices = torch.nonzero(image_patch_input_indices[batch_idx] >= 0, as_tuple=True)[0]
            
            
            src_indices = image_patch_input_indices[batch_idx][dst_indices]
            
            if src_indices.shape[0] > continuous_embeddings[batch_idx].shape[0]:
                raise ValueError(
                    f"Number of continuous embeddings {continuous_embeddings[batch_idx].shape=} does not match "
                    f"number of continuous token ids {src_indices.shape=} in batch element {batch_idx}."
                )
            output_embeddings[batch_idx, dst_indices] = continuous_embeddings[batch_idx][src_indices]
        return output_embeddings

    @add_start_docstrings_to_model_forward(FUYU_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        image_patches: torch.Tensor = None,  
        image_patches_indices: torch.Tensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        labels: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        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
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache

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

        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:
            batch_size, seq_length = input_ids.shape
        elif inputs_embeds is not None:
            batch_size, seq_length, _ = inputs_embeds.shape
        else:
            raise ValueError("You have to specify either input_is or inputs_embeds")

        seq_length_with_past = seq_length
        past_key_values_length = 0

        if past_key_values is not None:
            past_key_values_length = past_key_values[0][0].shape[2]
            seq_length_with_past = seq_length_with_past + past_key_values_length

        if position_ids is None:
            device = input_ids.device if input_ids is not None else inputs_embeds.device
            position_ids = torch.arange(
                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
            )
            position_ids = position_ids.unsqueeze(0)

        if inputs_embeds is None:
            inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
            if image_patches is not None and past_key_values is None:
                patch_embeddings = [
                    self.vision_embed_tokens(patch.to(self.vision_embed_tokens.weight.dtype)).squeeze(0)
                    for patch in image_patches
                ]
                inputs_embeds = self.gather_continuous_embeddings(
                    word_embeddings=inputs_embeds,
                    continuous_embeddings=patch_embeddings,
                    image_patch_input_indices=image_patches_indices,
                )

        outputs = self.language_model(
            inputs_embeds=inputs_embeds,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            labels=labels,
            use_cache=use_cache,
            return_dict=return_dict,
        )

        return outputs

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        image_patches=None,
        image_patches_indices=None,
        **kwargs,
    ):
        if past_key_values:
            input_ids = input_ids[:, -1:]

        position_ids = kwargs.get("position_ids", None)
        if attention_mask is not None and position_ids is None:
            
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -1].unsqueeze(-1)

        
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        if image_patches_indices is not None:
            model_inputs["image_patches_indices"] = image_patches_indices

        model_inputs.update(
            {
                "position_ids": position_ids,
                "past_key_values": past_key_values,
                "use_cache": kwargs.get("use_cache"),
                "attention_mask": attention_mask,
                "image_patches_indices": image_patches_indices if past_key_values is None else None,
                "image_patches": image_patches if past_key_values is None else None,
            }
        )
        return model_inputs