



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


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


SCREAMING_SNAKE_CASE				:     Dict          =	logging.get_logger(__name__)






def 	__UpperCAmelCase					(       snake_case_  :  str   ,	snake_case_  :  int   )							->	Dict:






   """simple docstring"""


   import tensorflow.compat.va as tf

   _lowerCAmelCase					     =		tf.data.TFRecordDataset(lowercase__   ,	buffer_size=512 * 1024 * 1024   )
   _lowerCAmelCase					     =		blocks_dataset.batch(lowercase__   ,	drop_remainder=lowercase__   )
   _lowerCAmelCase					     =		next(blocks_dataset.take(1   ).as_numpy_iterator()   )

   return np_record


class       __lowerCamelCase	:







   def __init__(self       ,	lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase=2       ,	lowerCamelCase=1_000       ,	lowerCamelCase=100       ,	lowerCamelCase=100_000       ,	):




      '''simple docstring'''




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

      _lowerCAmelCase					     =		Builder(db=lowerCamelCase       ,	num_neighbors=lowerCamelCase       ,	distance_measure="""dot_product"""				)
      _lowerCAmelCase					     =		builder.tree(
          num_leaves=lowerCamelCase       ,	num_leaves_to_search=lowerCamelCase       ,	training_sample_size=lowerCamelCase				)
      _lowerCAmelCase					     =		builder.score_ah(dimensions_per_block=lowerCamelCase				)

      _lowerCAmelCase					     =		builder.build()







   def 							A__						(self       ,	lowerCamelCase				):




      '''simple docstring'''




      _lowerCAmelCase					     =		self.searcher.search_batched(question_projection.detach().cpu()				)
      return retrieved_block_ids.astype("""int64"""				)





class       __lowerCamelCase	:







   def __init__(self       ,	lowerCamelCase       ,	lowerCamelCase				):




      '''simple docstring'''




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







   def __call__(self       ,	lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase       ,	lowerCamelCase=None       ,	lowerCamelCase="pt"				):




      '''simple docstring'''




      _lowerCAmelCase					     =		np.take(self.block_records       ,	indices=lowerCamelCase       ,	axis=0				)

      _lowerCAmelCase					     =		self.tokenizer.decode(question_input_ids[0]       ,	skip_special_tokens=lowerCamelCase				)

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

      _lowerCAmelCase					     =		self.tokenizer(
          lowerCamelCase       ,	lowerCamelCase       ,	padding=lowerCamelCase       ,	truncation=lowerCamelCase       ,	return_special_tokens_mask=lowerCamelCase       ,	max_length=lowerCamelCase				)
      _lowerCAmelCase					     =		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 							A__						(cls       ,	lowerCamelCase       ,	*lowerCamelCase       ,	**lowerCamelCase				):




      '''simple docstring'''




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

      _lowerCAmelCase					     =		AutoTokenizer.from_pretrained(lowerCamelCase       ,	*lowerCamelCase       ,	**lowerCamelCase				)

      return cls(lowerCamelCase       ,	lowerCamelCase				)







   def 							A__						(self       ,	lowerCamelCase				):




      '''simple docstring'''




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







   def 							A__						(self       ,	lowerCamelCase       ,	lowerCamelCase				):




      '''simple docstring'''




      _lowerCAmelCase					     =		[]
      _lowerCAmelCase					     =		[]
      _lowerCAmelCase					     =		[]
      _lowerCAmelCase					     =		0

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

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




"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


SCREAMING_SNAKE_CASE				:     Union[str, Any]          =	{
    '''configuration_align''': [
        '''ALIGN_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''AlignConfig''',
        '''AlignTextConfig''',
        '''AlignVisionConfig''',
    ],
    '''processing_align''': ['''AlignProcessor'''],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  SCREAMING_SNAKE_CASE				:     Any          =	[
      '''ALIGN_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''AlignModel''',
      '''AlignPreTrainedModel''',
      '''AlignTextModel''',
      '''AlignVisionModel''',
  ]

if TYPE_CHECKING:
  from .configuration_align import (
      ALIGN_PRETRAINED_CONFIG_ARCHIVE_MAP,
      AlignConfig,
      AlignTextConfig,
      AlignVisionConfig,
  )
  from .processing_align import AlignProcessor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_align import (
        ALIGN_PRETRAINED_MODEL_ARCHIVE_LIST,
        AlignModel,
        AlignPreTrainedModel,
        AlignTextModel,
        AlignVisionModel,
    )

else:
  import sys

  SCREAMING_SNAKE_CASE				:     Any          =	_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)