





import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, BatchEncoding, PreTrainedTokenizer
from ...utils import logging


__A				:    List[Any]       =   logging.get_logger(__name__)

__A				:    Optional[Any]       =   "▁"

__A				:    Optional[int]       =   {"vocab_file": "sentencepiece.bpe.model"}

__A				:    List[Any]       =   {
    "vocab_file": {
        "facebook/nllb-200-distilled-600M": (
            "https://huggingface.co/facebook/nllb-200-distilled-600M/blob/main/sentencepiece.bpe.model"
        ),
    }
}

__A				:    List[str]       =   {
    "facebook/nllb-200-distilled-600M": 1_024,
}

# fmt: off
__A				:    Dict       =   ["ace_Arab", "ace_Latn", "acm_Arab", "acq_Arab", "aeb_Arab", "afr_Latn", "ajp_Arab", "aka_Latn", "amh_Ethi", "apc_Arab", "arb_Arab", "ars_Arab", "ary_Arab", "arz_Arab", "asm_Beng", "ast_Latn", "awa_Deva", "ayr_Latn", "azb_Arab", "azj_Latn", "bak_Cyrl", "bam_Latn", "ban_Latn", "bel_Cyrl", "bem_Latn", "ben_Beng", "bho_Deva", "bjn_Arab", "bjn_Latn", "bod_Tibt", "bos_Latn", "bug_Latn", "bul_Cyrl", "cat_Latn", "ceb_Latn", "ces_Latn", "cjk_Latn", "ckb_Arab", "crh_Latn", "cym_Latn", "dan_Latn", "deu_Latn", "dik_Latn", "dyu_Latn", "dzo_Tibt", "ell_Grek", "eng_Latn", "epo_Latn", "est_Latn", "eus_Latn", "ewe_Latn", "fao_Latn", "pes_Arab", "fij_Latn", "fin_Latn", "fon_Latn", "fra_Latn", "fur_Latn", "fuv_Latn", "gla_Latn", "gle_Latn", "glg_Latn", "grn_Latn", "guj_Gujr", "hat_Latn", "hau_Latn", "heb_Hebr", "hin_Deva", "hne_Deva", "hrv_Latn", "hun_Latn", "hye_Armn", "ibo_Latn", "ilo_Latn", "ind_Latn", "isl_Latn", "ita_Latn", "jav_Latn", "jpn_Jpan", "kab_Latn", "kac_Latn", "kam_Latn", "kan_Knda", "kas_Arab", "kas_Deva", "kat_Geor", "knc_Arab", "knc_Latn", "kaz_Cyrl", "kbp_Latn", "kea_Latn", "khm_Khmr", "kik_Latn", "kin_Latn", "kir_Cyrl", "kmb_Latn", "kon_Latn", "kor_Hang", "kmr_Latn", "lao_Laoo", "lvs_Latn", "lij_Latn", "lim_Latn", "lin_Latn", "lit_Latn", "lmo_Latn", "ltg_Latn", "ltz_Latn", "lua_Latn", "lug_Latn", "luo_Latn", "lus_Latn", "mag_Deva", "mai_Deva", "mal_Mlym", "mar_Deva", "min_Latn", "mkd_Cyrl", "plt_Latn", "mlt_Latn", "mni_Beng", "khk_Cyrl", "mos_Latn", "mri_Latn", "zsm_Latn", "mya_Mymr", "nld_Latn", "nno_Latn", "nob_Latn", "npi_Deva", "nso_Latn", "nus_Latn", "nya_Latn", "oci_Latn", "gaz_Latn", "ory_Orya", "pag_Latn", "pan_Guru", "pap_Latn", "pol_Latn", "por_Latn", "prs_Arab", "pbt_Arab", "quy_Latn", "ron_Latn", "run_Latn", "rus_Cyrl", "sag_Latn", "san_Deva", "sat_Beng", "scn_Latn", "shn_Mymr", "sin_Sinh", "slk_Latn", "slv_Latn", "smo_Latn", "sna_Latn", "snd_Arab", "som_Latn", "sot_Latn", "spa_Latn", "als_Latn", "srd_Latn", "srp_Cyrl", "ssw_Latn", "sun_Latn", "swe_Latn", "swh_Latn", "szl_Latn", "tam_Taml", "tat_Cyrl", "tel_Telu", "tgk_Cyrl", "tgl_Latn", "tha_Thai", "tir_Ethi", "taq_Latn", "taq_Tfng", "tpi_Latn", "tsn_Latn", "tso_Latn", "tuk_Latn", "tum_Latn", "tur_Latn", "twi_Latn", "tzm_Tfng", "uig_Arab", "ukr_Cyrl", "umb_Latn", "urd_Arab", "uzn_Latn", "vec_Latn", "vie_Latn", "war_Latn", "wol_Latn", "xho_Latn", "ydd_Hebr", "yor_Latn", "yue_Hant", "zho_Hans", "zho_Hant", "zul_Latn"]

class 	A_       (a_			):
	UpperCAmelCase__					=			VOCAB_FILES_NAMES
	UpperCAmelCase__					=			PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
	UpperCAmelCase__					=			PRETRAINED_VOCAB_FILES_MAP
	UpperCAmelCase__					=			['''input_ids''', '''attention_mask''']

	UpperCAmelCase__					=			[]
	UpperCAmelCase__					=			[]
	def __init__(    self							,					_A							,					_A="<s>"							,					_A="</s>"							,					_A="</s>"							,					_A="<s>"							,					_A="<unk>"							,					_A="<pad>"							,					_A="<mask>"							,					_A=None							,					_A=None							,					_A=None							,					_A = None							,					_A=None							,					_A=False							,					**_A							,					):






				'''simple docstring'''







				UpperCAmelCase   		=				AddedToken(_A							,					lstrip=_A							,					rstrip=_A		) if isinstance(_A							,					_A		) else mask_token

				UpperCAmelCase   		=				{} if sp_model_kwargs is None else sp_model_kwargs
				UpperCAmelCase   		=				legacy_behaviour

				super().__init__(
				    bos_token=_A							,					eos_token=_A							,					unk_token=_A							,					sep_token=_A							,					cls_token=_A							,					pad_token=_A							,					mask_token=_A							,					tokenizer_file=_A							,					src_lang=_A							,					tgt_lang=_A							,					additional_special_tokens=_A							,					sp_model_kwargs=self.sp_model_kwargs							,					legacy_behaviour=_A							,					**_A							,					)

				UpperCAmelCase   		=				spm.SentencePieceProcessor(**self.sp_model_kwargs		)
				self.sp_model.Load(str(_A		)		)
				UpperCAmelCase   		=				vocab_file
				# Original fairseq vocab and spm vocab must be "aligned":
				# Vocab    |    0    |    1    |   2    |    3    |  4   |  5   |  6   |   7  |   8  |  9
				# -------- | ------- | ------- | ------ | ------- | ---- | ---- | ---- | ---- | ---- | ----
				# fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | 'an' | '▁n' | '▁m' | '▁t' | '▁k' | '▁a'
				# spm      | '<unk>' | '<s>'   | '</s>' | 'an'    | '▁n' | '▁m' | '▁t' | '▁k' | '▁a' | '▁s'

				# Mimic fairseq token-to-id alignment for the first 4 token
				UpperCAmelCase   		=				{'''<s>''': 0, '''<pad>''': 1, '''</s>''': 2, '''<unk>''': 3}

				# The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
				UpperCAmelCase   		=				1

				UpperCAmelCase   		=				len(self.sp_model		)
				UpperCAmelCase   		=				{
				    code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(_A		)
				}
				UpperCAmelCase   		=				{v: k for k, v in self.lang_code_to_id.items()}
				UpperCAmelCase   		=				len(self.sp_model		) + len(self.lang_code_to_id		) + self.fairseq_offset

				self.fairseq_tokens_to_ids.update(self.lang_code_to_id		)
				UpperCAmelCase   		=				{v: k for k, v in self.fairseq_tokens_to_ids.items()}
				UpperCAmelCase   		=				list(self.lang_code_to_id.keys()		)

				if additional_special_tokens is not None:
							# Only add those special tokens if they are not already there.
							self._additional_special_tokens.extend(
							    [t for t in additional_special_tokens if t not in self._additional_special_tokens]		)

				UpperCAmelCase   		=				src_lang if src_lang is not None else '''eng_Latn'''
				UpperCAmelCase   		=				self.lang_code_to_id[self._src_lang]
				UpperCAmelCase   		=				tgt_lang
				self.set_src_lang_special_tokens(self._src_lang		)
	def __getstate__(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.__dict__.copy()
				UpperCAmelCase   		=				None
				UpperCAmelCase   		=				self.sp_model.serialized_model_proto()
				return state
	def __setstate__(    self							,					_A		):






				'''simple docstring'''







				UpperCAmelCase   		=				d

				# for backward compatibility
				if not hasattr(self							,					'''sp_model_kwargs'''		):
							UpperCAmelCase   		=				{}

				UpperCAmelCase   		=				spm.SentencePieceProcessor(**self.sp_model_kwargs		)
				self.sp_model.LoadFromSerializedProto(self.sp_model_proto		)
	@property
	def      _lowercase				(    self		):






				'''simple docstring'''







				return len(self.sp_model		) + len(self.lang_code_to_id		) + self.fairseq_offset + 1  # Plus 1 for the mask token
	@property
	def      _lowercase				(    self		):






				'''simple docstring'''







				return self._src_lang
	@src_lang.setter
	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				UpperCAmelCase   		=				new_src_lang
				self.set_src_lang_special_tokens(self._src_lang		)
	def      _lowercase				(    self							,					_A							,					_A = None							,					_A = False		):






				'''simple docstring'''







				if already_has_special_tokens:
							return super().get_special_tokens_mask(
							    token_ids_a=_A							,					token_ids_a=_A							,					already_has_special_tokens=_A		)

				UpperCAmelCase   		=				[1] * len(self.prefix_tokens		)
				UpperCAmelCase   		=				[1] * len(self.suffix_tokens		)
				if token_ids_a is None:
							return prefix_ones + ([0] * len(_A		)) + suffix_ones
				return prefix_ones + ([0] * len(_A		)) + ([0] * len(_A		)) + suffix_ones
	def      _lowercase				(    self							,					_A							,					_A = None		):






				'''simple docstring'''







				if token_ids_a is None:
							return self.prefix_tokens + token_ids_a + self.suffix_tokens
				# We don't expect to process pairs, but leave the pair logic for API consistency
				return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens
	def      _lowercase				(    self							,					_A							,					_A = None		):






				'''simple docstring'''







				UpperCAmelCase   		=				[self.sep_token_id]
				UpperCAmelCase   		=				[self.cls_token_id]

				if token_ids_a is None:
							return len(cls + token_ids_a + sep		) * [0]
				return len(cls + token_ids_a + sep + sep + token_ids_a + sep		) * [0]
	def      _lowercase				(    self							,					_A							,					_A							,					_A							,					_A							,					**_A		):






				'''simple docstring'''







				if src_lang is None or tgt_lang is None:
							raise ValueError('''Translation requires a `src_lang` and a `tgt_lang` for this model'''		)
				UpperCAmelCase   		=				src_lang
				UpperCAmelCase   		=				self(_A							,					add_special_tokens=_A							,					return_tensors=_A							,					**_A		)
				UpperCAmelCase   		=				self.convert_tokens_to_ids(_A		)
				UpperCAmelCase   		=				tgt_lang_id
				return inputs
	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				{self.convert_ids_to_tokens(_A		): i for i in range(self.vocab_size		)}
				vocab.update(self.added_tokens_encoder		)
				return vocab
	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				return self.sp_model.encode(_A							,					out_type=_A		)
	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				if token in self.fairseq_tokens_to_ids:
							return self.fairseq_tokens_to_ids[token]
				UpperCAmelCase   		=				self.sp_model.PieceToId(_A		)

				# Need to return unknown token if the SP model returned 0
				return spm_id + self.fairseq_offset if spm_id else self.unk_token_id
	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				if index in self.fairseq_ids_to_tokens:
							return self.fairseq_ids_to_tokens[index]
				return self.sp_model.IdToPiece(index - self.fairseq_offset		)
	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				UpperCAmelCase   		=				''''''.join(_A		).replace(_A							,					''' '''		).strip()
				return out_string
	def      _lowercase				(    self							,					_A							,					_A = None		):






				'''simple docstring'''







				if not os.path.isdir(_A		):
							logger.error(F"""Vocabulary path ({save_directory}) should be a directory"""		)
							return
				UpperCAmelCase   		=				os.path.join(
				    _A							,					(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']		)

				if os.path.abspath(self.vocab_file		) != os.path.abspath(_A		) and os.path.isfile(self.vocab_file		):
							copyfile(self.vocab_file							,					_A		)
				elif not os.path.isfile(self.vocab_file		):
							with open(_A							,					'''wb'''		) as fi:
										UpperCAmelCase   		=				self.sp_model.serialized_model_proto()
										fi.write(_A		)

				return (out_vocab_file,)
	def      _lowercase				(    self							,					_A							,					_A = "eng_Latn"							,					_A = None							,					_A = "fra_Latn"							,					**_A							,					):






				'''simple docstring'''







				UpperCAmelCase   		=				src_lang
				UpperCAmelCase   		=				tgt_lang
				return super().prepare_seqaseq_batch(_A							,					_A							,					**_A		)
	def      _lowercase				(    self		):






				'''simple docstring'''







				return self.set_src_lang_special_tokens(self.src_lang		)
	def      _lowercase				(    self		):






				'''simple docstring'''







				return self.set_tgt_lang_special_tokens(self.tgt_lang		)
	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.lang_code_to_id[src_lang]
				if self.legacy_behaviour:
							UpperCAmelCase   		=				[]
							UpperCAmelCase   		=				[self.eos_token_id, self.cur_lang_code]
				else:
							UpperCAmelCase   		=				[self.cur_lang_code]
							UpperCAmelCase   		=				[self.eos_token_id]







	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.lang_code_to_id[lang]
				if self.legacy_behaviour:
							UpperCAmelCase   		=				[]
							UpperCAmelCase   		=				[self.eos_token_id, self.cur_lang_code]
				else:
							UpperCAmelCase   		=				[self.cur_lang_code]
							UpperCAmelCase   		=				[self.eos_token_id]










import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
				from sagemaker import Session, TrainingJobAnalytics
				from sagemaker.huggingface import HuggingFace

@pytest.mark.skipif(
    literal_eval(os.getenv('''TEST_SAGEMAKER'''     ,							'''False'''			)			) is not True     ,							reason='''Skipping test because should only be run when releasing minor transformers version'''     ,							)
@pytest.mark.usefixtures('''sm_env'''			)
@parameterized_class(
    [
        {
            '''framework''': '''pytorch''',
            '''script''': '''run_glue.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.g4dn.xlarge''',
            '''results''': {'''train_runtime''': 6_5_0, '''eval_accuracy''': 0.6, '''eval_loss''': 0.9},
        },
        {
            '''framework''': '''tensorflow''',
            '''script''': '''run_tf.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.g4dn.xlarge''',
            '''results''': {'''train_runtime''': 6_0_0, '''eval_accuracy''': 0.3, '''eval_loss''': 0.9},
        },
    ]			)
class 	A_       (unittest.TestCase			):
	def      _lowercase				(    self		):






				'''simple docstring'''







				if self.framework == "pytorch":
							subprocess.run(
							    F"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()							,					encoding='''utf-8'''							,					check=_A							,					)
				assert hasattr(self							,					'''env'''		)
	def      _lowercase				(    self							,					_A=1		):






				'''simple docstring'''







				return HuggingFace(
				    entry_point=self.script							,					source_dir=self.env.test_path							,					role=self.env.role							,					image_uri=self.env.image_uri							,					base_job_name=F"""{self.env.base_job_name}-single"""							,					instance_count=_A							,					instance_type=self.instance_type							,					debugger_hook_config=_A							,					hyperparameters={**self.env.hyperparameters, '''model_name_or_path''': self.model_name_or_path}							,					metric_definitions=self.env.metric_definitions							,					py_version='''py36'''							,					)
	def      _lowercase				(    self							,					_A		):






				'''simple docstring'''







				TrainingJobAnalytics(_A		).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv"""		)







	def      _lowercase				(    self		):






				'''simple docstring'''







				UpperCAmelCase   		=				self.create_estimator()

				# run training
				estimator.fit()

				# result dataframe
				UpperCAmelCase   		=				TrainingJobAnalytics(estimator.latest_training_job.name		).dataframe()

				# extract kpis
				UpperCAmelCase   		=				list(result_metrics_df[result_metrics_df.metric_name == '''eval_accuracy''']['''value''']		)
				UpperCAmelCase   		=				list(result_metrics_df[result_metrics_df.metric_name == '''eval_loss''']['''value''']		)
				# get train time from SageMaker job, this includes starting, preprocessing, stopping
				UpperCAmelCase   		=				(
				    Session().describe_training_job(estimator.latest_training_job.name		).get('''TrainingTimeInSeconds'''							,					9_9_9_9_9_9		)
				)

				# assert kpis
				assert train_runtime <= self.results["train_runtime"]
				assert all(t >= self.results['''eval_accuracy'''] for t in eval_accuracy		)
				assert all(t <= self.results['''eval_loss'''] for t in eval_loss		)

				# dump tests result into json file to share in PR
				with open(F"""{estimator.latest_training_job.name}.json"""							,					'''w'''		) as outfile:
							json.dump({'''train_time''': train_runtime, '''eval_accuracy''': eval_accuracy, '''eval_loss''': eval_loss}							,					_A		)



