




"""simple docstring"""

import argparse
import gc
import os
import threading

import evaluate
import psutil
import torch
from datasets import load_dataset
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
from torch.utils.data import DataLoader
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
    set_seed,
)

from accelerate import Accelerator, DistributedType, FullyShardedDataParallelPlugin


########################################################################
# This is a fully working simple example to use Accelerate
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#   - FSDP
#
# This example also demonstrates the checkpointing and sharding capabilities
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


UpperCamelCase__			:     str			       =				16
UpperCamelCase__			:     Union[str, Any]			       =				32




def    snake_case__						(							lowerCAmelCase :		Any		):
					return int(x / 2**20		)
class      A__   :




	def __enter__( self	:						Tuple						)       ->			List[Any]:
						gc.collect()
						torch.cuda.empty_cache()
						torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
						__snake_case        =						torch.cuda.memory_allocated()
						__snake_case        =						psutil.Process()

						__snake_case        =						self.cpu_mem_used()
						__snake_case        =						True
						__snake_case        =						threading.Thread(target=self.peak_monitor_func						)
						__snake_case        =						True
						peak_monitor_thread.start()
						return self




	def 				a_		( self	:						List[Any]						)       ->			Optional[Any]:
						return self.process.memory_info().rss




	def 				a_		( self	:						Any						)       ->			List[Any]:
						__snake_case        =						-1

						while True:
											__snake_case        =						max(self.cpu_mem_used()   ,     self.cpu_peak						)

											# can't sleep or will not catch the peak right (this comment is here on purpose)
											# time.sleep(0.001) # 1msec

											if not self.peak_monitoring:
																break




	def __exit__( self	:						Optional[int]   ,     *lowercase	:						int						)       ->			int:
						__snake_case        =						False

						gc.collect()
						torch.cuda.empty_cache()
						__snake_case        =						torch.cuda.memory_allocated()
						__snake_case        =						torch.cuda.max_memory_allocated()
						__snake_case        =						bamb(self.end - self.begin						)
						__snake_case        =						bamb(self.peak - self.begin						)

						__snake_case        =						self.cpu_mem_used()
						__snake_case        =						bamb(self.cpu_end - self.cpu_begin						)
						__snake_case        =						bamb(self.cpu_peak - self.cpu_begin						)
						# print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")


# For testing only
if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
		from accelerate.test_utils.training import mocked_dataloaders

		UpperCamelCase__			:     Tuple			       =				mocked_dataloaders  # noqa: F811




def    snake_case__						(							lowerCAmelCase :		str					,       lowerCAmelCase :		Optional[int]		):
					# For testing only
					if os.environ.get("""TESTING_MOCKED_DATALOADERS"""					,       lowerCAmelCase		) == "1":
										__snake_case        =						2

					# New Code #
					# Pass the advanced FSDP settings not part of the accelerate config by creating fsdp_plugin
					__snake_case        =						FullyShardedDataParallelPlugin(
					    state_dict_config=FullStateDictConfig(offload_to_cpu=lowerCAmelCase					,       ranka_only=lowerCAmelCase		)					,       optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=lowerCAmelCase					,       ranka_only=lowerCAmelCase		)					,       )

					# Initialize accelerator
					if args.with_tracking:
										__snake_case        =						Accelerator(
										    cpu=args.cpu					,       mixed_precision=args.mixed_precision					,       log_with="""wandb"""					,       project_dir=args.logging_dir					,       fsdp_plugin=lowerCAmelCase					,       )
					else:
										__snake_case        =						Accelerator(fsdp_plugin=lowerCAmelCase		)
					accelerator.print(accelerator.distributed_type		)

					if hasattr(args.checkpointing_steps					,       """isdigit"""		):
										if args.checkpointing_steps == "epoch":
															__snake_case        =						args.checkpointing_steps
										elif args.checkpointing_steps.isdigit():
															__snake_case        =						int(args.checkpointing_steps		)
										else:
															raise ValueError(
															    F'''Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.'''		)
					else:
										__snake_case        =						None
					# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
					__snake_case        =						config["""lr"""]
					__snake_case        =						int(config["""num_epochs"""]		)
					__snake_case        =						int(config["""seed"""]		)
					__snake_case        =						int(config["""batch_size"""]		)

					# We need to initialize the trackers we use, and also store our configuration
					if args.with_tracking:
										__snake_case        =						vars(lowerCAmelCase		)
										accelerator.init_trackers("""fsdp_glue_no_trainer"""					,       lowerCAmelCase		)

					__snake_case        =						AutoTokenizer.from_pretrained(args.model_name_or_path		)
					__snake_case        =						load_dataset("""glue"""					,       """mrpc"""		)
					__snake_case        =						evaluate.load("""glue"""					,       """mrpc"""		)

					def tokenize_function(lowerCAmelCase :		Any		):
										# max_length=None => use the model max length (it's actually the default)
										__snake_case        =						tokenizer(examples["""sentence1"""]					,       examples["""sentence2"""]					,       truncation=lowerCAmelCase					,       max_length=lowerCAmelCase		)
										return outputs

					# Apply the method we just defined to all the examples in all the splits of the dataset
					# starting with the main process first:
					with accelerator.main_process_first():
										__snake_case        =						datasets.map(
										    lowerCAmelCase					,       batched=lowerCAmelCase					,       remove_columns=["""idx""", """sentence1""", """sentence2"""]					,       )

					# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
					# transformers library
					__snake_case        =						tokenized_datasets.rename_column("""label"""					,       """labels"""		)

					# If the batch size is too big we use gradient accumulation
					__snake_case        =						1
					if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
										__snake_case        =						batch_size // MAX_GPU_BATCH_SIZE
										__snake_case        =						MAX_GPU_BATCH_SIZE

					def collate_fn(lowerCAmelCase :		List[str]		):
										# On TPU it's best to pad everything to the same length or training will be very slow.
										__snake_case        =						128 if accelerator.distributed_type == DistributedType.TPU else None
										# When using mixed precision we want round multiples of 8/16
										if accelerator.mixed_precision == "fp8":
															__snake_case        =						16
										elif accelerator.mixed_precision != "no":
															__snake_case        =						8
										else:
															__snake_case        =						None

										return tokenizer.pad(
										    lowerCAmelCase					,       padding="""longest"""					,       max_length=lowerCAmelCase					,       pad_to_multiple_of=lowerCAmelCase					,       return_tensors="""pt"""					,       )

					# Instantiate dataloaders.
					__snake_case        =						DataLoader(
					    tokenized_datasets["""train"""]					,       shuffle=lowerCAmelCase					,       collate_fn=lowerCAmelCase					,       batch_size=lowerCAmelCase		)
					__snake_case        =						DataLoader(
					    tokenized_datasets["""validation"""]					,       shuffle=lowerCAmelCase					,       collate_fn=lowerCAmelCase					,       batch_size=lowerCAmelCase		)

					set_seed(lowerCAmelCase		)

					# Instantiate the model (we build the model here so that the seed also control new weights initialization)
					__snake_case        =						AutoModelForSequenceClassification.from_pretrained(
					    args.model_name_or_path					,       return_dict=lowerCAmelCase					,       low_cpu_mem_usage=lowerCAmelCase		)

					# New Code #
					# For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
					__snake_case        =						accelerator.prepare(lowerCAmelCase		)
					accelerator.print(lowerCAmelCase		)

					# Instantiate optimizer
					# New Code #
					# For FSDP feature, at present it doesn't support multiple parameter groups,
					# so we need to create a single parameter group for the whole model
					__snake_case        =						torch.optim.AdamW(params=model.parameters()					,       lr=lowerCAmelCase					,       weight_decay=2e-4		)

					# Instantiate scheduler
					__snake_case        =						get_linear_schedule_with_warmup(
					    optimizer=lowerCAmelCase					,       num_warmup_steps=10					,       num_training_steps=(len(lowerCAmelCase		) * num_epochs) // gradient_accumulation_steps					,       )

					# New Code #
					# For FSDP feature, prepare everything except the model as we have already prepared the model
					# before creating the optimizer
					# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
					# prepare method.
					__snake_case   ,   __snake_case   ,   __snake_case   ,   __snake_case        =						accelerator.prepare(
					    lowerCAmelCase					,       lowerCAmelCase					,       lowerCAmelCase					,       lowerCAmelCase		)

					__snake_case        =						0

					# Potentially load in the weights and states from a previous save
					if args.resume_from_checkpoint:
										if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
															accelerator.print(F'''Resumed from checkpoint: {args.resume_from_checkpoint}'''		)
															accelerator.load_state(args.resume_from_checkpoint		)
															__snake_case        =						os.path.basename(args.resume_from_checkpoint		)
										else:
															# Get the most recent checkpoint
															__snake_case        =						[f.name for f in os.scandir(os.getcwd()		) if f.is_dir()]
															dirs.sort(key=os.path.getctime		)
															__snake_case        =						dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
										# Extract `epoch_{i}` or `step_{i}`
										__snake_case        =						os.path.splitext(lowerCAmelCase		)[0]

										if "epoch" in training_difference:
															num_epochs -= int(training_difference.replace("""epoch_"""					,       """"""		)		)
															__snake_case        =						None
										else:
															__snake_case        =						int(training_difference.replace("""step_"""					,       """"""		)		)
															num_epochs -= resume_step // len(lowerCAmelCase		)
															# If resuming by step, we also need to know exactly how far into the DataLoader we went
															__snake_case        =						(num_epochs * len(lowerCAmelCase		)) - resume_step

    # Now we train the model
					for epoch in range(lowerCAmelCase		):
										# New Code #
										# context manager to track the peak memory usage during the training epoch
										with TorchTracemalloc() as tracemalloc:
															model.train()
															if args.with_tracking:
																				__snake_case        =						0
															for step, batch in enumerate(lowerCAmelCase		):
																				# We need to skip steps until we reach the resumed step
																				if args.resume_from_checkpoint and epoch == 0:
																									if resume_step is not None and step < resume_step:
																														pass
                # We could avoid this line since we set the accelerator with `device_placement=True`.
																				batch.to(accelerator.device		)
																				__snake_case        =						model(**lowerCAmelCase		)
																				__snake_case        =						outputs.loss
																				__snake_case        =						loss / gradient_accumulation_steps
																				# We keep track of the loss at each epoch
																				if args.with_tracking:
																									total_loss += loss.detach().float()
																				accelerator.backward(lowerCAmelCase		)
																				if step % gradient_accumulation_steps == 0:
																									optimizer.step()
																									lr_scheduler.step()
																									optimizer.zero_grad()
																									# accelerator.print(lr_scheduler.get_lr())

																				overall_step += 1

																				if isinstance(lowerCAmelCase					,       lowerCAmelCase		):
																									__snake_case        =						F'''step_{overall_step}'''
																									if overall_step % checkpointing_steps == 0:
																														if args.output_dir is not None:
																																			__snake_case        =						os.path.join(args.output_dir					,       lowerCAmelCase		)
																														accelerator.save_state(lowerCAmelCase		)
        # New Code #
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
										accelerator.print("""Memory before entering the train : {}""".format(bamb(tracemalloc.begin		)		)		)
										accelerator.print("""Memory consumed at the end of the train (end-begin): {}""".format(tracemalloc.used		)		)
										accelerator.print("""Peak Memory consumed during the train (max-begin): {}""".format(tracemalloc.peaked		)		)
										accelerator.print(
										    """Total Peak Memory consumed during the train (max): {}""".format(
										        tracemalloc.peaked + bamb(tracemalloc.begin		)		)		)
										# Logging the peak memory usage of the GPU to the tracker
										if args.with_tracking:
															accelerator.log(
															    {
															        """train_total_peak_memory""": tracemalloc.peaked + bamb(tracemalloc.begin		),
															    }					,       step=lowerCAmelCase					,       )

										# New Code #
										# context manager to track the peak memory usage during the evaluation
										with TorchTracemalloc() as tracemalloc:
															model.eval()
															for step, batch in enumerate(lowerCAmelCase		):
																				# We could avoid this line since we set the accelerator with `device_placement=True`.
																				batch.to(accelerator.device		)
																				with torch.no_grad():
																									__snake_case        =						model(**lowerCAmelCase		)
																				__snake_case        =						outputs.logits.argmax(dim=-1		)
																				__snake_case   ,   __snake_case        =						accelerator.gather_for_metrics((predictions, batch["""labels"""])		)
																				metric.add_batch(
																				    predictions=lowerCAmelCase					,       references=lowerCAmelCase					,       )

															__snake_case        =						metric.compute()
															# Use accelerator.print to print only on the main process.
															accelerator.print(F'''epoch {epoch}:'''					,       lowerCAmelCase		)
															if args.with_tracking:
																				accelerator.log(
																				    {
																				        """accuracy""": eval_metric["""accuracy"""],
																				        """f1""": eval_metric["""f1"""],
																				        """train_loss""": total_loss.item() / len(lowerCAmelCase		),
																				    }					,       step=lowerCAmelCase					,       )

															if checkpointing_steps == "epoch":
																				__snake_case        =						F'''epoch_{epoch}'''
																				if args.output_dir is not None:
																									__snake_case        =						os.path.join(args.output_dir					,       lowerCAmelCase		)
																				accelerator.save_state(lowerCAmelCase		)
        # New Code #
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
										accelerator.print("""Memory before entering the eval : {}""".format(bamb(tracemalloc.begin		)		)		)
										accelerator.print("""Memory consumed at the end of the eval (end-begin): {}""".format(tracemalloc.used		)		)
										accelerator.print("""Peak Memory consumed during the eval (max-begin): {}""".format(tracemalloc.peaked		)		)
										accelerator.print(
										    """Total Peak Memory consumed during the eval (max): {}""".format(tracemalloc.peaked + bamb(tracemalloc.begin		)		)		)
										# Logging the peak memory usage of the GPU to the tracker
										if args.with_tracking:
															accelerator.log(
															    {
															        """eval_total_peak_memory""": tracemalloc.peaked + bamb(tracemalloc.begin		),
															    }					,       step=lowerCAmelCase					,       )

					if args.with_tracking:
										accelerator.end_training()




def    snake_case__						(							):
					__snake_case        =						argparse.ArgumentParser(description="""Simple example of training script."""		)
					parser.add_argument(
					    """--mixed_precision"""					,       type=lowerCAmelCase					,       default=lowerCAmelCase					,       choices=["""no""", """fp16""", """bf16""", """fp8"""]					,       help="""Whether to use mixed precision. Choose"""
					    """between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."""
					    """and an Nvidia Ampere GPU."""					,       )
					parser.add_argument("""--cpu"""					,       action="""store_true"""					,       help="""If passed, will train on the CPU."""		)
					parser.add_argument(
					    """--checkpointing_steps"""					,       type=lowerCAmelCase					,       default=lowerCAmelCase					,       help="""Whether the various states should be saved at the end of every n steps, or 'epoch' for each epoch."""					,       )
					parser.add_argument(
					    """--resume_from_checkpoint"""					,       type=lowerCAmelCase					,       default=lowerCAmelCase					,       help="""If the training should continue from a checkpoint folder."""					,       )
					parser.add_argument(
					    """--with_tracking"""					,       action="""store_true"""					,       help="""Whether to load in all available experiment trackers from the environment and use them for logging."""					,       )
					parser.add_argument(
					    """--output_dir"""					,       type=lowerCAmelCase					,       default="""."""					,       help="""Optional save directory where all checkpoint folders will be stored. Default is the current working directory."""					,       )
					parser.add_argument(
					    """--logging_dir"""					,       type=lowerCAmelCase					,       default="""logs"""					,       help="""Location on where to store experiment tracking logs`"""					,       )
					parser.add_argument(
					    """--model_name_or_path"""					,       type=lowerCAmelCase					,       help="""Path to pretrained model or model identifier from huggingface.co/models."""					,       required=lowerCAmelCase					,       )
					__snake_case        =						parser.parse_args()
					__snake_case        =						{"""lr""": 2e-5, """num_epochs""": 3, """seed""": 42, """batch_size""": 16}
					training_function(lowerCAmelCase					,       lowerCAmelCase		)


if __name__ == "__main__":
		main()








"""simple docstring"""

import tempfile
import unittest

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

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import DPRContextEncoder, DPRQuestionEncoder, DPRReader, DPRReaderTokenizer
		from transformers.models.dpr.modeling_dpr import (
		    DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
		    DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
		)
class      A__   :




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




	def 				a_		( self	:						str						)       ->			Optional[Any]:
						__snake_case        =						ids_tensor([self.batch_size, self.seq_length]   ,     self.vocab_size						)

						__snake_case        =						None
						if self.use_input_mask:
											__snake_case        =						random_attention_mask([self.batch_size, self.seq_length]						)

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

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

						__snake_case        =						self.get_config()

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




	def 				a_		( self	:						str						)       ->			str:
						return DPRConfig(
						    projection_dim=self.projection_dim   ,     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   ,     initializer_range=self.initializer_range   ,     )




	def 				a_		( self	:						Optional[Any]   ,     lowercase	:						str   ,     lowercase	:						Optional[Any]   ,     lowercase	:						str   ,     lowercase	:						Dict   ,     lowercase	:						Union[str, Any]   ,     lowercase	:						Any   ,     lowercase	:						List[str]						)       ->			Optional[Any]:
						__snake_case        =						DPRContextEncoder(config=lowercase						)
						model.to(lowercase						)
						model.eval()
						__snake_case        =						model(lowercase   ,     attention_mask=lowercase   ,     token_type_ids=lowercase						)
						__snake_case        =						model(lowercase   ,     token_type_ids=lowercase						)
						__snake_case        =						model(lowercase						)
						self.parent.assertEqual(result.pooler_output.shape   ,     (self.batch_size, self.projection_dim or self.hidden_size)						)




	def 				a_		( self	:						Tuple   ,     lowercase	:						Tuple   ,     lowercase	:						Union[str, Any]   ,     lowercase	:						List[str]   ,     lowercase	:						Optional[int]   ,     lowercase	:						Any   ,     lowercase	:						Tuple   ,     lowercase	:						Union[str, Any]						)       ->			Optional[int]:
						__snake_case        =						DPRQuestionEncoder(config=lowercase						)
						model.to(lowercase						)
						model.eval()
						__snake_case        =						model(lowercase   ,     attention_mask=lowercase   ,     token_type_ids=lowercase						)
						__snake_case        =						model(lowercase   ,     token_type_ids=lowercase						)
						__snake_case        =						model(lowercase						)
						self.parent.assertEqual(result.pooler_output.shape   ,     (self.batch_size, self.projection_dim or self.hidden_size)						)




	def 				a_		( self	:						str   ,     lowercase	:						List[Any]   ,     lowercase	:						int   ,     lowercase	:						Optional[Any]   ,     lowercase	:						Any   ,     lowercase	:						Tuple   ,     lowercase	:						Optional[Any]   ,     lowercase	:						int						)       ->			str:
						__snake_case        =						DPRReader(config=lowercase						)
						model.to(lowercase						)
						model.eval()
						__snake_case        =						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 				a_		( self	:						str						)       ->			Optional[Any]:
						__snake_case        =						self.prepare_config_and_inputs()
						(
						    (
						    __snake_case
						)   ,   (
						    __snake_case
						)   ,   (
						    __snake_case
						)   ,   (
						    __snake_case
						)   ,   (
						    __snake_case
						)   ,   (
						    __snake_case
						)   ,   (
						    __snake_case
						)   ,   
						)        =						config_and_inputs
						__snake_case        =						{"""input_ids""": input_ids}
						return config, inputs_dict

@require_torch
class      A__   (   lowercase_  ,						lowercase_  ,						unittest.TestCase  ):
	A_				:					Dict      =   (
	    (
	        DPRContextEncoder,
	        DPRQuestionEncoder,
	        DPRReader,
	    )
	    if is_torch_available()
	    else ()
	)
	A_				:					Any      =   {"""feature-extraction""": DPRQuestionEncoder} if is_torch_available() else {}

	A_				:					List[Any]      =   False
	A_				:					Union[str, Any]      =   False  # why?
	A_				:					Any      =   False
	A_				:					List[str]      =   False




	def 				a_		( self	:						Union[str, Any]						)       ->			Dict:
						__snake_case        =						DPRModelTester(self						)
						__snake_case        =						ConfigTester(self   ,     config_class=lowercase   ,     hidden_size=37						)




	def 				a_		( self	:						Optional[int]						)       ->			Optional[int]:
						self.config_tester.run_common_tests()




	def 				a_		( self	:						Dict						)       ->			List[str]:
						__snake_case        =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_context_encoder(*lowercase						)




	def 				a_		( self	:						Any						)       ->			Tuple:
						__snake_case        =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_question_encoder(*lowercase						)




	def 				a_		( self	:						Any						)       ->			Optional[int]:
						__snake_case        =						self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_reader(*lowercase						)




	def 				a_		( self	:						Tuple						)       ->			str:
						__snake_case        =						self.model_tester.prepare_config_and_inputs()[0]

						__snake_case        =						DPRQuestionEncoder(config=lowercase						)
						model.to(lowercase						)
						model.eval()

						with tempfile.TemporaryDirectory() as tmp_dirname:
											model.save_pretrained(lowercase						)
											__snake_case        =						DPRQuestionEncoder.from_pretrained(lowercase   ,     projection_dim=512						)

						self.assertIsNotNone(lowercase						)




	@slow
	def 				a_		( self	:						Any						)       ->			Optional[Any]:
						for model_name in DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__snake_case        =						DPRContextEncoder.from_pretrained(lowercase						)
											self.assertIsNotNone(lowercase						)

						for model_name in DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__snake_case        =						DPRContextEncoder.from_pretrained(lowercase						)
											self.assertIsNotNone(lowercase						)

						for model_name in DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__snake_case        =						DPRQuestionEncoder.from_pretrained(lowercase						)
											self.assertIsNotNone(lowercase						)

						for model_name in DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__snake_case        =						DPRReader.from_pretrained(lowercase						)
											self.assertIsNotNone(lowercase						)

@require_torch
class      A__   (   unittest.TestCase  ):




	@slow
	def 				a_		( self	:						List[str]						)       ->			Union[str, Any]:
						__snake_case        =						DPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base"""   ,     return_dict=lowercase						)
						model.to(lowercase						)

						__snake_case        =						torch.tensor(
						    [[101, 7_592, 1_010, 2_003, 2_026, 3_899, 10_140, 1_029, 102]]   ,     dtype=torch.long   ,     device=lowercase						)  # [CLS] hello, is my dog cute? [SEP]
						__snake_case        =						model(lowercase						)[0]  # embedding shape = (1, 768)
						# compare the actual values for a slice.
						__snake_case        =						torch.tensor(
						    [
						        [
						            0.03_23_62_53,
						            0.12_75_33_35,
						            0.16_81_85_09,
						            0.00_27_97_86,
						            0.3_89_69_33,
						            0.24_26_49_45,
						            0.2_17_89_71,
						            -0.02_33_52_27,
						            -0.08_48_19_59,
						            -0.14_32_41_17,
						        ]
						    ]   ,     dtype=torch.float   ,     device=lowercase   ,     )
						self.assertTrue(torch.allclose(output[:, :10]   ,     lowercase   ,     atol=1e-4						)						)




	@slow
	def 				a_		( self	:						Tuple						)       ->			int:
						__snake_case        =						DPRReaderTokenizer.from_pretrained("""facebook/dpr-reader-single-nq-base"""						)
						__snake_case        =						DPRReader.from_pretrained("""facebook/dpr-reader-single-nq-base"""						)
						model.to(lowercase						)

						__snake_case        =						tokenizer(
						    questions="""What is love ?"""   ,     titles="""Haddaway"""   ,     texts="""What Is Love is a song recorded by the artist Haddaway"""   ,     padding=lowercase   ,     return_tensors="""pt"""   ,     )
						encoded_inputs.to(lowercase						)

						__snake_case        =						model(**lowercase						)

						# compare the actual values for a slice.
						__snake_case        =						torch.tensor(
						    [[-10.30_05, -10.77_65, -11.48_72, -11.68_41, -11.93_12, -10.30_02, -9.85_44, -11.73_78, -12.08_21, -10.29_75]]   ,     dtype=torch.float   ,     device=lowercase   ,     )

						__snake_case        =						torch.tensor(
						    [[-11.06_84, -11.70_41, -11.53_97, -10.34_65, -10.87_91, -6.84_43, -11.99_59, -11.03_64, -10.00_96, -6.84_05]]   ,     dtype=torch.float   ,     device=lowercase   ,     )
						self.assertTrue(torch.allclose(outputs.start_logits[:, :10]   ,     lowercase   ,     atol=1e-4						)						)
						self.assertTrue(torch.allclose(outputs.end_logits[:, :10]   ,     lowercase   ,     atol=1e-4						)						)


