





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 						lowerCamelCase       (unittest.TestCase      ):







							"""simple docstring"""




							@slow
							def  A_     ( self     :					Union[str, Any]						)				->			str:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						FlaxMTaForConditionalGeneration.from_pretrained("google/mt5-small"						)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						AutoTokenizer.from_pretrained("google/mt5-small"						)

														SCREAMING_SNAKE_CASE__       :      Any     =						tokenizer("Hello there",			return_tensors="np"						).input_ids
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						tokenizer("Hi I am",			return_tensors="np"						).input_ids

														SCREAMING_SNAKE_CASE__       :      Any     =						shift_tokens_right(_UpperCAmelCase,			model.config.pad_token_id,			model.config.decoder_start_token_id						)

														SCREAMING_SNAKE_CASE__       :      Dict     =						model(_UpperCAmelCase,			decoder_input_ids=_UpperCAmelCase						).logits
														SCREAMING_SNAKE_CASE__       :      str     =						optax.softmax_cross_entropy(_UpperCAmelCase,			onehot(_UpperCAmelCase,			logits.shape[-1]						)						).mean()

														SCREAMING_SNAKE_CASE__       :      List[str]     =						-(labels.shape[-1] * loss.item())

														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						-84.9127
														self.assertTrue(abs(mtf_score - EXPECTED_SCORE						) < 1E-4						)













import pytest

from datasets.splits import SplitDict, SplitInfo
from datasets.utils.py_utils import asdict




@pytest.mark.parametrize(
    "split_dict"			,				[
        SplitDict(),
        SplitDict({"train": SplitInfo(name="train"			,				num_bytes=13_37			,				num_examples=42			,				dataset_name="my_dataset"       )}       ),
        SplitDict({"train": SplitInfo(name="train"			,				num_bytes=13_37			,				num_examples=42       )}       ),
        SplitDict({"train": SplitInfo()}       ),
    ]			,				)
def _a (  SCREAMING_SNAKE_CASE__      :  SplitDict       )	->    str:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      Dict     =						split_dict._to_yaml_list()
							assert len(SCREAMING_SNAKE_CASE__       ) == len(SCREAMING_SNAKE_CASE__       )
							SCREAMING_SNAKE_CASE__       :      List[Any]     =						SplitDict._from_yaml_list(SCREAMING_SNAKE_CASE__       )
							for split_name, split_info in split_dict.items():
														# dataset_name field is deprecated, and is therefore not part of the YAML dump
														SCREAMING_SNAKE_CASE__       :      str     =						None
														# the split name of split_dict takes over the name of the split info object
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						split_name
							assert split_dict == reloaded







@pytest.mark.parametrize(
    "split_info"			,				[SplitInfo(), SplitInfo(dataset_name=SCREAMING_SNAKE_CASE__       ), SplitInfo(dataset_name="my_dataset"       )]       )
def _a (  SCREAMING_SNAKE_CASE__      :  Any       )	->    Dict:

							'''simple docstring'''






							SCREAMING_SNAKE_CASE__       :      str     =						asdict(SplitDict({"train": split_info}       )       )
							assert "dataset_name" in split_dict_asdict["train"]
							assert split_dict_asdict["train"]["dataset_name"] == split_info.dataset_name






