




import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
					from sagemaker import Session, TrainingJobAnalytics
					from sagemaker.huggingface import HuggingFace
@pytest.mark.skipif(
    literal_eval(os.getenv("""TEST_SAGEMAKER"""			,    """False"""				)				) is not True			,    reason="""Skipping test because should only be run when releasing minor transformers version"""			,    )
@pytest.mark.usefixtures("""sm_env"""				)
@parameterized_class(
    [
        {
            """framework""": """pytorch""",
            """script""": """run_glue_model_parallelism.py""",
            """model_name_or_path""": """roberta-large""",
            """instance_type""": """ml.p3dn.24xlarge""",
            """results""": {"""train_runtime""": 1_6_0_0, """eval_accuracy""": 0.3, """eval_loss""": 1.2},
        },
        {
            """framework""": """pytorch""",
            """script""": """run_glue.py""",
            """model_name_or_path""": """roberta-large""",
            """instance_type""": """ml.p3dn.24xlarge""",
            """results""": {"""train_runtime""": 1_6_0_0, """eval_accuracy""": 0.3, """eval_loss""": 1.2},
        },
    ]				)
class       a__   ( unittest.TestCase				):






						def     __SCREAMING_SNAKE_CASE(       self      ):




											"""simple docstring"""



											if self.framework == "pytorch":
																subprocess.run(
																    f"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()		,      encoding="utf-8"		,      check=lowerCAmelCase__		,      )
											assert hasattr(self		,      "env"      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	{
											    "enabled": True,
											    "processes_per_host": 8,
											}
											__lowerCAmelCase           =	{
											    "enabled": True,
											    "parameters": {
											        "microbatches": 4,
											        "placement_strategy": "spread",
											        "pipeline": "interleaved",
											        "optimize": "speed",
											        "partitions": 4,
											        "ddp": True,
											    },
											}

											__lowerCAmelCase           =	{"smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options}

											__lowerCAmelCase           =	"trainer" if self.script == "run_glue.py" else "smtrainer"
											# creates estimator
											return HuggingFace(
											    entry_point=self.script		,      source_dir=self.env.test_path		,      role=self.env.role		,      image_uri=self.env.image_uri		,      base_job_name=f"""{self.env.base_job_name}-{instance_count}-smp-{name_extension}"""		,      instance_count=lowerCAmelCase__		,      instance_type=self.instance_type		,      debugger_hook_config=lowerCAmelCase__		,      hyperparameters={
											        **self.env.hyperparameters,
											        "model_name_or_path": self.model_name_or_path,
											        "max_steps": 5_0_0,
											    }		,      metric_definitions=self.env.metric_definitions		,      distribution=lowerCAmelCase__		,      py_version="py36"		,      )






						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											TrainingJobAnalytics(lowerCAmelCase__      ).export_csv(f"""{self.env.test_path}/{job_name}_metrics.csv"""      )






						@parameterized.expand([(1,)]      )
						def     __SCREAMING_SNAKE_CASE(       self		,      _A      ):




											"""simple docstring"""



											__lowerCAmelCase           =	self.create_estimator(lowerCAmelCase__      )

											# run training
											estimator.fit()

											# result dataframe
											__lowerCAmelCase           =	TrainingJobAnalytics(estimator.latest_training_job.name      ).dataframe()

											# extract kpis
											__lowerCAmelCase           =	list(result_metrics_df[result_metrics_df.metric_name == "eval_accuracy"]["value"]      )
											__lowerCAmelCase           =	list(result_metrics_df[result_metrics_df.metric_name == "eval_loss"]["value"]      )
											# get train time from SageMaker job, this includes starting, preprocessing, stopping
											__lowerCAmelCase           =	(
											    Session().describe_training_job(estimator.latest_training_job.name      ).get("TrainingTimeInSeconds"		,      9_9_9_9_9_9      )
											)

											# assert kpis
											assert train_runtime <= self.results["train_runtime"]
											assert all(t >= self.results["eval_accuracy"] for t in eval_accuracy      )
											assert all(t <= self.results["eval_loss"] for t in eval_loss      )

											# dump tests result into json file to share in PR
											with open(f"""{estimator.latest_training_job.name}.json"""		,      "w"      ) as outfile:
																json.dump({"train_time": train_runtime, "eval_accuracy": eval_accuracy, "eval_loss": eval_loss}		,      lowerCAmelCase__      )










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      _a				(						SCREAMING_SNAKE_CASE_							:    Any					):
					# 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      _a				(						SCREAMING_SNAKE_CASE_							:    Dict				,				SCREAMING_SNAKE_CASE_							:    Dict					):
					__lowerCAmelCase           =	{}

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

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

										__lowerCAmelCase           =	value.float()

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

					return upgrade







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

					__lowerCAmelCase           =	FlavaForPreTraining(SCREAMING_SNAKE_CASE_					).eval()

					__lowerCAmelCase           =	convert_dalle_checkpoint(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_				,				save_checkpoint=SCREAMING_SNAKE_CASE_					)

					if os.path.exists(SCREAMING_SNAKE_CASE_					):
										__lowerCAmelCase           =	torch.load(SCREAMING_SNAKE_CASE_				,				map_location="cpu"					)
					else:
										__lowerCAmelCase           =	torch.hub.load_state_dict_from_url(SCREAMING_SNAKE_CASE_				,				map_location="cpu"					)

					__lowerCAmelCase           =	upgrade_state_dict(SCREAMING_SNAKE_CASE_				,				SCREAMING_SNAKE_CASE_					)
					hf_model.load_state_dict(SCREAMING_SNAKE_CASE_					)
					__lowerCAmelCase           =	hf_model.state_dict()
					__lowerCAmelCase           =	count_parameters(SCREAMING_SNAKE_CASE_					)
					__lowerCAmelCase           =	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__":
					UpperCamelCase__         =      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""")
					UpperCamelCase__         =      parser.parse_args()

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




