



'''simple docstring'''

import gc
import unittest

from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
from transformers.pipelines import PipelineException
from transformers.testing_utils import (
    is_pipeline_test,
    is_torch_available,
    nested_simplify,
    require_tf,
    require_torch,
    require_torch_gpu,
    slow,
)

from .test_pipelines_common import ANY

@is_pipeline_test
class 							a__( unittest.TestCase     ):
			a_	: Optional[int]									=      MODEL_FOR_MASKED_LM_MAPPING
			a_	: Any									=      TF_MODEL_FOR_MASKED_LM_MAPPING






			def 						_lowercase  (     self						)							->							Union[str, Any]:
					super().tearDown()
					# clean-up as much as possible GPU memory occupied by PyTorch
					gc.collect()
					if is_torch_available():
							import torch

							torch.cuda.empty_cache()






			@require_tf
			def 						_lowercase  (     self						)							->							str:
					snake_case__	     =pipeline(task='fill-mask'					,   model='sshleifer/tiny-distilroberta-base'					,   top_k=2					,   framework='tf'						)
					snake_case__	     =unmasker('My name is <mask>'						)
					self.assertEqual(
					    nested_simplify(_UpperCAmelCase					,   decimals=6						)					,   [
					        {'sequence': 'My name is grouped', 'score': 2.1E-05, 'token': 3_8015, 'token_str': ' grouped'},
					        {'sequence': 'My name is accuser', 'score': 2.1E-05, 'token': 2_5506, 'token_str': ' accuser'},
					    ]					,   )

					snake_case__	     =unmasker('The largest city in France is <mask>'						)
					self.assertEqual(
					    nested_simplify(_UpperCAmelCase					,   decimals=6						)					,   [
					        {
					            'sequence': 'The largest city in France is grouped',
					            'score': 2.1E-05,
					            'token': 3_8015,
					            'token_str': ' grouped',
					        },
					        {
					            'sequence': 'The largest city in France is accuser',
					            'score': 2.1E-05,
					            'token': 2_5506,
					            'token_str': ' accuser',
					        },
					    ]					,   )

					snake_case__	     =unmasker('My name is <mask>'					,   targets=[' Patrick', ' Clara', ' Teven']					,   top_k=3						)
					self.assertEqual(
					    nested_simplify(_UpperCAmelCase					,   decimals=6						)					,   [
					        {'sequence': 'My name is Clara', 'score': 2E-05, 'token': 1_3606, 'token_str': ' Clara'},
					        {'sequence': 'My name is Patrick', 'score': 2E-05, 'token': 3499, 'token_str': ' Patrick'},
					        {'sequence': 'My name is Te', 'score': 1.9E-05, 'token': 2941, 'token_str': ' Te'},
					    ]					,   )






			@require_torch
			def 						_lowercase  (     self						)							->							Any:
					snake_case__	     =pipeline(task='fill-mask'					,   model='sshleifer/tiny-distilroberta-base'					,   top_k=2					,   framework='pt'						)

					snake_case__	     =unmasker('My name is <mask>'						)
					self.assertEqual(
					    nested_simplify(_UpperCAmelCase					,   decimals=6						)					,   [
					        {'sequence': 'My name is Maul', 'score': 2.2E-05, 'token': 3_5676, 'token_str': ' Maul'},
					        {'sequence': 'My name isELS', 'score': 2.2E-05, 'token': 1_6416, 'token_str': 'ELS'},
					    ]					,   )

					snake_case__	     =unmasker('The largest city in France is <mask>'						)
					self.assertEqual(
					    nested_simplify(_UpperCAmelCase					,   decimals=6						)					,   [
					        {
					            'sequence': 'The largest city in France is Maul',
					            'score': 2.2E-05,
					            'token': 3_5676,
					            'token_str': ' Maul',
					        },
					        {'sequence': 'The largest city in France isELS', 'score': 2.2E-05, 'token': 1_6416, 'token_str': 'ELS'},
					    ]					,   )

					snake_case__	     =unmasker('My name is <mask>'					,   targets=[' Patrick', ' Clara', ' Teven']					,   top_k=3						)
					self.assertEqual(
					    nested_simplify(_UpperCAmelCase					,   decimals=6						)					,   [
					        {'sequence': 'My name is Patrick', 'score': 2.1E-05, 'token': 3499, 'token_str': ' Patrick'},
					        {'sequence': 'My name is Te', 'score': 2E-05, 'token': 2941, 'token_str': ' Te'},
					        {'sequence': 'My name is Clara', 'score': 2E-05, 'token': 1_3606, 'token_str': ' Clara'},
					    ]					,   )

					snake_case__	     =unmasker('My name is <mask> <mask>'					,   top_k=2						)

					self.assertEqual(
					    nested_simplify(_UpperCAmelCase					,   decimals=6						)					,   [
					        [
					            {
					                'score': 2.2E-05,
					                'token': 3_5676,
					                'token_str': ' Maul',
					                'sequence': '<s>My name is Maul<mask></s>',
					            },
					            {'score': 2.2E-05, 'token': 1_6416, 'token_str': 'ELS', 'sequence': '<s>My name isELS<mask></s>'},
					        ],
					        [
					            {
					                'score': 2.2E-05,
					                'token': 3_5676,
					                'token_str': ' Maul',
					                'sequence': '<s>My name is<mask> Maul</s>',
					            },
					            {'score': 2.2E-05, 'token': 1_6416, 'token_str': 'ELS', 'sequence': '<s>My name is<mask>ELS</s>'},
					        ],
					    ]					,   )






			@require_torch_gpu
			def 						_lowercase  (     self						)							->							Optional[Any]:
					snake_case__	     =pipeline('fill-mask'					,   model='hf-internal-testing/tiny-random-distilbert'					,   device=0					,   framework='pt'						)

					# convert model to fp16
					pipe.model.half()

					snake_case__	     =pipe('Paris is the [MASK] of France.'						)
					# We actually don't care about the result, we just want to make sure
					# it works, meaning the float16 tensor got casted back to float32
					# for postprocessing.
					self.assertIsInstance(_UpperCAmelCase					,   _UpperCAmelCase						)






			@slow
			@require_torch
			def 						_lowercase  (     self						)							->							str:
					snake_case__	     =pipeline(task='fill-mask'					,   model='distilroberta-base'					,   top_k=2					,   framework='pt'						)
					self.run_large_test(_UpperCAmelCase						)






			@slow
			@require_tf
			def 						_lowercase  (     self						)							->							Any:
					snake_case__	     =pipeline(task='fill-mask'					,   model='distilroberta-base'					,   top_k=2					,   framework='tf'						)
					self.run_large_test(_UpperCAmelCase						)






			def 						_lowercase  (     self					,   _UpperCAmelCase						)							->							Any:
					snake_case__	     =unmasker('My name is <mask>'						)
					self.assertEqual(
					    nested_simplify(_UpperCAmelCase						)					,   [
					        {'sequence': 'My name is John', 'score': 0.008, 'token': 610, 'token_str': ' John'},
					        {'sequence': 'My name is Chris', 'score': 0.007, 'token': 1573, 'token_str': ' Chris'},
					    ]					,   )
					snake_case__	     =unmasker('The largest city in France is <mask>'						)
					self.assertEqual(
					    nested_simplify(_UpperCAmelCase						)					,   [
					        {
					            'sequence': 'The largest city in France is Paris',
					            'score': 0.251,
					            'token': 2201,
					            'token_str': ' Paris',
					        },
					        {
					            'sequence': 'The largest city in France is Lyon',
					            'score': 0.214,
					            'token': 1_2790,
					            'token_str': ' Lyon',
					        },
					    ]					,   )

					snake_case__	     =unmasker('My name is <mask>'					,   targets=[' Patrick', ' Clara', ' Teven']					,   top_k=3						)
					self.assertEqual(
					    nested_simplify(_UpperCAmelCase						)					,   [
					        {'sequence': 'My name is Patrick', 'score': 0.005, 'token': 3499, 'token_str': ' Patrick'},
					        {'sequence': 'My name is Clara', 'score': 0.000, 'token': 1_3606, 'token_str': ' Clara'},
					        {'sequence': 'My name is Te', 'score': 0.000, 'token': 2941, 'token_str': ' Te'},
					    ]					,   )






			@require_torch
			def 						_lowercase  (     self						)							->							int:
					snake_case__	     =pipeline(task='fill-mask'					,   model='sshleifer/tiny-distilroberta-base'					,   framework='pt'						)
					snake_case__	     =None
					snake_case__	     =None
					self.run_pipeline_test(_UpperCAmelCase					,   []						)






			@require_tf
			def 						_lowercase  (     self						)							->							Optional[int]:
					snake_case__	     =pipeline(task='fill-mask'					,   model='sshleifer/tiny-distilroberta-base'					,   framework='tf'						)
					snake_case__	     =None
					snake_case__	     =None
					self.run_pipeline_test(_UpperCAmelCase					,   []						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Dict:
					if tokenizer is None or tokenizer.mask_token_id is None:
							self.skipTest('The provided tokenizer has no mask token, (probably reformer or wav2vec2)'						)

					snake_case__	     =FillMaskPipeline(model=_UpperCAmelCase					,   tokenizer=_UpperCAmelCase						)
					snake_case__	     =[
					    f"""This is another {tokenizer.mask_token} test""",
					]
					return fill_masker, examples






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							List[Any]:
					snake_case__	     =fill_masker.tokenizer
					snake_case__	     =fill_masker.model

					snake_case__	     =fill_masker(
					    f"""This is a {tokenizer.mask_token}"""					,   )
					self.assertEqual(
					    _UpperCAmelCase					,   [
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					    ]					,   )

					snake_case__	     =fill_masker([f"""This is a {tokenizer.mask_token}"""]						)
					self.assertEqual(
					    _UpperCAmelCase					,   [
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					    ]					,   )

					snake_case__	     =fill_masker([f"""This is a {tokenizer.mask_token}""", f"""Another {tokenizer.mask_token} great test."""]						)
					self.assertEqual(
					    _UpperCAmelCase					,   [
					        [
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        ],
					        [
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        ],
					    ]					,   )

					with self.assertRaises(_UpperCAmelCase						):
							fill_masker([None]						)
					# No mask_token is not supported
					with self.assertRaises(_UpperCAmelCase						):
							fill_masker('This is'						)

					self.run_test_top_k(_UpperCAmelCase					,   _UpperCAmelCase						)
					self.run_test_targets(_UpperCAmelCase					,   _UpperCAmelCase						)
					self.run_test_top_k_targets(_UpperCAmelCase					,   _UpperCAmelCase						)
					self.fill_mask_with_duplicate_targets_and_top_k(_UpperCAmelCase					,   _UpperCAmelCase						)
					self.fill_mask_with_multiple_masks(_UpperCAmelCase					,   _UpperCAmelCase						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							str:
					snake_case__	     =tokenizer.get_vocab()
					snake_case__	     =sorted(vocab.keys()						)[:2]
					# Pipeline argument
					snake_case__	     =FillMaskPipeline(model=_UpperCAmelCase					,   tokenizer=_UpperCAmelCase					,   targets=_UpperCAmelCase						)
					snake_case__	     =fill_masker(f"""This is a {tokenizer.mask_token}"""						)
					self.assertEqual(
					    _UpperCAmelCase					,   [
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					    ]					,   )
					snake_case__	     ={vocab[el] for el in targets}
					self.assertEqual({el['token'] for el in outputs}					,   _UpperCAmelCase						)
					snake_case__	     =[tokenizer.decode([x]						) for x in target_ids]
					self.assertEqual({el['token_str'] for el in outputs}					,   set(_UpperCAmelCase						)						)

					# Call argument
					snake_case__	     =FillMaskPipeline(model=_UpperCAmelCase					,   tokenizer=_UpperCAmelCase						)
					snake_case__	     =fill_masker(f"""This is a {tokenizer.mask_token}"""					,   targets=_UpperCAmelCase						)
					self.assertEqual(
					    _UpperCAmelCase					,   [
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					    ]					,   )
					snake_case__	     ={vocab[el] for el in targets}
					self.assertEqual({el['token'] for el in outputs}					,   _UpperCAmelCase						)
					snake_case__	     =[tokenizer.decode([x]						) for x in target_ids]
					self.assertEqual({el['token_str'] for el in outputs}					,   set(_UpperCAmelCase						)						)

					# Score equivalence
					snake_case__	     =fill_masker(f"""This is a {tokenizer.mask_token}"""					,   targets=_UpperCAmelCase						)
					snake_case__	     =[top_mask['token_str'] for top_mask in outputs]
					snake_case__	     =[top_mask['score'] for top_mask in outputs]

					# For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
					if set(_UpperCAmelCase						) == set(_UpperCAmelCase						):
							snake_case__	     =fill_masker(f"""This is a {tokenizer.mask_token}"""					,   targets=_UpperCAmelCase						)
							snake_case__	     =[top_mask['score'] for top_mask in unmasked_targets]
							self.assertEqual(nested_simplify(_UpperCAmelCase						)					,   nested_simplify(_UpperCAmelCase						)						)

					# Raises with invalid
					with self.assertRaises(_UpperCAmelCase						):
							snake_case__	     =fill_masker(f"""This is a {tokenizer.mask_token}"""					,   targets=[]						)
					# For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
					if "" not in tokenizer.get_vocab():
							with self.assertRaises(_UpperCAmelCase						):
									snake_case__	     =fill_masker(f"""This is a {tokenizer.mask_token}"""					,   targets=['']						)
							with self.assertRaises(_UpperCAmelCase						):
									snake_case__	     =fill_masker(f"""This is a {tokenizer.mask_token}"""					,   targets=''						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Any:
					snake_case__	     =FillMaskPipeline(model=_UpperCAmelCase					,   tokenizer=_UpperCAmelCase					,   top_k=2						)
					snake_case__	     =fill_masker(f"""This is a {tokenizer.mask_token}"""						)
					self.assertEqual(
					    _UpperCAmelCase					,   [
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					    ]					,   )

					snake_case__	     =FillMaskPipeline(model=_UpperCAmelCase					,   tokenizer=_UpperCAmelCase						)
					snake_case__	     =fill_masker(f"""This is a {tokenizer.mask_token}"""					,   top_k=2						)
					self.assertEqual(
					    _UpperCAmelCase					,   [
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					    ]					,   )
					self.assertEqual(nested_simplify(_UpperCAmelCase						)					,   nested_simplify(_UpperCAmelCase						)						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							Union[str, Any]:
					snake_case__	     =tokenizer.get_vocab()
					snake_case__	     =FillMaskPipeline(model=_UpperCAmelCase					,   tokenizer=_UpperCAmelCase						)

					# top_k=2, ntargets=3
					snake_case__	     =sorted(vocab.keys()						)[:3]
					snake_case__	     =fill_masker(f"""This is a {tokenizer.mask_token}"""					,   top_k=2					,   targets=_UpperCAmelCase						)

					# If we use the most probably targets, and filter differently, we should still
					# have the same results
					snake_case__	     =[el['token_str'] for el in sorted(_UpperCAmelCase					,   key=lambda _UpperCAmelCase						: x["score"]					,   reverse=_UpperCAmelCase						)]
					# For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
					if set(_UpperCAmelCase						).issubset(_UpperCAmelCase						):
							snake_case__	     =fill_masker(f"""This is a {tokenizer.mask_token}"""					,   top_k=3					,   targets=_UpperCAmelCase						)
							# They should yield exactly the same result
							self.assertEqual(nested_simplify(_UpperCAmelCase						)					,   nested_simplify(_UpperCAmelCase						)						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							List[Any]:
					snake_case__	     =FillMaskPipeline(model=_UpperCAmelCase					,   tokenizer=_UpperCAmelCase						)
					snake_case__	     =tokenizer.get_vocab()
					# String duplicates + id duplicates
					snake_case__	     =sorted(vocab.keys()						)[:3]
					snake_case__	     =[targets[0], targets[1], targets[0], targets[2], targets[1]]
					snake_case__	     =fill_masker(f"""My name is {tokenizer.mask_token}"""					,   targets=_UpperCAmelCase					,   top_k=10						)

					# The target list contains duplicates, so we can't output more
					# than them
					self.assertEqual(len(_UpperCAmelCase						)					,   3						)






			def 						_lowercase  (     self					,   _UpperCAmelCase					,   _UpperCAmelCase						)							->							List[Any]:
					snake_case__	     =FillMaskPipeline(model=_UpperCAmelCase					,   tokenizer=_UpperCAmelCase						)

					snake_case__	     =fill_masker(
					    f"""This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}"""					,   top_k=2						)
					self.assertEqual(
					    _UpperCAmelCase					,   [
					        [
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        ],
					        [
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        ],
					        [
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					            {'sequence': ANY(_UpperCAmelCase						), 'score': ANY(_UpperCAmelCase						), 'token': ANY(_UpperCAmelCase						), 'token_str': ANY(_UpperCAmelCase						)},
					        ],
					    ]					,   )









'''simple docstring'''

import argparse
import json
import os
from collections import OrderedDict

import numpy as np
import tensorflow as tf
import torch






def 				a  (     UpperCamelCase_    :    Any			)       ->	List[str]:
		snake_case__	     =os.path.join(args.tf_model_dir      ,					'parameters.json'			)
		snake_case__	     =json.loads(open(UpperCamelCase_			).read()			)
		if not params:
				raise ValueError(
				    f"""It seems that the json file at {parameter_file} is empty. Make sure you have a correct json file."""			)
		if not args.output.endswith('.pt'			):
				snake_case__	     =args.output + '.pt'
		snake_case__	     =OrderedDict()
		with tf.device('/CPU:0'			):
				snake_case__	     =tf.train.load_checkpoint(args.tf_model_dir			)
				snake_case__	     =reader.get_variable_to_shape_map()
				for key_name in shapes.keys():
						snake_case__	     =reader.get_tensor(UpperCamelCase_			).astype(np.floataa			)
						if key_name.endswith('/adam_m'			) or key_name.endswith('/adam_v'			):
								continue
						if key_name.startswith('pasts/'			):
								if key_name.startswith('pasts/mlp'			):
										snake_case__	     =int(key_name[9]			)
								elif key_name.startswith('pasts/out'			):
										snake_case__	     =8
								snake_case__	     ='model.sqout.%d.weight' % (player * 2)  # enter to nn.Sequencial with Tanh, so 2 at a time
								snake_case__	     =vnp.transpose([1, 0]			).copy()  # Mesh-Tensorflow is a diagonal matrix
								snake_case__	     =torch.tensor(UpperCamelCase_			)
						elif key_name.startswith('model/moe'			):
								snake_case__	     =int(key_name[9:].split('/'			)[0]			)
								if key_name.endswith('/switch_gating/kernel'			):
										snake_case__	     ='model.blocks.%d.feed_forward.mlp.router.classifier.weight' % player
										snake_case__	     =vnp.transpose([1, 0]			).copy()  # Mesh-Tensorflow is a diagonal matrix
										snake_case__	     =torch.tensor(UpperCamelCase_			)
								elif key_name.endswith('/softmlp/kernel'			):
										snake_case__	     ='model.blocks.%d.feed_forward.soft_bypass_mlp.weight' % player
										snake_case__	     =vnp.transpose([1, 0]			).copy()  # Mesh-Tensorflow is a diagonal matrix
										snake_case__	     =torch.tensor(UpperCamelCase_			)
								elif key_name.endswith('/wo/kernel'			) or key_name.endswith('/wi/kernel'			):
										snake_case__	     =key_name[-9:-7]
										for i in range(16			):
												snake_case__	     ='model.blocks.%d.feed_forward.mlp.experts.expert_%d.%s.weight' % (player, i, nlayer)
												snake_case__	     =(
												    vnp[i].transpose([1, 0]			).copy()
												)  # In Mesh-Tensorflow, it is one array, so it is divided
												snake_case__	     =torch.tensor(UpperCamelCase_			)
						elif key_name.startswith('model/mlp'			):
								snake_case__	     =int(key_name[9:].split('/'			)[0]			)
								if key_name.endswith('/p1/kernel'			):
										snake_case__	     ='model.blocks.%d.feed_forward.mlp.wi.weight' % player
										snake_case__	     =vnp.transpose([1, 0]			).copy()  # Mesh-Tensorflow is a diagonal matrix
										snake_case__	     =torch.tensor(UpperCamelCase_			)
								elif key_name.endswith('/p1/bias'			):
										snake_case__	     ='model.blocks.%d.feed_forward.mlp.wi.bias' % player
										snake_case__	     =vnp.copy()  # same because it is one dimensional
										snake_case__	     =torch.tensor(UpperCamelCase_			)
								elif key_name.endswith('/p2/kernel'			):
										snake_case__	     ='model.blocks.%d.feed_forward.mlp.wo.weight' % player
										snake_case__	     =vnp.transpose([1, 0]			).copy()  # Mesh-Tensorflow is a diagonal matrix
										snake_case__	     =torch.tensor(UpperCamelCase_			)
								elif key_name.endswith('/p2/bias'			):
										snake_case__	     ='model.blocks.%d.feed_forward.mlp.wo.bias' % player
										snake_case__	     =vnp.copy()  # same because it is one dimensional
										snake_case__	     =torch.tensor(UpperCamelCase_			)
						elif key_name.startswith('model/ln'			):
								snake_case__	     =int(key_name[8:].split('/'			)[0]			)
								if key_name.endswith('/b'			):
										snake_case__	     ='model.blocks.%d.feed_forward.norm.bias' % player
										snake_case__	     =vnp.copy()  # same because it is one dimensional
										snake_case__	     =torch.tensor(UpperCamelCase_			)
								elif key_name.endswith('/g'			):
										snake_case__	     ='model.blocks.%d.feed_forward.norm.weight' % player
										snake_case__	     =vnp.copy()  # same because it is one dimensional
										snake_case__	     =torch.tensor(UpperCamelCase_			)
						elif key_name.startswith('model/att'			):
								snake_case__	     =int(key_name[9:].split('/'			)[0]			)
								if key_name.endswith('/qkv/kernel'			):
										snake_case__	     =vnp.copy()  # Compute same dimension as Mesh-tensorflow using einsum
										snake_case__	     =state[:, 0, :, :]
										snake_case__	     =state[:, 1, :, :]
										snake_case__	     =state[:, 2, :, :]
										snake_case__	     =(
										    state_q.reshape([state_q.shape[0], state_q.shape[1] * state_q.shape[2]]			)
										    .transpose([1, 0]			)
										    .copy()
										)  # Mesh-Tensorflow is a diagonal matrix
										snake_case__	     =(
										    state_k.reshape([state_k.shape[0], state_k.shape[1] * state_k.shape[2]]			)
										    .transpose([1, 0]			)
										    .copy()
										)  # Mesh-Tensorflow is a diagonal matrix
										snake_case__	     =(
										    state_v.reshape([state_v.shape[0], state_v.shape[1] * state_v.shape[2]]			)
										    .transpose([1, 0]			)
										    .copy()
										)  # Mesh-Tensorflow is a diagonal matrix
										snake_case__	     ='model.blocks.%d.self_attn.self_attn.q_proj.weight' % player
										snake_case__	     =torch.tensor(UpperCamelCase_			)
										snake_case__	     ='model.blocks.%d.self_attn.self_attn.k_proj.weight' % player
										snake_case__	     =torch.tensor(UpperCamelCase_			)
										snake_case__	     ='model.blocks.%d.self_attn.self_attn.v_proj.weight' % player
										snake_case__	     =torch.tensor(UpperCamelCase_			)
								elif key_name.endswith('/o/kernel'			):
										snake_case__	     ='model.blocks.%d.self_attn.self_attn.out_proj.weight' % player
										snake_case__	     =(
										    vnp.reshape([vnp.shape[0] * vnp.shape[1], vnp.shape[2]]			).transpose([1, 0]			).copy()
										)  # Mesh-Tensorflow is a diagonal matrix
										snake_case__	     =torch.tensor(UpperCamelCase_			)
						elif key_name.startswith('model/an'			):
								snake_case__	     =int(key_name[8:].split('/'			)[0]			)
								if key_name.endswith('/b'			):
										snake_case__	     ='model.blocks.%d.self_attn.norm.bias' % player
										snake_case__	     =vnp.copy()  # same because it is one dimensional
										snake_case__	     =torch.tensor(UpperCamelCase_			)
								elif key_name.endswith('/g'			):
										snake_case__	     ='model.blocks.%d.self_attn.norm.weight' % player
										snake_case__	     =vnp.copy()  # same because it is one dimensional
										snake_case__	     =torch.tensor(UpperCamelCase_			)
						elif (
						    key_name.startswith('model/wte'			)
						    or key_name.startswith('model/wpe'			)
						    or key_name.startswith('model/ete'			)
						):
								snake_case__	     ={'wte': 'embed_tokens', 'wpe': 'position_embeddings', 'ete': 'extra_position_embeddings'}[
								    key_name[-3:]
								]
								snake_case__	     ='model.%s.weight' % nlayer
								snake_case__	     =vnp.copy()  # same in embedded
								snake_case__	     =torch.tensor(UpperCamelCase_			)
								if key_name.startswith('model/wte'			):
										snake_case__	     ='lm_head.weight'
										snake_case__	     =vnp.copy()  # same in embedded
										snake_case__	     =torch.tensor(UpperCamelCase_			)
						elif key_name.startswith('model/wob'			):
								snake_case__	     ='final_logits_bias'
								snake_case__	     =vnp.copy()  # same in embedded
								snake_case__	     =state.reshape((1, -1)			)
								snake_case__	     =torch.tensor(UpperCamelCase_			)
						elif key_name == "model/dense/kernel":
								snake_case__	     ='model.last_project.weight'
								snake_case__	     =vnp.transpose([1, 0]			).copy()  # Mesh-Tensorflow is a diagonal matrix
								snake_case__	     =torch.tensor(UpperCamelCase_			)
						elif key_name == "model/dense_1/bias":
								snake_case__	     ='model.last_project.bias'
								snake_case__	     =vnp.copy()  # same because it is one dimensional
								snake_case__	     =torch.tensor(UpperCamelCase_			)
		torch.save(UpperCamelCase_      ,					args.output			)


if __name__ == "__main__":
			SCREAMING_SNAKE_CASE__     :						int	=  argparse.ArgumentParser(
			    description='''model converter.''', formatter_class=argparse.ArgumentDefaultsHelpFormatter
			)
			parser.add_argument('''--tf_model_dir''', metavar='''PATH''', type=str, required=True, help='''import model''')
			parser.add_argument('''--output''', metavar='''PATH''', type=str, required=True, help='''output model''')
			SCREAMING_SNAKE_CASE__     :						str	=  parser.parse_args()
			convert_tf_gptsan_to_pt(args)




