"""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


__lowerCamelCase									=	logging.get_logger(__name__)

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

__lowerCamelCase									=	{
    "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__(      __A			):





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

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

													A__						  =				3
													A__						  =				do_lower_case
													A__						  =				remove_space
													A__						  =				keep_accents
													A__						  =				vocab_file
													A__						  =				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.'  )
													A__						  =				jieba
													A__						  =				str.maketrans(' \n'		,'\u2582\u2583'  )





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase = None  )       ->					List[int]:
													A__						  =				[self.sep_token_id]
													A__						  =				[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		,__UpperCAmelCase		,__UpperCAmelCase = None  )       ->					List[int]:
													A__						  =				[self.sep_token_id]
													A__						  =				[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		,__UpperCAmelCase		,__UpperCAmelCase = 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(__UpperCAmelCase  ):
																				logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''  )
																				return
													A__						  =				os.path.join(
													    __UpperCAmelCase		,(filename_prefix + '-' if filename_prefix else '') + VOCAB_FILES_NAMES['vocab_file']  )

													if os.path.abspath(self.vocab_file  ) != os.path.abspath(__UpperCAmelCase  ):
																				copyfile(self.vocab_file		,__UpperCAmelCase  )

													return (out_vocab_file,)





						def 					snake_case__						(							self		,__UpperCAmelCase		,*__UpperCAmelCase		,**__UpperCAmelCase  )       ->					Union[str, Any]:
													A__						  =				[
													    ' '.join([x.translate(self.translator  ) for x in self.jieba.cut(__UpperCAmelCase		,cut_all=__UpperCAmelCase  )]  )
													    for text in batch_text_or_text_pairs
													]
													return super()._batch_encode_plus(__UpperCAmelCase		,*__UpperCAmelCase		,**__UpperCAmelCase  )





						def 					snake_case__						(							self		,*__UpperCAmelCase		,**__UpperCAmelCase  )       ->					List[str]:
													A__						  =				super()._decode(*__UpperCAmelCase		,**__UpperCAmelCase  )
													A__						  =				text.replace(' '		,''  ).replace('\u2582'		,' '  ).replace('\u2583'		,'\n'  )
													return text






"""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


__lowerCamelCase									=	"block_records.npy"


__lowerCamelCase									=	logging.get_logger(__name__)






def 	UpperCAmelCase       (						UpperCamelCase__							,     UpperCamelCase__					):






							"""simple docstring"""





							import tensorflow.compat.va as tf

							A__						  =				tf.data.TFRecordDataset(UpperCamelCase__							,     buffer_size=512 * 1_024 * 1_024					)
							A__						  =				blocks_dataset.batch(UpperCamelCase__							,     drop_remainder=UpperCamelCase__					)
							A__						  =				next(blocks_dataset.take(1					).as_numpy_iterator()					)

							return np_record
class      UpperCamelCase__:





						def __init__(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase=2		,__UpperCAmelCase=10_00		,__UpperCAmelCase=1_00		,__UpperCAmelCase=10_00_00		,)       ->					Optional[int]:
													from scann.scann_ops.py.scann_ops_pybind import builder as Builder

													A__						  =				Builder(db=__UpperCAmelCase		,num_neighbors=__UpperCAmelCase		,distance_measure='dot_product'  )
													A__						  =				builder.tree(
													    num_leaves=__UpperCAmelCase		,num_leaves_to_search=__UpperCAmelCase		,training_sample_size=__UpperCAmelCase  )
													A__						  =				builder.score_ah(dimensions_per_block=__UpperCAmelCase  )

													A__						  =				builder.build()





						def 					snake_case__						(							self		,__UpperCAmelCase  )       ->					Tuple:
													A__  ,	A__						  =				self.searcher.search_batched(question_projection.detach().cpu()  )
													return retrieved_block_ids.astype('int64'  )






class      UpperCamelCase__:





						def __init__(							self		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					List[str]:
													super().__init__()
													A__						  =				block_records
													A__						  =				tokenizer





						def __call__(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase=None		,__UpperCAmelCase="pt"  )       ->					str:
													A__						  =				np.take(self.block_records		,indices=__UpperCAmelCase		,axis=0  )

													A__						  =				self.tokenizer.decode(question_input_ids[0]		,skip_special_tokens=__UpperCAmelCase  )

													A__						  =				[]
													A__						  =				[]
													for retrieved_block in retrieved_blocks:
																				text.append(__UpperCAmelCase  )
																				text_pair.append(retrieved_block.decode()  )

													A__						  =				self.tokenizer(
													    __UpperCAmelCase		,__UpperCAmelCase		,padding=__UpperCAmelCase		,truncation=__UpperCAmelCase		,return_special_tokens_mask=__UpperCAmelCase		,max_length=__UpperCAmelCase  )
													A__						  =				concat_inputs.convert_to_tensors(__UpperCAmelCase  )

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





						@classmethod
						def 					snake_case__						(							cls		,__UpperCAmelCase		,*__UpperCAmelCase		,**__UpperCAmelCase  )       ->					List[Any]:
													if os.path.isdir(__UpperCAmelCase  ):
																				A__						  =				os.path.join(__UpperCAmelCase		,_REALM_BLOCK_RECORDS_FILENAME  )
													else:
																				A__						  =				hf_hub_download(
																				    repo_id=__UpperCAmelCase		,filename=_REALM_BLOCK_RECORDS_FILENAME		,**__UpperCAmelCase  )
													A__						  =				np.load(__UpperCAmelCase		,allow_pickle=__UpperCAmelCase  )

													A__						  =				AutoTokenizer.from_pretrained(__UpperCAmelCase		,*__UpperCAmelCase		,**__UpperCAmelCase  )

													return cls(__UpperCAmelCase		,__UpperCAmelCase  )





						def 					snake_case__						(							self		,__UpperCAmelCase  )       ->					Tuple:
													# save block records
													np.save(os.path.join(__UpperCAmelCase		,_REALM_BLOCK_RECORDS_FILENAME  )		,self.block_records  )
													# save tokenizer
													self.tokenizer.save_pretrained(__UpperCAmelCase  )





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					List[Any]:
													A__						  =				[]
													A__						  =				[]
													A__						  =				[]
													A__						  =				0

													for input_id in concat_inputs.input_ids:
																				A__						  =				input_id.tolist()
																				# Check answers between two [SEP] tokens
																				A__						  =				input_id_list.index(self.tokenizer.sep_token_id  )
																				A__						  =				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		,__UpperCAmelCase  ):
																																		if answer[0] == input_id_list[idx]:
																																									if input_id_list[idx : idx + len(__UpperCAmelCase  )] == answer:
																																																start_pos[-1].append(__UpperCAmelCase  )
																																																end_pos[-1].append(idx + len(__UpperCAmelCase  ) - 1  )

																				if len(start_pos[-1]  ) == 0:
																											has_answers.append(__UpperCAmelCase  )
																				else:
																											has_answers.append(__UpperCAmelCase  )
																											if len(start_pos[-1]  ) > max_answers:
																																		A__						  =				len(start_pos[-1]  )

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





