

from copy import deepcopy

import torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader

from accelerate.accelerator import Accelerator
from accelerate.state import GradientState
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import DistributedType, is_torch_version, set_seed

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Dict  ,      lowerCAmelCase:  List[Any]  ,      lowerCAmelCase:  int  ,      lowerCAmelCase:  Dict )   ->       Any:
				for param, grad_param in zip(model_a.parameters()  ,      model_b.parameters() ):
								if not param.requires_grad:
												continue
								if not did_step:
												# Grads should not be in sync
												assert (
												    torch.allclose(param.grad  ,      grad_param.grad ) is False
												), F'Gradients in sync when they should not be at iteration {iteration}:\nmodel_a grad ({param.grad}) == model_b grad ({grad_param.grad})'
								else:
												# Grads should be in sync
												assert (
												    torch.allclose(param.grad  ,      grad_param.grad ) is True
												), F'Gradients not in sync when they should be at iteration {iteration}:\nmodel_a grad ({param.grad}) != model_b grad ({grad_param.grad})'

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Any  ,      lowerCAmelCase:  List[Any]  ,      lowerCAmelCase:  List[Any]  ,      lowerCAmelCase:  Optional[int]  ,      lowerCAmelCase:  int=True )   ->       Optional[int]:
				model.train()
				_UpperCAmelCase       :		str 		=    model(lowerCAmelCase )
				_UpperCAmelCase       :		Tuple 		=    F.mse_loss(lowerCAmelCase  ,      target.to(output.device ) )
				if not do_backward:
								loss /= accelerator.gradient_accumulation_steps
								loss.backward()
				else:
								accelerator.backward(lowerCAmelCase )

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  int  ,      lowerCAmelCase:  Tuple=False )   ->       Union[str, Any]:
				set_seed(42 )
				_UpperCAmelCase       :		Dict 		=    RegressionModel()
				_UpperCAmelCase       :		Any 		=    deepcopy(lowerCAmelCase )
				_UpperCAmelCase       :		str 		=    RegressionDataset(length=80 )
				_UpperCAmelCase       :		Optional[int] 		=    DataLoader(lowerCAmelCase  ,      batch_size=16 )
				model.to(accelerator.device )
				if sched:
								_UpperCAmelCase       :		Dict 		=    AdamW(params=model.parameters()  ,      lr=1E-3 )
								_UpperCAmelCase       :		Optional[Any] 		=    AdamW(params=ddp_model.parameters()  ,      lr=1E-3 )
								_UpperCAmelCase       :		Union[str, Any] 		=    LambdaLR(lowerCAmelCase  ,      lr_lambda=lambda lowerCAmelCase : epoch**0.65 )
								_UpperCAmelCase       :		Optional[Any] 		=    LambdaLR(lowerCAmelCase  ,      lr_lambda=lambda lowerCAmelCase : epoch**0.65 )
				# Make a copy of `model`
				if sched:
								_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase       :		Union[str, Any] 		=    accelerator.prepare(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )
				else:
								_UpperCAmelCase		,							_UpperCAmelCase       :		int 		=    accelerator.prepare(lowerCAmelCase  ,      lowerCAmelCase )
				if sched:
								return (model, opt, sched, dataloader, ddp_model, ddp_opt, ddp_sched)
				return model, ddp_model, dataloader

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Optional[int] )   ->       str:
				# Test when on a single CPU or GPU that the context manager does nothing
				_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase       :		Union[str, Any] 		=    get_training_setup(lowerCAmelCase )
				# Use a single batch
				_UpperCAmelCase		,							_UpperCAmelCase       :		Any 		=    next(iter(lowerCAmelCase ) ).values()
				for iteration in range(3 ):
								# Gather the distributed inputs and targs for the base model
								_UpperCAmelCase		,							_UpperCAmelCase       :		List[str] 		=    accelerator.gather((ddp_input, ddp_target) )
								_UpperCAmelCase		,							_UpperCAmelCase       :		Optional[int] 		=    input.to(accelerator.device ), target.to(accelerator.device )
								# Perform our initial ground truth step in non "DDP"
								step_model(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )
								# Do "gradient accumulation" (noop)
								if iteration % 2 == 0:
												# Accumulate grads locally
												with accelerator.no_sync(lowerCAmelCase ):
																step_model(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )
								else:
												# Sync grads
												step_model(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )

								# Since `no_sync` is a noop, `ddp_model` and `model` grads should always be in sync
								check_model_parameters(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )
								for param, ddp_param in zip(model.parameters()  ,      ddp_model.parameters() ):
												if not param.requires_grad:
																continue
												assert torch.allclose(
												    param.grad  ,      ddp_param.grad ), F'Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'

								# Shuffle ddp_input on each iteration
								torch.manual_seed(1337 + iteration )
								_UpperCAmelCase       :		Dict 		=    ddp_input[torch.randperm(len(lowerCAmelCase ) )]

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  List[Any] )   ->       Dict:
				# Test on distributed setup that context manager behaves properly
				_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase       :		List[str] 		=    get_training_setup(lowerCAmelCase )
				# Use a single batch
				_UpperCAmelCase		,							_UpperCAmelCase       :		Union[str, Any] 		=    next(iter(lowerCAmelCase ) ).values()
				for iteration in range(3 ):
								# Gather the distributed inputs and targs for the base model
								_UpperCAmelCase		,							_UpperCAmelCase       :		Optional[int] 		=    accelerator.gather((ddp_input, ddp_target) )
								_UpperCAmelCase		,							_UpperCAmelCase       :		Dict 		=    input.to(accelerator.device ), target.to(accelerator.device )
								# Perform our initial ground truth step in non "DDP"
								step_model(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )
								# Do "gradient accumulation" (noop)
								if iteration % 2 == 0:
												# Accumulate grads locally
												with accelerator.no_sync(lowerCAmelCase ):
																step_model(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )
								else:
												# Sync grads
												step_model(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )

								# DDP model and model should only be in sync when not (iteration % 2 == 0)
								for param, ddp_param in zip(model.parameters()  ,      ddp_model.parameters() ):
												if not param.requires_grad:
																continue
												if iteration % 2 == 0:
																# Grads should not be in sync
																assert (
																    torch.allclose(param.grad  ,      ddp_param.grad ) is False
																), F'Gradients in sync when they should not be:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})'
												else:
																# Grads should be in sync
																assert (
																    torch.allclose(param.grad  ,      ddp_param.grad ) is True
																), F'Gradients not in sync when they should be:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'

        # Shuffle ddp_input on each iteration
								torch.manual_seed(1337 + iteration )
								_UpperCAmelCase       :		Tuple 		=    ddp_input[torch.randperm(len(lowerCAmelCase ) )]

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  List[Any]=False  ,      lowerCAmelCase:  Optional[int]=False )   ->       Union[str, Any]:
				_UpperCAmelCase       :		str 		=    Accelerator(
				    split_batches=lowerCAmelCase  ,      dispatch_batches=lowerCAmelCase  ,      gradient_accumulation_steps=2 )
				# Test that context manager behaves properly
				_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase       :		int 		=    get_training_setup(lowerCAmelCase )
				for iteration, batch in enumerate(lowerCAmelCase ):
								_UpperCAmelCase		,							_UpperCAmelCase       :		List[str] 		=    batch.values()
								# Gather the distributed inputs and targs for the base model
								_UpperCAmelCase		,							_UpperCAmelCase       :		Tuple 		=    accelerator.gather((ddp_input, ddp_target) )
								_UpperCAmelCase		,							_UpperCAmelCase       :		Dict 		=    input.to(accelerator.device ), target.to(accelerator.device )
								# Perform our initial ground truth step in non "DDP"
								step_model(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )
								# Do "gradient accumulation" (noop)
								with accelerator.accumulate(lowerCAmelCase ):
												step_model(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )

								# DDP model and model should only be in sync when not (iteration % 2 == 0)
								for param, ddp_param in zip(model.parameters()  ,      ddp_model.parameters() ):
												if not param.requires_grad:
																continue
												if ((iteration + 1) % 2 == 0) or (iteration == len(lowerCAmelCase ) - 1):
																# Grads should be in sync
																assert (
																    torch.allclose(param.grad  ,      ddp_param.grad ) is True
																), F'Gradients not in sync when they should be at iteration {iteration}:\nModel grad ({param.grad}) != DDP grad ({ddp_param.grad})'
												else:
																# Grads should not be in sync
																assert (
																    torch.allclose(param.grad  ,      ddp_param.grad ) is False
																), F'Gradients in sync when they should not be at iteration {iteration}:\nModel grad ({param.grad}) == DDP grad ({ddp_param.grad})'

        # Shuffle ddp_input on each iteration
								torch.manual_seed(1337 + iteration )
								_UpperCAmelCase       :		int 		=    ddp_input[torch.randperm(len(lowerCAmelCase ) )]
				GradientState._reset_state()

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  List[str]=False  ,      lowerCAmelCase:  Optional[Any]=False )   ->       Optional[int]:
				_UpperCAmelCase       :		Optional[int] 		=    Accelerator(
				    split_batches=lowerCAmelCase  ,      dispatch_batches=lowerCAmelCase  ,      gradient_accumulation_steps=2 )
				# Test that context manager behaves properly
				_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase       :		Optional[Any] 		=    get_training_setup(lowerCAmelCase  ,      lowerCAmelCase )
				for iteration, batch in enumerate(lowerCAmelCase ):
								_UpperCAmelCase		,							_UpperCAmelCase       :		Optional[int] 		=    batch.values()
								# Gather the distributed inputs and targs for the base model
								_UpperCAmelCase		,							_UpperCAmelCase       :		int 		=    accelerator.gather((ddp_input, ddp_target) )
								_UpperCAmelCase		,							_UpperCAmelCase       :		Dict 		=    input.to(accelerator.device ), target.to(accelerator.device )
								# Perform our initial ground truth step in non "DDP"
								model.train()
								ddp_model.train()
								step_model(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )
								opt.step()

								if ((iteration + 1) % 2 == 0) or ((iteration + 1) == len(lowerCAmelCase )):
												if split_batches:
																sched.step()
												else:
																for _ in range(accelerator.num_processes ):
																				sched.step()
								opt.zero_grad()
								# Perform gradient accumulation under wrapper
								with accelerator.accumulate(lowerCAmelCase ):
												step_model(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )
												ddp_opt.step()
												ddp_sched.step()
												ddp_opt.zero_grad()

								# Learning rates should be the same
								assert (
								    opt.param_groups[0]["lr"] == ddp_opt.param_groups[0]["lr"]
								), F'Learning rates found in each optimizer did not align\nopt: {opt.param_groups[0]["lr"]}\nDDP opt: {ddp_opt.param_groups[0]["lr"]}\n'
								_UpperCAmelCase       :		Tuple 		=    (((iteration + 1) % 2) == 0) or ((iteration + 1) == len(lowerCAmelCase ))
								if accelerator.num_processes > 1:
												check_model_parameters(lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase  ,      lowerCAmelCase )
								# Shuffle ddp_input on each iteration
								torch.manual_seed(1337 + iteration )
				GradientState._reset_state()

def 							__SCREAMING_SNAKE_CASE							(				)   ->       Any:
				_UpperCAmelCase       :		Optional[int] 		=    Accelerator()

				_UpperCAmelCase       :		Dict 		=    RegressionDataset(length=80 )
				_UpperCAmelCase       :		Optional[Any] 		=    DataLoader(lowerCAmelCase  ,      batch_size=16 )
				_UpperCAmelCase       :		Any 		=    RegressionDataset(length=96 )
				_UpperCAmelCase       :		Union[str, Any] 		=    DataLoader(lowerCAmelCase  ,      batch_size=16 )
				_UpperCAmelCase		,							_UpperCAmelCase       :		List[Any] 		=    accelerator.prepare(lowerCAmelCase  ,      lowerCAmelCase )
				assert accelerator.gradient_state.active_dataloader is None
				for iteration, _ in enumerate(lowerCAmelCase ):
								assert id(accelerator.gradient_state.active_dataloader ) == id(lowerCAmelCase )
								if iteration < len(lowerCAmelCase ) - 1:
												assert not accelerator.gradient_state.end_of_dataloader
												if iteration == 1:
																for batch_num, _ in enumerate(lowerCAmelCase ):
																				assert id(accelerator.gradient_state.active_dataloader ) == id(lowerCAmelCase )
																				if batch_num < len(lowerCAmelCase ) - 1:
																								assert not accelerator.gradient_state.end_of_dataloader
																				else:
																								assert accelerator.gradient_state.end_of_dataloader
								else:
												assert accelerator.gradient_state.end_of_dataloader
				assert accelerator.gradient_state.active_dataloader is None

def 							__SCREAMING_SNAKE_CASE							(				)   ->       Dict:
				_UpperCAmelCase       :		Optional[int] 		=    Accelerator()
				_UpperCAmelCase       :		str 		=    accelerator.state
				if state.local_process_index == 0:
								print("**Test `accumulate` gradient accumulation with dataloader break**" )
				test_dataloader_break()
				if state.distributed_type == DistributedType.NO:
								if state.local_process_index == 0:
												print("**Test NOOP `no_sync` context manager**" )
								test_noop_sync(lowerCAmelCase )
				if state.distributed_type in (DistributedType.MULTI_GPU, DistributedType.MULTI_CPU):
								if state.local_process_index == 0:
												print("**Test Distributed `no_sync` context manager**" )
								test_distributed_sync(lowerCAmelCase )
				if state.distributed_type == DistributedType.MULTI_GPU:
								for split_batch in [True, False]:
												for dispatch_batches in [True, False]:
																if state.local_process_index == 0:
																				print(
																				    "**Test `accumulate` gradient accumulation, "  ,      F'`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**'  ,      )
																test_gradient_accumulation(lowerCAmelCase  ,      lowerCAmelCase )

    # Currently will break on torch 2.0 +, need to investigate why
				if is_torch_version("<"  ,      "2.0" ) or state.distributed_type == DistributedType.NO:
								if state.local_process_index == 0:
												print(
												    "**Test `accumulate` gradient accumulation with optimizer and scheduler, "  ,      "`split_batches=False`, `dispatch_batches=False`**"  ,      )
								test_gradient_accumulation_with_opt_and_scheduler()
								if state.distributed_type == DistributedType.MULTI_GPU:
												for split_batch in [True, False]:
																for dispatch_batches in [True, False]:
																				if not split_batch and not dispatch_batches:
																								continue
																				if state.local_process_index == 0:
																								print(
																								    "**Test `accumulate` gradient accumulation with optimizer and scheduler, "  ,      F'`split_batches={split_batch}` and `dispatch_batches={dispatch_batches}`**'  ,      )
																				test_gradient_accumulation_with_opt_and_scheduler(lowerCAmelCase  ,      lowerCAmelCase )

def 							__SCREAMING_SNAKE_CASE							(				lowerCAmelCase:  Dict )   ->       Union[str, Any]:
				# For xla_spawn (TPUs)
				main()


if __name__ == "__main__":
					main()





import gc
import unittest

from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
from transformers.pipelines import PipelineException
from transformers.testing_utils import (
    is_pipeline_test,
    is_torch_available,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_gpu,
    slow,
)

from .test_pipelines_common import ANY

@is_pipeline_test
class 				a  (	unittest.TestCase				):
						_lowercase    						=			MODEL_FOR_MASKED_LM_MAPPING
						_lowercase    						=			TF_MODEL_FOR_MASKED_LM_MAPPING


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										super().tearDown()
										# clean-up as much as possible GPU memory occupied by PyTorch
										gc.collect()
										if is_torch_available():
														import torch

														torch.cuda.empty_cache()


						@require_tf
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    pipeline(task="fill-mask"						,					model="sshleifer/tiny-distilroberta-base"						,					top_k=2						,					framework="tf"	)
										_UpperCAmelCase       :		Union[str, Any] 		=    unmasker("My name is <mask>"	)
										self.assertEqual(
										    nested_simplify(A_						,					decimals=6	)						,					[
										        {"sequence": "My name is grouped", "score": 2.1e-05, "token": 38015, "token_str": " grouped"},
										        {"sequence": "My name is accuser", "score": 2.1e-05, "token": 25506, "token_str": " accuser"},
										    ]						,					)

										_UpperCAmelCase       :		Optional[int] 		=    unmasker("The largest city in France is <mask>"	)
										self.assertEqual(
										    nested_simplify(A_						,					decimals=6	)						,					[
										        {
										            "sequence": "The largest city in France is grouped",
										            "score": 2.1e-05,
										            "token": 38015,
										            "token_str": " grouped",
										        },
										        {
										            "sequence": "The largest city in France is accuser",
										            "score": 2.1e-05,
										            "token": 25506,
										            "token_str": " accuser",
										        },
										    ]						,					)

										_UpperCAmelCase       :		Optional[Any] 		=    unmasker("My name is <mask>"						,					targets=[" Patrick", " Clara", " Teven"]						,					top_k=3	)
										self.assertEqual(
										    nested_simplify(A_						,					decimals=6	)						,					[
										        {"sequence": "My name is Clara", "score": 2e-05, "token": 13606, "token_str": " Clara"},
										        {"sequence": "My name is Patrick", "score": 2e-05, "token": 3499, "token_str": " Patrick"},
										        {"sequence": "My name is Te", "score": 1.9e-05, "token": 2941, "token_str": " Te"},
										    ]						,					)


						@require_torch
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[Any] 		=    pipeline(task="fill-mask"						,					model="sshleifer/tiny-distilroberta-base"						,					top_k=2						,					framework="pt"	)

										_UpperCAmelCase       :		List[str] 		=    unmasker("My name is <mask>"	)
										self.assertEqual(
										    nested_simplify(A_						,					decimals=6	)						,					[
										        {"sequence": "My name is Maul", "score": 2.2e-05, "token": 35676, "token_str": " Maul"},
										        {"sequence": "My name isELS", "score": 2.2e-05, "token": 16416, "token_str": "ELS"},
										    ]						,					)

										_UpperCAmelCase       :		str 		=    unmasker("The largest city in France is <mask>"	)
										self.assertEqual(
										    nested_simplify(A_						,					decimals=6	)						,					[
										        {
										            "sequence": "The largest city in France is Maul",
										            "score": 2.2e-05,
										            "token": 35676,
										            "token_str": " Maul",
										        },
										        {"sequence": "The largest city in France isELS", "score": 2.2e-05, "token": 16416, "token_str": "ELS"},
										    ]						,					)

										_UpperCAmelCase       :		List[Any] 		=    unmasker("My name is <mask>"						,					targets=[" Patrick", " Clara", " Teven"]						,					top_k=3	)
										self.assertEqual(
										    nested_simplify(A_						,					decimals=6	)						,					[
										        {"sequence": "My name is Patrick", "score": 2.1e-05, "token": 3499, "token_str": " Patrick"},
										        {"sequence": "My name is Te", "score": 2e-05, "token": 2941, "token_str": " Te"},
										        {"sequence": "My name is Clara", "score": 2e-05, "token": 13606, "token_str": " Clara"},
										    ]						,					)

										_UpperCAmelCase       :		Optional[Any] 		=    unmasker("My name is <mask> <mask>"						,					top_k=2	)

										self.assertEqual(
										    nested_simplify(A_						,					decimals=6	)						,					[
										        [
										            {
										                "score": 2.2e-05,
										                "token": 35676,
										                "token_str": " Maul",
										                "sequence": "<s>My name is Maul<mask></s>",
										            },
										            {"score": 2.2e-05, "token": 16416, "token_str": "ELS", "sequence": "<s>My name isELS<mask></s>"},
										        ],
										        [
										            {
										                "score": 2.2e-05,
										                "token": 35676,
										                "token_str": " Maul",
										                "sequence": "<s>My name is<mask> Maul</s>",
										            },
										            {"score": 2.2e-05, "token": 16416, "token_str": "ELS", "sequence": "<s>My name is<mask>ELS</s>"},
										        ],
										    ]						,					)


						@require_torch_gpu
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Union[str, Any] 		=    pipeline("fill-mask"						,					model="hf-internal-testing/tiny-random-distilbert"						,					device=0						,					framework="pt"	)

										# convert model to fp16
										pipe.model.half()

										_UpperCAmelCase       :		int 		=    pipe("Paris is the [MASK] of France."	)
										# We actually don't care about the result, we just want to make sure
										# it works, meaning the float16 tensor got casted back to float32
										# for postprocessing.
										self.assertIsInstance(A_						,					A_	)


						@slow
						@require_torch
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[Any] 		=    pipeline(task="fill-mask"						,					model="distilroberta-base"						,					top_k=2						,					framework="pt"	)
										self.run_large_test(A_	)


						@slow
						@require_tf
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    pipeline(task="fill-mask"						,					model="distilroberta-base"						,					top_k=2						,					framework="tf"	)
										self.run_large_test(A_	)


						def        _UpperCAmelCase				(       self						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    unmasker("My name is <mask>"	)
										self.assertEqual(
										    nested_simplify(A_	)						,					[
										        {"sequence": "My name is John", "score": 0.0_08, "token": 610, "token_str": " John"},
										        {"sequence": "My name is Chris", "score": 0.0_07, "token": 1573, "token_str": " Chris"},
										    ]						,					)
										_UpperCAmelCase       :		Optional[int] 		=    unmasker("The largest city in France is <mask>"	)
										self.assertEqual(
										    nested_simplify(A_	)						,					[
										        {
										            "sequence": "The largest city in France is Paris",
										            "score": 0.2_51,
										            "token": 2201,
										            "token_str": " Paris",
										        },
										        {
										            "sequence": "The largest city in France is Lyon",
										            "score": 0.2_14,
										            "token": 12790,
										            "token_str": " Lyon",
										        },
										    ]						,					)

										_UpperCAmelCase       :		Any 		=    unmasker("My name is <mask>"						,					targets=[" Patrick", " Clara", " Teven"]						,					top_k=3	)
										self.assertEqual(
										    nested_simplify(A_	)						,					[
										        {"sequence": "My name is Patrick", "score": 0.0_05, "token": 3499, "token_str": " Patrick"},
										        {"sequence": "My name is Clara", "score": 0.0_00, "token": 13606, "token_str": " Clara"},
										        {"sequence": "My name is Te", "score": 0.0_00, "token": 2941, "token_str": " Te"},
										    ]						,					)


						@require_torch
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    pipeline(task="fill-mask"						,					model="sshleifer/tiny-distilroberta-base"						,					framework="pt"	)
										_UpperCAmelCase       :		Optional[int] 		=    None
										_UpperCAmelCase       :		List[Any] 		=    None
										self.run_pipeline_test(A_						,					[]	)


						@require_tf
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    pipeline(task="fill-mask"						,					model="sshleifer/tiny-distilroberta-base"						,					framework="tf"	)
										_UpperCAmelCase       :		Union[str, Any] 		=    None
										_UpperCAmelCase       :		str 		=    None
										self.run_pipeline_test(A_						,					[]	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										if tokenizer is None or tokenizer.mask_token_id is None:
														self.skipTest("The provided tokenizer has no mask token, (probably reformer or wav2vec2)"	)

										_UpperCAmelCase       :		Tuple 		=    FillMaskPipeline(model=A_						,					tokenizer=A_	)
										_UpperCAmelCase       :		Tuple 		=    [
										    f'This is another {tokenizer.mask_token} test',
										]
										return fill_masker, examples


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Dict 		=    fill_masker.tokenizer
										_UpperCAmelCase       :		Tuple 		=    fill_masker.model

										_UpperCAmelCase       :		Union[str, Any] 		=    fill_masker(
										    f'This is a {tokenizer.mask_token}'						,					)
										self.assertEqual(
										    A_						,					[
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										    ]						,					)

										_UpperCAmelCase       :		Optional[Any] 		=    fill_masker([f'This is a {tokenizer.mask_token}']	)
										self.assertEqual(
										    A_						,					[
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										    ]						,					)

										_UpperCAmelCase       :		List[str] 		=    fill_masker([f'This is a {tokenizer.mask_token}', f'Another {tokenizer.mask_token} great test.']	)
										self.assertEqual(
										    A_						,					[
										        [
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        ],
										        [
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        ],
										    ]						,					)

										with self.assertRaises(A_	):
														fill_masker([None]	)
										# No mask_token is not supported
										with self.assertRaises(A_	):
														fill_masker("This is"	)

										self.run_test_top_k(A_						,					A_	)
										self.run_test_targets(A_						,					A_	)
										self.run_test_top_k_targets(A_						,					A_	)
										self.fill_mask_with_duplicate_targets_and_top_k(A_						,					A_	)
										self.fill_mask_with_multiple_masks(A_						,					A_	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Any 		=    tokenizer.get_vocab()
										_UpperCAmelCase       :		int 		=    sorted(vocab.keys()	)[:2]
										# Pipeline argument
										_UpperCAmelCase       :		Any 		=    FillMaskPipeline(model=A_						,					tokenizer=A_						,					targets=A_	)
										_UpperCAmelCase       :		Optional[int] 		=    fill_masker(f'This is a {tokenizer.mask_token}'	)
										self.assertEqual(
										    A_						,					[
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										    ]						,					)
										_UpperCAmelCase       :		Tuple 		=    {vocab[el] for el in targets}
										self.assertEqual({el["token"] for el in outputs}						,					A_	)
										_UpperCAmelCase       :		List[Any] 		=    [tokenizer.decode([x]	) for x in target_ids]
										self.assertEqual({el["token_str"] for el in outputs}						,					set(A_	)	)

										# Call argument
										_UpperCAmelCase       :		Union[str, Any] 		=    FillMaskPipeline(model=A_						,					tokenizer=A_	)
										_UpperCAmelCase       :		List[Any] 		=    fill_masker(f'This is a {tokenizer.mask_token}'						,					targets=A_	)
										self.assertEqual(
										    A_						,					[
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										    ]						,					)
										_UpperCAmelCase       :		Optional[Any] 		=    {vocab[el] for el in targets}
										self.assertEqual({el["token"] for el in outputs}						,					A_	)
										_UpperCAmelCase       :		Any 		=    [tokenizer.decode([x]	) for x in target_ids]
										self.assertEqual({el["token_str"] for el in outputs}						,					set(A_	)	)

										# Score equivalence
										_UpperCAmelCase       :		int 		=    fill_masker(f'This is a {tokenizer.mask_token}'						,					targets=A_	)
										_UpperCAmelCase       :		Tuple 		=    [top_mask["token_str"] for top_mask in outputs]
										_UpperCAmelCase       :		Union[str, Any] 		=    [top_mask["score"] for top_mask in outputs]

										# For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
										if set(A_	) == set(A_	):
														_UpperCAmelCase       :		List[str] 		=    fill_masker(f'This is a {tokenizer.mask_token}'						,					targets=A_	)
														_UpperCAmelCase       :		List[str] 		=    [top_mask["score"] for top_mask in unmasked_targets]
														self.assertEqual(nested_simplify(A_	)						,					nested_simplify(A_	)	)

										# Raises with invalid
										with self.assertRaises(A_	):
														_UpperCAmelCase       :		Optional[int] 		=    fill_masker(f'This is a {tokenizer.mask_token}'						,					targets=[]	)
										# For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
										if "" not in tokenizer.get_vocab():
														with self.assertRaises(A_	):
																		_UpperCAmelCase       :		Dict 		=    fill_masker(f'This is a {tokenizer.mask_token}'						,					targets=[""]	)
														with self.assertRaises(A_	):
																		_UpperCAmelCase       :		List[str] 		=    fill_masker(f'This is a {tokenizer.mask_token}'						,					targets=""	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Union[str, Any] 		=    FillMaskPipeline(model=A_						,					tokenizer=A_						,					top_k=2	)
										_UpperCAmelCase       :		Union[str, Any] 		=    fill_masker(f'This is a {tokenizer.mask_token}'	)
										self.assertEqual(
										    A_						,					[
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										    ]						,					)

										_UpperCAmelCase       :		int 		=    FillMaskPipeline(model=A_						,					tokenizer=A_	)
										_UpperCAmelCase       :		Optional[Any] 		=    fill_masker(f'This is a {tokenizer.mask_token}'						,					top_k=2	)
										self.assertEqual(
										    A_						,					[
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										    ]						,					)
										self.assertEqual(nested_simplify(A_	)						,					nested_simplify(A_	)	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    tokenizer.get_vocab()
										_UpperCAmelCase       :		Union[str, Any] 		=    FillMaskPipeline(model=A_						,					tokenizer=A_	)

										# top_k=2, ntargets=3
										_UpperCAmelCase       :		Tuple 		=    sorted(vocab.keys()	)[:3]
										_UpperCAmelCase       :		Union[str, Any] 		=    fill_masker(f'This is a {tokenizer.mask_token}'						,					top_k=2						,					targets=A_	)

										# If we use the most probably targets, and filter differently, we should still
										# have the same results
										_UpperCAmelCase       :		Tuple 		=    [el["token_str"] for el in sorted(A_						,					key=lambda A_	: x["score"]						,					reverse=A_	)]
										# For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
										if set(A_	).issubset(A_	):
														_UpperCAmelCase       :		int 		=    fill_masker(f'This is a {tokenizer.mask_token}'						,					top_k=3						,					targets=A_	)
														# They should yield exactly the same result
														self.assertEqual(nested_simplify(A_	)						,					nested_simplify(A_	)	)


						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    FillMaskPipeline(model=A_						,					tokenizer=A_	)
										_UpperCAmelCase       :		Dict 		=    tokenizer.get_vocab()
										# String duplicates + id duplicates
										_UpperCAmelCase       :		Dict 		=    sorted(vocab.keys()	)[:3]
										_UpperCAmelCase       :		Tuple 		=    [targets[0], targets[1], targets[0], targets[2], targets[1]]
										_UpperCAmelCase       :		Dict 		=    fill_masker(f'My name is {tokenizer.mask_token}'						,					targets=A_						,					top_k=10	)

										# The target list contains duplicates, so we can't output more
										# than them
										self.assertEqual(len(A_	)						,					3	)





						def        _UpperCAmelCase				(       self						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    FillMaskPipeline(model=A_						,					tokenizer=A_	)

										_UpperCAmelCase       :		List[str] 		=    fill_masker(
										    f'This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}'						,					top_k=2	)
										self.assertEqual(
										    A_						,					[
										        [
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        ],
										        [
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        ],
										        [
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										            {"sequence": ANY(A_	), "score": ANY(A_	), "token": ANY(A_	), "token_str": ANY(A_	)},
										        ],
										    ]						,					)


