
'''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						=      "block_records.npy"


__SCREAMING_SNAKE_CASE						=      logging.get_logger(__name__)
def 					__magic_name__(       UpperCAmelCase__ :       str							,				UpperCAmelCase__ :       int	):
						import tensorflow.compat.va as tf

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

						return np_record




class     __snake_case			:




		def __init__(  self					:Union[str, Any]			,    _SCREAMING_SNAKE_CASE					:int			,    _SCREAMING_SNAKE_CASE					:Dict			,    _SCREAMING_SNAKE_CASE					:int=2			,    _SCREAMING_SNAKE_CASE					:Optional[int]=1_0_0_0			,    _SCREAMING_SNAKE_CASE					:Union[str, Any]=1_0_0			,    _SCREAMING_SNAKE_CASE					:int=1_0_0_0_0_0			,    )					->			List[str]:






								'''simple docstring'''






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

								UpperCamelCase_     :	List[Any]					     =     Builder(db=_SCREAMING_SNAKE_CASE			,    num_neighbors=_SCREAMING_SNAKE_CASE			,    distance_measure="""dot_product""")
								UpperCamelCase_     :	Union[str, Any]					     =     builder.tree(
								    num_leaves=_SCREAMING_SNAKE_CASE			,    num_leaves_to_search=_SCREAMING_SNAKE_CASE			,    training_sample_size=_SCREAMING_SNAKE_CASE)
								UpperCamelCase_     :	Tuple					     =     builder.score_ah(dimensions_per_block=_SCREAMING_SNAKE_CASE)

								UpperCamelCase_     :	Optional[Any]					     =     builder.build()





		def 	_snake_case	(  self					:List[str]			,    _SCREAMING_SNAKE_CASE					:Dict)					->			List[Any]:






								'''simple docstring'''






								UpperCamelCase_     ,UpperCamelCase_     :	str					     =     self.searcher.search_batched(question_projection.detach().cpu())
								return retrieved_block_ids.astype("""int64""")







class     __snake_case			:




		def __init__(  self					:Optional[Any]			,    _SCREAMING_SNAKE_CASE					:Optional[int]			,    _SCREAMING_SNAKE_CASE					:Dict)					->			int:






								'''simple docstring'''






								super().__init__()
								UpperCamelCase_     :	int					     =     block_records
								UpperCamelCase_     :	str					     =     tokenizer




		def __call__(  self					:Dict			,    _SCREAMING_SNAKE_CASE					:Optional[Any]			,    _SCREAMING_SNAKE_CASE					:Tuple			,    _SCREAMING_SNAKE_CASE					:Optional[int]			,    _SCREAMING_SNAKE_CASE					:List[Any]=None			,    _SCREAMING_SNAKE_CASE					:List[Any]="pt")					->			Any:






								'''simple docstring'''






								UpperCamelCase_     :	List[str]					     =     np.take(self.block_records			,    indices=_SCREAMING_SNAKE_CASE			,    axis=0)

								UpperCamelCase_     :	Optional[int]					     =     self.tokenizer.decode(question_input_ids[0]			,    skip_special_tokens=_SCREAMING_SNAKE_CASE)

								UpperCamelCase_     :	List[Any]					     =     []
								UpperCamelCase_     :	Dict					     =     []
								for retrieved_block in retrieved_blocks:
														text.append(_SCREAMING_SNAKE_CASE)
														text_pair.append(retrieved_block.decode())

								UpperCamelCase_     :	Optional[int]					     =     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)
								UpperCamelCase_     :	str					     =     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 	_snake_case	(  cls					:Optional[Any]			,    _SCREAMING_SNAKE_CASE					:Optional[Union[str, os.PathLike]]			,    *_SCREAMING_SNAKE_CASE					:List[str]			,    **_SCREAMING_SNAKE_CASE					:List[Any])					->			Any:






								'''simple docstring'''






								if os.path.isdir(_SCREAMING_SNAKE_CASE):
														UpperCamelCase_     :	List[str]					     =     os.path.join(_SCREAMING_SNAKE_CASE			,    _REALM_BLOCK_RECORDS_FILENAME)
								else:
														UpperCamelCase_     :	int					     =     hf_hub_download(
														    repo_id=_SCREAMING_SNAKE_CASE			,    filename=_REALM_BLOCK_RECORDS_FILENAME			,    **_SCREAMING_SNAKE_CASE)
								UpperCamelCase_     :	int					     =     np.load(_SCREAMING_SNAKE_CASE			,    allow_pickle=_SCREAMING_SNAKE_CASE)

								UpperCamelCase_     :	Dict					     =     AutoTokenizer.from_pretrained(_SCREAMING_SNAKE_CASE			,    *_SCREAMING_SNAKE_CASE			,    **_SCREAMING_SNAKE_CASE)

								return cls(_SCREAMING_SNAKE_CASE			,    _SCREAMING_SNAKE_CASE)




		def 	_snake_case	(  self					:int			,    _SCREAMING_SNAKE_CASE					:int)					->			Tuple:






								'''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 	_snake_case	(  self					:List[str]			,    _SCREAMING_SNAKE_CASE					:Optional[Any]			,    _SCREAMING_SNAKE_CASE					:Optional[int])					->			str:






								'''simple docstring'''






								UpperCamelCase_     :	Tuple					     =     []
								UpperCamelCase_     :	Union[str, Any]					     =     []
								UpperCamelCase_     :	Optional[int]					     =     []
								UpperCamelCase_     :	List[str]					     =     0

								for input_id in concat_inputs.input_ids:
														UpperCamelCase_     :	Union[str, Any]					     =     input_id.tolist()
														# Check answers between two [SEP] tokens
														UpperCamelCase_     :	Optional[int]					     =     input_id_list.index(self.tokenizer.sep_token_id)
														UpperCamelCase_     :	str					     =     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:
																										UpperCamelCase_     :	List[Any]					     =     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:
																				UpperCamelCase_     :	Optional[Any]					     =     [-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_tokenizers_available, is_torch_available


__SCREAMING_SNAKE_CASE						=      {
    "configuration_markuplm": ["MARKUPLM_PRETRAINED_CONFIG_ARCHIVE_MAP", "MarkupLMConfig"],
    "feature_extraction_markuplm": ["MarkupLMFeatureExtractor"],
    "processing_markuplm": ["MarkupLMProcessor"],
    "tokenization_markuplm": ["MarkupLMTokenizer"],
}

try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__SCREAMING_SNAKE_CASE						=      ["MarkupLMTokenizerFast"]

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__SCREAMING_SNAKE_CASE						=      [
	    "MARKUPLM_PRETRAINED_MODEL_ARCHIVE_LIST",
	    "MarkupLMForQuestionAnswering",
	    "MarkupLMForSequenceClassification",
	    "MarkupLMForTokenClassification",
	    "MarkupLMModel",
	    "MarkupLMPreTrainedModel",
	]


if TYPE_CHECKING:
	from .configuration_markuplm import MARKUPLM_PRETRAINED_CONFIG_ARCHIVE_MAP, MarkupLMConfig
	from .feature_extraction_markuplm import MarkupLMFeatureExtractor
	from .processing_markuplm import MarkupLMProcessor
	from .tokenization_markuplm import MarkupLMTokenizer

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_markuplm_fast import MarkupLMTokenizerFast

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_markuplm import (
		    MARKUPLM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    MarkupLMForQuestionAnswering,
		    MarkupLMForSequenceClassification,
		    MarkupLMForTokenClassification,
		    MarkupLMModel,
		    MarkupLMPreTrainedModel,
		)


else:
	import sys

	__SCREAMING_SNAKE_CASE						=      _LazyModule(__name__, globals()["__file__"], _import_structure)




