



import tempfile
import unittest
from pathlib import Path
from shutil import copyfile

from transformers import BatchEncoding, MarianTokenizer
from transformers.testing_utils import get_tests_dir, require_sentencepiece, slow
from transformers.utils import is_sentencepiece_available, is_tf_available, is_torch_available


if is_sentencepiece_available():
							from transformers.models.marian.tokenization_marian import VOCAB_FILES_NAMES, save_json

from ...test_tokenization_common import TokenizerTesterMixin


__snake_case       :	int            =get_tests_dir('fixtures/test_sentencepiece.model')

__snake_case       :	Optional[Any]            ={'target_lang': 'fi', 'source_lang': 'en'}
__snake_case       :	Union[str, Any]            ='>>zh<<'
__snake_case       :	Dict            ='Helsinki-NLP/'

if is_torch_available():
							__snake_case       :	Union[str, Any]            ='pt'
elif is_tf_available():
							__snake_case       :	Union[str, Any]            ='tf'
else:
							__snake_case       :	Optional[Any]            ='jax'



@require_sentencepiece
class  lowerCamelCase__    (							lowerCamelCase__   ,       unittest.TestCase):




		'''simple docstring'''



		snake_case_			=MarianTokenizer
		snake_case_			=False
		snake_case_			=True



		def 	lowerCAmelCase__   (self				)  ->    Union[str, Any]:

						"""simple docstring"""






						super().setUp()
						lowerCAmelCase__  :       Union[str, Any]				=  ['''</s>''', '''<unk>''', '''▁This''', '''▁is''', '''▁a''', '''▁t''', '''est''', '''\u0120''', '''<pad>''']
						lowerCAmelCase__  :       Union[str, Any]				=  dict(zip(__lowerCamelCase	,range(len(__lowerCamelCase				)				)				)				)
						lowerCAmelCase__  :       Any				=  Path(self.tmpdirname				)
						save_json(__lowerCamelCase	,save_dir / VOCAB_FILES_NAMES['''vocab''']				)
						save_json(__lowerCamelCase	,save_dir / VOCAB_FILES_NAMES['''tokenizer_config_file''']				)
						if not (save_dir / VOCAB_FILES_NAMES["source_spm"]).exists():
										copyfile(__lowerCamelCase	,save_dir / VOCAB_FILES_NAMES['''source_spm''']				)
										copyfile(__lowerCamelCase	,save_dir / VOCAB_FILES_NAMES['''target_spm''']				)

						lowerCAmelCase__  :       List[str]				=  MarianTokenizer.from_pretrained(self.tmpdirname				)
						tokenizer.save_pretrained(self.tmpdirname				)



		def 	lowerCAmelCase__   (self	,**__lowerCamelCase				)  ->    MarianTokenizer:

						"""simple docstring"""






						return MarianTokenizer.from_pretrained(self.tmpdirname	,**__lowerCamelCase				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase				)  ->    str:

						"""simple docstring"""






						return (
						    "This is a test",
						    "This is a test",
						)



		def 	lowerCAmelCase__   (self				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       str				=  '''</s>'''
						lowerCAmelCase__  :       List[str]				=  0

						self.assertEqual(self.get_tokenizer()._convert_token_to_id(__lowerCamelCase				)	,__lowerCamelCase				)
						self.assertEqual(self.get_tokenizer()._convert_id_to_token(__lowerCamelCase				)	,__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       List[Any]				=  list(self.get_tokenizer().get_vocab().keys()				)

						self.assertEqual(vocab_keys[0]	,'''</s>'''				)
						self.assertEqual(vocab_keys[1]	,'''<unk>'''				)
						self.assertEqual(vocab_keys[-1]	,'''<pad>'''				)
						self.assertEqual(len(__lowerCamelCase				)	,9				)



		def 	lowerCAmelCase__   (self				)  ->    Optional[Any]:

						"""simple docstring"""






						self.assertEqual(self.get_tokenizer().vocab_size	,9				)



		def 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       str				=  MarianTokenizer.from_pretrained(f"""{ORG_NAME}opus-mt-en-de"""				)
						lowerCAmelCase__  :       List[str]				=  en_de_tokenizer(['''I am a small frog''']	,return_tensors=__lowerCamelCase				)
						self.assertIsInstance(__lowerCamelCase	,__lowerCamelCase				)
						lowerCAmelCase__  :       int				=  [38, 1_21, 14, 6_97, 3_88_48, 0]
						self.assertListEqual(__lowerCamelCase	,batch.input_ids[0]				)

						lowerCAmelCase__  :       Tuple				=  tempfile.mkdtemp()
						en_de_tokenizer.save_pretrained(__lowerCamelCase				)
						lowerCAmelCase__  :       Optional[int]				=  [x.name for x in Path(__lowerCamelCase				).glob('''*'''				)]
						self.assertIn('''source.spm'''	,__lowerCamelCase				)
						MarianTokenizer.from_pretrained(__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  self.get_tokenizer()

						lowerCAmelCase__  :       Union[str, Any]				=  tok(
						    ['''I am a small frog''' * 10_00, '''I am a small frog''']	,padding=__lowerCamelCase	,truncation=__lowerCamelCase	,return_tensors=__lowerCamelCase				)
						self.assertIsInstance(__lowerCamelCase	,__lowerCamelCase				)
						self.assertEqual(batch.input_ids.shape	,(2, 5_12)				)



		def 	lowerCAmelCase__   (self				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Union[str, Any]				=  self.get_tokenizer()
						lowerCAmelCase__  :       Dict				=  tok(['''I am a tiny frog''', '''I am a small frog''']	,padding=__lowerCamelCase	,return_tensors=__lowerCamelCase				)
						self.assertIsInstance(__lowerCamelCase	,__lowerCamelCase				)
						self.assertEqual(batch_smaller.input_ids.shape	,(2, 10)				)



		@slow
		def 	lowerCAmelCase__   (self				)  ->    List[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  {'''input_ids''': [[4_34_95, 4_62, 20, 4_21_64, 13_69, 52, 4_64, 1_32, 17_03, 4_92, 13, 74_91, 3_89_99, 6, 8, 4_64, 1_32, 17_03, 4_92, 13, 46_69, 3_78_67, 13, 75_25, 27, 15_93, 9_88, 13, 3_39_72, 70_29, 6, 20, 82_51, 3_83, 2, 2_70, 58_66, 37_88, 2, 23_53, 82_51, 1_23_38, 2, 1_39_58, 3_87, 2, 36_29, 69_53, 1_88, 29_00, 2, 1_39_58, 80_11, 1_15_01, 23, 84_60, 40_73, 3_40_09, 20, 4_35, 1_14_39, 27, 8, 84_60, 40_73, 60_04, 20, 99_88, 3_75, 27, 33, 2_66, 19_45, 10_76, 13_50, 3_78_67, 32_88, 5, 5_77, 10_76, 43_74, 8, 50_82, 5, 2_64_53, 2_57, 5_56, 4_03, 2, 2_42, 1_32, 3_83, 3_16, 4_92, 8, 1_07_67, 6, 3_16, 3_04, 42_39, 3, 0], [1_48, 1_57_22, 19, 18_39, 12, 13_50, 13, 2_23_27, 50_82, 54_18, 4_75_67, 3_59_38, 59, 3_18, 1_95_52, 1_08, 21_83, 54, 1_49_76, 48_35, 32, 5_47, 11_14, 8, 3_15, 24_17, 5, 92, 1_90_88, 3, 0, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00], [36, 63_95, 1_25_70, 3_91_47, 1_15_97, 6, 2_66, 4, 4_54_05, 72_96, 3, 0, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00, 5_81_00]], '''attention_mask''': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}  # noqa: E501
						# fmt: on

						self.tokenizer_integration_test_util(
						    expected_encoding=__lowerCamelCase	,model_name='''Helsinki-NLP/opus-mt-en-de'''	,revision='''1a8c2263da11e68e50938f97e10cd57820bd504c'''	,decode_kwargs={'''use_source_tokenizer''': True}	,)



		def 	lowerCAmelCase__   (self				)  ->    List[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       str				=  MarianTokenizer.from_pretrained('''hf-internal-testing/test-marian-two-vocabs'''				)

						lowerCAmelCase__  :       Any				=  '''Tämä on testi'''
						lowerCAmelCase__  :       List[Any]				=  '''This is a test'''

						lowerCAmelCase__  :       str				=  [76, 7, 20_47, 2]
						lowerCAmelCase__  :       Dict				=  [69, 12, 11, 9_40, 2]

						lowerCAmelCase__  :       Union[str, Any]				=  tokenizer(__lowerCamelCase				).input_ids
						self.assertListEqual(__lowerCamelCase	,__lowerCamelCase				)

						lowerCAmelCase__  :       Dict				=  tokenizer(text_target=__lowerCamelCase				).input_ids
						self.assertListEqual(__lowerCamelCase	,__lowerCamelCase				)

						lowerCAmelCase__  :       int				=  tokenizer.decode(__lowerCamelCase	,skip_special_tokens=__lowerCamelCase				)
						self.assertEqual(__lowerCamelCase	,__lowerCamelCase				)






from typing import Dict

import numpy as np
import torch

from . import residue_constants as rc
from .tensor_utils import tensor_tree_map, tree_map






def 			lowerCAmelCase__	(					lowerCamelCase_	:	Dict[str, torch.Tensor]):





				'''simple docstring'''

				lowerCAmelCase__  :       List[Any]				=  []
				lowerCAmelCase__  :       int				=  []
				lowerCAmelCase__  :       Optional[int]				=  []

				for rt in rc.restypes:
								lowerCAmelCase__  :       int				=  rc.restype_name_to_atomaa_names[rc.restype_atoa[rt]]
								restype_atomaa_to_atomaa_list.append([(rc.atom_order[name] if name else 0) for name in atom_names])
								lowerCAmelCase__  :       Any				=  {name: i for i, name in enumerate(lowerCamelCase_)}
								restype_atomaa_to_atomaa_list.append(
								    [(atom_name_to_idxaa[name] if name in atom_name_to_idxaa else 0) for name in rc.atom_types])

								restype_atomaa_mask_list.append([(1.0 if name else 0.0) for name in atom_names])

				# Add dummy mapping for restype 'UNK'
				restype_atomaa_to_atomaa_list.append([0] * 14)
				restype_atomaa_to_atomaa_list.append([0] * 37)
				restype_atomaa_mask_list.append([0.0] * 14)

				lowerCAmelCase__  :       Optional[Any]				=  torch.tensor(
				    lowerCamelCase_ ,dtype=torch.intaa ,device=protein['''aatype'''].device ,)
				lowerCAmelCase__  :       Optional[int]				=  torch.tensor(
				    lowerCamelCase_ ,dtype=torch.intaa ,device=protein['''aatype'''].device ,)
				lowerCAmelCase__  :       List[str]				=  torch.tensor(
				    lowerCamelCase_ ,dtype=torch.floataa ,device=protein['''aatype'''].device ,)
				lowerCAmelCase__  :       Optional[Any]				=  protein['''aatype'''].to(torch.long)

				# create the mapping for (residx, atom14) --> atom37, i.e. an array
				# with shape (num_res, 14) containing the atom37 indices for this protein
				lowerCAmelCase__  :       Optional[int]				=  restype_atomaa_to_atomaa[protein_aatype]
				lowerCAmelCase__  :       List[str]				=  restype_atomaa_mask[protein_aatype]

				lowerCAmelCase__  :       Tuple				=  residx_atomaa_mask
				lowerCAmelCase__  :       Dict				=  residx_atomaa_to_atomaa.long()

				# create the gather indices for mapping back
				lowerCAmelCase__  :       Union[str, Any]				=  restype_atomaa_to_atomaa[protein_aatype]
				lowerCAmelCase__  :       int				=  residx_atomaa_to_atomaa.long()

				# create the corresponding mask
				lowerCAmelCase__  :       List[str]				=  torch.zeros([21, 37] ,dtype=torch.floataa ,device=protein['''aatype'''].device)
				for restype, restype_letter in enumerate(rc.restypes):
								lowerCAmelCase__  :       List[str]				=  rc.restype_atoa[restype_letter]
								lowerCAmelCase__  :       List[str]				=  rc.residue_atoms[restype_name]
								for atom_name in atom_names:
												lowerCAmelCase__  :       List[Any]				=  rc.atom_order[atom_name]
												lowerCAmelCase__  :       Any				=  1

				lowerCAmelCase__  :       int				=  restype_atomaa_mask[protein_aatype]
				lowerCAmelCase__  :       int				=  residx_atomaa_mask

				return protein






def 			lowerCAmelCase__	(					lowerCamelCase_	:	Dict[str, torch.Tensor]):





				'''simple docstring'''

				lowerCAmelCase__  :       Optional[Any]				=  tree_map(lambda lowerCamelCase_: torch.tensor(lowerCamelCase_ ,device=batch['''aatype'''].device) ,lowerCamelCase_ ,np.ndarray)
				lowerCAmelCase__  :       Union[str, Any]				=  tensor_tree_map(lambda lowerCamelCase_: np.array(lowerCamelCase_) ,make_atomaa_masks(lowerCamelCase_))
				return out

