


'''simple docstring'''



from typing import Optional

import numpy as np
import torch
from torch import nn
from transformers import GPTaConfig, GPTaLMHeadModel
from transformers.modeling_utils import ModuleUtilsMixin

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin


class 						_snake_case   (       a__    ,				a__    ,				a__		):
					lowerCAmelCase :Dict						=   [R'''h\.\d+\.attn\.bias''', R'''h\.\d+\.attn\.masked_bias''']




					@register_to_config
					def __init__( self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase = None							,							_lowerCamelCase = 5_0257							,							_lowerCamelCase = 1024							,							_lowerCamelCase = 768							,							_lowerCamelCase = 12							,							_lowerCamelCase = 12							,							_lowerCamelCase = None							,							_lowerCamelCase = "gelu_new"							,							_lowerCamelCase = 0.1							,							_lowerCamelCase = 0.1							,							_lowerCamelCase = 0.1							,							_lowerCamelCase = 1e-5							,							_lowerCamelCase = 0.02							,							_lowerCamelCase = True							,							_lowerCamelCase = True							,							_lowerCamelCase = False							,							_lowerCamelCase = False							,							):
								super().__init__()

								UpperCAmelCase__					:   str       				=       prefix_length

								if prefix_inner_dim != n_embd and prefix_hidden_dim is None:
											raise ValueError(
											    f'''`prefix_hidden_dim` cannot be `None` when `prefix_inner_dim`: {prefix_hidden_dim} and'''
											    f''' `n_embd`: {n_embd} are not equal.''')

								UpperCAmelCase__					:   int       				=       prefix_inner_dim
								UpperCAmelCase__					:   Optional[Any]       				=       prefix_hidden_dim

								UpperCAmelCase__					:   int       				=       (
								    nn.Linear(self.prefix_inner_dim							,							self.prefix_hidden_dim)
								    if self.prefix_hidden_dim is not None
								    else nn.Identity()
								)
								UpperCAmelCase__					:   Tuple       				=       (
								    nn.Linear(self.prefix_hidden_dim							,							_lowerCamelCase) if self.prefix_hidden_dim is not None else nn.Identity()
								)

								UpperCAmelCase__					:   Optional[Any]       				=       GPTaConfig(
								    vocab_size=_lowerCamelCase							,							n_positions=_lowerCamelCase							,							n_embd=_lowerCamelCase							,							n_layer=_lowerCamelCase							,							n_head=_lowerCamelCase							,							n_inner=_lowerCamelCase							,							activation_function=_lowerCamelCase							,							resid_pdrop=_lowerCamelCase							,							embd_pdrop=_lowerCamelCase							,							attn_pdrop=_lowerCamelCase							,							layer_norm_epsilon=_lowerCamelCase							,							initializer_range=_lowerCamelCase							,							scale_attn_weights=_lowerCamelCase							,							use_cache=_lowerCamelCase							,							scale_attn_by_inverse_layer_idx=_lowerCamelCase							,							reorder_and_upcast_attn=_lowerCamelCase							,							)
								UpperCAmelCase__					:   int       				=       GPTaLMHeadModel(_lowerCamelCase)




					def 			snake_case__  ( self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase = None							,							_lowerCamelCase = None							,							):
								UpperCAmelCase__					:   Union[str, Any]       				=       self.transformer.transformer.wte(_lowerCamelCase)
								UpperCAmelCase__					:   int       				=       self.encode_prefix(_lowerCamelCase)
								UpperCAmelCase__					:   List[str]       				=       self.decode_prefix(_lowerCamelCase)
								UpperCAmelCase__					:   str       				=       torch.cat((prefix_embeds, embedding_text)							,							dim=1)

								if labels is not None:
											UpperCAmelCase__					:   List[Any]       				=       self.get_dummy_token(input_ids.shape[0]							,							input_ids.device)
											UpperCAmelCase__					:   Optional[int]       				=       torch.cat((dummy_token, input_ids)							,							dim=1)
								UpperCAmelCase__					:   List[Any]       				=       self.transformer(inputs_embeds=_lowerCamelCase							,							labels=_lowerCamelCase							,							attention_mask=_lowerCamelCase)
								if self.prefix_hidden_dim is not None:
											return out, hidden
								else:
											return out




					def 			snake_case__  ( self							,							_lowerCamelCase							,							_lowerCamelCase):
								return torch.zeros(_lowerCamelCase							,							self.prefix_length							,							dtype=torch.intaa							,							device=_lowerCamelCase)




					def 			snake_case__  ( self							,							_lowerCamelCase):
								return self.encode_prefix(_lowerCamelCase)




					@torch.no_grad()
					def 			snake_case__  ( self							,							_lowerCamelCase							,							_lowerCamelCase							,							_lowerCamelCase):
								UpperCAmelCase__					:   int       				=       torch.split(_lowerCamelCase							,							1							,							dim=0)
								UpperCAmelCase__					:   Optional[int]       				=       []
								UpperCAmelCase__					:   Union[str, Any]       				=       []
								for feature in features:
											UpperCAmelCase__					:   Dict       				=       self.decode_prefix(feature.to(_lowerCamelCase))  # back to the clip feature
											# Only support beam search for now
											UpperCAmelCase__					:   int       				=       self.generate_beam(
											    input_embeds=_lowerCamelCase							,							device=_lowerCamelCase							,							eos_token_id=_lowerCamelCase)
											generated_tokens.append(output_tokens[0])
											generated_seq_lengths.append(seq_lengths[0])
								UpperCAmelCase__					:   Dict       				=       torch.stack(_lowerCamelCase)
								UpperCAmelCase__					:   Any       				=       torch.stack(_lowerCamelCase)
								return generated_tokens, generated_seq_lengths




					@torch.no_grad()
					def 			snake_case__  ( self							,							_lowerCamelCase=None							,							_lowerCamelCase=None							,							_lowerCamelCase=None							,							_lowerCamelCase = 5							,							_lowerCamelCase = 67							,							_lowerCamelCase = 1.0							,							_lowerCamelCase = None							,							):
								UpperCAmelCase__					:   Union[str, Any]       				=       eos_token_id
								UpperCAmelCase__					:   Dict       				=       None
								UpperCAmelCase__					:   Any       				=       None
								UpperCAmelCase__					:   Dict       				=       torch.ones(_lowerCamelCase							,							device=_lowerCamelCase							,							dtype=torch.int)
								UpperCAmelCase__					:   List[str]       				=       torch.zeros(_lowerCamelCase							,							device=_lowerCamelCase							,							dtype=torch.bool)

								if input_embeds is not None:
											UpperCAmelCase__					:   str       				=       input_embeds
								else:
											UpperCAmelCase__					:   int       				=       self.transformer.transformer.wte(_lowerCamelCase)

								for i in range(_lowerCamelCase):
											UpperCAmelCase__					:   str       				=       self.transformer(inputs_embeds=_lowerCamelCase)
											UpperCAmelCase__					:   Tuple       				=       outputs.logits
											UpperCAmelCase__					:   List[str]       				=       logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
											UpperCAmelCase__					:   Any       				=       logits.softmax(-1).log()

											if scores is None:
														UpperCAmelCase__					:   List[Any]       				=       logits.topk(_lowerCamelCase							,							-1)
														UpperCAmelCase__					:   int       				=       generated.expand(_lowerCamelCase							,							*generated.shape[1:])
														UpperCAmelCase__					:   str       				=       next_tokens.permute(1							,							0), scores.squeeze(0)
														if tokens is None:
																	UpperCAmelCase__					:   Union[str, Any]       				=       next_tokens
														else:
																	UpperCAmelCase__					:   List[Any]       				=       tokens.expand(_lowerCamelCase							,							*tokens.shape[1:])
																	UpperCAmelCase__					:   Optional[Any]       				=       torch.cat((tokens, next_tokens)							,							dim=1)
											else:
														UpperCAmelCase__					:   Optional[int]       				=       -float(np.inf)
														UpperCAmelCase__					:   Optional[int]       				=       0
														UpperCAmelCase__					:   Union[str, Any]       				=       scores[:, None] + logits
														seq_lengths[~is_stopped] += 1
														UpperCAmelCase__					:   Optional[Any]       				=       scores_sum / seq_lengths[:, None]
														UpperCAmelCase__					:   Dict       				=       scores_sum_average.view(-1).topk(_lowerCamelCase							,							-1)
														UpperCAmelCase__					:   Optional[int]       				=       next_tokens // scores_sum.shape[1]
														UpperCAmelCase__					:   int       				=       seq_lengths[next_tokens_source]
														UpperCAmelCase__					:   Union[str, Any]       				=       next_tokens % scores_sum.shape[1]
														UpperCAmelCase__					:   Any       				=       next_tokens.unsqueeze(1)
														UpperCAmelCase__					:   str       				=       tokens[next_tokens_source]
														UpperCAmelCase__					:   Dict       				=       torch.cat((tokens, next_tokens)							,							dim=1)
														UpperCAmelCase__					:   Union[str, Any]       				=       generated[next_tokens_source]
														UpperCAmelCase__					:   Tuple       				=       scores_sum_average * seq_lengths
														UpperCAmelCase__					:   int       				=       is_stopped[next_tokens_source]

											UpperCAmelCase__					:   List[str]       				=       self.transformer.transformer.wte(next_tokens.squeeze()).view(generated.shape[0]							,							1							,							-1)
											UpperCAmelCase__					:   Dict       				=       torch.cat((generated, next_token_embed)							,							dim=1)
											UpperCAmelCase__					:   List[Any]       				=       is_stopped + next_tokens.eq(_lowerCamelCase).squeeze()
											if is_stopped.all():
														break

								UpperCAmelCase__					:   Dict       				=       scores / seq_lengths
								UpperCAmelCase__					:   Dict       				=       scores.argsort(descending=_lowerCamelCase)
								# tokens tensors are already padded to max_seq_length
								UpperCAmelCase__					:   List[str]       				=       [tokens[i] for i in order]
								UpperCAmelCase__					:   Optional[Any]       				=       torch.stack(_lowerCamelCase							,							dim=0)
								UpperCAmelCase__					:   Tuple       				=       torch.tensor([seq_lengths[i] for i in order]							,							dtype=seq_lengths.dtype)
								return output_texts, seq_lengths



'''simple docstring'''



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

import sentencepiece as spm

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


__A		=logging.get_logger(__name__)

__A		={'vocab_file': 'spiece.model'}

__A		={
    'vocab_file': {
        'TsinghuaAI/CPM-Generate': 'https://huggingface.co/TsinghuaAI/CPM-Generate/resolve/main/spiece.model',
    }
}


class 						_snake_case   (       a__		):




					def __init__( self							,							_lowerCamelCase							,							_lowerCamelCase=False							,							_lowerCamelCase=True							,							_lowerCamelCase=False							,							_lowerCamelCase="<s>"							,							_lowerCamelCase="</s>"							,							_lowerCamelCase="<unk>"							,							_lowerCamelCase="<sep>"							,							_lowerCamelCase="<pad>"							,							_lowerCamelCase="<cls>"							,							_lowerCamelCase="<mask>"							,							_lowerCamelCase=["<eop>", "<eod>"]							,							_lowerCamelCase = None							,							**_lowerCamelCase							,							):
								UpperCAmelCase__					:   Dict       				=       AddedToken(_lowerCamelCase							,							lstrip=_lowerCamelCase							,							rstrip=_lowerCamelCase) if isinstance(_lowerCamelCase							,							_lowerCamelCase) else mask_token

								UpperCAmelCase__					:   Any       				=       {} if sp_model_kwargs is None else sp_model_kwargs

								super().__init__(
								    do_lower_case=_lowerCamelCase							,							remove_space=_lowerCamelCase							,							keep_accents=_lowerCamelCase							,							bos_token=_lowerCamelCase							,							eos_token=_lowerCamelCase							,							unk_token=_lowerCamelCase							,							sep_token=_lowerCamelCase							,							pad_token=_lowerCamelCase							,							cls_token=_lowerCamelCase							,							mask_token=_lowerCamelCase							,							additional_special_tokens=_lowerCamelCase							,							sp_model_kwargs=self.sp_model_kwargs							,							**_lowerCamelCase							,							)

								UpperCAmelCase__					:   List[str]       				=       3

								UpperCAmelCase__					:   str       				=       do_lower_case
								UpperCAmelCase__					:   Tuple       				=       remove_space
								UpperCAmelCase__					:   List[str]       				=       keep_accents
								UpperCAmelCase__					:   Any       				=       vocab_file

								UpperCAmelCase__					:   str       				=       spm.SentencePieceProcessor(**self.sp_model_kwargs)
								self.sp_model.Load(_lowerCamelCase)

								try:
											import jieba
								except ModuleNotFoundError as error:
											raise error.__class__(
											    """You need to install jieba to use CpmTokenizer or CpmTokenizerFast. """
											    """See https://pypi.org/project/jieba/ for installation.""")
								UpperCAmelCase__					:   Union[str, Any]       				=       jieba
								UpperCAmelCase__					:   int       				=       str.maketrans(""" \n"""							,							"""\u2582\u2583""")




					@property
					# Copied from transformers.models.xlnet.tokenization_xlnet.XLNetTokenizer.vocab_size
					def 			snake_case__  ( self):
								return len(self.sp_model)




					def 			snake_case__  ( self):
								UpperCAmelCase__					:   Optional[Any]       				=       {self.convert_ids_to_tokens(_lowerCamelCase): i for i in range(self.vocab_size)}
								vocab.update(self.added_tokens_encoder)
								return vocab




					def __getstate__( self):
								UpperCAmelCase__					:   List[str]       				=       self.__dict__.copy()
								UpperCAmelCase__					:   Optional[int]       				=       None
								return state




					def __setstate__( self							,							_lowerCamelCase):
								UpperCAmelCase__					:   Optional[Any]       				=       d

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

								UpperCAmelCase__					:   Union[str, Any]       				=       spm.SentencePieceProcessor(**self.sp_model_kwargs)
								self.sp_model.Load(self.vocab_file)




					def 			snake_case__  ( self							,							_lowerCamelCase):
								if self.remove_space:
											UpperCAmelCase__					:   Tuple       				=       """ """.join(inputs.strip().split())
								else:
											UpperCAmelCase__					:   str       				=       inputs
								UpperCAmelCase__					:   List[Any]       				=       outputs.replace("""``"""							,							"""\"""").replace("""''"""							,							"""\"""")

								if not self.keep_accents:
											UpperCAmelCase__					:   str       				=       unicodedata.normalize("""NFKD"""							,							_lowerCamelCase)
											UpperCAmelCase__					:   Any       				=       """""".join([c for c in outputs if not unicodedata.combining(_lowerCamelCase)])
								if self.do_lower_case:
											UpperCAmelCase__					:   Optional[int]       				=       outputs.lower()

								return outputs




					def 			snake_case__  ( self							,							_lowerCamelCase):
								UpperCAmelCase__					:   List[str]       				=       self.preprocess_text(_lowerCamelCase)
								UpperCAmelCase__					:   Dict       				=       self.sp_model.encode(_lowerCamelCase							,							out_type=_lowerCamelCase)
								UpperCAmelCase__					:   List[Any]       				=       []
								for piece in pieces:
											if len(_lowerCamelCase) > 1 and piece[-1] == str(""",""") and piece[-2].isdigit():
														UpperCAmelCase__					:   str       				=       self.sp_model.EncodeAsPieces(piece[:-1].replace(_lowerCamelCase							,							""""""))
														if piece[0] != SPIECE_UNDERLINE and cur_pieces[0][0] == SPIECE_UNDERLINE:
																	if len(cur_pieces[0]) == 1:
																				UpperCAmelCase__					:   Optional[int]       				=       cur_pieces[1:]
																	else:
																				UpperCAmelCase__					:   Optional[int]       				=       cur_pieces[0][1:]
														cur_pieces.append(piece[-1])
														new_pieces.extend(_lowerCamelCase)
											else:
														new_pieces.append(_lowerCamelCase)

								return new_pieces




					def 			snake_case__  ( self							,							_lowerCamelCase):
								return self.sp_model.PieceToId(_lowerCamelCase)




					def 			snake_case__  ( self							,							_lowerCamelCase):
								return self.sp_model.IdToPiece(_lowerCamelCase)




					def 			snake_case__  ( self							,							_lowerCamelCase):
								UpperCAmelCase__					:   str       				=       """""".join(_lowerCamelCase).replace(_lowerCamelCase							,							""" """).strip()
								return out_string




					def 			snake_case__  ( self							,							_lowerCamelCase							,							_lowerCamelCase = None):
								UpperCAmelCase__					:   Optional[Any]       				=       [self.sep_token_id]
								UpperCAmelCase__					:   Optional[int]       				=       [self.cls_token_id]
								if token_ids_a is None:
											return token_ids_a + sep + cls
								return token_ids_a + sep + token_ids_a + sep + cls




					def 			snake_case__  ( self							,							_lowerCamelCase							,							_lowerCamelCase = None							,							_lowerCamelCase = False):

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

								if token_ids_a is not None:
											return ([0] * len(_lowerCamelCase)) + [1] + ([0] * len(_lowerCamelCase)) + [1, 1]
								return ([0] * len(_lowerCamelCase)) + [1, 1]




					def 			snake_case__  ( self							,							_lowerCamelCase							,							_lowerCamelCase = None):
								UpperCAmelCase__					:   int       				=       [self.sep_token_id]
								UpperCAmelCase__					:   Tuple       				=       [2]

								if token_ids_a is None:
											return len(token_ids_a + sep) * [0] + cls_segment_id
								return len(token_ids_a + sep) * [0] + len(token_ids_a + sep) * [1] + cls_segment_id




					def 			snake_case__  ( self							,							_lowerCamelCase							,							_lowerCamelCase = None):
								if not os.path.isdir(_lowerCamelCase):
											logger.error(f'''Vocabulary path ({save_directory}) should be a directory''')
											return
								UpperCAmelCase__					:   str       				=       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__					:   str       				=       self.sp_model.serialized_model_proto()
														fi.write(_lowerCamelCase)

								return (out_vocab_file,)




					def 			snake_case__  ( self							,							*_lowerCamelCase							,							**_lowerCamelCase):
								UpperCAmelCase__					:   Union[str, Any]       				=       super()._decode(*_lowerCamelCase							,							**_lowerCamelCase)
								UpperCAmelCase__					:   Any       				=       text.replace(""" """							,							"""""").replace("""\u2582"""							,							""" """).replace("""\u2583"""							,							"""\n""")
								return text