



import unittest

from transformers import is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device


if is_torch_available():
				import torch

				from transformers import AutoModelForImageClassification

if is_vision_available():
				from transformers import AutoImageProcessor

@require_torch
@require_vision
class 				snake_case__(unittest.TestCase      ):

					"""simple docstring"""






					@slow
					def 			snake_case					(							self :     List[Any]					):
									lowercase__					:   Tuple													=					AutoImageProcessor.from_pretrained("microsoft/dit-base-finetuned-rvlcdip"					)
									lowercase__					:   str													=					AutoModelForImageClassification.from_pretrained("microsoft/dit-base-finetuned-rvlcdip"					)
									model.to(SCREAMING_SNAKE_CASE					)

									from datasets import load_dataset

									lowercase__					:   str													=					load_dataset("nielsr/rvlcdip-demo"					)

									lowercase__					:   List[Any]													=					dataset["train"][0]["image"].convert("RGB"					)

									lowercase__					:   str													=					image_processor(SCREAMING_SNAKE_CASE				,       return_tensors="pt"					).to(SCREAMING_SNAKE_CASE					)

									# forward pass
									with torch.no_grad():
													lowercase__					:   Optional[Any]													=					model(**SCREAMING_SNAKE_CASE					)
													lowercase__					:   Optional[int]													=					outputs.logits

									lowercase__					:   Tuple													=					torch.Size((1, 16)					)
									self.assertEqual(logits.shape				,       SCREAMING_SNAKE_CASE					)

									lowercase__					:   Dict													=					torch.tensor(
									    [-0.4_158, -0.4_092, -0.4_347]				,       device=SCREAMING_SNAKE_CASE				,       dtype=torch.float				,       )
									self.assertTrue(torch.allclose(logits[0, :3]				,       SCREAMING_SNAKE_CASE				,       atol=1E-4					)					)









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 :     Optional[int]					):
									lowercase__					:   Dict													=					FlaxMTaForConditionalGeneration.from_pretrained("google/mt5-small"					)
									lowercase__					:   Dict													=					AutoTokenizer.from_pretrained("google/mt5-small"					)

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

									lowercase__					:   int													=					shift_tokens_right(SCREAMING_SNAKE_CASE				,       model.config.pad_token_id				,       model.config.decoder_start_token_id					)

									lowercase__					:   Tuple													=					model(SCREAMING_SNAKE_CASE				,       decoder_input_ids=SCREAMING_SNAKE_CASE					).logits
									lowercase__					:   Dict													=					optax.softmax_cross_entropy(SCREAMING_SNAKE_CASE				,       onehot(SCREAMING_SNAKE_CASE				,       logits.shape[-1]					)					).mean()

									lowercase__					:   Union[str, Any]													=					-(labels.shape[-1] * loss.item())

									lowercase__					:   Union[str, Any]													=					-84.9_127
									self.assertTrue(abs(mtf_score - EXPECTED_SCORE					) < 1E-4					)




