
"""simple docstring"""






import os
import re
import unicodedata
from shutil import copyfile
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union

import sentencepiece as spm

from ...tokenization_utils import PreTrainedTokenizer
from ...utils import is_torch_available, logging


if is_torch_available():
					import torch


if TYPE_CHECKING:
					from transformers.pipelines.conversational import Conversation


UpperCAmelCase        =				logging.get_logger(__name__)
UpperCAmelCase        =				{"""vocab_file""": """spiece.model"""}

UpperCAmelCase        =				{
    """vocab_file""": {
        """AI-Sweden/gpt-sw3-126m""": """https://huggingface.co/AI-Sweden/gpt-sw3-126m/resolve/main/spiece.model""",
        """AI-Sweden/gpt-sw3-350m""": """https://huggingface.co/AI-Sweden/gpt-sw3-350m/resolve/main/spiece.model""",
        """AI-Sweden/gpt-sw3-1.6b""": """https://huggingface.co/AI-Sweden/gpt-sw3-1.6b/resolve/main/spiece.model""",
        """AI-Sweden/gpt-sw3-6.7b""": """https://huggingface.co/AI-Sweden/gpt-sw3-6.7b/resolve/main/spiece.model""",
        """AI-Sweden/gpt-sw3-20b""": """https://huggingface.co/AI-Sweden/gpt-sw3-20b/resolve/main/spiece.model""",
    }
}

UpperCAmelCase        =				{
    """AI-Sweden/gpt-sw3-126m""": 2_0_4_8,
    """AI-Sweden/gpt-sw3-350m""": 2_0_4_8,
    """AI-Sweden/gpt-sw3-1.6b""": 2_0_4_8,
    """AI-Sweden/gpt-sw3-6.7b""": 2_0_4_8,
    """AI-Sweden/gpt-sw3-20b""": 2_0_4_8,
}



class 			lowercase					( lowercase__		):
			lowercase        =     VOCAB_FILES_NAMES
			lowercase        =     PRETRAINED_VOCAB_FILES_MAP
			lowercase        =     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			lowercase        =     ['''input_ids''', '''attention_mask''']






			def __init__(self  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	Dict      ,SCREAMING_SNAKE_CASE_  :	int=False      ,SCREAMING_SNAKE_CASE_  :	List[Any]=False      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=False      ,SCREAMING_SNAKE_CASE_  :	List[str]=None      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=None      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]=None      ,SCREAMING_SNAKE_CASE_  :	int=None      ,SCREAMING_SNAKE_CASE_  :	Optional[Dict[str, Any]] = None      ,**SCREAMING_SNAKE_CASE_  :	Union[str, Any]      ,)       ->	None:





				"""simple docstring"""


				lowerCAmelCase												=  {} if sp_model_kwargs is None else sp_model_kwargs

				lowerCAmelCase												=  kwargs.get('''name_or_path'''			)
				if name_or_path is None:
					logger.warning(
					    '''name_or_path not provided, will work for all GPTSw3 models except gpt-sw3-7b,'''
					    ''' you are testing the model, this can safely be ignored'''			)
					lowerCAmelCase												=  '''None'''

				# Default definitions for our 2 tokenizer versions, with None-checks to enable proper testing
				lowerCAmelCase												=  '''<|endoftext|>''' if eos_token is None else eos_token
				lowerCAmelCase												=  '''<unk>''' if unk_token is None else unk_token
				if "gpt-sw3-7b" in name_or_path:
					lowerCAmelCase												=  unk_token if pad_token is None else pad_token
					lowerCAmelCase												=  eos_token if bos_token is None else bos_token
				else:
					lowerCAmelCase												=  '''<pad>''' if pad_token is None else pad_token
					lowerCAmelCase												=  '''<s>''' if bos_token is None else bos_token

				super().__init__(
				    do_lower_case=SCREAMING_SNAKE_CASE_      ,remove_space=SCREAMING_SNAKE_CASE_      ,keep_accents=SCREAMING_SNAKE_CASE_      ,bos_token=SCREAMING_SNAKE_CASE_      ,eos_token=SCREAMING_SNAKE_CASE_      ,unk_token=SCREAMING_SNAKE_CASE_      ,pad_token=SCREAMING_SNAKE_CASE_      ,sp_model_kwargs=self.sp_model_kwargs      ,**SCREAMING_SNAKE_CASE_      ,)

				lowerCAmelCase												=  do_lower_case
				lowerCAmelCase												=  remove_space
				lowerCAmelCase												=  keep_accents
				lowerCAmelCase												=  vocab_file

				lowerCAmelCase												=  spm.SentencePieceProcessor(**self.sp_model_kwargs			)
				self.sp_model.Load(SCREAMING_SNAKE_CASE_			)

				# Used for whitespace normalization in input texts
				# fmt : off
				lowerCAmelCase												=  {''' ''', ''' ''', ''' ''', ''' ''', ''' ''', '''　''', ''' ''', ''' ''', ''' ''', ''' ''', '''￼''', ''''''}
				# fmt : on

				# Regular expression to remove non-printing characters (e.g. some unicode control chars) in preprocessing
				lowerCAmelCase												=  re.compile(
				    F"""[{"".join(map(SCREAMING_SNAKE_CASE_      ,list(range(0      ,9			)			) + list(range(11      ,32			)			) + list(range(127      ,160			)			) + [160, 173, 8_203]			)			)}]"""			)






			def __getstate__(self  :	Tuple			)       ->	Any:





				"""simple docstring"""


				lowerCAmelCase												=  self.__dict__.copy()
				lowerCAmelCase												=  None
				return state






			def __setstate__(self  :	Dict      ,SCREAMING_SNAKE_CASE_  :	Dict			)       ->	List[str]:





				"""simple docstring"""


				lowerCAmelCase												=  d

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

				lowerCAmelCase												=  spm.SentencePieceProcessor(**self.sp_model_kwargs			)
				self.sp_model.Load(self.vocab_file			)






			@property
			# Copied from transformers.models.albert.tokenization_albert.AlbertTokenizer.vocab_size
			def 					UpperCAmelCase							(self  :	int			)       ->	int:





				"""simple docstring"""


				return len(self.sp_model			)






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





				"""simple docstring"""


				lowerCAmelCase												=  self.non_printing_characters_re.sub(''''''      ,SCREAMING_SNAKE_CASE_			)

				# Normalize whitespaces
				lowerCAmelCase												=  ''''''.join([char if char not in self.whitespaces else ''' ''' for char in text]			)

				# NFC Unicode normalization
				lowerCAmelCase												=  unicodedata.normalize('''NFC'''      ,SCREAMING_SNAKE_CASE_			)
				return text






			def 					UpperCAmelCase							(self  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	str      ,**SCREAMING_SNAKE_CASE_  :	List[str]			)       ->	List[str]:





				"""simple docstring"""


				lowerCAmelCase												=  self.preprocess_text(SCREAMING_SNAKE_CASE_			)
				return self.sp_model.encode(SCREAMING_SNAKE_CASE_      ,out_type=SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	str			)       ->	int:





				"""simple docstring"""


				return self.sp_model.PieceToId(SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	int			)       ->	str:





				"""simple docstring"""


				return self.sp_model.IdToPiece(SCREAMING_SNAKE_CASE_			)






			@staticmethod
			def 					UpperCAmelCase							(SCREAMING_SNAKE_CASE_  :	str			)       ->	str:





				"""simple docstring"""


				return out_string






			def 					UpperCAmelCase							(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	List[str]			)       ->	str:





				"""simple docstring"""


				lowerCAmelCase												=  []
				lowerCAmelCase												=  ''''''
				lowerCAmelCase												=  False
				for token in tokens:
					# make sure that special tokens are not decoded using sentencepiece model
					if token in self.all_special_tokens:
						# TODO: Check if this is needed, as it ensures that decode(encode(doc)) != doc by adding extra whitespace in the decoded document
						if not prev_is_special:
							out_string += " "

						out_string += self.sp_model.decode(SCREAMING_SNAKE_CASE_			) + token
						lowerCAmelCase												=  True
						lowerCAmelCase												=  []
					else:
						current_sub_tokens.append(SCREAMING_SNAKE_CASE_			)
						lowerCAmelCase												=  False
				out_string += self.sp_model.decode(SCREAMING_SNAKE_CASE_			)

				return out_string






			def 					UpperCAmelCase							(self  :	Dict			)       ->	Dict[str, int]:





				"""simple docstring"""


				lowerCAmelCase												=  {self.convert_ids_to_tokens(SCREAMING_SNAKE_CASE_			): i for i in range(self.vocab_size			)}
				vocab.update(self.added_tokens_encoder			)
				return vocab






			def 					UpperCAmelCase							(self  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Optional[str] = None			)       ->	Tuple[str]:





				"""simple docstring"""


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

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

				return (out_vocab_file,)






			def 					UpperCAmelCase							(self  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	Union[str, List[str]]      ,SCREAMING_SNAKE_CASE_  :	Union[str, bool] = False			)       ->	Union[List[int], List[List[int]], "torch.Tensor"]:





				"""simple docstring"""


				if isinstance(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			):
					lowerCAmelCase												=  self.preprocess_text(SCREAMING_SNAKE_CASE_			)
					lowerCAmelCase												=  self.sp_model.encode(SCREAMING_SNAKE_CASE_			)
				else:
					lowerCAmelCase												=  [self.preprocess_text(SCREAMING_SNAKE_CASE_			) for t in text]
					lowerCAmelCase												=  self.sp_model.encode(SCREAMING_SNAKE_CASE_			)

				if return_tensors is True or return_tensors == "pt":
					lowerCAmelCase												=  torch.tensor(SCREAMING_SNAKE_CASE_			)

				return token_ids






			def 					UpperCAmelCase							(self  :	str      ,SCREAMING_SNAKE_CASE_  :	Union[int, List[int]]			)       ->	str:





				"""simple docstring"""


				return self.sp_model.decode(SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	"Conversation"			)       ->	List[int]:





				"""simple docstring"""


				lowerCAmelCase												=  [F"""User: {text}""" if is_user else F"""Bot: {text}""" for is_user, text in conversation.iter_texts()]
				lowerCAmelCase												=  (
				    F"""{self.eos_token}{self.bos_token}""" + F"""{self.bos_token}""".join(SCREAMING_SNAKE_CASE_			) + F"""{self.bos_token}Bot:"""
				)
				return self.encode(text=SCREAMING_SNAKE_CASE_			)




"""simple docstring"""






from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_torch_dist_unique_port,
    require_torch_multi_gpu,
    require_torch_neuroncore,
)
from transformers.training_args import ParallelMode
from transformers.utils import logging


UpperCAmelCase        =				logging.get_logger(__name__)


if is_torch_available():
					import torch
					from torch import nn
					from torch.utils.data import Dataset

					from transformers import Trainer



					class 			lowercase					( lowercase__		):






								def __init__(self  :	int      ,SCREAMING_SNAKE_CASE_  :	int = 101			)       ->	List[str]:





									"""simple docstring"""


									lowerCAmelCase												=  length






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





									"""simple docstring"""


									return self.length






								def __getitem__(self  :	str      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]			)       ->	int:





									"""simple docstring"""


									return i





					class 			lowercase					:






								def __call__(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	List[Any]			)       ->	Optional[Any]:





									"""simple docstring"""


									return {"input_ids": torch.tensor(SCREAMING_SNAKE_CASE_			), "labels": torch.tensor(SCREAMING_SNAKE_CASE_			)}





					class 			lowercase					( nn.Module		):






								def __init__(self  :	List[str]			)       ->	List[Any]:





									"""simple docstring"""


									super().__init__()
									# Add some (unused) params otherwise DDP will complain.
									lowerCAmelCase												=  nn.Linear(120      ,80			)






								def 					UpperCAmelCase							(self  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=None			)       ->	Union[str, Any]:





									"""simple docstring"""


									if labels is not None:
										return torch.tensor(0.0      ,device=input_ids.device			), input_ids
									else:
										return input_ids



class 			lowercase					( lowercase__		):






			@require_torch_neuroncore
			def 					UpperCAmelCase							(self  :	Union[str, Any]			)       ->	Union[str, Any]:





				"""simple docstring"""


				lowerCAmelCase												=  F"""--nproc_per_node=2
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
				lowerCAmelCase												=  self.get_auto_remove_tmp_dir()
				lowerCAmelCase												=  F"""--output_dir {output_dir}""".split()
				lowerCAmelCase												=  ['''torchrun'''] + distributed_args + args
				execute_subprocess_async(SCREAMING_SNAKE_CASE_      ,env=self.get_env()			)
				# successful return here == success - any errors would have caused an error in the sub-call





class 			lowercase					( lowercase__		):






			@require_torch_multi_gpu
			def 					UpperCAmelCase							(self  :	List[str]			)       ->	List[str]:





				"""simple docstring"""


				lowerCAmelCase												=  F"""--nproc_per_node={torch.cuda.device_count()}
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
				lowerCAmelCase												=  self.get_auto_remove_tmp_dir()
				lowerCAmelCase												=  F"""--output_dir {output_dir}""".split()
				lowerCAmelCase												=  ['''torchrun'''] + distributed_args + args
				execute_subprocess_async(SCREAMING_SNAKE_CASE_      ,env=self.get_env()			)
				# successful return here == success - any errors would have caused an error in the sub-call


if __name__ == "__main__":
					# The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
					#
					# PYTHONPATH="src" python -m torch.distributed.run --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py

					UpperCAmelCase        =				HfArgumentParser((TrainingArguments,))
					UpperCAmelCase        =				parser.parse_args_into_dataclasses()[0]

					logger.warning(
					    f"""Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, """
					    f"""distributed training: {training_args.parallel_mode != ParallelMode.NOT_DISTRIBUTED}"""
					)

					# Essentially, what we want to verify in the distributed case is that we get all samples back,
					# in the right order. (this is crucial for prediction for instance)
					for dataset_length in [1_0_1, 4_0, 7]:
										UpperCAmelCase        =				DummyDataset(dataset_length)




										def __magic_name__	(				_lowerCamelCase:		EvalPrediction			)					->						Dict:

											'''simple docstring'''




											lowerCAmelCase												=  list(range(len(_lowerCamelCase			)			)			)
											lowerCAmelCase												=  p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
											if not success and training_args.local_rank == 0:
												logger.warning(
												    '''Predictions and/or labels do not match expected results:\n  - predictions: '''
												    F"""{p.predictions.tolist()}\n  - labels: {p.label_ids.tolist()}\n  - expected: {sequential}"""			)
											return {"success": success}

										UpperCAmelCase        =				Trainer(
										    model=DummyModel(),
										    args=training_args,
										    data_collator=DummyDataCollator(),
										    eval_dataset=dataset,
										    compute_metrics=compute_metrics,
										)
										UpperCAmelCase        =				trainer.evaluate()
										logger.info(metrics)
										if metrics["eval_success"] is not True:
															logger.error(metrics)
															exit(1)

										UpperCAmelCase        =				trainer.predict(dataset)
										logger.info(p.metrics)
										if p.metrics["test_success"] is not True:
															logger.error(p.metrics)
															exit(1)

										UpperCAmelCase        =				2

										UpperCAmelCase        =				trainer.evaluate()
										logger.info(metrics)
										if metrics["eval_success"] is not True:
															logger.error(metrics)
															exit(1)

										UpperCAmelCase        =				trainer.predict(dataset)
										logger.info(p.metrics)
										if p.metrics["test_success"] is not True:
															logger.error(p.metrics)
															exit(1)



										UpperCAmelCase        =				None


