

import os
import re
import shutil
from argparse import ArgumentParser, Namespace

from datasets.commands import BaseDatasetsCLICommand
from datasets.utils.logging import get_logger


A											=     '<<<<<<< This should probably be modified because it mentions: '

A											=     '=======\n>>>>>>>\n'

A											=     [
    'TextEncoderConfig',
    'ByteTextEncoder',
    'SubwordTextEncoder',
    'encoder_config',
    'maybe_build_from_corpus',
    'manual_dir',
]

A											=     [
    # (pattern, replacement)
    # Order is important here for some replacements
    (R'tfds\.core', R'datasets'),
    (R'tf\.io\.gfile\.GFile', R'open'),
    (R'tf\.([\w\d]+)', R'datasets.Value(\'\1\')'),
    (R'tfds\.features\.Text\(\)', R'datasets.Value(\'string\')'),
    (R'tfds\.features\.Text\(', R'datasets.Value(\'string\'),'),
    (R'features\s*=\s*tfds.features.FeaturesDict\(', R'features=datasets.Features('),
    (R'tfds\.features\.FeaturesDict\(', R'dict('),
    (R'The TensorFlow Datasets Authors', R'The TensorFlow Datasets Authors and the HuggingFace Datasets Authors'),
    (R'tfds\.', R'datasets.'),
    (R'dl_manager\.manual_dir', R'self.config.data_dir'),
    (R'self\.builder_config', R'self.config'),
]





def a(lowercase__	):





	'''simple docstring'''
	return ConvertCommand(args.tfds_path			, args.datasets_directory	)






class SCREAMING_SNAKE_CASE  (		__snake_case							):






			"""simple docstring"""




			@staticmethod
			def 			__lowerCAmelCase (    __UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= parser.add_parser(
				    'convert' ,				help='Convert a TensorFlow Datasets dataset to a HuggingFace Datasets dataset.' ,				)
				train_parser.add_argument(
				    '--tfds_path' ,				type=__UpperCamelCase ,				required=__UpperCamelCase ,				help='Path to a TensorFlow Datasets folder to convert or a single tfds file to convert.' ,				)
				train_parser.add_argument(
				    '--datasets_directory' ,				type=__UpperCamelCase ,				required=__UpperCamelCase ,				help='Path to the HuggingFace Datasets folder.'							)
				train_parser.set_defaults(func=__UpperCamelCase							)


			def __init__(    self ,				__UpperCamelCase ,				__UpperCamelCase ,				*__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= get_logger('datasets-cli/converting'							)

				snake_case_      		= tfds_path
				snake_case_      		= datasets_directory







			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				if os.path.isdir(self._tfds_path							):
					snake_case_      		= os.path.abspath(self._tfds_path							)
				elif os.path.isfile(self._tfds_path							):
					snake_case_      		= os.path.dirname(self._tfds_path							)
				else:
					raise ValueError('--tfds_path is neither a directory nor a file. Please check path.'							)

				snake_case_      		= os.path.abspath(self._datasets_directory							)

				self._logger.info(f"""Converting datasets from {abs_tfds_path} to {abs_datasets_path}"""							)

				snake_case_      		= []
				snake_case_      		= []
				snake_case_      		= {}

				if os.path.isdir(self._tfds_path							):
					snake_case_      		= os.listdir(__UpperCamelCase							)
				else:
					snake_case_      		= [os.path.basename(self._tfds_path							)]

				for f_name in file_names:
					self._logger.info(f"""Looking at file {f_name}"""							)
					snake_case_      		= os.path.join(__UpperCamelCase ,				__UpperCamelCase							)
					snake_case_      		= os.path.join(__UpperCamelCase ,				__UpperCamelCase							)

					if not os.path.isfile(__UpperCamelCase							) or "__init__" in f_name or "_test" in f_name or ".py" not in f_name:
						self._logger.info('Skipping file'							)
						continue

					with open(__UpperCamelCase ,				encoding='utf-8'							) as f:
						snake_case_      		= f.readlines()

					snake_case_      		= []
					snake_case_      		= False
					snake_case_      		= False
					snake_case_      		= []
					for line in lines:
						snake_case_      		= line

						# Convert imports
						if "import tensorflow.compat.v2 as tf" in out_line:
							continue
						elif "@tfds.core" in out_line:
							continue
						elif "builder=self" in out_line:
							continue
						elif "import tensorflow_datasets.public_api as tfds" in out_line:
							snake_case_      		= 'import datasets\n'
						elif "import tensorflow" in out_line:
							# order is important here
							snake_case_      		= ''
							continue
						elif "from absl import logging" in out_line:
							snake_case_      		= 'from datasets import logging\n'
						elif "getLogger" in out_line:
							snake_case_      		= out_line.replace('getLogger' ,				'get_logger'							)
						elif any(expression in out_line for expression in TO_HIGHLIGHT							):
							snake_case_      		= True
							snake_case_      		= list(filter(lambda __UpperCamelCase							: e in out_line ,				__UpperCamelCase							)							)
							out_lines.append(HIGHLIGHT_MESSAGE_PRE + str(__UpperCamelCase							) + '\n'							)
							out_lines.append(__UpperCamelCase							)
							out_lines.append(__UpperCamelCase							)
							continue
						else:
							for pattern, replacement in TO_CONVERT:
								snake_case_      		= re.sub(__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							)

                # Take care of saving utilities (to later move them together with main script)
						if "tensorflow_datasets" in out_line:
							snake_case_      		= re.match(r'from\stensorflow_datasets.*import\s([^\.\r\n]+)' ,				__UpperCamelCase							)
							tfds_imports.extend(imp.strip() for imp in match.group(1							).split(','							)							)
							snake_case_      		= 'from . import ' + match.group(1							)

						# Check we have not forget anything
						if "tf." in out_line or "tfds." in out_line or "tensorflow_datasets" in out_line:
							raise ValueError(f"""Error converting {out_line.strip()}"""							)

						if "GeneratorBasedBuilder" in out_line or "BeamBasedBuilder" in out_line:
							snake_case_      		= True
						out_lines.append(__UpperCamelCase							)

					if is_builder or "wmt" in f_name:
						# We create a new directory for each dataset
						snake_case_      		= f_name.replace('.py' ,				''							)
						snake_case_      		= os.path.join(__UpperCamelCase ,				__UpperCamelCase							)
						snake_case_      		= os.path.join(__UpperCamelCase ,				__UpperCamelCase							)
						os.makedirs(__UpperCamelCase ,				exist_ok=__UpperCamelCase							)
						self._logger.info(f"""Adding directory {output_dir}"""							)
						imports_to_builder_map.update({imp: output_dir for imp in tfds_imports}							)
					else:
						# Utilities will be moved at the end
						utils_files.append(__UpperCamelCase							)

					if needs_manual_update:
						with_manual_update.append(__UpperCamelCase							)

					with open(__UpperCamelCase ,				'w' ,				encoding='utf-8'							) as f:
						f.writelines(__UpperCamelCase							)
					self._logger.info(f"""Converted in {output_file}"""							)

				for utils_file in utils_files:
					try:
						snake_case_      		= os.path.basename(__UpperCamelCase							)
						snake_case_      		= imports_to_builder_map[f_name.replace('.py' ,				''							)]
						self._logger.info(f"""Moving {dest_folder} to {utils_file}"""							)
						shutil.copy(__UpperCamelCase ,				__UpperCamelCase							)
					except KeyError:
						self._logger.error(f"""Cannot find destination folder for {utils_file}. Please copy manually."""							)

				if with_manual_update:
					for file_path in with_manual_update:
						self._logger.warning(
						    f"""You need to manually update file {file_path} to remove configurations using 'TextEncoderConfig'."""							)



import unittest

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

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
		import torch

		from transformers import (
		    MraForMaskedLM,
		    MraForMultipleChoice,
		    MraForQuestionAnswering,
		    MraForSequenceClassification,
		    MraForTokenClassification,
		    MraModel,
		)
		from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST


class SCREAMING_SNAKE_CASE  :






			"""simple docstring"""




			def __init__(    self ,				__UpperCamelCase ,				__UpperCamelCase=2 ,				__UpperCamelCase=8 ,				__UpperCamelCase=True ,				__UpperCamelCase=True ,				__UpperCamelCase=True ,				__UpperCamelCase=True ,				__UpperCamelCase=99 ,				__UpperCamelCase=16 ,				__UpperCamelCase=5 ,				__UpperCamelCase=2 ,				__UpperCamelCase=36 ,				__UpperCamelCase="gelu" ,				__UpperCamelCase=0.0 ,				__UpperCamelCase=0.0 ,				__UpperCamelCase=5_12 ,				__UpperCamelCase=16 ,				__UpperCamelCase=2 ,				__UpperCamelCase=0.02 ,				__UpperCamelCase=3 ,				__UpperCamelCase=4 ,				__UpperCamelCase=None ,				):

				"""simple docstring"""


				snake_case_      		= parent
				snake_case_      		= batch_size
				snake_case_      		= seq_length
				snake_case_      		= is_training
				snake_case_      		= use_input_mask
				snake_case_      		= use_token_type_ids
				snake_case_      		= use_labels
				snake_case_      		= vocab_size
				snake_case_      		= hidden_size
				snake_case_      		= num_hidden_layers
				snake_case_      		= num_attention_heads
				snake_case_      		= intermediate_size
				snake_case_      		= hidden_act
				snake_case_      		= hidden_dropout_prob
				snake_case_      		= attention_probs_dropout_prob
				snake_case_      		= max_position_embeddings
				snake_case_      		= type_vocab_size
				snake_case_      		= type_sequence_label_size
				snake_case_      		= initializer_range
				snake_case_      		= num_labels
				snake_case_      		= num_choices
				snake_case_      		= scope


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= ids_tensor([self.batch_size, self.seq_length] ,				self.vocab_size							)

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

				snake_case_      		= None
				if self.use_token_type_ids:
					snake_case_      		= ids_tensor([self.batch_size, self.seq_length] ,				self.type_vocab_size							)

				snake_case_      		= None
				snake_case_      		= None
				snake_case_      		= None
				if self.use_labels:
					snake_case_      		= ids_tensor([self.batch_size] ,				self.type_sequence_label_size							)
					snake_case_      		= ids_tensor([self.batch_size, self.seq_length] ,				self.num_labels							)
					snake_case_      		= ids_tensor([self.batch_size] ,				self.num_choices							)

				snake_case_      		= self.get_config()

				return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				return MraConfig(
				    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 ,				type_vocab_size=self.type_vocab_size ,				is_decoder=__UpperCamelCase ,				initializer_range=self.initializer_range ,				)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.get_config()
				snake_case_      		= 3_00
				return config


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				(
				    (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       
				)      		= self.prepare_config_and_inputs()

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

				return (
				    config,
				    input_ids,
				    token_type_ids,
				    input_mask,
				    sequence_labels,
				    token_labels,
				    choice_labels,
				    encoder_hidden_states,
				    encoder_attention_mask,
				)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= MraModel(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase							)
				snake_case_      		= model(__UpperCamelCase ,				token_type_ids=__UpperCamelCase							)
				snake_case_      		= model(__UpperCamelCase							)
				self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.hidden_size)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				):

				"""simple docstring"""


				snake_case_      		= True
				snake_case_      		= MraModel(__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(
				    __UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				encoder_hidden_states=__UpperCamelCase ,				encoder_attention_mask=__UpperCamelCase ,				)
				snake_case_      		= model(
				    __UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				encoder_hidden_states=__UpperCamelCase ,				)
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase							)
				self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.hidden_size)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= MraForMaskedLM(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				labels=__UpperCamelCase							)
				self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.vocab_size)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= MraForQuestionAnswering(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(
				    __UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				start_positions=__UpperCamelCase ,				end_positions=__UpperCamelCase ,				)
				self.parent.assertEqual(result.start_logits.shape ,				(self.batch_size, self.seq_length)							)
				self.parent.assertEqual(result.end_logits.shape ,				(self.batch_size, self.seq_length)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= self.num_labels
				snake_case_      		= MraForSequenceClassification(__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				labels=__UpperCamelCase							)
				self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_labels)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= self.num_labels
				snake_case_      		= MraForTokenClassification(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= model(__UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				labels=__UpperCamelCase							)
				self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.num_labels)							)


			def 			__lowerCAmelCase (    self ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase ,				__UpperCamelCase							):

				"""simple docstring"""


				snake_case_      		= self.num_choices
				snake_case_      		= MraForMultipleChoice(config=__UpperCamelCase							)
				model.to(__UpperCamelCase							)
				model.eval()
				snake_case_      		= input_ids.unsqueeze(1							).expand(-1 ,				self.num_choices ,				-1							).contiguous()
				snake_case_      		= token_type_ids.unsqueeze(1							).expand(-1 ,				self.num_choices ,				-1							).contiguous()
				snake_case_      		= input_mask.unsqueeze(1							).expand(-1 ,				self.num_choices ,				-1							).contiguous()
				snake_case_      		= model(
				    __UpperCamelCase ,				attention_mask=__UpperCamelCase ,				token_type_ids=__UpperCamelCase ,				labels=__UpperCamelCase ,				)
				self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_choices)							)







			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.prepare_config_and_inputs()
				(
				    (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       (
				    snake_case_
				)		,       
				)      		= config_and_inputs
				snake_case_      		= {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
				return config, inputs_dict







@require_torch
class SCREAMING_SNAKE_CASE  (		__snake_case							,						unittest.TestCase							):






			"""simple docstring"""




			__A						  =  (
			    (
			        MraModel,
			        MraForMaskedLM,
			        MraForMultipleChoice,
			        MraForQuestionAnswering,
			        MraForSequenceClassification,
			        MraForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			__A						  =  False
			__A						  =  False
			__A						  =  False
			__A						  =  False

			__A						  =  ()


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= MraModelTester(self							)
				snake_case_      		= ConfigTester(self ,				config_class=__UpperCamelCase ,				hidden_size=37							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				self.config_tester.run_common_tests()


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				for type in ["absolute", "relative_key", "relative_key_query"]:
					snake_case_      		= type
					self.model_tester.create_and_check_model(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_multiple_choice(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_question_answering(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_sequence_classification(*__UpperCamelCase							)


			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*__UpperCamelCase							)


			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					snake_case_      		= MraModel.from_pretrained(__UpperCamelCase							)
					self.assertIsNotNone(__UpperCamelCase							)







			@unittest.skip(reason='MRA does not output attentions'							)
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				return







@require_torch
class SCREAMING_SNAKE_CASE  (		unittest.TestCase							):






			"""simple docstring"""




			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= MraModel.from_pretrained('uw-madison/mra-base-512-4'							)
				snake_case_      		= torch.arange(2_56							).unsqueeze(0							)

				with torch.no_grad():
					snake_case_      		= model(__UpperCamelCase							)[0]

				snake_case_      		= torch.Size((1, 2_56, 7_68)							)
				self.assertEqual(output.shape ,				__UpperCamelCase							)

				snake_case_      		= torch.tensor(
				    [[[-0.0140, 0.0830, -0.0381], [0.1546, 0.1402, 0.0220], [0.1162, 0.0851, 0.0165]]]							)

				self.assertTrue(torch.allclose(output[:, :3, :3] ,				__UpperCamelCase ,				atol=1E-4							)							)


			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= MraForMaskedLM.from_pretrained('uw-madison/mra-base-512-4'							)
				snake_case_      		= torch.arange(2_56							).unsqueeze(0							)

				with torch.no_grad():
					snake_case_      		= model(__UpperCamelCase							)[0]

				snake_case_      		= 5_02_65

				snake_case_      		= torch.Size((1, 2_56, vocab_size)							)
				self.assertEqual(output.shape ,				__UpperCamelCase							)

				snake_case_      		= torch.tensor(
				    [[[9.2595, -3.6038, 11.8819], [9.3869, -3.2693, 11.0956], [11.8524, -3.4938, 13.1210]]]							)

				self.assertTrue(torch.allclose(output[:, :3, :3] ,				__UpperCamelCase ,				atol=1E-4							)							)







			@slow
			def 			__lowerCAmelCase (    self							):

				"""simple docstring"""


				snake_case_      		= MraForMaskedLM.from_pretrained('uw-madison/mra-base-4096-8-d3'							)
				snake_case_      		= torch.arange(40_96							).unsqueeze(0							)

				with torch.no_grad():
					snake_case_      		= model(__UpperCamelCase							)[0]

				snake_case_      		= 5_02_65

				snake_case_      		= torch.Size((1, 40_96, vocab_size)							)
				self.assertEqual(output.shape ,				__UpperCamelCase							)

				snake_case_      		= torch.tensor(
				    [[[5.4789, -2.3564, 7.5064], [7.9067, -1.3369, 9.9668], [9.0712, -1.8106, 7.0380]]]							)

				self.assertTrue(torch.allclose(output[:, :3, :3] ,				__UpperCamelCase ,				atol=1E-4							)							)
