




'''simple docstring'''



import unittest
import warnings

from datasets.utils import experimental







@experimental
def 	_UpperCamelCase       ()->      Any:




     '''simple docstring'''





     return "success"







class 			lowerCAmelCase      (       unittest.TestCase):


       def 					lowerCAmelCase					(      self     )					->     Optional[Any]:



            '''simple docstring'''




            with warnings.catch_warnings(record=__SCREAMING_SNAKE_CASE     ) as w:
                 warnings.simplefilter('''always'''     )
                 self.assertEqual(dummy_function()			,     '''success'''     )
            self.assertEqual(len(__SCREAMING_SNAKE_CASE     )			,     1     )












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


UpperCAmelCase_				:     Optional[int]     			=						'''block_records.npy'''


UpperCAmelCase_				:     Any     			=						logging.get_logger(__name__)







def 	_UpperCamelCase       (_lowerCamelCase							:   str      ,		_lowerCamelCase							:   int		)->      np.ndarray:




     '''simple docstring'''





     import tensorflow.compat.va as tf

     __snake_case			    =     tf.data.TFRecordDataset(_lowerCamelCase      ,		buffer_size=5_12 * 10_24 * 10_24		)
     __snake_case			    =     blocks_dataset.batch(_lowerCamelCase      ,		drop_remainder=_lowerCamelCase		)
     __snake_case			    =     next(blocks_dataset.take(1		).as_numpy_iterator()		)

     return np_record





class 			lowerCAmelCase      :


       def __init__(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE=2			,     __SCREAMING_SNAKE_CASE=1000			,     __SCREAMING_SNAKE_CASE=100			,     __SCREAMING_SNAKE_CASE=10_0000			,     )					->     Optional[int]:



            '''simple docstring'''




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

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

            __snake_case			    =     builder.build()




       def 					lowerCAmelCase					(      self			,     __SCREAMING_SNAKE_CASE     )					->     Any:



            '''simple docstring'''




            __snake_case    ,			__snake_case			    =     self.searcher.search_batched(question_projection.detach().cpu()     )
            return retrieved_block_ids.astype('''int64'''     )





class 			lowerCAmelCase      :


       def __init__(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE     )					->     List[Any]:



            '''simple docstring'''




            super().__init__()
            __snake_case			    =     block_records
            __snake_case			    =     tokenizer


       def __call__(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE=None			,     __SCREAMING_SNAKE_CASE="pt"     )					->     Union[str, Any]:



            '''simple docstring'''




            __snake_case			    =     np.take(self.block_records			,     indices=__SCREAMING_SNAKE_CASE			,     axis=0     )

            __snake_case			    =     self.tokenizer.decode(question_input_ids[0]			,     skip_special_tokens=__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     []
            __snake_case			    =     []
            for retrieved_block in retrieved_blocks:
                 text.append(__SCREAMING_SNAKE_CASE     )
                 text_pair.append(retrieved_block.decode()     )

            __snake_case			    =     self.tokenizer(
                __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE			,     padding=__SCREAMING_SNAKE_CASE			,     truncation=__SCREAMING_SNAKE_CASE			,     return_special_tokens_mask=__SCREAMING_SNAKE_CASE			,     max_length=__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     concat_inputs.convert_to_tensors(__SCREAMING_SNAKE_CASE     )

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


       @classmethod
       def 					lowerCAmelCase					(      cls			,     __SCREAMING_SNAKE_CASE			,     *__SCREAMING_SNAKE_CASE			,     **__SCREAMING_SNAKE_CASE     )					->     Optional[int]:



            '''simple docstring'''




            if os.path.isdir(__SCREAMING_SNAKE_CASE     ):
                 __snake_case			    =     os.path.join(__SCREAMING_SNAKE_CASE			,     _REALM_BLOCK_RECORDS_FILENAME     )
            else:
                 __snake_case			    =     hf_hub_download(
                     repo_id=__SCREAMING_SNAKE_CASE			,     filename=_REALM_BLOCK_RECORDS_FILENAME			,     **__SCREAMING_SNAKE_CASE     )
            __snake_case			    =     np.load(__SCREAMING_SNAKE_CASE			,     allow_pickle=__SCREAMING_SNAKE_CASE     )

            __snake_case			    =     AutoTokenizer.from_pretrained(__SCREAMING_SNAKE_CASE			,     *__SCREAMING_SNAKE_CASE			,     **__SCREAMING_SNAKE_CASE     )

            return cls(__SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE     )


       def 					lowerCAmelCase					(      self			,     __SCREAMING_SNAKE_CASE     )					->     Any:



            '''simple docstring'''




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




       def 					lowerCAmelCase					(      self			,     __SCREAMING_SNAKE_CASE			,     __SCREAMING_SNAKE_CASE     )					->     Union[str, Any]:



            '''simple docstring'''




            __snake_case			    =     []
            __snake_case			    =     []
            __snake_case			    =     []
            __snake_case			    =     0

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

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

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






