



'''simple docstring'''

import unittest

from transformers import is_flax_available
from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, require_torch, slow


if is_flax_available():
				import optax
				from flax.training.common_utils import onehot

				from transformers import AutoTokenizer, FlaxMTaForConditionalGeneration
				from transformers.models.ta.modeling_flax_ta import shift_tokens_right

@require_torch
@require_sentencepiece
@require_tokenizers
@require_flax
class __snake_case(				unittest.TestCase ):


						'''simple docstring'''



						@slow
						def     __snake_case   ( self				)     ->     Union[str, Any]:
							lowerCAmelCase   =     FlaxMTaForConditionalGeneration.from_pretrained("""google/mt5-small"""				)
							lowerCAmelCase   =     AutoTokenizer.from_pretrained("""google/mt5-small"""				)

							lowerCAmelCase   =     tokenizer("""Hello there"""     ,	return_tensors="""np"""				).input_ids
							lowerCAmelCase   =     tokenizer("""Hi I am"""     ,	return_tensors="""np"""				).input_ids

							lowerCAmelCase   =     shift_tokens_right(A_     ,	model.config.pad_token_id     ,	model.config.decoder_start_token_id				)

							lowerCAmelCase   =     model(A_     ,	decoder_input_ids=A_				).logits
							lowerCAmelCase   =     optax.softmax_cross_entropy(A_     ,	onehot(A_     ,	logits.shape[-1]				)				).mean()

							lowerCAmelCase   =     -(labels.shape[-1] * loss.item())

							lowerCAmelCase   =     -8_4.9_1_2_7
							self.assertTrue(abs(mtf_score - EXPECTED_SCORE				) < 1e-4				)




'''simple docstring'''

import time
import unittest

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

from ..test_modeling_common import ids_tensor


if is_torch_available():
				import torch

				from transformers.generation import (
				    MaxLengthCriteria,
				    MaxNewTokensCriteria,
				    MaxTimeCriteria,
				    StoppingCriteriaList,
				    validate_stopping_criteria,
				)

@require_torch
class __snake_case(				unittest.TestCase ):


						'''simple docstring'''



						def     __snake_case   ( self     ,	A_				)     ->     str:
							lowerCAmelCase   =     3
							lowerCAmelCase   =     250

							lowerCAmelCase   =     ids_tensor((batch_size, length)     ,	A_				)
							lowerCAmelCase   =     torch.ones((batch_size, length)     ,	device=A_     ,	dtype=torch.float				) / length
							return input_ids, scores


						def     __snake_case   ( self				)     ->     Any:
							lowerCAmelCase,       lowerCAmelCase   =     self._get_tensors(5				)

							lowerCAmelCase   =     StoppingCriteriaList(
							    [
							        MaxLengthCriteria(max_length=10				),
							        MaxTimeCriteria(max_time=0.1				),
							    ]				)

							self.assertFalse(criteria(A_     ,	A_				)				)

							lowerCAmelCase,       lowerCAmelCase   =     self._get_tensors(9				)
							self.assertFalse(criteria(A_     ,	A_				)				)

							lowerCAmelCase,       lowerCAmelCase   =     self._get_tensors(10				)
							self.assertTrue(criteria(A_     ,	A_				)				)


						def     __snake_case   ( self				)     ->     Optional[int]:
							lowerCAmelCase   =     MaxLengthCriteria(max_length=10				)

							lowerCAmelCase,       lowerCAmelCase   =     self._get_tensors(5				)
							self.assertFalse(criteria(A_     ,	A_				)				)

							lowerCAmelCase,       lowerCAmelCase   =     self._get_tensors(9				)
							self.assertFalse(criteria(A_     ,	A_				)				)

							lowerCAmelCase,       lowerCAmelCase   =     self._get_tensors(10				)
							self.assertTrue(criteria(A_     ,	A_				)				)


						def     __snake_case   ( self				)     ->     List[Any]:
							lowerCAmelCase   =     MaxNewTokensCriteria(start_length=5     ,	max_new_tokens=5				)

							lowerCAmelCase,       lowerCAmelCase   =     self._get_tensors(5				)
							self.assertFalse(criteria(A_     ,	A_				)				)

							lowerCAmelCase,       lowerCAmelCase   =     self._get_tensors(9				)
							self.assertFalse(criteria(A_     ,	A_				)				)

							lowerCAmelCase,       lowerCAmelCase   =     self._get_tensors(10				)
							self.assertTrue(criteria(A_     ,	A_				)				)

							lowerCAmelCase   =     StoppingCriteriaList([criteria]				)
							self.assertEqual(criteria_list.max_length     ,	10				)


						def     __snake_case   ( self				)     ->     List[str]:
							lowerCAmelCase,       lowerCAmelCase   =     self._get_tensors(5				)

							lowerCAmelCase   =     MaxTimeCriteria(max_time=0.1				)
							self.assertFalse(criteria(A_     ,	A_				)				)

							lowerCAmelCase   =     MaxTimeCriteria(max_time=0.1     ,	initial_timestamp=time.time() - 0.2				)
							self.assertTrue(criteria(A_     ,	A_				)				)




						def     __snake_case   ( self				)     ->     Optional[int]:
							validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10				)]				)     ,	10				)

							with self.assertWarns(A_				):
								validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10				)]				)     ,	11				)

							lowerCAmelCase   =     validate_stopping_criteria(StoppingCriteriaList()     ,	11				)

							self.assertEqual(len(A_				)     ,	1				)