






"""simple docstring"""
import unittest

from transformers import MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING, is_vision_available
from transformers.pipelines import pipeline
from transformers.testing_utils import (
    is_pipeline_test,
    nested_simplify,
    require_tf,
    require_torch,
    require_vision,
    slow,
)

from .test_pipelines_common import ANY


if is_vision_available():
							from PIL import Image
else:



							class      SCREAMING_SNAKE_CASE  :





								"""simple docstring"""




								@staticmethod
								def    lowerCamelCase(*lowerCAmelCase_     ,     **lowerCAmelCase_							):
									pass



@is_pipeline_test
@require_torch
@require_vision
class      SCREAMING_SNAKE_CASE  (				unittest.TestCase						):





	"""simple docstring"""




	_A    =				MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING





	def    lowerCamelCase(self     ,     lowerCAmelCase_     ,     lowerCAmelCase_     ,     lowerCAmelCase_							):
		A_		:						Union[str, Any]	    =   pipeline("""visual-question-answering"""     ,     model="""hf-internal-testing/tiny-vilt-random-vqa"""							)
		A_		:						Union[str, Any]	    =   [
		    {
		        """image""": Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""							),
		        """question""": """How many cats are there?""",
		    },
		    {
		        """image""": """./tests/fixtures/tests_samples/COCO/000000039769.png""",
		        """question""": """How many cats are there?""",
		    },
		]
		return vqa_pipeline, examples





	def    lowerCamelCase(self     ,     lowerCAmelCase_     ,     lowerCAmelCase_							):
		A_		:						str	    =   vqa_pipeline(lowerCAmelCase_     ,     top_k=1							)
		self.assertEqual(
		    lowerCAmelCase_     ,     [
		        [{"""score""": ANY(lowerCAmelCase_							), """answer""": ANY(lowerCAmelCase_							)}],
		        [{"""score""": ANY(lowerCAmelCase_							), """answer""": ANY(lowerCAmelCase_							)}],
		    ]     ,     )





	@require_torch
	def    lowerCamelCase(self							):
		A_		:						Tuple	    =   pipeline("""visual-question-answering"""     ,     model="""hf-internal-testing/tiny-vilt-random-vqa"""							)
		A_		:						Any	    =   """./tests/fixtures/tests_samples/COCO/000000039769.png"""
		A_		:						Any	    =   """How many cats are there?"""

		A_		:						Optional[int]	    =   vqa_pipeline(image=lowerCAmelCase_     ,     question="""How many cats are there?"""     ,     top_k=2							)
		self.assertEqual(
		    lowerCAmelCase_     ,     [{"""score""": ANY(lowerCAmelCase_							), """answer""": ANY(lowerCAmelCase_							)}, {"""score""": ANY(lowerCAmelCase_							), """answer""": ANY(lowerCAmelCase_							)}]							)

		A_		:						Any	    =   vqa_pipeline({"""image""": image, """question""": question}     ,     top_k=2							)
		self.assertEqual(
		    lowerCAmelCase_     ,     [{"""score""": ANY(lowerCAmelCase_							), """answer""": ANY(lowerCAmelCase_							)}, {"""score""": ANY(lowerCAmelCase_							), """answer""": ANY(lowerCAmelCase_							)}]							)





	@slow
	@require_torch
	def    lowerCamelCase(self							):
		A_		:						int	    =   pipeline("""visual-question-answering"""     ,     model="""dandelin/vilt-b32-finetuned-vqa"""							)
		A_		:						Tuple	    =   """./tests/fixtures/tests_samples/COCO/000000039769.png"""
		A_		:						int	    =   """How many cats are there?"""

		A_		:						List[Any]	    =   vqa_pipeline(image=lowerCAmelCase_     ,     question=lowerCAmelCase_     ,     top_k=2							)
		self.assertEqual(
		    nested_simplify(lowerCAmelCase_     ,     decimals=4							)     ,     [{"""score""": 0.8799, """answer""": """2"""}, {"""score""": 0.296, """answer""": """1"""}]							)

		A_		:						Any	    =   vqa_pipeline({"""image""": image, """question""": question}     ,     top_k=2							)
		self.assertEqual(
		    nested_simplify(lowerCAmelCase_     ,     decimals=4							)     ,     [{"""score""": 0.8799, """answer""": """2"""}, {"""score""": 0.296, """answer""": """1"""}]							)

		A_		:						Tuple	    =   vqa_pipeline(
		    [{"""image""": image, """question""": question}, {"""image""": image, """question""": question}]     ,     top_k=2							)
		self.assertEqual(
		    nested_simplify(lowerCAmelCase_     ,     decimals=4							)     ,     [[{"""score""": 0.8799, """answer""": """2"""}, {"""score""": 0.296, """answer""": """1"""}]] * 2     ,     )





	@require_tf
	@unittest.skip("""Visual question answering not implemented in TF"""							)
	def    lowerCamelCase(self							):
		pass









"""simple docstring"""
import os

import torch

from ..logging import get_logger
from .constants import FSDP_PYTORCH_VERSION, MODEL_NAME, OPTIMIZER_NAME
from .versions import is_torch_version


if is_torch_version(">=", FSDP_PYTORCH_VERSION):
							import torch.distributed.checkpoint as dist_cp
							from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
							from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
							from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
							from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType


_lowerCAmelCase         =							get_logger(__name__)






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__=0					):
	os.makedirs(snake_case__   ,   exist_ok=snake_case__					)
	with FSDP.state_dict_type(
	    snake_case__   ,   fsdp_plugin.state_dict_type   ,   fsdp_plugin.state_dict_config   ,   fsdp_plugin.optim_state_dict_config					):
		A_		:						Optional[int]	    =   model.state_dict()
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			A_		:						int	    =   F"""{MODEL_NAME}.bin""" if model_index == 0 else F"""{MODEL_NAME}_{model_index}.bin"""
			A_		:						Union[str, Any]	    =   os.path.join(snake_case__   ,   snake_case__					)
			if accelerator.process_index == 0:
				logger.info(F"""Saving model to {output_model_file}"""					)
				torch.save(snake_case__   ,   snake_case__					)
				logger.info(F"""Model saved to {output_model_file}"""					)
		elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
			A_		:						List[str]	    =   (
			    F"""{MODEL_NAME}_rank{accelerator.process_index}.bin"""
			    if model_index == 0
			    else F"""{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"""
			)
			A_		:						List[str]	    =   os.path.join(snake_case__   ,   snake_case__					)
			logger.info(F"""Saving model to {output_model_file}"""					)
			torch.save(snake_case__   ,   snake_case__					)
			logger.info(F"""Model saved to {output_model_file}"""					)
		elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
			A_		:						str	    =   os.path.join(snake_case__   ,   F"""{MODEL_NAME}_{model_index}"""					)
			os.makedirs(snake_case__   ,   exist_ok=snake_case__					)
			logger.info(F"""Saving model to {ckpt_dir}"""					)
			A_		:						Dict	    =   {"""model""": state_dict}

			dist_cp.save_state_dict(
			    state_dict=snake_case__   ,   storage_writer=dist_cp.FileSystemWriter(snake_case__					)   ,   planner=DefaultSavePlanner()   ,   )
			logger.info(F"""Model saved to {ckpt_dir}"""					)






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__=0					):
	accelerator.wait_for_everyone()
	with FSDP.state_dict_type(
	    snake_case__   ,   fsdp_plugin.state_dict_type   ,   fsdp_plugin.state_dict_config   ,   fsdp_plugin.optim_state_dict_config					):
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			if type(snake_case__					) != FSDP and accelerator.process_index != 0:
				if not fsdp_plugin.sync_module_states:
					raise ValueError(
					    """Set the `sync_module_states` flag to `True` so that model states are synced across processes when """
					    """initializing FSDP object"""					)
				return
			A_		:						Tuple	    =   F"""{MODEL_NAME}.bin""" if model_index == 0 else F"""{MODEL_NAME}_{model_index}.bin"""
			A_		:						Any	    =   os.path.join(snake_case__   ,   snake_case__					)
			logger.info(F"""Loading model from {input_model_file}"""					)
			A_		:						List[Any]	    =   torch.load(snake_case__					)
			logger.info(F"""Model loaded from {input_model_file}"""					)
		elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
			A_		:						List[str]	    =   (
			    F"""{MODEL_NAME}_rank{accelerator.process_index}.bin"""
			    if model_index == 0
			    else F"""{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"""
			)
			A_		:						Any	    =   os.path.join(snake_case__   ,   snake_case__					)
			logger.info(F"""Loading model from {input_model_file}"""					)
			A_		:						List[str]	    =   torch.load(snake_case__					)
			logger.info(F"""Model loaded from {input_model_file}"""					)
		elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
			A_		:						Optional[Any]	    =   (
			    os.path.join(snake_case__   ,   F"""{MODEL_NAME}_{model_index}"""					)
			    if F"""{MODEL_NAME}""" not in input_dir
			    else input_dir
			)
			logger.info(F"""Loading model from {ckpt_dir}"""					)
			A_		:						List[str]	    =   {"""model""": model.state_dict()}
			dist_cp.load_state_dict(
			    state_dict=snake_case__   ,   storage_reader=dist_cp.FileSystemReader(snake_case__					)   ,   planner=DefaultLoadPlanner()   ,   )
			A_		:						Tuple	    =   state_dict["""model"""]
			logger.info(F"""Model loaded from {ckpt_dir}"""					)
		model.load_state_dict(snake_case__					)






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__=0					):
	os.makedirs(snake_case__   ,   exist_ok=snake_case__					)
	with FSDP.state_dict_type(
	    snake_case__   ,   fsdp_plugin.state_dict_type   ,   fsdp_plugin.state_dict_config   ,   fsdp_plugin.optim_state_dict_config					):
		A_		:						Optional[Any]	    =   FSDP.optim_state_dict(snake_case__   ,   snake_case__					)
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			if accelerator.process_index == 0:
				A_		:						Any	    =   (
				    F"""{OPTIMIZER_NAME}.bin""" if optimizer_index == 0 else F"""{OPTIMIZER_NAME}_{optimizer_index}.bin"""
				)
				A_		:						Optional[int]	    =   os.path.join(snake_case__   ,   snake_case__					)
				logger.info(F"""Saving Optimizer state to {output_optimizer_file}"""					)
				torch.save(snake_case__   ,   snake_case__					)
				logger.info(F"""Optimizer state saved in {output_optimizer_file}"""					)
		else:
			A_		:						Union[str, Any]	    =   os.path.join(snake_case__   ,   F"""{OPTIMIZER_NAME}_{optimizer_index}"""					)
			os.makedirs(snake_case__   ,   exist_ok=snake_case__					)
			logger.info(F"""Saving Optimizer state to {ckpt_dir}"""					)
			dist_cp.save_state_dict(
			    state_dict={"""optimizer""": optim_state}   ,   storage_writer=dist_cp.FileSystemWriter(snake_case__					)   ,   planner=DefaultSavePlanner()   ,   )
			logger.info(F"""Optimizer state saved in {ckpt_dir}"""					)






def 			__UpperCamelCase	(		snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__   ,   snake_case__=0					):
	accelerator.wait_for_everyone()
	with FSDP.state_dict_type(
	    snake_case__   ,   fsdp_plugin.state_dict_type   ,   fsdp_plugin.state_dict_config   ,   fsdp_plugin.optim_state_dict_config					):
		if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
			A_		:						List[str]	    =   None
			# below check should work but currently it isn't working (mostly opytorch issue),
			# in the meantime disabling it at the cost of excess memory usage
			# if accelerator.process_index == 0 or not fsdp_plugin.optim_state_dict_config.rank0_only:
			A_		:						Tuple	    =   (
			    F"""{OPTIMIZER_NAME}.bin""" if optimizer_index == 0 else F"""{OPTIMIZER_NAME}_{optimizer_index}.bin"""
			)
			A_		:						int	    =   os.path.join(snake_case__   ,   snake_case__					)
			logger.info(F"""Loading Optimizer state from {input_optimizer_file}"""					)
			A_		:						Tuple	    =   torch.load(snake_case__					)
			logger.info(F"""Optimizer state loaded from {input_optimizer_file}"""					)
		else:
			A_		:						str	    =   (
			    os.path.join(snake_case__   ,   F"""{OPTIMIZER_NAME}_{optimizer_index}"""					)
			    if F"""{OPTIMIZER_NAME}""" not in input_dir
			    else input_dir
			)
			logger.info(F"""Loading Optimizer from {ckpt_dir}"""					)
			A_		:						Union[str, Any]	    =   load_sharded_optimizer_state_dict(
			    model_state_dict=model.state_dict()   ,   optimizer_key="""optimizer"""   ,   storage_reader=dist_cp.FileSystemReader(snake_case__					)   ,   )
			A_		:						Tuple	    =   optim_state["""optimizer"""]
			logger.info(F"""Optimizer loaded from {ckpt_dir}"""					)
		A_		:						List[Any]	    =   FSDP.optim_state_dict_to_load(snake_case__   ,   snake_case__   ,   snake_case__					)
		optimizer.load_state_dict(snake_case__					)

