usrbinenv python codingutf8 2022 the huggingface 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 create a visionencoderdecodermodel instance from pretrained encoderdecoder models the crossattention will be randomly initialized arguments pertaining to which modelconfigtokenizer we are going to finetune or train from scratch load pretrained model and tokenizer use explicit specified encoder config use pretrained encoder model s config use explicit specified decoder config use pretrained decoder model s config necessary for fromencoderdecoderpretrained when decoderconfig is passed gpt2 only has boseos tokens but not decoderstartpad tokens this is necessary to make flax s generate work usr bin env python coding utf 8 2022 the huggingface 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 create a visionencoderdecodermodel instance from pretrained encoder decoder models the cross attention will be randomly initialized arguments pertaining to which model config tokenizer we are going to fine tune or train from scratch load pretrained model and tokenizer use explicit specified encoder config use pretrained encoder model s config use explicit specified decoder config use pretrained decoder model s config necessary for from_encoder_decoder_pretrained when decoder_config is passed gpt2 only has bos eos tokens but not decoder_start pad tokens this is necessary to make flax s generate work
from dataclasses import dataclass, field
from typing import Optional

from transformers import AutoConfig, AutoImageProcessor, AutoTokenizer, FlaxVisionEncoderDecoderModel, HfArgumentParser


@dataclass
class ModelArguments:
    

    output_dir: str = field(
        metadata={"help": "The output directory where the model will be written."},
    )
    encoder_model_name_or_path: str = field(
        metadata={
            "help": (
                "The encoder model checkpoint for weights initialization. "
                "Don't set if you want to train an encoder model from scratch."
            )
        },
    )
    decoder_model_name_or_path: str = field(
        metadata={
            "help": (
                "The decoder model checkpoint for weights initialization. "
                "Don't set if you want to train a decoder model from scratch."
            )
        },
    )
    encoder_config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained encoder config name or path if not the same as encoder_model_name"}
    )
    decoder_config_name: Optional[str] = field(
        default=None, metadata={"help": "Pretrained decoder config name or path if not the same as decoder_model_name"}
    )


def main():
    parser = HfArgumentParser((ModelArguments,))
    (model_args,) = parser.parse_args_into_dataclasses()

    

    
    if model_args.encoder_config_name:
        encoder_config = AutoConfig.from_pretrained(model_args.encoder_config_name)
    
    else:
        encoder_config = AutoConfig.from_pretrained(model_args.encoder_model_name_or_path)

    
    if model_args.decoder_config_name:
        decoder_config = AutoConfig.from_pretrained(model_args.decoder_config_name)
    
    else:
        decoder_config = AutoConfig.from_pretrained(model_args.decoder_model_name_or_path)

    
    decoder_config.is_decoder = True
    decoder_config.add_cross_attention = True

    model = FlaxVisionEncoderDecoderModel.from_encoder_decoder_pretrained(
        encoder_pretrained_model_name_or_path=model_args.encoder_model_name_or_path,
        decoder_pretrained_model_name_or_path=model_args.decoder_model_name_or_path,
        encoder_config=encoder_config,
        decoder_config=decoder_config,
    )

    
    decoder_start_token_id = decoder_config.decoder_start_token_id
    pad_token_id = decoder_config.pad_token_id
    if decoder_start_token_id is None:
        decoder_start_token_id = decoder_config.bos_token_id
    if pad_token_id is None:
        pad_token_id = decoder_config.eos_token_id

    
    model.config.eos_token_id = decoder_config.eos_token_id
    model.config.decoder_start_token_id = decoder_start_token_id
    model.config.pad_token_id = pad_token_id

    image_processor = AutoImageProcessor.from_pretrained(model_args.encoder_model_name_or_path)

    tokenizer = AutoTokenizer.from_pretrained(model_args.decoder_model_name_or_path)
    tokenizer.pad_token = tokenizer.convert_ids_to_tokens(model.config.pad_token_id)

    model.save_pretrained(model_args.output_dir)
    image_processor.save_pretrained(model_args.output_dir)
    tokenizer.save_pretrained(model_args.output_dir)


if __name__ == "__main__":
    main()