'''simple docstring'''





from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
			import numpy
			import tensorflow as tf

			from transformers import (
			    TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
			    TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
			    TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
			    BertConfig,
			    DPRConfig,
			    TFDPRContextEncoder,
			    TFDPRQuestionEncoder,
			    TFDPRReader,
			)




class 		SCREAMING_SNAKE_CASE     :







				'''simple docstring'''





				def __init__(						self					:						List[Any]       ,		lowercase__					:						Dict       ,		lowercase__					:						Union[str, Any]=13       ,		lowercase__					:						Any=7       ,		lowercase__					:						Any=True       ,		lowercase__					:						Optional[int]=True       ,		lowercase__					:						Optional[int]=True       ,		lowercase__					:						int=True       ,		lowercase__					:						List[str]=99       ,		lowercase__					:						Union[str, Any]=32       ,		lowercase__					:						int=2       ,		lowercase__					:						int=4       ,		lowercase__					:						Optional[Any]=37       ,		lowercase__					:						Dict="gelu"       ,		lowercase__					:						Any=0.1       ,		lowercase__					:						Any=0.1       ,		lowercase__					:						Union[str, Any]=512       ,		lowercase__					:						Union[str, Any]=16       ,		lowercase__					:						List[str]=2       ,		lowercase__					:						Union[str, Any]=0.0_2       ,		lowercase__					:						Tuple=3       ,		lowercase__					:						List[str]=4       ,		lowercase__					:						Optional[int]=None       ,		lowercase__					:						Optional[int]=0       ,		)							->Any:


										'''simple docstring'''




										_UpperCamelCase			:			Optional[Any]   			=						parent
										_UpperCamelCase			:			Optional[int]   			=						batch_size
										_UpperCamelCase			:			str   			=						seq_length
										_UpperCamelCase			:			Union[str, Any]   			=						is_training
										_UpperCamelCase			:			Optional[int]   			=						use_input_mask
										_UpperCamelCase			:			Union[str, Any]   			=						use_token_type_ids
										_UpperCamelCase			:			List[Any]   			=						use_labels
										_UpperCamelCase			:			int   			=						vocab_size
										_UpperCamelCase			:			Union[str, Any]   			=						hidden_size
										_UpperCamelCase			:			Tuple   			=						num_hidden_layers
										_UpperCamelCase			:			Any   			=						num_attention_heads
										_UpperCamelCase			:			Optional[int]   			=						intermediate_size
										_UpperCamelCase			:			str   			=						hidden_act
										_UpperCamelCase			:			int   			=						hidden_dropout_prob
										_UpperCamelCase			:			Optional[int]   			=						attention_probs_dropout_prob
										_UpperCamelCase			:			Union[str, Any]   			=						max_position_embeddings
										_UpperCamelCase			:			str   			=						type_vocab_size
										_UpperCamelCase			:			Optional[int]   			=						type_sequence_label_size
										_UpperCamelCase			:			Union[str, Any]   			=						initializer_range
										_UpperCamelCase			:			Tuple   			=						num_labels
										_UpperCamelCase			:			Optional[Any]   			=						num_choices
										_UpperCamelCase			:			Optional[Any]   			=						scope
										_UpperCamelCase			:			Optional[int]   			=						projection_dim





				def  snake_case__				(						self					:						Any					)							->List[Any]:


										'''simple docstring'''




										_UpperCamelCase			:			Optional[Any]   			=						ids_tensor([self.batch_size, self.seq_length]       ,		self.vocab_size					)

										_UpperCamelCase			:			Dict   			=						None
										if self.use_input_mask:
																# follow test_modeling_tf_ctrl.py
																_UpperCamelCase			:			Any   			=						random_attention_mask([self.batch_size, self.seq_length]					)

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

										_UpperCamelCase			:			int   			=						None
										_UpperCamelCase			:			Dict   			=						None
										_UpperCamelCase			:			List[str]   			=						None
										if self.use_labels:
																_UpperCamelCase			:			Union[str, Any]   			=						ids_tensor([self.batch_size]       ,		self.type_sequence_label_size					)
																_UpperCamelCase			:			Optional[int]   			=						ids_tensor([self.batch_size, self.seq_length]       ,		self.num_labels					)
																_UpperCamelCase			:			int   			=						ids_tensor([self.batch_size]       ,		self.num_choices					)

										_UpperCamelCase			:			Optional[Any]   			=						BertConfig(
										    vocab_size=self.vocab_size       ,		hidden_size=self.hidden_size       ,		num_hidden_layers=self.num_hidden_layers       ,		num_attention_heads=self.num_attention_heads       ,		intermediate_size=self.intermediate_size       ,		hidden_act=self.hidden_act       ,		hidden_dropout_prob=self.hidden_dropout_prob       ,		attention_probs_dropout_prob=self.attention_probs_dropout_prob       ,		max_position_embeddings=self.max_position_embeddings       ,		type_vocab_size=self.type_vocab_size       ,		is_decoder=lowercase__       ,		initializer_range=self.initializer_range       ,		)
										_UpperCamelCase			:			Union[str, Any]   			=						DPRConfig(projection_dim=self.projection_dim       ,		**config.to_dict()					)

										return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels





				def  snake_case__				(						self					:						List[str]       ,		lowercase__					:						Tuple       ,		lowercase__					:						str       ,		lowercase__					:						Tuple       ,		lowercase__					:						List[Any]       ,		lowercase__					:						List[Any]       ,		lowercase__					:						Tuple       ,		lowercase__					:						Union[str, Any]					)							->Tuple:


										'''simple docstring'''




										_UpperCamelCase			:			Optional[int]   			=						TFDPRContextEncoder(config=lowercase__					)
										_UpperCamelCase			:			str   			=						model(lowercase__       ,		attention_mask=lowercase__       ,		token_type_ids=lowercase__					)
										_UpperCamelCase			:			Any   			=						model(lowercase__       ,		token_type_ids=lowercase__					)
										_UpperCamelCase			:			Tuple   			=						model(lowercase__					)
										self.parent.assertEqual(result.pooler_output.shape       ,		(self.batch_size, self.projection_dim or self.hidden_size)					)





				def  snake_case__				(						self					:						Optional[int]       ,		lowercase__					:						Tuple       ,		lowercase__					:						int       ,		lowercase__					:						Dict       ,		lowercase__					:						Any       ,		lowercase__					:						List[str]       ,		lowercase__					:						Optional[Any]       ,		lowercase__					:						int					)							->Dict:


										'''simple docstring'''




										_UpperCamelCase			:			Optional[Any]   			=						TFDPRQuestionEncoder(config=lowercase__					)
										_UpperCamelCase			:			Any   			=						model(lowercase__       ,		attention_mask=lowercase__       ,		token_type_ids=lowercase__					)
										_UpperCamelCase			:			Dict   			=						model(lowercase__       ,		token_type_ids=lowercase__					)
										_UpperCamelCase			:			Dict   			=						model(lowercase__					)
										self.parent.assertEqual(result.pooler_output.shape       ,		(self.batch_size, self.projection_dim or self.hidden_size)					)





				def  snake_case__				(						self					:						str       ,		lowercase__					:						Tuple       ,		lowercase__					:						Union[str, Any]       ,		lowercase__					:						Dict       ,		lowercase__					:						str       ,		lowercase__					:						Union[str, Any]       ,		lowercase__					:						Optional[int]       ,		lowercase__					:						List[str]					)							->Any:


										'''simple docstring'''




										_UpperCamelCase			:			Tuple   			=						TFDPRReader(config=lowercase__					)
										_UpperCamelCase			:			Optional[Any]   			=						model(lowercase__       ,		attention_mask=lowercase__					)

										self.parent.assertEqual(result.start_logits.shape       ,		(self.batch_size, self.seq_length)					)
										self.parent.assertEqual(result.end_logits.shape       ,		(self.batch_size, self.seq_length)					)
										self.parent.assertEqual(result.relevance_logits.shape       ,		(self.batch_size,)					)







				def  snake_case__				(						self					:						Optional[Any]					)							->str:


										'''simple docstring'''




										_UpperCamelCase			:			Union[str, Any]   			=						self.prepare_config_and_inputs()
										(
										    (
										    _UpperCamelCase
										)			,       (
										    _UpperCamelCase
										)			,       (
										    _UpperCamelCase
										)			,       (
										    _UpperCamelCase
										)			,       (
										    _UpperCamelCase
										)			,       (
										    _UpperCamelCase
										)			,       (
										    _UpperCamelCase
										)			,       
										)			:			str   			=						config_and_inputs
										_UpperCamelCase			:			List[Any]   			=						{"input_ids": input_ids}
										return config, inputs_dict





@require_tf
class 		SCREAMING_SNAKE_CASE     (					SCREAMING_SNAKE_CASE   , SCREAMING_SNAKE_CASE   , unittest.TestCase	):







				'''simple docstring'''
				UpperCAmelCase__							=							(
				    (
				        TFDPRContextEncoder,
				        TFDPRQuestionEncoder,
				        TFDPRReader,
				    )
				    if is_tf_available()
				    else ()
				)
				UpperCAmelCase__							=							{'''feature-extraction''': TFDPRQuestionEncoder} if is_tf_available() else {}

				UpperCAmelCase__							=							False
				UpperCAmelCase__							=							False
				UpperCAmelCase__							=							False
				UpperCAmelCase__							=							False
				UpperCAmelCase__							=							False





				def  snake_case__				(						self					:						Tuple					)							->Optional[Any]:


										'''simple docstring'''




										_UpperCamelCase			:			Optional[int]   			=						TFDPRModelTester(self					)
										_UpperCamelCase			:			Tuple   			=						ConfigTester(self       ,		config_class=lowercase__       ,		hidden_size=37					)





				def  snake_case__				(						self					:						Any					)							->str:


										'''simple docstring'''




										self.config_tester.run_common_tests()





				def  snake_case__				(						self					:						List[str]					)							->Dict:


										'''simple docstring'''




										_UpperCamelCase			:			str   			=						self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_dpr_context_encoder(*lowercase__					)





				def  snake_case__				(						self					:						Any					)							->Optional[int]:


										'''simple docstring'''




										_UpperCamelCase			:			Optional[Any]   			=						self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_dpr_question_encoder(*lowercase__					)





				def  snake_case__				(						self					:						Union[str, Any]					)							->Union[str, Any]:


										'''simple docstring'''




										_UpperCamelCase			:			Optional[int]   			=						self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_dpr_reader(*lowercase__					)







				@slow
				def  snake_case__				(						self					:						Tuple					)							->Tuple:


										'''simple docstring'''




										for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																_UpperCamelCase			:			Tuple   			=						TFDPRContextEncoder.from_pretrained(lowercase__					)
																self.assertIsNotNone(lowercase__					)

										for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																_UpperCamelCase			:			Optional[Any]   			=						TFDPRContextEncoder.from_pretrained(lowercase__					)
																self.assertIsNotNone(lowercase__					)

										for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																_UpperCamelCase			:			Union[str, Any]   			=						TFDPRQuestionEncoder.from_pretrained(lowercase__					)
																self.assertIsNotNone(lowercase__					)

										for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																_UpperCamelCase			:			str   			=						TFDPRReader.from_pretrained(lowercase__					)
																self.assertIsNotNone(lowercase__					)





@require_tf
class 		SCREAMING_SNAKE_CASE     (					unittest.TestCase	):







				'''simple docstring'''





				@slow
				def  snake_case__				(						self					:						Optional[Any]					)							->int:


										'''simple docstring'''




										_UpperCamelCase			:			Dict   			=						TFDPRQuestionEncoder.from_pretrained("facebook/dpr-question_encoder-single-nq-base"					)

										_UpperCamelCase			:			List[str]   			=						tf.constant(
										    [[101, 7_592, 1_010, 2_003, 2_026, 3_899, 10_140, 1_029, 102]]					)  # [CLS] hello, is my dog cute? [SEP]
										_UpperCamelCase			:			Optional[Any]   			=						model(lowercase__					)[0]  # embedding shape = (1, 768)
										# compare the actual values for a slice.
										_UpperCamelCase			:			Optional[Any]   			=						tf.constant(
										    [
										        [
										            0.0_3_2_3_6_2_5_3,
										            0.1_2_7_5_3_3_3_5,
										            0.1_6_8_1_8_5_0_9,
										            0.0_0_2_7_9_7_8_6,
										            0.3_8_9_6_9_3_3,
										            0.2_4_2_6_4_9_4_5,
										            0.2_1_7_8_9_7_1,
										            -0.0_2_3_3_5_2_2_7,
										            -0.0_8_4_8_1_9_5_9,
										            -0.1_4_3_2_4_1_1_7,
										        ]
										    ]					)
										self.assertTrue(numpy.allclose(output[:, :10].numpy()       ,		expected_slice.numpy()       ,		atol=1e-4					)					)



'''simple docstring'''





import math
import os
from copy import deepcopy

import datasets
import evaluate
import torch
import transformers
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from accelerate import Accelerator
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import is_tpu_available, set_seed


lowerCAmelCase_	:		Tuple			=     """true"""
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase=8_2     ,UpperCAmelCase=1_6					)     ->			Union[str, Any]:




						'''simple docstring'''
						set_seed(4_2					)
						_UpperCamelCase			:			List[Any]   			=						RegressionModel()
						_UpperCamelCase			:			Any   			=						deepcopy(UpperCAmelCase					)
						_UpperCamelCase			:			Tuple   			=						RegressionDataset(length=UpperCAmelCase					)
						_UpperCamelCase			:			Union[str, Any]   			=						DataLoader(UpperCAmelCase     ,batch_size=UpperCAmelCase					)
						model.to(accelerator.device					)
						_UpperCamelCase			,       _UpperCamelCase			:			Dict   			=						accelerator.prepare(UpperCAmelCase     ,UpperCAmelCase					)
						return model, ddp_model, dataloader
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase=False					)     ->			List[str]:




						'''simple docstring'''
						_UpperCamelCase			:			str   			=						AutoTokenizer.from_pretrained("hf-internal-testing/mrpc-bert-base-cased"					)
						_UpperCamelCase			:			Optional[Any]   			=						load_dataset("glue"     ,"mrpc"     ,split="validation"					)

						def tokenize_function(UpperCAmelCase					):
												_UpperCamelCase			:			Tuple   			=						tokenizer(examples["sentence1"]     ,examples["sentence2"]     ,truncation=UpperCAmelCase     ,max_length=UpperCAmelCase					)
												return outputs

						with accelerator.main_process_first():
												_UpperCamelCase			:			str   			=						dataset.map(
												    UpperCAmelCase     ,batched=UpperCAmelCase     ,remove_columns=["idx", "sentence1", "sentence2"]     ,)

						_UpperCamelCase			:			Optional[int]   			=						tokenized_datasets.rename_column("label"     ,"labels"					)

						def collate_fn(UpperCAmelCase					):
												if use_longest:
																		return tokenizer.pad(UpperCAmelCase     ,padding="longest"     ,return_tensors="pt"					)
												return tokenizer.pad(UpperCAmelCase     ,padding="max_length"     ,max_length=1_2_8     ,return_tensors="pt"					)

						return DataLoader(UpperCAmelCase     ,shuffle=UpperCAmelCase     ,collate_fn=UpperCAmelCase     ,batch_size=1_6					)
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase					)     ->			Dict:




						'''simple docstring'''
						_UpperCamelCase			:			str   			=						Accelerator(dispatch_batches=UpperCAmelCase     ,split_batches=UpperCAmelCase					)
						_UpperCamelCase			:			Union[str, Any]   			=						get_dataloader(UpperCAmelCase     ,not dispatch_batches					)
						_UpperCamelCase			:			Optional[int]   			=						AutoModelForSequenceClassification.from_pretrained(
						    "hf-internal-testing/mrpc-bert-base-cased"     ,return_dict=UpperCAmelCase					)
						_UpperCamelCase			,       _UpperCamelCase			:			Union[str, Any]   			=						accelerator.prepare(UpperCAmelCase     ,UpperCAmelCase					)
						return {"ddp": [ddp_model, ddp_dataloader, "cuda:0"], "no": [model, dataloader, accelerator.device]}, accelerator
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase					)     ->			int:




						'''simple docstring'''
						_UpperCamelCase			:			List[str]   			=						[]
						for batch in dataloader:
												_UpperCamelCase			,       _UpperCamelCase			:			int   			=						batch.values()
												with torch.no_grad():
																		_UpperCamelCase			:			Tuple   			=						model(UpperCAmelCase					)
																		_UpperCamelCase			,       _UpperCamelCase			:			List[Any]   			=						accelerator.gather_for_metrics((logit, target)					)
																		logits_and_targets.append((logit, target)					)
						_UpperCamelCase			,       _UpperCamelCase			:			Optional[Any]   			=						[], []
						for logit, targ in logits_and_targets:
												logits.append(UpperCAmelCase					)
												targs.append(UpperCAmelCase					)
						_UpperCamelCase			,       _UpperCamelCase			:			int   			=						torch.cat(UpperCAmelCase					), torch.cat(UpperCAmelCase					)
						return logits, targs
def 						__A							(     UpperCAmelCase     ,UpperCAmelCase=8_2     ,UpperCAmelCase=False     ,UpperCAmelCase=False     ,UpperCAmelCase=1_6					)     ->			Any:




						'''simple docstring'''
						_UpperCamelCase			,       _UpperCamelCase			,       _UpperCamelCase			:			Union[str, Any]   			=						get_basic_setup(UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase					)
						_UpperCamelCase			,       _UpperCamelCase			:			Tuple   			=						generate_predictions(UpperCAmelCase     ,UpperCAmelCase     ,UpperCAmelCase					)
						assert (
						    len(UpperCAmelCase					) == num_samples
						), f'''Unexpected number of inputs:\n    Expected: {num_samples}\n    Actual: {len(UpperCAmelCase					)}'''
def 						__A							(     UpperCAmelCase = False     ,UpperCAmelCase = False					)     ->			Tuple:




						'''simple docstring'''
						_UpperCamelCase			:			int   			=						evaluate.load("glue"     ,"mrpc"					)
						_UpperCamelCase			,       _UpperCamelCase			:			Any   			=						get_mrpc_setup(UpperCAmelCase     ,UpperCAmelCase					)
						# First do baseline
						_UpperCamelCase			,       _UpperCamelCase			,       _UpperCamelCase			:			Tuple   			=						setup["no"]
						model.to(UpperCAmelCase					)
						model.eval()
						for batch in dataloader:
												batch.to(UpperCAmelCase					)
												with torch.inference_mode():
																		_UpperCamelCase			:			Optional[Any]   			=						model(**UpperCAmelCase					)
												_UpperCamelCase			:			Any   			=						outputs.logits.argmax(dim=-1					)
												metric.add_batch(predictions=UpperCAmelCase     ,references=batch["labels"]					)
						_UpperCamelCase			:			List[str]   			=						metric.compute()

						# Then do distributed
						_UpperCamelCase			,       _UpperCamelCase			,       _UpperCamelCase			:			List[Any]   			=						setup["ddp"]
						model.eval()
						for batch in dataloader:
												with torch.inference_mode():
																		_UpperCamelCase			:			int   			=						model(**UpperCAmelCase					)
												_UpperCamelCase			:			Tuple   			=						outputs.logits.argmax(dim=-1					)
												_UpperCamelCase			:			List[str]   			=						batch["labels"]
												_UpperCamelCase			,       _UpperCamelCase			:			List[str]   			=						accelerator.gather_for_metrics((preds, references)					)
												metric.add_batch(predictions=UpperCAmelCase     ,references=UpperCAmelCase					)
						_UpperCamelCase			:			int   			=						metric.compute()

						for key in "accuracy f1".split():
												assert math.isclose(
												    baseline[key]     ,distributed[key]					), f'''Baseline and Distributed are not the same for key {key}:\n\tBaseline: {baseline[key]}\n\tDistributed: {distributed[key]}\n'''
def 						__A							(     )     ->			Union[str, Any]:




						'''simple docstring'''
						_UpperCamelCase			:			Dict   			=						Accelerator(split_batches=UpperCAmelCase     ,dispatch_batches=UpperCAmelCase					)
						if accelerator.is_local_main_process:
												datasets.utils.logging.set_verbosity_warning()
												transformers.utils.logging.set_verbosity_warning()
						else:
												datasets.utils.logging.set_verbosity_error()
												transformers.utils.logging.set_verbosity_error()
						# These are a bit slower so they should only be ran on the GPU or TPU
						if torch.cuda.is_available() or is_tpu_available():
												if accelerator.is_local_main_process:
																		print("**Testing gather_for_metrics**"					)
												for split_batches in [True, False]:
																		for dispatch_batches in [True, False]:
																								if accelerator.is_local_main_process:
																														print(f'''With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`'''					)
																								test_mrpc(UpperCAmelCase     ,UpperCAmelCase					)
																								accelerator.state._reset_state()
						if accelerator.is_local_main_process:
												print("**Test torch metrics**"					)
						for split_batches in [True, False]:
												for dispatch_batches in [True, False]:
																		_UpperCamelCase			:			int   			=						Accelerator(split_batches=UpperCAmelCase     ,dispatch_batches=UpperCAmelCase					)
																		if accelerator.is_local_main_process:
																								print(f'''With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`, length=99'''					)
																		test_torch_metrics(UpperCAmelCase     ,9_9					)
																		accelerator.state._reset_state()
						if accelerator.is_local_main_process:
												print("**Test last batch is not dropped when perfectly divisible**"					)
						_UpperCamelCase			:			int   			=						Accelerator()
						test_torch_metrics(UpperCAmelCase     ,5_1_2					)
						accelerator.state._reset_state()
def 						__A							(     UpperCAmelCase					)     ->			List[str]:




						'''simple docstring'''
						# For xla_spawn (TPUs)
						main()


if __name__ == "__main__":
			main()


