



import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers import (
							    OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
							    OpenAIGPTConfig,
							    OpenAIGPTDoubleHeadsModel,
							    OpenAIGPTForSequenceClassification,
							    OpenAIGPTLMHeadModel,
							    OpenAIGPTModel,
							)
class 					_snake_case      :




				def __init__(    self       :    List[str],       __lowercase       :    int,       __lowercase       :    List[str]=13,       __lowercase       :    str=7,       __lowercase       :    int=True,       __lowercase       :    Any=True,       __lowercase       :    Any=True,       __lowercase       :    Any=99,       __lowercase       :    Optional[Any]=32,       __lowercase       :    Tuple=5,       __lowercase       :    List[Any]=4,       __lowercase       :    str=37,       __lowercase       :    Optional[Any]="gelu",       __lowercase       :    Optional[int]=0.1,       __lowercase       :    Tuple=0.1,       __lowercase       :    List[str]=512,       __lowercase       :    List[str]=16,       __lowercase       :    List[Any]=2,       __lowercase       :    Optional[int]=0.02,       __lowercase       :    Any=3,       __lowercase       :    int=4,       __lowercase       :    Tuple=None,       ):
							lowercase__       	=					parent
							lowercase__       	=					batch_size
							lowercase__       	=					seq_length
							lowercase__       	=					is_training
							lowercase__       	=					use_token_type_ids
							lowercase__       	=					use_labels
							lowercase__       	=					vocab_size
							lowercase__       	=					hidden_size
							lowercase__       	=					num_hidden_layers
							lowercase__       	=					num_attention_heads
							lowercase__       	=					intermediate_size
							lowercase__       	=					hidden_act
							lowercase__       	=					hidden_dropout_prob
							lowercase__       	=					attention_probs_dropout_prob
							lowercase__       	=					max_position_embeddings
							lowercase__       	=					type_vocab_size
							lowercase__       	=					type_sequence_label_size
							lowercase__       	=					initializer_range
							lowercase__       	=					num_labels
							lowercase__       	=					num_choices
							lowercase__       	=					scope
							lowercase__       	=					self.vocab_size - 1




				def 				A__   (    self       :    Union[str, Any]    ):
							lowercase__       	=					ids_tensor([self.batch_size, self.seq_length],       self.vocab_size    )

							lowercase__       	=					None
							if self.use_token_type_ids:
										lowercase__       	=					ids_tensor([self.batch_size, self.seq_length],       self.type_vocab_size    )

							lowercase__       	=					None
							lowercase__       	=					None
							lowercase__       	=					None
							if self.use_labels:
										lowercase__       	=					ids_tensor([self.batch_size],       self.type_sequence_label_size    )
										lowercase__       	=					ids_tensor([self.batch_size, self.seq_length],       self.num_labels    )
										lowercase__       	=					ids_tensor([self.batch_size],       self.num_choices    )

							lowercase__       	=					OpenAIGPTConfig(
							    vocab_size=self.vocab_size,       n_embd=self.hidden_size,       n_layer=self.num_hidden_layers,       n_head=self.num_attention_heads,       n_positions=self.max_position_embeddings,       pad_token_id=self.pad_token_id,       )

							lowercase__       	=					ids_tensor([self.num_hidden_layers, self.num_attention_heads],       2    )

							return (
							    config,
							    input_ids,
							    head_mask,
							    token_type_ids,
							    sequence_labels,
							    token_labels,
							    choice_labels,
							)




				def 				A__   (    self       :    Tuple,       __lowercase       :    int,       __lowercase       :    Any,       __lowercase       :    int,       __lowercase       :    List[Any],       *__lowercase       :    int    ):
							lowercase__       	=					OpenAIGPTModel(config=a_    )
							model.to(a_    )
							model.eval()

							lowercase__       	=					model(a_,       token_type_ids=a_,       head_mask=a_    )
							lowercase__       	=					model(a_,       token_type_ids=a_    )
							lowercase__       	=					model(a_    )

							self.parent.assertEqual(result.last_hidden_state.shape,       (self.batch_size, self.seq_length, self.hidden_size)    )




				def 				A__   (    self       :    Any,       __lowercase       :    str,       __lowercase       :    Optional[Any],       __lowercase       :    Dict,       __lowercase       :    str,       *__lowercase       :    Union[str, Any]    ):
							lowercase__       	=					OpenAIGPTLMHeadModel(a_    )
							model.to(a_    )
							model.eval()

							lowercase__       	=					model(a_,       token_type_ids=a_,       labels=a_    )
							self.parent.assertEqual(result.loss.shape,       ()    )
							self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.seq_length, self.vocab_size)    )




				def 				A__   (    self       :    List[str],       __lowercase       :    int,       __lowercase       :    Optional[Any],       __lowercase       :    Union[str, Any],       __lowercase       :    Optional[Any],       *__lowercase       :    List[str]    ):
							lowercase__       	=					OpenAIGPTDoubleHeadsModel(a_    )
							model.to(a_    )
							model.eval()

							lowercase__       	=					model(a_,       token_type_ids=a_,       labels=a_    )
							self.parent.assertEqual(result.loss.shape,       ()    )
							self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.seq_length, self.vocab_size)    )




				def 				A__   (    self       :    Optional[int],       __lowercase       :    List[str],       __lowercase       :    Dict,       __lowercase       :    Tuple,       __lowercase       :    Optional[Any],       *__lowercase       :    Union[str, Any]    ):
							lowercase__       	=					self.num_labels
							lowercase__       	=					OpenAIGPTForSequenceClassification(a_    )
							model.to(a_    )
							model.eval()

							lowercase__       	=					ids_tensor([self.batch_size],       self.type_sequence_label_size    )
							lowercase__       	=					model(a_,       token_type_ids=a_,       labels=a_    )
							self.parent.assertEqual(result.logits.shape,       (self.batch_size, self.num_labels)    )




				def 				A__   (    self       :    Optional[int]    ):
							lowercase__       	=					self.prepare_config_and_inputs()
							(
							    (
							    lowercase__   
							)	,							(
							    lowercase__   
							)	,							(
							    lowercase__   
							)	,							(
							    lowercase__   
							)	,							(
							    lowercase__   
							)	,							(
							    lowercase__   
							)	,							(
							    lowercase__   
							)	,							   
							)       	=					config_and_inputs
							lowercase__       	=					{
							    "input_ids": input_ids,
							    "token_type_ids": token_type_ids,
							    "head_mask": head_mask,
							}

							return config, inputs_dict
@require_torch
class 					_snake_case      (       lowercase__     ,						lowercase__     ,						lowercase__     ,						unittest.TestCase):
				UpperCamelCase__					:       Union[str, Any]  		=(
				    (OpenAIGPTModel, OpenAIGPTLMHeadModel, OpenAIGPTDoubleHeadsModel, OpenAIGPTForSequenceClassification)
				    if is_torch_available()
				    else ()
				)
				UpperCamelCase__					:       Tuple  		=(
				    (OpenAIGPTLMHeadModel,) if is_torch_available() else ()
				)  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
				UpperCamelCase__					:       str  		=(
				    {
				        "feature-extraction": OpenAIGPTModel,
				        "text-classification": OpenAIGPTForSequenceClassification,
				        "text-generation": OpenAIGPTLMHeadModel,
				        "zero-shot": OpenAIGPTForSequenceClassification,
				    }
				    if is_torch_available()
				    else {}
				)




				def 				A__   (    self       :    str,       __lowercase       :    Union[str, Any],       __lowercase       :    int,       __lowercase       :    Optional[Any],       __lowercase       :    Dict,       __lowercase       :    List[Any]    ):
							if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
										# Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
										# `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
										# tiny config could not be created.
										return True

							return False




				def 				A__   (    self       :    List[Any],       __lowercase       :    Optional[Any],       __lowercase       :    List[str],       __lowercase       :    List[str]=False    ):
							lowercase__       	=					super()._prepare_for_class(a_,       a_,       return_labels=a_    )

							if return_labels:
										if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
													lowercase__       	=					torch.zeros(
													    (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length),       dtype=torch.long,       device=a_,       )
													lowercase__       	=					inputs_dict["labels"]
													lowercase__       	=					inputs_dict["labels"]
													lowercase__       	=					torch.zeros(
													    (self.model_tester.batch_size, self.model_tester.num_choices),       dtype=torch.long,       device=a_,       )
													lowercase__       	=					torch.zeros(
													    self.model_tester.batch_size,       dtype=torch.long,       device=a_    )
							return inputs_dict




				def 				A__   (    self       :    Dict    ):
							lowercase__       	=					OpenAIGPTModelTester(self    )
							lowercase__       	=					ConfigTester(self,       config_class=a_,       n_embd=37    )




				def 				A__   (    self       :    Tuple    ):
							self.config_tester.run_common_tests()




				def 				A__   (    self       :    Dict    ):
							lowercase__       	=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_openai_gpt_model(*a_    )




				def 				A__   (    self       :    Dict    ):
							lowercase__       	=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_lm_head_model(*a_    )




				def 				A__   (    self       :    List[str]    ):
							lowercase__       	=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_double_lm_head_model(*a_    )




				def 				A__   (    self       :    Union[str, Any]    ):
							lowercase__       	=					self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*a_    )




				@slow
				def 				A__   (    self       :    Union[str, Any]    ):
							for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowercase__       	=					OpenAIGPTModel.from_pretrained(a_    )
										self.assertIsNotNone(a_    )
@require_torch
class 					_snake_case      (       unittest.TestCase):




				@slow
				def 				A__   (    self       :    Dict    ):
							lowercase__       	=					OpenAIGPTLMHeadModel.from_pretrained("openai-gpt"    )
							model.to(a_    )
							lowercase__       	=					torch.tensor([[481, 4735, 544]],       dtype=torch.long,       device=a_    )  # the president is
							lowercase__       	=					[
							    481,
							    4735,
							    544,
							    246,
							    963,
							    870,
							    762,
							    239,
							    244,
							    4_0477,
							    244,
							    249,
							    719,
							    881,
							    487,
							    544,
							    240,
							    244,
							    603,
							    481,
							]  # the president is a very good man. " \n " i\'m sure he is, " said the

							lowercase__       	=					model.generate(a_,       do_sample=a_    )
							self.assertListEqual(output_ids[0].tolist(),       a_    )







'''simple docstring'''




import argparse
import os

import torch

from transformers import FlavaConfig, FlavaForPreTraining
from transformers.models.flava.convert_dalle_to_flava_codebook import convert_dalle_checkpoint


def    lowerCamelCase__		(    SCREAMING_SNAKE_CASE   : str    ):
  # encoder.embeddings are double copied in original FLAVA
  return sum(param.float().sum() if 'encoder.embeddings' not in key else 0 for key, param in state_dict.items()    )


def    lowerCamelCase__		(    SCREAMING_SNAKE_CASE   : int							, SCREAMING_SNAKE_CASE   : Optional[int]    ):
  UpperCAmelCase    				=						{}

  for key, value in state_dict.items():
    if "text_encoder.embeddings" in key or "image_encoder.embeddings" in key:
      continue

    UpperCAmelCase    				=						key.replace('heads.cmd.mim_head.cls.predictions'							, 'mmm_image_head'    )
    UpperCAmelCase    				=						key.replace('heads.cmd.mlm_head.cls.predictions'							, 'mmm_text_head'    )
    UpperCAmelCase    				=						key.replace('heads.cmd.itm_head.cls'							, 'itm_head'    )
    UpperCAmelCase    				=						key.replace('heads.cmd.itm_head.pooler'							, 'itm_head.pooler'    )
    UpperCAmelCase    				=						key.replace('heads.cmd.clip_head.logit_scale'							, 'flava.logit_scale'    )
    UpperCAmelCase    				=						key.replace('heads.fairseq_mlm.cls.predictions'							, 'mlm_head'    )
    UpperCAmelCase    				=						key.replace('heads.imagenet.mim_head.cls.predictions'							, 'mim_head'    )
    UpperCAmelCase    				=						key.replace('mm_text_projection'							, 'flava.text_to_mm_projection'    )
    UpperCAmelCase    				=						key.replace('mm_image_projection'							, 'flava.image_to_mm_projection'    )
    UpperCAmelCase    				=						key.replace('image_encoder.module'							, 'flava.image_model'    )
    UpperCAmelCase    				=						key.replace('text_encoder.module'							, 'flava.text_model'    )
    UpperCAmelCase    				=						key.replace('mm_encoder.module.encoder.cls_token'							, 'flava.multimodal_model.cls_token'    )
    UpperCAmelCase    				=						key.replace('mm_encoder.module'							, 'flava.multimodal_model'    )
    UpperCAmelCase    				=						key.replace('text_projection'							, 'flava.text_projection'    )
    UpperCAmelCase    				=						key.replace('image_projection'							, 'flava.image_projection'    )

    UpperCAmelCase    				=						value.float()

  for key, value in codebook_state_dict.items():
    UpperCAmelCase    				=						value

  return upgrade


@torch.no_grad()
def    lowerCamelCase__		(    SCREAMING_SNAKE_CASE   : List[str]							, SCREAMING_SNAKE_CASE   : Optional[Any]							, SCREAMING_SNAKE_CASE   : Optional[int]							, SCREAMING_SNAKE_CASE   : Any=None    ):
  if config_path is not None:
    UpperCAmelCase    				=						FlavaConfig.from_pretrained(SCREAMING_SNAKE_CASE    )
  else:
    UpperCAmelCase    				=						FlavaConfig()

  UpperCAmelCase    				=						FlavaForPreTraining(SCREAMING_SNAKE_CASE    ).eval()

  UpperCAmelCase    				=						convert_dalle_checkpoint(SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE							, save_checkpoint=SCREAMING_SNAKE_CASE    )

  if os.path.exists(SCREAMING_SNAKE_CASE    ):
    UpperCAmelCase    				=						torch.load(SCREAMING_SNAKE_CASE							, map_location='cpu'    )
  else:
    UpperCAmelCase    				=						torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE							, map_location='cpu'    )

  UpperCAmelCase    				=						upgrade_state_dict(SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE    )
  hf_model.load_state_dict(SCREAMING_SNAKE_CASE    )
  UpperCAmelCase    				=						hf_model.state_dict()
  UpperCAmelCase    				=						count_parameters(SCREAMING_SNAKE_CASE    )
  UpperCAmelCase    				=						count_parameters(SCREAMING_SNAKE_CASE    ) + count_parameters(SCREAMING_SNAKE_CASE    )

  assert torch.allclose(SCREAMING_SNAKE_CASE							, SCREAMING_SNAKE_CASE							, atol=1E-3    )

  hf_model.save_pretrained(SCREAMING_SNAKE_CASE    )


if __name__ == "__main__":
  _a					:			str							=    argparse.ArgumentParser()
  parser.add_argument('--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model.')
  parser.add_argument('--checkpoint_path', default=None, type=str, help='Path to flava checkpoint')
  parser.add_argument('--codebook_path', default=None, type=str, help='Path to flava codebook checkpoint')
  parser.add_argument('--config_path', default=None, type=str, help='Path to hf config.json of model to convert')
  _a					:			str							=    parser.parse_args()

  convert_flava_checkpoint(args.checkpoint_path, args.codebook_path, args.pytorch_dump_folder_path, args.config_path)


