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


__lowerCamelCase									=	"block_records.npy"


__lowerCamelCase									=	logging.get_logger(__name__)






def 	UpperCAmelCase       (						UpperCamelCase__							,     UpperCamelCase__					):






							"""simple docstring"""





							import tensorflow.compat.va as tf

							A__						  =				tf.data.TFRecordDataset(UpperCamelCase__							,     buffer_size=512 * 1_024 * 1_024					)
							A__						  =				blocks_dataset.batch(UpperCamelCase__							,     drop_remainder=UpperCamelCase__					)
							A__						  =				next(blocks_dataset.take(1					).as_numpy_iterator()					)

							return np_record
class      UpperCamelCase__:





						def __init__(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase=2		,__UpperCAmelCase=10_00		,__UpperCAmelCase=1_00		,__UpperCAmelCase=10_00_00		,)       ->					Optional[int]:
													from scann.scann_ops.py.scann_ops_pybind import builder as Builder

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

													A__						  =				builder.build()





						def 					snake_case__						(							self		,__UpperCAmelCase  )       ->					Tuple:
													A__  ,	A__						  =				self.searcher.search_batched(question_projection.detach().cpu()  )
													return retrieved_block_ids.astype('int64'  )






class      UpperCamelCase__:





						def __init__(							self		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					List[str]:
													super().__init__()
													A__						  =				block_records
													A__						  =				tokenizer





						def __call__(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase=None		,__UpperCAmelCase="pt"  )       ->					str:
													A__						  =				np.take(self.block_records		,indices=__UpperCAmelCase		,axis=0  )

													A__						  =				self.tokenizer.decode(question_input_ids[0]		,skip_special_tokens=__UpperCAmelCase  )

													A__						  =				[]
													A__						  =				[]
													for retrieved_block in retrieved_blocks:
																				text.append(__UpperCAmelCase  )
																				text_pair.append(retrieved_block.decode()  )

													A__						  =				self.tokenizer(
													    __UpperCAmelCase		,__UpperCAmelCase		,padding=__UpperCAmelCase		,truncation=__UpperCAmelCase		,return_special_tokens_mask=__UpperCAmelCase		,max_length=__UpperCAmelCase  )
													A__						  =				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 					snake_case__						(							cls		,__UpperCAmelCase		,*__UpperCAmelCase		,**__UpperCAmelCase  )       ->					List[Any]:
													if os.path.isdir(__UpperCAmelCase  ):
																				A__						  =				os.path.join(__UpperCAmelCase		,_REALM_BLOCK_RECORDS_FILENAME  )
													else:
																				A__						  =				hf_hub_download(
																				    repo_id=__UpperCAmelCase		,filename=_REALM_BLOCK_RECORDS_FILENAME		,**__UpperCAmelCase  )
													A__						  =				np.load(__UpperCAmelCase		,allow_pickle=__UpperCAmelCase  )

													A__						  =				AutoTokenizer.from_pretrained(__UpperCAmelCase		,*__UpperCAmelCase		,**__UpperCAmelCase  )

													return cls(__UpperCAmelCase		,__UpperCAmelCase  )





						def 					snake_case__						(							self		,__UpperCAmelCase  )       ->					Tuple:
													# save block records
													np.save(os.path.join(__UpperCAmelCase		,_REALM_BLOCK_RECORDS_FILENAME  )		,self.block_records  )
													# save tokenizer
													self.tokenizer.save_pretrained(__UpperCAmelCase  )





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					List[Any]:
													A__						  =				[]
													A__						  =				[]
													A__						  =				[]
													A__						  =				0

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

        # Pad -1 to max_answers
													for start_pos_, end_pos_ in zip(__UpperCAmelCase		,__UpperCAmelCase  ):
																				if len(start_pos_  ) < max_answers:
																											A__						  =				[-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_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__lowerCamelCase									=	{
    "configuration_layoutlm": ["LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP", "LayoutLMConfig", "LayoutLMOnnxConfig"],
    "tokenization_layoutlm": ["LayoutLMTokenizer"],
}

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

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__lowerCamelCase									=	[
				    "LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST",
				    "LayoutLMForMaskedLM",
				    "LayoutLMForSequenceClassification",
				    "LayoutLMForTokenClassification",
				    "LayoutLMForQuestionAnswering",
				    "LayoutLMModel",
				    "LayoutLMPreTrainedModel",
				]

try:
				if not is_tf_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				__lowerCamelCase									=	[
				    "TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST",
				    "TFLayoutLMForMaskedLM",
				    "TFLayoutLMForSequenceClassification",
				    "TFLayoutLMForTokenClassification",
				    "TFLayoutLMForQuestionAnswering",
				    "TFLayoutLMMainLayer",
				    "TFLayoutLMModel",
				    "TFLayoutLMPreTrainedModel",
				]


if TYPE_CHECKING:
				from .configuration_layoutlm import LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP, LayoutLMConfig, LayoutLMOnnxConfig
				from .tokenization_layoutlm import LayoutLMTokenizer

				try:
								if not is_tokenizers_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .tokenization_layoutlm_fast import LayoutLMTokenizerFast

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_layoutlm import (
								    LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
								    LayoutLMForMaskedLM,
								    LayoutLMForQuestionAnswering,
								    LayoutLMForSequenceClassification,
								    LayoutLMForTokenClassification,
								    LayoutLMModel,
								    LayoutLMPreTrainedModel,
								)
				try:
								if not is_tf_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_tf_layoutlm import (
								    TF_LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
								    TFLayoutLMForMaskedLM,
								    TFLayoutLMForQuestionAnswering,
								    TFLayoutLMForSequenceClassification,
								    TFLayoutLMForTokenClassification,
								    TFLayoutLMMainLayer,
								    TFLayoutLMModel,
								    TFLayoutLMPreTrainedModel,
								)

else:
				import sys

				__lowerCamelCase									=	_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





