



'''simple docstring'''






import json
import os
from functools import lru_cache
from typing import List, Optional, Tuple

import regex as re

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


A_				:					Dict							=      logging.get_logger(__name__)


A_				:					List[str]							=      {"""vocab_file""": """vocab.json""", """merges_file""": """merges.txt"""}

# See all MVP models at https://huggingface.co/models?filter=mvp
A_				:					Tuple							=      {
    """vocab_file""": {
        """RUCAIBox/mvp""": """https://huggingface.co/RUCAIBox/mvp/resolve/main/vocab.json""",
    },
    """added_tokens.json""": {
        """RUCAIBox/mvp""": """https://huggingface.co/RUCAIBox/mvp/resolve/main/added_tokens.json""",
    },
    """merges_file""": {
        """RUCAIBox/mvp""": """https://huggingface.co/RUCAIBox/mvp/resolve/main/merges.txt""",
    },
}

A_				:					Dict							=      {
    """RUCAIBox/mvp""": 1_0_2_4,
}




@lru_cache()
def        snake_case_   (	)->						int:





       '''simple docstring'''


       _UpperCAmelCase    :		Union[str, Any]     	=					(
           list(range(ord("""!"""						)			,  ord("""~"""						) + 1						)						) + list(range(ord("""¡"""						)			,  ord("""¬"""						) + 1						)						) + list(range(ord("""®"""						)			,  ord("""ÿ"""						) + 1						)						)
       )
       _UpperCAmelCase    :		List[str]     	=					bs[:]
       _UpperCAmelCase    :		Optional[Any]     	=					0
       for b in range(2**8						):
              if b not in bs:
                     bs.append(lowerCAmelCase_						)
                     cs.append(2**8 + n						)
                     n += 1
       _UpperCAmelCase    :		Any     	=					[chr(lowerCAmelCase_						) for n in cs]
       return dict(zip(lowerCAmelCase_			,  lowerCAmelCase_						)						)




def        snake_case_   (	lowerCAmelCase_						)->						str:





       '''simple docstring'''


       _UpperCAmelCase    :		Dict     	=					set()
       _UpperCAmelCase    :		int     	=					word[0]
       for char in word[1:]:
              pairs.add((prev_char, char)						)
              _UpperCAmelCase    :		Union[str, Any]     	=					char
       return pairs



class 				lowercase			(    _lowerCamelCase							):


 """simple docstring"""


 UpperCAmelCase		      =    VOCAB_FILES_NAMES
 UpperCAmelCase		      =    PRETRAINED_VOCAB_FILES_MAP
 UpperCAmelCase		      =    PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 UpperCAmelCase		      =    ["""input_ids""", """attention_mask"""]







 def __init__(				self    ,a_    ,a_    ,a_="replace"    ,a_="<s>"    ,a_="</s>"    ,a_="</s>"    ,a_="<s>"    ,a_="<unk>"    ,a_="<pad>"    ,a_="<mask>"    ,a_=False    ,**a_    ,)						->  Any:
        _UpperCAmelCase    :		Optional[Any]     	=					AddedToken(a_    ,lstrip=a_    ,rstrip=a_				) if isinstance(a_    ,a_				) else bos_token
        _UpperCAmelCase    :		List[Any]     	=					AddedToken(a_    ,lstrip=a_    ,rstrip=a_				) if isinstance(a_    ,a_				) else eos_token
        _UpperCAmelCase    :		Union[str, Any]     	=					AddedToken(a_    ,lstrip=a_    ,rstrip=a_				) if isinstance(a_    ,a_				) else sep_token
        _UpperCAmelCase    :		List[str]     	=					AddedToken(a_    ,lstrip=a_    ,rstrip=a_				) if isinstance(a_    ,a_				) else cls_token
        _UpperCAmelCase    :		Any     	=					AddedToken(a_    ,lstrip=a_    ,rstrip=a_				) if isinstance(a_    ,a_				) else unk_token
        _UpperCAmelCase    :		Optional[int]     	=					AddedToken(a_    ,lstrip=a_    ,rstrip=a_				) if isinstance(a_    ,a_				) else pad_token

        # Mask token behave like a normal word, i.e. include the space before it
        _UpperCAmelCase    :		str     	=					AddedToken(a_    ,lstrip=a_    ,rstrip=a_				) if isinstance(a_    ,a_				) else mask_token

        super().__init__(
            errors=a_    ,bos_token=a_    ,eos_token=a_    ,unk_token=a_    ,sep_token=a_    ,cls_token=a_    ,pad_token=a_    ,mask_token=a_    ,add_prefix_space=a_    ,**a_    ,)
        with open(a_    ,encoding="""utf-8"""				) as vocab_handle:
               _UpperCAmelCase    :		Any     	=					json.load(a_				)
        _UpperCAmelCase    :		Optional[Any]     	=					{v: k for k, v in self.encoder.items()}
        _UpperCAmelCase    :		Any     	=					errors  # how to handle errors in decoding
        _UpperCAmelCase    :		Optional[Any]     	=					bytes_to_unicode()
        _UpperCAmelCase    :		Optional[Any]     	=					{v: k for k, v in self.byte_encoder.items()}
        with open(a_    ,encoding="""utf-8"""				) as merges_handle:
               _UpperCAmelCase    :		Optional[int]     	=					merges_handle.read().split("""\n"""				)[1:-1]
        _UpperCAmelCase    :		List[str]     	=					[tuple(merge.split()				) for merge in bpe_merges]
        _UpperCAmelCase    :		Optional[Any]     	=					dict(zip(a_    ,range(len(a_				)				)				)				)
        _UpperCAmelCase    :		List[Any]     	=					{}
        _UpperCAmelCase    :		Union[str, Any]     	=					add_prefix_space

        # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
        _UpperCAmelCase    :		Optional[Any]     	=					re.compile(r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""				)







 @property
 def _snake_case					(				self				)						->  str:
        return len(self.encoder				)







 def _snake_case					(				self				)						->  Dict:
        return dict(self.encoder    ,**self.added_tokens_encoder				)







 def _snake_case					(				self    ,a_				)						->  Optional[Any]:
        if token in self.cache:
               return self.cache[token]
        _UpperCAmelCase    :		Any     	=					tuple(a_				)
        _UpperCAmelCase    :		int     	=					get_pairs(a_				)

        if not pairs:
               return token

        while True:
               _UpperCAmelCase    :		Optional[Any]     	=					min(a_    ,key=lambda a_				: self.bpe_ranks.get(a_    ,float("""inf"""				)				)				)
               if bigram not in self.bpe_ranks:
                      break
               _UpperCAmelCase	,_UpperCAmelCase    :		Optional[Any]     	=					bigram
               _UpperCAmelCase    :		Optional[Any]     	=					[]
               _UpperCAmelCase    :		Optional[int]     	=					0
               while i < len(a_				):
                      try:
                             _UpperCAmelCase    :		int     	=					word.index(a_    ,a_				)
                      except ValueError:
                             new_word.extend(word[i:]				)
                             break
                      else:
                             new_word.extend(word[i:j]				)
                             _UpperCAmelCase    :		List[Any]     	=					j

                      if word[i] == first and i < len(a_				) - 1 and word[i + 1] == second:
                             new_word.append(first + second				)
                             i += 2
                      else:
                             new_word.append(word[i]				)
                             i += 1
               _UpperCAmelCase    :		str     	=					tuple(a_				)
               _UpperCAmelCase    :		List[str]     	=					new_word
               if len(a_				) == 1:
                      break
               else:
                      _UpperCAmelCase    :		Any     	=					get_pairs(a_				)
        _UpperCAmelCase    :		List[str]     	=					""" """.join(a_				)
        _UpperCAmelCase    :		List[Any]     	=					word
        return word







 def _snake_case					(				self    ,a_				)						->  Optional[int]:
        _UpperCAmelCase    :		str     	=					[]
        for token in re.findall(self.pat    ,a_				):
               _UpperCAmelCase    :		List[str]     	=					"""""".join(
                   self.byte_encoder[b] for b in token.encode("""utf-8"""				)				)  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
               bpe_tokens.extend(bpe_token for bpe_token in self.bpe(a_				).split(""" """				)				)
        return bpe_tokens







 def _snake_case					(				self    ,a_				)						->  Union[str, Any]:
        return self.encoder.get(a_    ,self.encoder.get(self.unk_token				)				)







 def _snake_case					(				self    ,a_				)						->  Optional[int]:
        return self.decoder.get(a_				)







 def _snake_case					(				self    ,a_				)						->  str:
        _UpperCAmelCase    :		Dict     	=					"""""".join(a_				)
        _UpperCAmelCase    :		str     	=					bytearray([self.byte_decoder[c] for c in text]				).decode("""utf-8"""    ,errors=self.errors				)
        return text







 def _snake_case					(				self    ,a_    ,a_ = None				)						->  Tuple[str]:
        if not os.path.isdir(a_				):
               logger.error(f'''Vocabulary path ({save_directory}) should be a directory'''				)
               return
        _UpperCAmelCase    :		Union[str, Any]     	=					os.path.join(
            a_    ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]				)
        _UpperCAmelCase    :		Any     	=					os.path.join(
            a_    ,(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""merges_file"""]				)

        with open(a_    ,"""w"""    ,encoding="""utf-8"""				) as f:
               f.write(json.dumps(self.encoder    ,indent=2    ,sort_keys=a_    ,ensure_ascii=a_				) + """\n"""				)

        _UpperCAmelCase    :		Optional[Any]     	=					0
        with open(a_    ,"""w"""    ,encoding="""utf-8"""				) as writer:
               writer.write("""#version: 0.2\n"""				)
               for bpe_tokens, token_index in sorted(self.bpe_ranks.items()    ,key=lambda a_				: kv[1]				):
                      if index != token_index:
                             logger.warning(
                                 f'''Saving vocabulary to {merge_file}: BPE merge indices are not consecutive.'''
                                 """ Please check that the tokenizer is not corrupted!"""				)
                             _UpperCAmelCase    :		Optional[Any]     	=					token_index
                      writer.write(""" """.join(a_				) + """\n"""				)
                      index += 1

        return vocab_file, merge_file







 def _snake_case					(				self    ,a_    ,a_ = None				)						->  List[int]:
        if token_ids_a is None:
               return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
        _UpperCAmelCase    :		Union[str, Any]     	=					[self.cls_token_id]
        _UpperCAmelCase    :		Optional[int]     	=					[self.sep_token_id]
        return cls + token_ids_a + sep + sep + token_ids_a + sep







 def _snake_case					(				self    ,a_    ,a_ = None    ,a_ = False				)						->  List[int]:
        if already_has_special_tokens:
               return super().get_special_tokens_mask(
                   token_ids_a=a_    ,token_ids_a=a_    ,already_has_special_tokens=a_				)

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







 def _snake_case					(				self    ,a_    ,a_ = None				)						->  List[int]:
        _UpperCAmelCase    :		str     	=					[self.sep_token_id]
        _UpperCAmelCase    :		Union[str, Any]     	=					[self.cls_token_id]

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







 def _snake_case					(				self    ,a_    ,a_=False    ,**a_				)						->  int:
        _UpperCAmelCase    :		Any     	=					kwargs.pop("""add_prefix_space"""    ,self.add_prefix_space				)
        if (is_split_into_words or add_prefix_space) and (len(a_				) > 0 and not text[0].isspace()):
               _UpperCAmelCase    :		Dict     	=					""" """ + text
        return (text, kwargs)





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


_A				=					'block_records.npy'


_A				=					logging.get_logger(__name__)
def 						UpperCAmelCase   (						a_,		a_     ):





		'''simple docstring'''


		import tensorflow.compat.va as tf

		lowerCamelCase	:     Union[str, Any]            =       tf.data.TFRecordDataset(a_,		buffer_size=512 * 1024 * 1024     )
		lowerCamelCase	:     Union[str, Any]            =       blocks_dataset.batch(a_,		drop_remainder=a_     )
		lowerCamelCase	:     Union[str, Any]            =       next(blocks_dataset.take(1     ).as_numpy_iterator()     )

		return np_record






class 			_lowercase  :




	def __init__(					self						,					UpperCAmelCase_						,					UpperCAmelCase_						,					UpperCAmelCase_=2						,					UpperCAmelCase_=1000						,					UpperCAmelCase_=100						,					UpperCAmelCase_=100000						,					)						->				List[Any]:
			from scann.scann_ops.py.scann_ops_pybind import builder as Builder

			lowerCamelCase	:     int            =       Builder(db=UpperCAmelCase_						,					num_neighbors=UpperCAmelCase_						,					distance_measure='dot_product'   )
			lowerCamelCase	:     Optional[int]            =       builder.tree(
			    num_leaves=UpperCAmelCase_						,					num_leaves_to_search=UpperCAmelCase_						,					training_sample_size=UpperCAmelCase_   )
			lowerCamelCase	:     Union[str, Any]            =       builder.score_ah(dimensions_per_block=UpperCAmelCase_   )

			lowerCamelCase	:     List[Any]            =       builder.build()




	def      _UpperCamelCase       (					self						,					UpperCAmelCase_   )						->				Optional[int]:
			lowerCamelCase      ,       lowerCamelCase	:     List[Any]            =       self.searcher.search_batched(question_projection.detach().cpu()   )
			return retrieved_block_ids.astype('int64'   )






class 			_lowercase  :




	def __init__(					self						,					UpperCAmelCase_						,					UpperCAmelCase_   )						->				Optional[Any]:
			super().__init__()
			lowerCamelCase	:     Tuple            =       block_records
			lowerCamelCase	:     Any            =       tokenizer




	def __call__(					self						,					UpperCAmelCase_						,					UpperCAmelCase_						,					UpperCAmelCase_						,					UpperCAmelCase_=None						,					UpperCAmelCase_="pt"   )						->				Tuple:
			lowerCamelCase	:     List[str]            =       np.take(self.block_records						,					indices=UpperCAmelCase_						,					axis=0   )

			lowerCamelCase	:     Any            =       self.tokenizer.decode(question_input_ids[0]						,					skip_special_tokens=UpperCAmelCase_   )

			lowerCamelCase	:     int            =       []
			lowerCamelCase	:     Optional[int]            =       []
			for retrieved_block in retrieved_blocks:
					text.append(UpperCAmelCase_   )
					text_pair.append(retrieved_block.decode()   )

			lowerCamelCase	:     Union[str, Any]            =       self.tokenizer(
			    UpperCAmelCase_						,					UpperCAmelCase_						,					padding=UpperCAmelCase_						,					truncation=UpperCAmelCase_						,					return_special_tokens_mask=UpperCAmelCase_						,					max_length=UpperCAmelCase_   )
			lowerCamelCase	:     str            =       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      _UpperCamelCase       (					cls						,					UpperCAmelCase_						,					*UpperCAmelCase_						,					**UpperCAmelCase_   )						->				Dict:
			if os.path.isdir(UpperCAmelCase_   ):
					lowerCamelCase	:     str            =       os.path.join(UpperCAmelCase_						,					_REALM_BLOCK_RECORDS_FILENAME   )
			else:
					lowerCamelCase	:     Tuple            =       hf_hub_download(
					    repo_id=UpperCAmelCase_						,					filename=_REALM_BLOCK_RECORDS_FILENAME						,					**UpperCAmelCase_   )
			lowerCamelCase	:     Dict            =       np.load(UpperCAmelCase_						,					allow_pickle=UpperCAmelCase_   )

			lowerCamelCase	:     Union[str, Any]            =       AutoTokenizer.from_pretrained(UpperCAmelCase_						,					*UpperCAmelCase_						,					**UpperCAmelCase_   )

			return cls(UpperCAmelCase_						,					UpperCAmelCase_   )




	def      _UpperCamelCase       (					self						,					UpperCAmelCase_   )						->				Union[str, Any]:
			# save block records
			np.save(os.path.join(UpperCAmelCase_						,					_REALM_BLOCK_RECORDS_FILENAME   )						,					self.block_records   )
			# save tokenizer
			self.tokenizer.save_pretrained(UpperCAmelCase_   )




	def      _UpperCamelCase       (					self						,					UpperCAmelCase_						,					UpperCAmelCase_   )						->				Tuple:
			lowerCamelCase	:     List[Any]            =       []
			lowerCamelCase	:     Any            =       []
			lowerCamelCase	:     int            =       []
			lowerCamelCase	:     Union[str, Any]            =       0

			for input_id in concat_inputs.input_ids:
					lowerCamelCase	:     Optional[Any]            =       input_id.tolist()
					# Check answers between two [SEP] tokens
					lowerCamelCase	:     Tuple            =       input_id_list.index(self.tokenizer.sep_token_id   )
					lowerCamelCase	:     Any            =       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:
									lowerCamelCase	:     List[str]            =       len(start_pos[-1]   )

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






