





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_048,
    "AI-Sweden/gpt-sw3-350m": 2_048,
    "AI-Sweden/gpt-sw3-1.6b": 2_048,
    "AI-Sweden/gpt-sw3-6.7b": 2_048,
    "AI-Sweden/gpt-sw3-20b": 2_048,
}

class     __SCREAMING_SNAKE_CASE   (       __lowerCAmelCase    ):
							snake_case       : str         =    VOCAB_FILES_NAMES
							snake_case       : List[Any]         =    PRETRAINED_VOCAB_FILES_MAP
							snake_case       : Optional[Any]         =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							snake_case       : Tuple         =    ['''input_ids''', '''attention_mask''']






							def __init__(  self					,	__lowerCAmelCase					,	__lowerCAmelCase=False					,	__lowerCAmelCase=False					,	__lowerCAmelCase=False					,	__lowerCAmelCase=None					,	__lowerCAmelCase=None					,	__lowerCAmelCase=None					,	__lowerCAmelCase=None					,	__lowerCAmelCase = None					,	**__lowerCAmelCase					,	):
														UpperCamelCase__			       =     {} if sp_model_kwargs is None else sp_model_kwargs

														UpperCamelCase__			       =     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"""					)
																					UpperCamelCase__			       =     """None"""

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

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

														UpperCamelCase__			       =     do_lower_case
														UpperCamelCase__			       =     remove_space
														UpperCamelCase__			       =     keep_accents
														UpperCamelCase__			       =     vocab_file

														UpperCamelCase__			       =     spm.SentencePieceProcessor(**self.sp_model_kwargs					)
														self.sp_model.Load(lowerCAmelCase_					)

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

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






							def __getstate__(  self					):
														UpperCamelCase__			       =     self.__dict__.copy()
														UpperCamelCase__			       =     None
														return state






							def __setstate__(  self					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     d

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

														UpperCamelCase__			       =     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 			_lowerCamelCase						(  self					):
														return len(self.sp_model					)






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     self.non_printing_characters_re.sub(""""""					,	lowerCAmelCase_					)

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

														# NFC Unicode normalization
														UpperCamelCase__			       =     unicodedata.normalize("""NFC"""					,	lowerCAmelCase_					)
														return text






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	**__lowerCAmelCase					):
														UpperCamelCase__			       =     self.preprocess_text(lowerCAmelCase_					)
														return self.sp_model.encode(lowerCAmelCase_					,	out_type=lowerCAmelCase_					)






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														return self.sp_model.PieceToId(lowerCAmelCase_					)






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														return self.sp_model.IdToPiece(lowerCAmelCase_					)






							@staticmethod
							def 			_lowerCamelCase						(  __lowerCAmelCase					):
														return out_string






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     []
														UpperCamelCase__			       =     """"""
														UpperCamelCase__			       =     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(lowerCAmelCase_					) + token
																												UpperCamelCase__			       =     True
																												UpperCamelCase__			       =     []
																					else:
																												current_sub_tokens.append(lowerCAmelCase_					)
																												UpperCamelCase__			       =     False
														out_string += self.sp_model.decode(lowerCAmelCase_					)

														return out_string






							def 			_lowerCamelCase						(  self					):
														UpperCamelCase__			       =     {self.convert_ids_to_tokens(lowerCAmelCase_					): i for i in range(self.vocab_size					)}
														vocab.update(self.added_tokens_encoder					)
														return vocab






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase = None					):
														if not os.path.isdir(lowerCAmelCase_					):
																					logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""					)
																					return
														UpperCamelCase__			       =     os.path.join(
														    lowerCAmelCase_					,	(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]					)

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

														return (out_vocab_file,)






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					,	__lowerCAmelCase = False					):

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

														if return_tensors is True or return_tensors == "pt":
																					UpperCamelCase__			       =     torch.tensor(lowerCAmelCase_					)

														return token_ids






							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														return self.sp_model.decode(lowerCAmelCase_					)







							def 			_lowerCamelCase						(  self					,	__lowerCAmelCase					):
														UpperCamelCase__			       =     [f"""User: {text}""" if is_user else f"""Bot: {text}""" for is_user, text in conversation.iter_texts()]
														UpperCamelCase__			       =     (
														    f"""{self.eos_token}{self.bos_token}""" + f"""{self.bos_token}""".join(lowerCAmelCase_					) + f"""{self.bos_token}Bot:"""
														)
														return self.encode(text=lowerCAmelCase_					)











from typing import Dict

import numpy as np
import torch

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





def       _UpperCamelCase       (a__				:Dict[str, torch.Tensor]    ):






							"""simple docstring"""





							UpperCamelCase__			       =     []
							UpperCamelCase__			       =     []
							UpperCamelCase__			       =     []

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

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

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

							UpperCamelCase__			       =     torch.tensor(
							    a__				,     dtype=torch.intaa				,     device=protein["""aatype"""].device				,     )
							UpperCamelCase__			       =     torch.tensor(
							    a__				,     dtype=torch.intaa				,     device=protein["""aatype"""].device				,     )
							UpperCamelCase__			       =     torch.tensor(
							    a__				,     dtype=torch.floataa				,     device=protein["""aatype"""].device				,     )
							UpperCamelCase__			       =     protein["""aatype"""].to(torch.long    )

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

							UpperCamelCase__			       =     residx_atomaa_mask
							UpperCamelCase__			       =     residx_atomaa_to_atomaa.long()

							# create the gather indices for mapping back
							UpperCamelCase__			       =     restype_atomaa_to_atomaa[protein_aatype]
							UpperCamelCase__			       =     residx_atomaa_to_atomaa.long()

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

							UpperCamelCase__			       =     restype_atomaa_mask[protein_aatype]
							UpperCamelCase__			       =     residx_atomaa_mask

							return protein





def       _UpperCamelCase       (a__				:Dict[str, torch.Tensor]    ):






							"""simple docstring"""





							UpperCamelCase__			       =     tree_map(lambda a__    : torch.tensor(a__				,     device=batch["""aatype"""].device    )				,     a__				,     np.ndarray    )
							UpperCamelCase__			       =     tensor_tree_map(lambda a__    : np.array(a__    )				,     make_atomaa_masks(a__    )    )
							return out




