

import argparse
import logging
import os
from datetime import datetime

import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, RandomSampler, TensorDataset
from tqdm import tqdm

from transformers import GPTaLMHeadModel


A_       : Optional[Any]         =			logging.getLogger(__name__)

def 						UpperCamelCase					(lowercase_: Tuple    ,      lowercase_: Optional[int]					)     ->      Optional[int]:
	# save results
	if os.path.exists(lowercase_					):
		if os.path.exists(os.path.join(lowercase_    ,      """config.json"""					)					) and os.path.isfile(
		    os.path.join(lowercase_    ,      """config.json"""					)					):
			os.remove(os.path.join(lowercase_    ,      """config.json"""					)					)
		if os.path.exists(os.path.join(lowercase_    ,      """pytorch_model.bin"""					)					) and os.path.isfile(
		    os.path.join(lowercase_    ,      """pytorch_model.bin"""					)					):
			os.remove(os.path.join(lowercase_    ,      """pytorch_model.bin"""					)					)
	else:
		os.makedirs(lowercase_					)
	model.save_pretrained(lowercase_					)

def 						UpperCamelCase					(lowercase_: Any    ,      lowercase_: Dict=False					)     ->      Dict:
	A__							:    int        =		2
	if unlogit:
		A__							:    Tuple        =		torch.pow(lowercase_    ,      lowercase_					)
	A__							:    Tuple        =		p * torch.log(lowercase_					)
	A__							:    Tuple        =		0
	return -plogp.sum(dim=-1					)

def 						UpperCamelCase					(lowercase_: Optional[Any]					)     ->      List[str]:
	logger.info("""lv, h >\t""" + """\t""".join(f"""{x + 1}""" for x in range(len(lowercase_					)					)					)					)
	for row in range(len(lowercase_					)					):
		if tensor.dtype != torch.long:
			logger.info(f"""layer {row + 1}:\t""" + """\t""".join(f"""{x:.5f}""" for x in tensor[row].cpu().data					)					)
		else:
			logger.info(f"""layer {row + 1}:\t""" + """\t""".join(f"""{x:d}""" for x in tensor[row].cpu().data					)					)

def 						UpperCamelCase					(lowercase_: int    ,      lowercase_: Union[str, Any]    ,      lowercase_: List[Any]    ,      lowercase_: List[Any]=True    ,      lowercase_: Any=True    ,      lowercase_: int=None    ,      lowercase_: int=False					)     ->      Optional[Any]:
	A__   ,   A__							:    Optional[Any]        =		model.config.num_hidden_layers, model.config.num_attention_heads
	A__							:    Optional[int]        =		torch.zeros(lowercase_    ,      lowercase_					).to(args.device					)
	A__							:    List[Any]        =		torch.zeros(lowercase_    ,      lowercase_					).to(args.device					)

	if head_mask is None:
		A__							:    int        =		torch.ones(lowercase_    ,      lowercase_					).to(args.device					)

	head_mask.requires_grad_(requires_grad=lowercase_					)
	# If actually pruned attention multi-head, set head mask to None to avoid shape mismatch
	if actually_pruned:
		A__							:    List[str]        =		None

	A__							:    int        =		0.0
	A__							:    Dict        =		0.0
	for step, inputs in enumerate(tqdm(lowercase_    ,      desc="""Iteration"""    ,      disable=args.local_rank not in [-1, 0]					)					):
		A__							:    Optional[int]        =		tuple(t.to(args.device					) for t in inputs					)
		((A__)   ,   )							:    int        =		inputs

		# Do a forward pass (not with torch.no_grad() since we need gradients for importance score - see below)
		A__							:    Tuple        =		model(lowercase_    ,      labels=lowercase_    ,      head_mask=lowercase_					)
		#  (loss), lm_logits, presents, (all hidden_states), (attentions)
		A__   ,   A__   ,   A__							:    Union[str, Any]        =		(
		    outputs[0],
		    outputs[1],
		    outputs[-1],
		)  # Loss and logits are the first, attention the last
		loss.backward()  # Backpropagate to populate the gradients in the head mask
		total_loss += loss.detach().cpu().numpy()
		if compute_entropy:
			for layer, attn in enumerate(lowercase_					):
				A__							:    List[Any]        =		entropy(attn.detach()    ,      lowercase_					)
				attn_entropy[layer] += masked_entropy.sum(-1					).sum(0					).sum(0					).detach()

		if compute_importance:
			head_importance += head_mask.grad.abs().detach()
		tot_tokens += torch.ones_like(lowercase_					).float().detach().sum().data

	# Normalize
	attn_entropy /= tot_tokens
	head_importance /= tot_tokens
	# Layerwise importance normalization
	if not args.dont_normalize_importance_by_layer:
		A__							:    Union[str, Any]        =		2
		A__							:    List[Any]        =		torch.pow(torch.pow(lowercase_    ,      lowercase_					).sum(-1					)    ,      1 / exponent					)
		head_importance /= norm_by_layer.unsqueeze(-1					) + 1E-20

	if not args.dont_normalize_global_importance:
		A__							:    Any        =		(head_importance - head_importance.min()) / (head_importance.max() - head_importance.min())

	# Print matrices
	if compute_entropy:
		logger.info("""Attention entropies"""					)
		print_ad_tensor(lowercase_					)
	if compute_importance:
		logger.info("""Head importance scores"""					)
		print_ad_tensor(lowercase_					)
	logger.info("""Head ranked by importance scores"""					)
	A__							:    List[Any]        =		torch.zeros(head_importance.numel()    ,      dtype=torch.long    ,      device=args.device					)
	A__							:    str        =		torch.arange(
	    head_importance.numel()    ,      device=args.device					)
	A__							:    str        =		head_ranks.view_as(lowercase_					)
	print_ad_tensor(lowercase_					)
	return attn_entropy, head_importance, total_loss

def 						UpperCamelCase					(lowercase_: Dict    ,      lowercase_: List[Any]    ,      lowercase_: Optional[int]					)     ->      Optional[int]:
	A__   ,   A__   ,   A__							:    Dict        =		compute_heads_importance(lowercase_    ,      lowercase_    ,      lowercase_    ,      compute_entropy=lowercase_					)
	A__							:    Tuple        =		1 / loss  # instead of downsteam score use the LM loss
	logger.info("""Pruning: original score: %f, threshold: %f"""    ,      lowercase_    ,      original_score * args.masking_threshold					)

	A__							:    int        =		torch.ones_like(lowercase_					)
	A__							:    int        =		max(1    ,      int(new_head_mask.numel() * args.masking_amount					)					)

	A__							:    List[Any]        =		original_score
	while current_score >= original_score * args.masking_threshold:
		A__							:    Dict        =		new_head_mask.clone().detach()  # save current head mask
		# heads from least important to most - keep only not-masked heads
		A__							:    List[Any]        =		float("""Inf"""					)
		A__							:    List[str]        =		head_importance.view(-1					).sort()[1]

		if len(lowercase_					) <= num_to_mask:
			print("""BREAK BY num_to_mask"""					)
			break

		# mask heads
		A__							:    Union[str, Any]        =		current_heads_to_mask[:num_to_mask]
		logger.info("""Heads to mask: %s"""    ,      str(current_heads_to_mask.tolist()					)					)
		A__							:    List[Any]        =		new_head_mask.view(-1					)
		A__							:    int        =		0.0
		A__							:    Any        =		new_head_mask.view_as(lowercase_					)
		A__							:    Union[str, Any]        =		new_head_mask.clone().detach()
		print_ad_tensor(lowercase_					)

		# Compute metric and head importance again
		A__   ,   A__   ,   A__							:    Tuple        =		compute_heads_importance(
		    lowercase_    ,      lowercase_    ,      lowercase_    ,      compute_entropy=lowercase_    ,      head_mask=lowercase_					)
		A__							:    str        =		1 / loss
		logger.info(
		    """Masking: current score: %f, remaining heads %d (%.1f percents)"""    ,      lowercase_    ,      new_head_mask.sum()    ,      new_head_mask.sum() / new_head_mask.numel() * 100    ,      )

	logger.info("""Final head mask"""					)
	print_ad_tensor(lowercase_					)
	np.save(os.path.join(args.output_dir    ,      """head_mask.npy"""					)    ,      head_mask.detach().cpu().numpy()					)

	return head_mask

def 						UpperCamelCase					(lowercase_: str    ,      lowercase_: Any    ,      lowercase_: Optional[int]    ,      lowercase_: List[str]					)     ->      Optional[int]:
	A__							:    List[Any]        =		datetime.now()
	A__   ,   A__   ,   A__							:    Optional[int]        =		compute_heads_importance(
	    lowercase_    ,      lowercase_    ,      lowercase_    ,      compute_entropy=lowercase_    ,      compute_importance=lowercase_    ,      head_mask=lowercase_					)
	A__							:    Tuple        =		1 / loss
	A__							:    Any        =		datetime.now() - before_time

	A__							:    Dict        =		sum(p.numel() for p in model.parameters()					)
	A__							:    Tuple        =		{
	    layer: (1 - head_mask[layer].long()).nonzero().squeeze().tolist() for layer in range(len(lowercase_					)					)
	}

	for k, v in heads_to_prune.items():
		if isinstance(lowercase_    ,      lowercase_					):
			A__							:    List[str]        =		[
			    v,
			]

	assert sum(len(lowercase_					) for h in heads_to_prune.values()					) == (1 - head_mask.long()).sum().item()
	model.prune_heads(lowercase_					)
	A__							:    List[Any]        =		sum(p.numel() for p in model.parameters()					)

	A__							:    List[Any]        =		datetime.now()
	A__   ,   A__   ,   A__							:    Optional[Any]        =		compute_heads_importance(
	    lowercase_    ,      lowercase_    ,      lowercase_    ,      compute_entropy=lowercase_    ,      compute_importance=lowercase_    ,      head_mask=lowercase_    ,      actually_pruned=lowercase_    ,      )

	A__							:    Tuple        =		1 / loss
	A__							:    List[str]        =		datetime.now() - before_time

	logger.info(
	    """Pruning: original num of params: %.2e, after pruning %.2e (%.1f percents)"""    ,      lowercase_    ,      lowercase_    ,      pruned_num_params / original_num_params * 100    ,      )
	logger.info("""Pruning: score with masking: %f score with pruning: %f"""    ,      lowercase_    ,      lowercase_					)
	logger.info("""Pruning: speed ratio (original timing / new timing): %f percents"""    ,      original_time / new_time * 100					)
	save_model(lowercase_    ,      args.output_dir					)

def 						UpperCamelCase					()     ->      Any:
	A__							:    List[Any]        =		argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    """--data_dir"""    ,      default=lowercase_    ,      type=lowercase_    ,      required=lowercase_    ,      help="""The input data dir. Should contain the .tsv files (or other data files) for the task."""    ,      )
	parser.add_argument(
	    """--model_name_or_path"""    ,      default=lowercase_    ,      type=lowercase_    ,      required=lowercase_    ,      help="""Path to pretrained model or model identifier from huggingface.co/models"""    ,      )
	parser.add_argument(
	    """--output_dir"""    ,      default=lowercase_    ,      type=lowercase_    ,      required=lowercase_    ,      help="""The output directory where the model predictions and checkpoints will be written."""    ,      )

	# Other parameters
	parser.add_argument(
	    """--config_name"""    ,      default=""""""    ,      type=lowercase_    ,      help="""Pretrained config name or path if not the same as model_name_or_path"""    ,      )
	parser.add_argument(
	    """--tokenizer_name"""    ,      default=""""""    ,      type=lowercase_    ,      help="""Pretrained tokenizer name or path if not the same as model_name_or_path"""    ,      )
	parser.add_argument(
	    """--cache_dir"""    ,      default=lowercase_    ,      type=lowercase_    ,      help="""Where do you want to store the pre-trained models downloaded from s3"""    ,      )
	parser.add_argument(
	    """--data_subset"""    ,      type=lowercase_    ,      default=-1    ,      help="""If > 0: limit the data to a subset of data_subset instances."""					)
	parser.add_argument(
	    """--overwrite_output_dir"""    ,      action="""store_true"""    ,      help="""Whether to overwrite data in output directory"""					)
	parser.add_argument(
	    """--overwrite_cache"""    ,      action="""store_true"""    ,      help="""Overwrite the cached training and evaluation sets"""					)

	parser.add_argument(
	    """--dont_normalize_importance_by_layer"""    ,      action="""store_true"""    ,      help="""Don't normalize importance score by layers"""					)
	parser.add_argument(
	    """--dont_normalize_global_importance"""    ,      action="""store_true"""    ,      help="""Don't normalize all importance scores between 0 and 1"""    ,      )

	parser.add_argument(
	    """--try_masking"""    ,      action="""store_true"""    ,      help="""Whether to try to mask head until a threshold of accuracy."""					)
	parser.add_argument(
	    """--masking_threshold"""    ,      default=0.9    ,      type=lowercase_    ,      help="""masking threshold in term of metrics (stop masking when metric < threshold * original metric value)."""    ,      )
	parser.add_argument(
	    """--masking_amount"""    ,      default=0.1    ,      type=lowercase_    ,      help="""Amount to heads to masking at each masking step."""					)
	parser.add_argument("""--metric_name"""    ,      default="""acc"""    ,      type=lowercase_    ,      help="""Metric to use for head masking."""					)

	parser.add_argument(
	    """--max_seq_length"""    ,      default=128    ,      type=lowercase_    ,      help=(
	        """The maximum total input sequence length after WordPiece tokenization. \n"""
	        """Sequences longer than this will be truncated, sequences shorter padded."""
	    )    ,      )
	parser.add_argument("""--batch_size"""    ,      default=1    ,      type=lowercase_    ,      help="""Batch size."""					)

	parser.add_argument("""--seed"""    ,      type=lowercase_    ,      default=42					)
	parser.add_argument("""--local_rank"""    ,      type=lowercase_    ,      default=-1    ,      help="""local_rank for distributed training on gpus"""					)
	parser.add_argument("""--no_cuda"""    ,      action="""store_true"""    ,      help="""Whether not to use CUDA when available"""					)
	parser.add_argument("""--server_ip"""    ,      type=lowercase_    ,      default=""""""    ,      help="""Can be used for distant debugging."""					)
	parser.add_argument("""--server_port"""    ,      type=lowercase_    ,      default=""""""    ,      help="""Can be used for distant debugging."""					)
	A__							:    List[Any]        =		parser.parse_args()

	if args.server_ip and args.server_port:
		# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
		import ptvsd

		print("""Waiting for debugger attach"""					)
		ptvsd.enable_attach(address=(args.server_ip, args.server_port)    ,      redirect_output=lowercase_					)
		ptvsd.wait_for_attach()

	# Setup devices and distributed training
	if args.local_rank == -1 or args.no_cuda:
		A__							:    List[Any]        =		torch.device("""cuda""" if torch.cuda.is_available() and not args.no_cuda else """cpu"""					)
		A__							:    int        =		0 if args.no_cuda else torch.cuda.device_count()
	else:
		torch.cuda.set_device(args.local_rank					)
		A__							:    Union[str, Any]        =		torch.device("""cuda"""    ,      args.local_rank					)
		A__							:    Dict        =		1
		torch.distributed.init_process_group(backend="""nccl"""					)  # Initializes the distributed backend

	# Setup logging
	logging.basicConfig(level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN					)
	logger.info("""device: {} n_gpu: {}, distributed: {}""".format(args.device    ,      args.n_gpu    ,      bool(args.local_rank != -1					)					)					)

	A__							:    List[Any]        =		GPTaLMHeadModel.from_pretrained(args.model_name_or_path					)

	# Distributed and parallel training
	model.to(args.device					)
	if args.local_rank != -1:
		A__							:    int        =		nn.parallel.DistributedDataParallel(
		    lowercase_    ,      device_ids=[args.local_rank]    ,      output_device=args.local_rank    ,      find_unused_parameters=lowercase_					)
	elif args.n_gpu > 1:
		A__							:    Dict        =		nn.DataParallel(lowercase_					)

	# Print/save training arguments
	os.makedirs(args.output_dir    ,      exist_ok=lowercase_					)
	torch.save(lowercase_    ,      os.path.join(args.output_dir    ,      """run_args.bin"""					)					)
	logger.info("""Training/evaluation parameters %s"""    ,      lowercase_					)

	# Prepare dataset
	A__							:    Optional[Any]        =		np.concatenate(
	    [
	        np.loadtxt(args.data_dir    ,      dtype=np.intaa					),
	    ]					)
	A__							:    Any        =		(torch.from_numpy(lowercase_					),)
	A__							:    Any        =		TensorDataset(*lowercase_					)
	A__							:    Dict        =		RandomSampler(lowercase_					)
	A__							:    List[Any]        =		DataLoader(lowercase_    ,      sampler=lowercase_    ,      batch_size=args.batch_size					)

	# Compute head entropy and importance score
	compute_heads_importance(lowercase_    ,      lowercase_    ,      lowercase_					)

	# Try head masking (set heads to zero until the score goes under a threshole)
	# and head pruning (remove masked heads and see the effect on the network)
	if args.try_masking and args.masking_threshold > 0.0 and args.masking_threshold < 1.0:
		A__							:    Dict        =		mask_heads(lowercase_    ,      lowercase_    ,      lowercase_					)
		prune_heads(lowercase_    ,      lowercase_    ,      lowercase_    ,      lowercase_					)


if __name__ == "__main__":
	main()






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

from transformers import MaMaaaTokenizer, is_torch_available
from transformers.testing_utils import (
    get_tests_dir,
    nested_simplify,
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
)
from transformers.utils import is_sentencepiece_available


if is_sentencepiece_available():
	from transformers.models.mam_aaa.tokenization_mam_aaa import VOCAB_FILES_NAMES, save_json

from ...test_tokenization_common import TokenizerTesterMixin


if is_sentencepiece_available():
	A_       : Tuple         =			get_tests_dir('fixtures/test_sentencepiece.model')


if is_torch_available():
	from transformers.models.mam_aaa.modeling_mam_aaa import shift_tokens_right

A_       : Tuple         =			12_8022
A_       : Optional[Any]         =			12_8028



@require_sentencepiece
class     _a						(__magic_name__		,   unittest.TestCase				):


					'''simple docstring'''






					UpperCAmelCase__:   Tuple  							=       MaMaaaTokenizer
					UpperCAmelCase__:   List[Any]  							=       False
					UpperCAmelCase__:   Any  							=       False
					UpperCAmelCase__:   Optional[Any]  							=       True






					def   __A			( self			):
						super().setUp()

						A__							:    Union[str, Any]        =		["""</s>""", """<unk>""", """▁This""", """▁is""", """▁a""", """▁t""", """est""", """\u0120""", """<pad>"""]
						A__							:    Optional[Any]        =		dict(zip(A__    ,       range(len(A__			)			)			)			)
						A__							:    Optional[int]        =		Path(self.tmpdirname			)
						save_json(A__    ,       save_dir / VOCAB_FILES_NAMES["""vocab_file"""]			)
						if not (save_dir / VOCAB_FILES_NAMES["spm_file"]).exists():
							copyfile(A__    ,       save_dir / VOCAB_FILES_NAMES["""spm_file"""]			)

						A__							:    Tuple        =		MaMaaaTokenizer.from_pretrained(self.tmpdirname			)
						tokenizer.save_pretrained(self.tmpdirname			)






					def   __A			( self    ,       **A__			):
						return MaMaaaTokenizer.from_pretrained(self.tmpdirname    ,       **A__			)






					def   __A			( self    ,       A__			):
						return (
						    "This is a test",
						    "This is a test",
						)






					def   __A			( self			):
						A__							:    Any        =		"""</s>"""
						A__							:    str        =		0

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






					def   __A			( self			):
						A__							:    str        =		self.get_tokenizer()
						A__							:    Dict        =		list(tokenizer.get_vocab().keys()			)

						self.assertEqual(vocab_keys[0]    ,       """</s>"""			)
						self.assertEqual(vocab_keys[1]    ,       """<unk>"""			)
						self.assertEqual(vocab_keys[-1]    ,       """<s>"""			)
						self.assertEqual(len(A__			)    ,       tokenizer.vocab_size + len(tokenizer.get_added_vocab()			)			)






					@unittest.skip("""Skip this test while all models are still to be uploaded."""			)
					def   __A			( self			):
						pass






					def   __A			( self			):
						A__							:    Optional[int]        =		self.get_tokenizer()

						A__							:    int        =		tokenizer.tokenize("""This is a test"""			)
						self.assertListEqual(A__    ,       ["""▁This""", """▁is""", """▁a""", """▁t""", """est"""]			)

						self.assertListEqual(
						    tokenizer.convert_tokens_to_ids(A__			)    ,       [2, 3, 4, 5, 6]    ,       )

						A__							:    Dict        =		tokenizer.convert_ids_to_tokens([2, 3, 4, 5, 6]			)
						self.assertListEqual(A__    ,       ["""▁This""", """▁is""", """▁a""", """▁t""", """est"""]			)

						A__							:    Any        =		tokenizer.convert_tokens_to_string(A__			)
						self.assertEqual(A__    ,       """This is a test"""			)






					@slow
					def   __A			( self			):
						# fmt: off
						A__							:    int        =		{"""input_ids""": [[12_8022, 11_0108, 397, 11, 3_8272, 2247, 12_4811, 285, 1_8105, 1586, 207, 7, 3_9534, 4428, 397, 1019, 1_8105, 1586, 207, 7, 4_1337, 1_6786, 241, 7, 2_0214, 17, 12_5690, 1_0398, 7, 4_4378, 5_8069, 6_8342, 7798, 7343, 11, 299, 3_3310, 4, 158, 3_7350, 9_4077, 4569, 299, 3_3310, 90, 4, 5_2840, 290, 4, 3_1270, 112, 299, 682, 4, 5_2840, 3_9953, 1_4079, 193, 5_2519, 9_0894, 1_7894, 12_0697, 11, 4_0445, 551, 17, 1019, 5_2519, 9_0894, 1_7756, 963, 11, 4_0445, 480, 17, 9792, 1120, 5173, 1393, 6240, 1_6786, 241, 12_0996, 28, 1245, 1393, 11_8240, 1_1123, 1019, 9_3612, 2691, 1_0618, 9_8058, 12_0409, 1928, 279, 4, 4_0683, 367, 178, 207, 1019, 103, 10_3121, 506, 6_5296, 5, 2], [12_8022, 2_1217, 367, 117, 12_5450, 128, 719, 7, 7308, 40, 9_3612, 1_2669, 1116, 1_6704, 71, 1_7785, 3699, 1_5592, 35, 144, 9584, 241, 1_1943, 713, 950, 799, 2247, 8_8427, 150, 149, 11_8813, 12_0706, 1019, 10_6906, 8_1518, 28, 1224, 2_2799, 397, 5, 2, 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], [12_8022, 1658, 12_3311, 5155, 5578, 4722, 279, 1_4947, 2366, 1120, 1197, 14, 1348, 9232, 5, 2, 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]], """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, 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], [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, 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=A__    ,       model_name="""facebook/m2m100_418M"""    ,       revision="""c168bae485c864188cf9aa0e4108b0b6934dc91e"""    ,       )



@require_torch
@require_sentencepiece
@require_tokenizers
class     _a						(unittest.TestCase				):


					'''simple docstring'''






					UpperCAmelCase__:   Optional[int]  							=       '''facebook/m2m100_418M'''
					UpperCAmelCase__:   Any  							=       [
					    '''In my opinion, there are two levels of response from the French government.''',
					    '''NSA Affair Emphasizes Complete Lack of Debate on Intelligence''',
					]
					UpperCAmelCase__:   Any  							=       [
					    '''Selon moi, il y a deux niveaux de réponse de la part du gouvernement français.''',
					    '''L\'affaire NSA souligne l\'absence totale de débat sur le renseignement''',
					]

					# fmt: off
					UpperCAmelCase__:   List[str]  							=       [EN_CODE, 5_93, 19_49, 11_57_81, 4, 7_15_86, 42_34, 6_06_33, 12_62_33, 4_32, 12_38_08, 1_55_92, 11_97, 11_71_32, 12_06_18, 5, 2]






					@classmethod
					def   __A			( cls			):
						A__							:    MaMaaaTokenizer        =		MaMaaaTokenizer.from_pretrained(
						    cls.checkpoint_name    ,       src_lang="""en"""    ,       tgt_lang="""fr"""			)
						A__							:    int        =		1
						return cls






					def   __A			( self			):
						self.assertEqual(self.tokenizer.get_lang_id("""ar"""			)    ,       12_8006			)
						self.assertEqual(self.tokenizer.get_lang_id("""en"""			)    ,       12_8022			)
						self.assertEqual(self.tokenizer.get_lang_id("""ro"""			)    ,       12_8076			)
						self.assertEqual(self.tokenizer.get_lang_id("""mr"""			)    ,       12_8063			)






					def   __A			( self			):
						A__							:    Optional[Any]        =		self.tokenizer.get_vocab()
						self.assertEqual(len(A__			)    ,       self.tokenizer.vocab_size			)
						self.assertEqual(vocab["""<unk>"""]    ,       3			)
						self.assertIn(self.tokenizer.get_lang_token("""en"""			)    ,       A__			)






					def   __A			( self			):
						A__							:    List[Any]        =		"""en"""
						A__							:    str        =		self.tokenizer.batch_encode_plus(self.src_text			).input_ids[0]
						self.assertListEqual(self.expected_src_tokens    ,       A__			)






					def   __A			( self			):
						self.assertIn(A__    ,       self.tokenizer.all_special_ids			)
						# fmt: off
						A__							:    Dict        =		[FR_CODE, 5364, 82, 8642, 4, 294, 47, 8, 1_4028, 136, 3286, 9706, 6, 9_0797, 6, 14_4012, 162, 8_8128, 3_0061, 5, 2]
						# fmt: on
						A__							:    Dict        =		self.tokenizer.decode(A__    ,       skip_special_tokens=A__			)
						A__							:    List[str]        =		self.tokenizer.decode(generated_ids[1:]    ,       skip_special_tokens=A__			)
						self.assertEqual(A__    ,       A__			)
						self.assertNotIn(self.tokenizer.eos_token    ,       A__			)






					def   __A			( self			):
						A__							:    str        =		tempfile.mkdtemp()
						A__							:    Dict        =		self.tokenizer.lang_token_to_id
						self.tokenizer.save_pretrained(A__			)
						A__							:    List[Any]        =		MaMaaaTokenizer.from_pretrained(A__			)
						self.assertDictEqual(new_tok.lang_token_to_id    ,       A__			)






					@require_torch
					def   __A			( self			):
						A__							:    List[str]        =		"""en"""
						A__							:    List[str]        =		"""fr"""

						A__							:    int        =		self.tokenizer(self.src_text    ,       text_target=self.tgt_text    ,       padding=A__    ,       return_tensors="""pt"""			)

						A__							:    int        =		shift_tokens_right(
						    batch["""labels"""]    ,       self.tokenizer.pad_token_id    ,       self.tokenizer.eos_token_id			)

						for k in batch:
							A__							:    Any        =		batch[k].tolist()
						# batch = {k: v.tolist() for k,v in batch.items()}
						# fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
						# batch.decoder_inputs_ids[0][0] ==
						assert batch.input_ids[1][0] == EN_CODE
						assert batch.input_ids[1][-1] == 2
						assert batch.labels[1][0] == FR_CODE
						assert batch.labels[1][-1] == 2
						assert batch.decoder_input_ids[1][:2] == [2, FR_CODE]






					@require_torch
					def   __A			( self			):
						A__							:    List[str]        =		"""mr"""
						self.assertListEqual(self.tokenizer.prefix_tokens    ,       [self.tokenizer.get_lang_id("""mr"""			)]			)
						self.assertListEqual(self.tokenizer.suffix_tokens    ,       [self.tokenizer.eos_token_id]			)

						A__							:    Any        =		"""zh"""
						self.assertListEqual(self.tokenizer.prefix_tokens    ,       [self.tokenizer.get_lang_id("""zh"""			)]			)
						self.assertListEqual(self.tokenizer.suffix_tokens    ,       [self.tokenizer.eos_token_id]			)






					@require_torch
					def   __A			( self			):
						A__							:    Optional[int]        =		"""mr"""
						self.tokenizer._switch_to_target_mode()
						self.assertListEqual(self.tokenizer.prefix_tokens    ,       [self.tokenizer.get_lang_id("""mr"""			)]			)
						self.assertListEqual(self.tokenizer.suffix_tokens    ,       [self.tokenizer.eos_token_id]			)
						self.tokenizer._switch_to_input_mode()
						self.assertListEqual(self.tokenizer.prefix_tokens    ,       [self.tokenizer.get_lang_id(self.tokenizer.src_lang			)]			)

						A__							:    Any        =		"""zh"""
						self.tokenizer._switch_to_target_mode()
						self.assertListEqual(self.tokenizer.prefix_tokens    ,       [self.tokenizer.get_lang_id("""zh"""			)]			)
						self.assertListEqual(self.tokenizer.suffix_tokens    ,       [self.tokenizer.eos_token_id]			)
						self.tokenizer._switch_to_input_mode()
						self.assertListEqual(self.tokenizer.prefix_tokens    ,       [self.tokenizer.get_lang_id(self.tokenizer.src_lang			)]			)






					@require_torch
					def   __A			( self			):
						A__							:    Any        =		self.tokenizer._build_translation_inputs("""A test"""    ,       return_tensors="""pt"""    ,       src_lang="""en"""    ,       tgt_lang="""ar"""			)

						self.assertEqual(
						    nested_simplify(A__			)    ,       {
						        # en_XX, A, test, EOS
						        """input_ids""": [[12_8022, 58, 4183, 2]],
						        """attention_mask""": [[1, 1, 1, 1]],
						        # ar_AR
						        """forced_bos_token_id""": 12_8006,
						    }    ,       )



