




import unittest

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

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
			import torch

			from transformers import BertGenerationDecoder, BertGenerationEncoder






class  __magic_name__						:
				'''simple docstring'''







				def __init__(							self:   List[Any]		,					_lowerCamelCase:   Any		,					_lowerCamelCase:   Optional[Any]=13		,					_lowerCamelCase:   Tuple=7		,					_lowerCamelCase:   Union[str, Any]=True		,					_lowerCamelCase:   int=True		,					_lowerCamelCase:   str=99		,					_lowerCamelCase:   List[str]=32		,					_lowerCamelCase:   Tuple=5		,					_lowerCamelCase:   Optional[Any]=4		,					_lowerCamelCase:   Tuple=37		,					_lowerCamelCase:   Dict="gelu"		,					_lowerCamelCase:   Optional[int]=0.1		,					_lowerCamelCase:   Optional[Any]=0.1		,					_lowerCamelCase:   int=50		,					_lowerCamelCase:   Optional[int]=0.02		,					_lowerCamelCase:   Any=True		,					_lowerCamelCase:   List[Any]=None		,					):
											SCREAMING_SNAKE_CASE_   	=    parent
											SCREAMING_SNAKE_CASE_   	=    batch_size
											SCREAMING_SNAKE_CASE_   	=    seq_length
											SCREAMING_SNAKE_CASE_   	=    is_training
											SCREAMING_SNAKE_CASE_   	=    use_input_mask
											SCREAMING_SNAKE_CASE_   	=    vocab_size
											SCREAMING_SNAKE_CASE_   	=    hidden_size
											SCREAMING_SNAKE_CASE_   	=    num_hidden_layers
											SCREAMING_SNAKE_CASE_   	=    num_attention_heads
											SCREAMING_SNAKE_CASE_   	=    intermediate_size
											SCREAMING_SNAKE_CASE_   	=    hidden_act
											SCREAMING_SNAKE_CASE_   	=    hidden_dropout_prob
											SCREAMING_SNAKE_CASE_   	=    attention_probs_dropout_prob
											SCREAMING_SNAKE_CASE_   	=    max_position_embeddings
											SCREAMING_SNAKE_CASE_   	=    initializer_range
											SCREAMING_SNAKE_CASE_   	=    use_labels
											SCREAMING_SNAKE_CASE_   	=    scope


				def     _A							(							self:   int		):
											SCREAMING_SNAKE_CASE_   	=    ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size		)

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

											if self.use_labels:
																		SCREAMING_SNAKE_CASE_   	=    ids_tensor([self.batch_size, self.seq_length]		,					self.vocab_size		)

											SCREAMING_SNAKE_CASE_   	=    self.get_config()

											return config, input_ids, input_mask, token_labels


				def     _A							(							self:   List[str]		):
											return BertGenerationConfig(
											    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		,					is_decoder=_lowerCamelCase		,					initializer_range=self.initializer_range		,					)


				def     _A							(							self:   Tuple		):
											(
											    (
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	
											)   	=    self.prepare_config_and_inputs()

											SCREAMING_SNAKE_CASE_   	=    True
											SCREAMING_SNAKE_CASE_   	=    floats_tensor([self.batch_size, self.seq_length, self.hidden_size]		)
											SCREAMING_SNAKE_CASE_   	=    ids_tensor([self.batch_size, self.seq_length]		,					vocab_size=2		)

											return (
											    config,
											    input_ids,
											    input_mask,
											    token_labels,
											    encoder_hidden_states,
											    encoder_attention_mask,
											)


				def     _A							(							self:   str		,					_lowerCamelCase:   Tuple		,					_lowerCamelCase:   Tuple		,					_lowerCamelCase:   Tuple		,					_lowerCamelCase:   Tuple		,					**_lowerCamelCase:   Optional[int]		,					):
											SCREAMING_SNAKE_CASE_   	=    BertGenerationEncoder(config=_lowerCamelCase		)
											model.to(_lowerCamelCase		)
											model.eval()
											SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		,					attention_mask=_lowerCamelCase		)
											SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		)
											self.parent.assertEqual(result.last_hidden_state.shape		,					(self.batch_size, self.seq_length, self.hidden_size)		)


				def     _A							(							self:   Any		,					_lowerCamelCase:   Tuple		,					_lowerCamelCase:   Any		,					_lowerCamelCase:   List[Any]		,					_lowerCamelCase:   Union[str, Any]		,					_lowerCamelCase:   Optional[int]		,					_lowerCamelCase:   Tuple		,					**_lowerCamelCase:   Any		,					):
											SCREAMING_SNAKE_CASE_   	=    True
											SCREAMING_SNAKE_CASE_   	=    BertGenerationEncoder(config=_lowerCamelCase		)
											model.to(_lowerCamelCase		)
											model.eval()
											SCREAMING_SNAKE_CASE_   	=    model(
											    _lowerCamelCase		,					attention_mask=_lowerCamelCase		,					encoder_hidden_states=_lowerCamelCase		,					encoder_attention_mask=_lowerCamelCase		,					)
											SCREAMING_SNAKE_CASE_   	=    model(
											    _lowerCamelCase		,					attention_mask=_lowerCamelCase		,					encoder_hidden_states=_lowerCamelCase		,					)
											self.parent.assertEqual(result.last_hidden_state.shape		,					(self.batch_size, self.seq_length, self.hidden_size)		)


				def     _A							(							self:   Optional[Any]		,					_lowerCamelCase:   List[Any]		,					_lowerCamelCase:   List[Any]		,					_lowerCamelCase:   Optional[int]		,					_lowerCamelCase:   List[str]		,					_lowerCamelCase:   Union[str, Any]		,					_lowerCamelCase:   str		,					**_lowerCamelCase:   Dict		,					):
											SCREAMING_SNAKE_CASE_   	=    True
											SCREAMING_SNAKE_CASE_   	=    True
											SCREAMING_SNAKE_CASE_   	=    BertGenerationDecoder(config=_lowerCamelCase		).to(_lowerCamelCase		).eval()

											# first forward pass
											SCREAMING_SNAKE_CASE_   	=    model(
											    _lowerCamelCase		,					attention_mask=_lowerCamelCase		,					encoder_hidden_states=_lowerCamelCase		,					encoder_attention_mask=_lowerCamelCase		,					use_cache=_lowerCamelCase		,					)
											SCREAMING_SNAKE_CASE_   	=    outputs.past_key_values

											# create hypothetical multiple next token and extent to next_input_ids
											SCREAMING_SNAKE_CASE_   	=    ids_tensor((self.batch_size, 3)		,					config.vocab_size		)
											SCREAMING_SNAKE_CASE_   	=    ids_tensor((self.batch_size, 3)		,					vocab_size=2		)

											# append to next input_ids and
											SCREAMING_SNAKE_CASE_   	=    torch.cat([input_ids, next_tokens]		,					dim=-1		)
											SCREAMING_SNAKE_CASE_   	=    torch.cat([input_mask, next_mask]		,					dim=-1		)

											SCREAMING_SNAKE_CASE_   	=    model(
											    _lowerCamelCase		,					attention_mask=_lowerCamelCase		,					encoder_hidden_states=_lowerCamelCase		,					encoder_attention_mask=_lowerCamelCase		,					output_hidden_states=_lowerCamelCase		,					)['''hidden_states'''][0]
											SCREAMING_SNAKE_CASE_   	=    model(
											    _lowerCamelCase		,					attention_mask=_lowerCamelCase		,					encoder_hidden_states=_lowerCamelCase		,					encoder_attention_mask=_lowerCamelCase		,					past_key_values=_lowerCamelCase		,					output_hidden_states=_lowerCamelCase		,					)['''hidden_states'''][0]

											# select random slice
											SCREAMING_SNAKE_CASE_   	=    ids_tensor((1,)		,					output_from_past.shape[-1]		).item()
											SCREAMING_SNAKE_CASE_   	=    output_from_no_past[:, -3:, random_slice_idx].detach()
											SCREAMING_SNAKE_CASE_   	=    output_from_past[:, :, random_slice_idx].detach()

											self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]		)

											# test that outputs are equal for slice
											self.parent.assertTrue(torch.allclose(_lowerCamelCase		,					_lowerCamelCase		,					atol=1E-3		)		)


				def     _A							(							self:   List[Any]		,					_lowerCamelCase:   Any		,					_lowerCamelCase:   int		,					_lowerCamelCase:   Tuple		,					_lowerCamelCase:   Optional[Any]		,					*_lowerCamelCase:   List[Any]		,					):
											SCREAMING_SNAKE_CASE_   	=    BertGenerationDecoder(_lowerCamelCase		)
											model.to(_lowerCamelCase		)
											model.eval()
											SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		,					attention_mask=_lowerCamelCase		,					labels=_lowerCamelCase		)
											self.parent.assertEqual(result.logits.shape		,					(self.batch_size, self.seq_length, self.vocab_size)		)





				def     _A							(							self:   str		):
											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    self.prepare_config_and_inputs()
											SCREAMING_SNAKE_CASE_   	=    {'''input_ids''': input_ids, '''attention_mask''': input_mask}
											return config, inputs_dict






@require_torch
class  __magic_name__						(  __UpperCAmelCase    ,  __UpperCAmelCase    ,  __UpperCAmelCase    ,  unittest.TestCase):
				'''simple docstring'''







				SCREAMING_SNAKE_CASE__					:							Optional[Any]	      =					(BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
				SCREAMING_SNAKE_CASE__					:							Any	      =					(BertGenerationDecoder,) if is_torch_available() else ()
				SCREAMING_SNAKE_CASE__					:							Dict	      =					(
				    {"feature-extraction": BertGenerationEncoder, "text-generation": BertGenerationDecoder}
				    if is_torch_available()
				    else {}
				)


				def     _A							(							self:   Union[str, Any]		):
											SCREAMING_SNAKE_CASE_   	=    BertGenerationEncoderTester(self		)
											SCREAMING_SNAKE_CASE_   	=    ConfigTester(self		,					config_class=_lowerCamelCase		,					hidden_size=37		)


				def     _A							(							self:   Any		):
											self.config_tester.run_common_tests()


				def     _A							(							self:   int		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_lowerCamelCase		)


				def     _A							(							self:   str		):
											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs()
											SCREAMING_SNAKE_CASE_   	=    '''bert'''
											self.model_tester.create_and_check_model(_lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		)


				def     _A							(							self:   Optional[int]		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(*_lowerCamelCase		)


				def     _A							(							self:   Dict		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_decoder_model_past_large_inputs(*_lowerCamelCase		)


				def     _A							(							self:   Optional[Any]		):
											# This regression test was failing with PyTorch < 1.3
											(
											    (
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	(
											    SCREAMING_SNAKE_CASE_
											)	,	
											)   	=    self.model_tester.prepare_config_and_inputs_for_decoder()

											SCREAMING_SNAKE_CASE_   	=    None

											self.model_tester.create_and_check_model_as_decoder(
											    _lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		,					_lowerCamelCase		,					)


				def     _A							(							self:   Optional[Any]		):
											SCREAMING_SNAKE_CASE_   	=    self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_for_causal_lm(*_lowerCamelCase		)





				@slow
				def     _A							(							self:   str		):
											SCREAMING_SNAKE_CASE_   	=    BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''		)
											self.assertIsNotNone(_lowerCamelCase		)






@require_torch
class  __magic_name__						(  unittest.TestCase):
				'''simple docstring'''







				@slow
				def     _A							(							self:   Dict		):
											SCREAMING_SNAKE_CASE_   	=    BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''		)
											SCREAMING_SNAKE_CASE_   	=    torch.tensor([[1_01, 75_92, 10_10, 20_26, 38_99, 20_03, 1_01_40, 1_02]]		)
											with torch.no_grad():
																		SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		)[0]
											SCREAMING_SNAKE_CASE_   	=    torch.Size([1, 8, 10_24]		)
											self.assertEqual(output.shape		,					_lowerCamelCase		)
											SCREAMING_SNAKE_CASE_   	=    torch.tensor(
											    [[[0.17_75, 0.00_83, -0.03_21], [1.60_02, 0.12_87, 0.39_12], [2.14_73, 0.57_91, 0.60_66]]]		)
											self.assertTrue(torch.allclose(output[:, :3, :3]		,					_lowerCamelCase		,					atol=1E-4		)		)






@require_torch
class  __magic_name__						(  unittest.TestCase):
				'''simple docstring'''







				@slow
				def     _A							(							self:   Optional[int]		):
											SCREAMING_SNAKE_CASE_   	=    BertGenerationDecoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''		)
											SCREAMING_SNAKE_CASE_   	=    torch.tensor([[1_01, 75_92, 10_10, 20_26, 38_99, 20_03, 1_01_40, 1_02]]		)
											with torch.no_grad():
																		SCREAMING_SNAKE_CASE_   	=    model(_lowerCamelCase		)[0]
											SCREAMING_SNAKE_CASE_   	=    torch.Size([1, 8, 5_03_58]		)
											self.assertEqual(output.shape		,					_lowerCamelCase		)
											SCREAMING_SNAKE_CASE_   	=    torch.tensor(
											    [[[-0.57_88, -2.59_94, -3.70_54], [0.04_38, 4.79_97, 1.87_95], [1.58_62, 6.64_09, 4.46_38]]]		)
											self.assertTrue(torch.allclose(output[:, :3, :3]		,					_lowerCamelCase		,					atol=1E-4		)		)








import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
			import torch

			from transformers.generation import DisjunctiveConstraint






@require_torch
class  __magic_name__						(  unittest.TestCase):
				'''simple docstring'''







				def     _A							(							self:   str		):
											# For consistency across different places the DisjunctiveConstraint is called,
											# dc.token_ids is a list of integers. It is also initialized only by integers.

											SCREAMING_SNAKE_CASE_   	=    [[1, 2, 4], [1, 2, 3, 4]]
											SCREAMING_SNAKE_CASE_   	=    DisjunctiveConstraint(_lowerCamelCase		)
											self.assertTrue(isinstance(dc.token_ids		,					_lowerCamelCase		)		)

											with self.assertRaises(_lowerCamelCase		):
																		DisjunctiveConstraint(torch.LongTensor([[1, 2, 4], [1, 2, 3]]		)		)

											with self.assertRaises(_lowerCamelCase		):
																		DisjunctiveConstraint([torch.LongTensor([1, 2, 4]		), torch.LongTensor([1, 2, 3, 4, 5]		)]		)


				def     _A							(							self:   Union[str, Any]		):
											# We can't have constraints that are complete subsets of another. This leads to a preverse
											# interpretation of "constraint fulfillment": does generating [1,2,3] fulfill the constraint?
											# It would mean that it generated [1,2] which fulfills it, but it's in the middle of potentially
											# fulfilling [1,2,3,4]. If we believe that [1,2,3] does fulfill the constraint, then the algorithm
											# will necessarily never reach [1,2,3,4], giving users a false sense of control (better to just not allow it).
											SCREAMING_SNAKE_CASE_   	=    [[1, 2], [1, 2, 3, 4]]

											with self.assertRaises(_lowerCamelCase		):
																		DisjunctiveConstraint(_lowerCamelCase		)  # fails here


				def     _A							(							self:   Dict		):
											SCREAMING_SNAKE_CASE_   	=    [[1, 2, 3], [1, 2, 4]]

											SCREAMING_SNAKE_CASE_   	=    DisjunctiveConstraint(_lowerCamelCase		)

											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    dc.update(1		)
											SCREAMING_SNAKE_CASE_   	=    stepped is True and completed is False and reset is False
											self.assertTrue(_lowerCamelCase		)
											self.assertTrue(not dc.completed		)
											self.assertTrue(dc.current_seq == [1]		)

											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    dc.update(2		)
											SCREAMING_SNAKE_CASE_   	=    stepped is True and completed is False and reset is False
											self.assertTrue(_lowerCamelCase		)
											self.assertTrue(not dc.completed		)
											self.assertTrue(dc.current_seq == [1, 2]		)

											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    dc.update(3		)
											SCREAMING_SNAKE_CASE_   	=    stepped is True and completed is True and reset is False
											self.assertTrue(_lowerCamelCase		)
											self.assertTrue(dc.completed		)  # Completed!
											self.assertTrue(dc.current_seq == [1, 2, 3]		)





				def     _A							(							self:   Optional[int]		):
											SCREAMING_SNAKE_CASE_   	=    [[1, 2, 3], [1, 2, 4, 5], [1, 2, 5]]

											SCREAMING_SNAKE_CASE_   	=    DisjunctiveConstraint(_lowerCamelCase		)

											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    dc.update(1		)
											self.assertTrue(not dc.completed		)
											self.assertTrue(dc.current_seq == [1]		)

											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    dc.update(2		)
											self.assertTrue(not dc.completed		)
											self.assertTrue(dc.current_seq == [1, 2]		)

											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    dc.update(4		)
											self.assertTrue(not dc.completed		)
											self.assertTrue(dc.current_seq == [1, 2, 4]		)

											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    dc.update(5		)
											self.assertTrue(dc.completed		)  # Completed!
											self.assertTrue(dc.current_seq == [1, 2, 4, 5]		)

											dc.reset()

											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    dc.update(1		)
											self.assertTrue(not dc.completed		)
											self.assertTrue(dc.remaining() == 3		)
											self.assertTrue(dc.current_seq == [1]		)

											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    dc.update(2		)
											self.assertTrue(not dc.completed		)
											self.assertTrue(dc.remaining() == 2		)
											self.assertTrue(dc.current_seq == [1, 2]		)

											SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_	,	SCREAMING_SNAKE_CASE_   	=    dc.update(5		)
											self.assertTrue(dc.completed		)  # Completed!
											self.assertTrue(dc.remaining() == 0		)
											self.assertTrue(dc.current_seq == [1, 2, 5]		)


