


"""simple docstring"""






import os
from typing import Optional, Union

import numpy as np
from huggingface_hub import hf_hub_download

from ... import AutoTokenizer
from ...utils import logging


lowercase  		=			"""block_records.npy"""


lowercase  		=			logging.get_logger(__name__)



def 			A__		(					_UpperCAmelCase    :      str				,    _UpperCAmelCase    :      int	)							->	np.ndarray:

	'''simple docstring'''


	import tensorflow.compat.va as tf

	snake_case__  :						str		=      tf.data.TFRecordDataset(_UpperCAmelCase				,    buffer_size=5_12 * 10_24 * 10_24	)
	snake_case__  :						str		=      blocks_dataset.batch(_UpperCAmelCase				,    drop_remainder=_UpperCAmelCase	)
	snake_case__  :						List[Any]		=      next(blocks_dataset.take(1	).as_numpy_iterator()	)

	return np_record


class 			SCREAMING_SNAKE_CASE_  :




							'''simple docstring'''





							def __init__(						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__=2	,					lowerCamelCase__=1_000	,					lowerCamelCase__=100	,					lowerCamelCase__=100_000	,					)   ->	str:







								'''simple docstring'''




								from scann.scann_ops.py.scann_ops_pybind import builder as Builder

								snake_case__  :						Optional[Any]		=      Builder(db=lowerCamelCase__	,					num_neighbors=lowerCamelCase__	,					distance_measure="dot_product")
								snake_case__  :						Optional[int]		=      builder.tree(
								    num_leaves=lowerCamelCase__	,					num_leaves_to_search=lowerCamelCase__	,					training_sample_size=lowerCamelCase__)
								snake_case__  :						str		=      builder.score_ah(dimensions_per_block=lowerCamelCase__)

								snake_case__  :						Union[str, Any]		=      builder.build()







							def        UpperCAmelCase (						self	,					lowerCamelCase__)   ->	str:







								'''simple docstring'''




								snake_case__,					snake_case__  :						str		=      self.searcher.search_batched(question_projection.detach().cpu())
								return retrieved_block_ids.astype("int64")


class 			SCREAMING_SNAKE_CASE_  :




							'''simple docstring'''





							def __init__(						self	,					lowerCamelCase__	,					lowerCamelCase__)   ->	str:







								'''simple docstring'''




								super().__init__()
								snake_case__  :						Tuple		=      block_records
								snake_case__  :						int		=      tokenizer




							def __call__(						self	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__	,					lowerCamelCase__=None	,					lowerCamelCase__="pt")   ->	Optional[Any]:







								'''simple docstring'''




								snake_case__  :						str		=      np.take(self.block_records	,					indices=lowerCamelCase__	,					axis=0)

								snake_case__  :						Dict		=      self.tokenizer.decode(question_input_ids[0]	,					skip_special_tokens=lowerCamelCase__)

								snake_case__  :						List[Any]		=      []
								snake_case__  :						Optional[int]		=      []
								for retrieved_block in retrieved_blocks:
									text.append(lowerCamelCase__)
									text_pair.append(retrieved_block.decode())

								snake_case__  :						str		=      self.tokenizer(
								    lowerCamelCase__	,					lowerCamelCase__	,					padding=lowerCamelCase__	,					truncation=lowerCamelCase__	,					return_special_tokens_mask=lowerCamelCase__	,					max_length=lowerCamelCase__)
								snake_case__  :						Tuple		=      concat_inputs.convert_to_tensors(lowerCamelCase__)

								if answer_ids is not None:
									return self.block_has_answer(lowerCamelCase__	,					lowerCamelCase__) + (concat_inputs_tensors,)
								else:
									return (None, None, None, concat_inputs_tensors)




							@classmethod
							def        UpperCAmelCase (						cls	,					lowerCamelCase__	,					*lowerCamelCase__	,					**lowerCamelCase__)   ->	Dict:







								'''simple docstring'''




								if os.path.isdir(lowerCamelCase__):
									snake_case__  :						Tuple		=      os.path.join(lowerCamelCase__	,					_REALM_BLOCK_RECORDS_FILENAME)
								else:
									snake_case__  :						Any		=      hf_hub_download(
									    repo_id=lowerCamelCase__	,					filename=_REALM_BLOCK_RECORDS_FILENAME	,					**lowerCamelCase__)
								snake_case__  :						int		=      np.load(lowerCamelCase__	,					allow_pickle=lowerCamelCase__)

								snake_case__  :						Union[str, Any]		=      AutoTokenizer.from_pretrained(lowerCamelCase__	,					*lowerCamelCase__	,					**lowerCamelCase__)

								return cls(lowerCamelCase__	,					lowerCamelCase__)




							def        UpperCAmelCase (						self	,					lowerCamelCase__)   ->	Tuple:







								'''simple docstring'''




								np.save(os.path.join(lowerCamelCase__	,					_REALM_BLOCK_RECORDS_FILENAME)	,					self.block_records)
								# save tokenizer
								self.tokenizer.save_pretrained(lowerCamelCase__)







							def        UpperCAmelCase (						self	,					lowerCamelCase__	,					lowerCamelCase__)   ->	int:







								'''simple docstring'''




								snake_case__  :						Optional[Any]		=      []
								snake_case__  :						int		=      []
								snake_case__  :						int		=      []
								snake_case__  :						Any		=      0

								for input_id in concat_inputs.input_ids:
									snake_case__  :						List[Any]		=      input_id.tolist()
									# Check answers between two [SEP] tokens
									snake_case__  :						Optional[int]		=      input_id_list.index(self.tokenizer.sep_token_id)
									snake_case__  :						Optional[int]		=      first_sep_idx + 1 + input_id_list[first_sep_idx + 1 :].index(self.tokenizer.sep_token_id)

									start_pos.append([])
									end_pos.append([])
									for answer in answer_ids:
										for idx in range(first_sep_idx + 1	,					lowerCamelCase__):
											if answer[0] == input_id_list[idx]:
												if input_id_list[idx : idx + len(lowerCamelCase__)] == answer:
													start_pos[-1].append(lowerCamelCase__)
													end_pos[-1].append(idx + len(lowerCamelCase__) - 1)

									if len(start_pos[-1]) == 0:
										has_answers.append(lowerCamelCase__)
									else:
										has_answers.append(lowerCamelCase__)
										if len(start_pos[-1]) > max_answers:
											snake_case__  :						List[str]		=      len(start_pos[-1])

        # Pad -1 to max_answers
								for start_pos_, end_pos_ in zip(lowerCamelCase__	,					lowerCamelCase__):
									if len(start_pos_) < max_answers:
										snake_case__  :						Union[str, Any]		=      [-1] * (max_answers - len(start_pos_))
										start_pos_ += padded
										end_pos_ += padded
								return has_answers, start_pos, end_pos










'''simple docstring'''




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

from ...tokenization_utils_fast import AddedToken, PreTrainedTokenizerFast
from ...utils import logging


a  :     Any       =       logging.get_logger(__name__)

a  :     int       =       {"""vocab_file""": """spiece.model""", """tokenizer_file""": """tokenizer.json"""}

a  :     List[Any]       =       {
    """vocab_file""": {
        """TsinghuaAI/CPM-Generate""": """https://huggingface.co/TsinghuaAI/CPM-Generate/resolve/main/spiece.model""",
    },
    """tokenizer_file""": {
        """TsinghuaAI/CPM-Generate""": """https://huggingface.co/TsinghuaAI/CPM-Generate/resolve/main/tokenizer.json""",
    },
}

class      UpperCamelCase_					(				__magic_name__     ):





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

								super().__init__(
								    vocab_file=A						,						tokenizer_file=A						,						do_lower_case=A						,						remove_space=A						,						keep_accents=A						,						bos_token=A						,						eos_token=A						,						unk_token=A						,						sep_token=A						,						pad_token=A						,						cls_token=A						,						mask_token=A						,						additional_special_tokens=A						,						**A						,						)

								UpperCAmelCase	:						List[Any]					=						3
								UpperCAmelCase	:						Tuple					=						do_lower_case
								UpperCAmelCase	:						Union[str, Any]					=						remove_space
								UpperCAmelCase	:						Tuple					=						keep_accents
								UpperCAmelCase	:						Any					=						vocab_file
								UpperCAmelCase	:						Tuple					=						False if not self.vocab_file else True

								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	:						int					=						jieba
								UpperCAmelCase	:						Optional[int]					=						str.maketrans(""" \n"""						,						"""\u2582\u2583"""       )





	def 					_lowercase(    self						,						A						,						A = None       )     ->     List[int]:
								UpperCAmelCase	:						Optional[int]					=						[self.sep_token_id]
								UpperCAmelCase	:						str					=						[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 					_lowercase(    self						,						A						,						A = None       )     ->     List[int]:
								UpperCAmelCase	:						Union[str, Any]					=						[self.sep_token_id]
								UpperCAmelCase	:						str					=						[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 					_lowercase(    self						,						A						,						A = None       )     ->     Tuple[str]:
								if not self.can_save_slow_tokenizer:
															raise ValueError(
															    """Your fast tokenizer does not have the necessary information to save the vocabulary for a slow """
															    """tokenizer."""       )

								if not os.path.isdir(A       ):
															logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''       )
															return
								UpperCAmelCase	:						Tuple					=						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       ):
															copyfile(self.vocab_file						,						A       )

								return (out_vocab_file,)





	def 					_lowercase(    self						,						A						,						*A						,						**A       )     ->     List[Any]:
								UpperCAmelCase	:						Optional[Any]					=						[
								    """ """.join([x.translate(self.translator       ) for x in self.jieba.cut(A						,						cut_all=A       )]       )
								    for text in batch_text_or_text_pairs
								]
								return super()._batch_encode_plus(A						,						*A						,						**A       )






	def 					_lowercase(    self						,						*A						,						**A       )     ->     str:
								UpperCAmelCase	:						Optional[int]					=						super()._decode(*A						,						**A       )
								UpperCAmelCase	:						int					=						text.replace(""" """						,						""""""       ).replace("""\u2582"""						,						""" """       ).replace("""\u2583"""						,						"""\n"""       )
								return text

