


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


SCREAMING_SNAKE_CASE__				:    str         =    'block_records.npy'


SCREAMING_SNAKE_CASE__				:    str         =    logging.get_logger(__name__)




def       a__						(		snake_case__				:       str	,      snake_case__				:       int    ):
  import tensorflow.compat.va as tf

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

  return np_record


class 		_SCREAMING_SNAKE_CASE    :


     def __init__(     self       ,				A_       ,				A_       ,				A_=2       ,				A_=10_00       ,				A_=1_00       ,				A_=10_00_00       ,				):
       from scann.scann_ops.py.scann_ops_pybind import builder as Builder

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

       _UpperCAmelCase		:       Any       =    builder.build()




     def 					__snake_case(     self       ,				A_						):
       _UpperCAmelCase		:       List[Any]       =    self.searcher.search_batched(question_projection.detach().cpu()						)
       return retrieved_block_ids.astype("""int64"""						)


class 		_SCREAMING_SNAKE_CASE    :


     def __init__(     self       ,				A_       ,				A_						):
       super().__init__()
       _UpperCAmelCase		:       Any       =    block_records
       _UpperCAmelCase		:       Union[str, Any]       =    tokenizer


     def __call__(     self       ,				A_       ,				A_       ,				A_       ,				A_=None       ,				A_="pt"						):
       _UpperCAmelCase		:       List[Any]       =    np.take(self.block_records       ,				indices=A_       ,				axis=0						)

       _UpperCAmelCase		:       str       =    self.tokenizer.decode(question_input_ids[0]       ,				skip_special_tokens=A_						)

       _UpperCAmelCase		:       List[Any]       =    []
       _UpperCAmelCase		:       Any       =    []
       for retrieved_block in retrieved_blocks:
         text.append(A_						)
         text_pair.append(retrieved_block.decode()						)

       _UpperCAmelCase		:       List[str]       =    self.tokenizer(
           A_       ,				A_       ,				padding=A_       ,				truncation=A_       ,				return_special_tokens_mask=A_       ,				max_length=A_						)
       _UpperCAmelCase		:       List[Any]       =    concat_inputs.convert_to_tensors(A_						)

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


     @classmethod
     def 					__snake_case(     cls       ,				A_       ,				*A_       ,				**A_						):
       if os.path.isdir(A_						):
         _UpperCAmelCase		:       List[Any]       =    os.path.join(A_       ,				_REALM_BLOCK_RECORDS_FILENAME						)
       else:
         _UpperCAmelCase		:       int       =    hf_hub_download(
             repo_id=A_       ,				filename=_REALM_BLOCK_RECORDS_FILENAME       ,				**A_						)
       _UpperCAmelCase		:       int       =    np.load(A_       ,				allow_pickle=A_						)

       _UpperCAmelCase		:       Dict       =    AutoTokenizer.from_pretrained(A_       ,				*A_       ,				**A_						)

       return cls(A_       ,				A_						)


     def 					__snake_case(     self       ,				A_						):
       # save block records
       np.save(os.path.join(A_       ,				_REALM_BLOCK_RECORDS_FILENAME						)       ,				self.block_records						)
       # save tokenizer
       self.tokenizer.save_pretrained(A_						)




     def 					__snake_case(     self       ,				A_       ,				A_						):
       _UpperCAmelCase		:       Optional[Any]       =    []
       _UpperCAmelCase		:       Tuple       =    []
       _UpperCAmelCase		:       str       =    []
       _UpperCAmelCase		:       List[str]       =    0

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

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

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









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

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


if is_sentencepiece_available():
 from .tokenization_xlm_roberta import XLMRobertaTokenizer
else:
 SCREAMING_SNAKE_CASE__				:    Optional[Any]         =    None


SCREAMING_SNAKE_CASE__				:    Dict         =    logging.get_logger(__name__)

SCREAMING_SNAKE_CASE__				:    List[Any]         =    {'vocab_file': 'sentencepiece.bpe.model', 'tokenizer_file': 'tokenizer.json'}

SCREAMING_SNAKE_CASE__				:    Any         =    {
    'vocab_file': {
        'xlm-roberta-base': 'https://huggingface.co/xlm-roberta-base/resolve/main/sentencepiece.bpe.model',
        'xlm-roberta-large': 'https://huggingface.co/xlm-roberta-large/resolve/main/sentencepiece.bpe.model',
        'xlm-roberta-large-finetuned-conll02-dutch': (
            'https://huggingface.co/xlm-roberta-large-finetuned-conll02-dutch/resolve/main/sentencepiece.bpe.model'
        ),
        'xlm-roberta-large-finetuned-conll02-spanish': (
            'https://huggingface.co/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/sentencepiece.bpe.model'
        ),
        'xlm-roberta-large-finetuned-conll03-english': (
            'https://huggingface.co/xlm-roberta-large-finetuned-conll03-english/resolve/main/sentencepiece.bpe.model'
        ),
        'xlm-roberta-large-finetuned-conll03-german': (
            'https://huggingface.co/xlm-roberta-large-finetuned-conll03-german/resolve/main/sentencepiece.bpe.model'
        ),
    },
    'tokenizer_file': {
        'xlm-roberta-base': 'https://huggingface.co/xlm-roberta-base/resolve/main/tokenizer.json',
        'xlm-roberta-large': 'https://huggingface.co/xlm-roberta-large/resolve/main/tokenizer.json',
        'xlm-roberta-large-finetuned-conll02-dutch': (
            'https://huggingface.co/xlm-roberta-large-finetuned-conll02-dutch/resolve/main/tokenizer.json'
        ),
        'xlm-roberta-large-finetuned-conll02-spanish': (
            'https://huggingface.co/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/tokenizer.json'
        ),
        'xlm-roberta-large-finetuned-conll03-english': (
            'https://huggingface.co/xlm-roberta-large-finetuned-conll03-english/resolve/main/tokenizer.json'
        ),
        'xlm-roberta-large-finetuned-conll03-german': (
            'https://huggingface.co/xlm-roberta-large-finetuned-conll03-german/resolve/main/tokenizer.json'
        ),
    },
}

SCREAMING_SNAKE_CASE__				:    Optional[int]         =    {
    'xlm-roberta-base': 512,
    'xlm-roberta-large': 512,
    'xlm-roberta-large-finetuned-conll02-dutch': 512,
    'xlm-roberta-large-finetuned-conll02-spanish': 512,
    'xlm-roberta-large-finetuned-conll03-english': 512,
    'xlm-roberta-large-finetuned-conll03-german': 512,
}


class 		_SCREAMING_SNAKE_CASE    (	A    ):
     __SCREAMING_SNAKE_CASE    				=						VOCAB_FILES_NAMES
     __SCREAMING_SNAKE_CASE    				=						PRETRAINED_VOCAB_FILES_MAP
     __SCREAMING_SNAKE_CASE    				=						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
     __SCREAMING_SNAKE_CASE    				=						['''input_ids''', '''attention_mask''']
     __SCREAMING_SNAKE_CASE    				=						XLMRobertaTokenizer


     def __init__(     self       ,				A_=None       ,				A_=None       ,				A_="<s>"       ,				A_="</s>"       ,				A_="</s>"       ,				A_="<s>"       ,				A_="<unk>"       ,				A_="<pad>"       ,				A_="<mask>"       ,				**A_       ,				):
       # Mask token behave like a normal word, i.e. include the space before it
       _UpperCAmelCase		:       int       =    AddedToken(A_       ,				lstrip=A_       ,				rstrip=A_						) if isinstance(A_       ,				A_						) else mask_token

       super().__init__(
           A_       ,				tokenizer_file=A_       ,				bos_token=A_       ,				eos_token=A_       ,				sep_token=A_       ,				cls_token=A_       ,				unk_token=A_       ,				pad_token=A_       ,				mask_token=A_       ,				**A_       ,				)

       _UpperCAmelCase		:       Union[str, Any]       =    vocab_file
       _UpperCAmelCase		:       List[str]       =    False if not self.vocab_file else True


     def 					__snake_case(     self       ,				A_       ,				A_ = None						):

       if token_ids_a is None:
         return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
       _UpperCAmelCase		:       Optional[Any]       =    [self.cls_token_id]
       _UpperCAmelCase		:       Union[str, Any]       =    [self.sep_token_id]
       return cls + token_ids_a + sep + sep + token_ids_a + sep


     def 					__snake_case(     self       ,				A_       ,				A_ = None						):
       _UpperCAmelCase		:       Tuple       =    [self.sep_token_id]
       _UpperCAmelCase		:       Optional[int]       =    [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_ = None						):
       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		:       List[str]       =    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,)





