"""simple docstring"""




from collections import OrderedDict
from typing import Any, Mapping, Optional

from ... import PreTrainedTokenizer
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig, OnnxConfigWithPast, OnnxSeqaSeqConfigWithPast
from ...onnx.utils import compute_effective_axis_dimension
from ...utils import TensorType, is_torch_available, logging


__lowerCamelCase									=	logging.get_logger(__name__)

__lowerCamelCase									=	{
    "facebook/mbart-large-cc25": "https://huggingface.co/facebook/mbart-large-cc25/resolve/main/config.json",
    # See all MBART models at https://huggingface.co/models?filter=mbart
}
class      UpperCamelCase__(      __A			):
						lowerCAmelCase__	:    Dict    				=		'mbart'
						lowerCAmelCase__	:    Optional[Any]    				=		['past_key_values']
						lowerCAmelCase__	:    Tuple    				=		{'num_attention_heads': 'encoder_attention_heads', 'hidden_size': 'd_model'}





						def __init__(							self		,__UpperCAmelCase=5_02_65		,__UpperCAmelCase=10_24		,__UpperCAmelCase=12		,__UpperCAmelCase=40_96		,__UpperCAmelCase=16		,__UpperCAmelCase=12		,__UpperCAmelCase=40_96		,__UpperCAmelCase=16		,__UpperCAmelCase=0.0		,__UpperCAmelCase=0.0		,__UpperCAmelCase=True		,__UpperCAmelCase=True		,__UpperCAmelCase="gelu"		,__UpperCAmelCase=10_24		,__UpperCAmelCase=0.1		,__UpperCAmelCase=0.0		,__UpperCAmelCase=0.0		,__UpperCAmelCase=0.0_2		,__UpperCAmelCase=0.0		,__UpperCAmelCase=False		,__UpperCAmelCase=1		,__UpperCAmelCase=0		,__UpperCAmelCase=2		,__UpperCAmelCase=2		,**__UpperCAmelCase		,)       ->					Optional[Any]:
													A__						  =				vocab_size
													A__						  =				max_position_embeddings
													A__						  =				d_model
													A__						  =				encoder_ffn_dim
													A__						  =				encoder_layers
													A__						  =				encoder_attention_heads
													A__						  =				decoder_ffn_dim
													A__						  =				decoder_layers
													A__						  =				decoder_attention_heads
													A__						  =				dropout
													A__						  =				attention_dropout
													A__						  =				activation_dropout
													A__						  =				activation_function
													A__						  =				init_std
													A__						  =				encoder_layerdrop
													A__						  =				decoder_layerdrop
													A__						  =				classifier_dropout
													A__						  =				use_cache
													A__						  =				encoder_layers
													A__						  =				scale_embedding  # scale factor will be sqrt(d_model) if True
													super().__init__(
													    pad_token_id=__UpperCAmelCase		,bos_token_id=__UpperCAmelCase		,eos_token_id=__UpperCAmelCase		,is_encoder_decoder=__UpperCAmelCase		,forced_eos_token_id=__UpperCAmelCase		,**__UpperCAmelCase		,)






class      UpperCamelCase__(      __A			):





						@property
						def 					snake_case__						(							self  )       ->					Mapping[str, Mapping[int, str]]:
													if self.task in ["default", "seq2seq-lm"]:
																				A__						  =				OrderedDict(
																				    [
																				        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
																				        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
																				    ]  )

																				if self.use_past:
																											A__						  =				{0: 'batch'}
																											A__						  =				{0: 'batch', 1: 'past_decoder_sequence + sequence'}
																				else:
																											A__						  =				{0: 'batch', 1: 'decoder_sequence'}
																											A__						  =				{0: 'batch', 1: 'decoder_sequence'}

																				if self.use_past:
																											self.fill_with_past_key_values_(__UpperCAmelCase		,direction='inputs'  )
													elif self.task == "causal-lm":
																				# TODO: figure this case out.
																				A__						  =				OrderedDict(
																				    [
																				        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
																				        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
																				    ]  )
																				if self.use_past:
																											A__  ,	A__						  =				self.num_layers
																											for i in range(__UpperCAmelCase  ):
																																		A__						  =				{0: 'batch', 2: 'past_sequence + sequence'}
																																		A__						  =				{0: 'batch', 2: 'past_sequence + sequence'}
													else:
																				A__						  =				OrderedDict(
																				    [
																				        ('input_ids', {0: 'batch', 1: 'encoder_sequence'}),
																				        ('attention_mask', {0: 'batch', 1: 'encoder_sequence'}),
																				        ('decoder_input_ids', {0: 'batch', 1: 'decoder_sequence'}),
																				        ('decoder_attention_mask', {0: 'batch', 1: 'decoder_sequence'}),
																				    ]  )

													return common_inputs





						@property
						def 					snake_case__						(							self  )       ->					Mapping[str, Mapping[int, str]]:
													if self.task in ["default", "seq2seq-lm"]:
																				A__						  =				super().outputs
													else:
																				A__						  =				super(__UpperCAmelCase		,self  ).outputs
																				if self.use_past:
																											A__  ,	A__						  =				self.num_layers
																											for i in range(__UpperCAmelCase  ):
																																		A__						  =				{0: 'batch', 2: 'past_sequence + sequence'}
																																		A__						  =				{0: 'batch', 2: 'past_sequence + sequence'}
													return common_outputs





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase = -1		,__UpperCAmelCase = -1		,__UpperCAmelCase = False		,__UpperCAmelCase = None		,)       ->					Mapping[str, Any]:
													A__						  =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
													    __UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )

													# Generate decoder inputs
													A__						  =				seq_length if not self.use_past else 1
													A__						  =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
													    __UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )
													A__						  =				{f'''decoder_{name}''': tensor for name, tensor in decoder_inputs.items()}
													A__						  =				dict(**__UpperCAmelCase		,**__UpperCAmelCase  )

													if self.use_past:
																				if not is_torch_available():
																											raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'  )
																				else:
																											import torch
																				A__  ,	A__						  =				common_inputs['input_ids'].shape
																				A__						  =				common_inputs['decoder_input_ids'].shape[1]
																				A__  ,	A__						  =				self.num_attention_heads
																				A__						  =				(
																				    batch,
																				    num_encoder_attention_heads,
																				    encoder_seq_length,
																				    self._config.hidden_size // num_encoder_attention_heads,
																				)
																				A__						  =				decoder_seq_length + 3
																				A__						  =				(
																				    batch,
																				    num_decoder_attention_heads,
																				    decoder_past_length,
																				    self._config.hidden_size // num_decoder_attention_heads,
																				)

																				A__						  =				torch.cat(
																				    [common_inputs['decoder_attention_mask'], torch.ones(__UpperCAmelCase		,__UpperCAmelCase  )]		,dim=1  )

																				A__						  =				[]
																				# If the number of encoder and decoder layers are present in the model configuration, both are considered
																				A__  ,	A__						  =				self.num_layers
																				A__						  =				min(__UpperCAmelCase		,__UpperCAmelCase  )
																				A__						  =				max(__UpperCAmelCase		,__UpperCAmelCase  ) - min_num_layers
																				A__						  =				'encoder' if num_encoder_layers > num_decoder_layers else 'decoder'

																				for _ in range(__UpperCAmelCase  ):
																											common_inputs["past_key_values"].append(
																											    (
																											        torch.zeros(__UpperCAmelCase  ),
																											        torch.zeros(__UpperCAmelCase  ),
																											        torch.zeros(__UpperCAmelCase  ),
																											        torch.zeros(__UpperCAmelCase  ),
																											    )  )
																				# TODO: test this.
																				A__						  =				encoder_shape if remaining_side_name == 'encoder' else decoder_shape
																				for _ in range(__UpperCAmelCase		,__UpperCAmelCase  ):
																											common_inputs["past_key_values"].append((torch.zeros(__UpperCAmelCase  ), torch.zeros(__UpperCAmelCase  ))  )
													return common_inputs





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase = -1		,__UpperCAmelCase = -1		,__UpperCAmelCase = False		,__UpperCAmelCase = None		,)       ->					Mapping[str, Any]:
													A__						  =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
													    __UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )

													if self.use_past:
																				if not is_torch_available():
																											raise ValueError('Cannot generate dummy past_keys inputs without PyTorch installed.'  )
																				else:
																											import torch
																				A__  ,	A__						  =				common_inputs['input_ids'].shape
																				# Not using the same length for past_key_values
																				A__						  =				seqlen + 2
																				A__  ,	A__						  =				self.num_layers
																				A__  ,	A__						  =				self.num_attention_heads
																				A__						  =				(
																				    batch,
																				    num_encoder_attention_heads,
																				    past_key_values_length,
																				    self._config.hidden_size // num_encoder_attention_heads,
																				)

																				A__						  =				common_inputs['attention_mask'].dtype
																				A__						  =				torch.cat(
																				    [common_inputs['attention_mask'], torch.ones(__UpperCAmelCase		,__UpperCAmelCase		,dtype=__UpperCAmelCase  )]		,dim=1  )
																				A__						  =				[
																				    (torch.zeros(__UpperCAmelCase  ), torch.zeros(__UpperCAmelCase  )) for _ in range(__UpperCAmelCase  )
																				]
													return common_inputs





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase = -1		,__UpperCAmelCase = -1		,__UpperCAmelCase = False		,__UpperCAmelCase = None		,)       ->					Mapping[str, Any]:
													# Copied from OnnxConfig.generate_dummy_inputs
													# Did not use super(OnnxConfigWithPast, self).generate_dummy_inputs for code clarity.
													# If dynamic axis (-1) we forward with a fixed dimension of 2 samples to avoid optimizations made by ONNX
													A__						  =				compute_effective_axis_dimension(
													    __UpperCAmelCase		,fixed_dimension=OnnxConfig.default_fixed_batch		,num_token_to_add=0  )

													# If dynamic axis (-1) we forward with a fixed dimension of 8 tokens to avoid optimizations made by ONNX
													A__						  =				tokenizer.num_special_tokens_to_add(__UpperCAmelCase  )
													A__						  =				compute_effective_axis_dimension(
													    __UpperCAmelCase		,fixed_dimension=OnnxConfig.default_fixed_sequence		,num_token_to_add=__UpperCAmelCase  )

													# Generate dummy inputs according to compute batch and sequence
													A__						  =				[' '.join([tokenizer.unk_token]  ) * seq_length] * batch_size
													A__						  =				dict(tokenizer(__UpperCAmelCase		,return_tensors=__UpperCAmelCase  )  )
													return common_inputs





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase = -1		,__UpperCAmelCase = -1		,__UpperCAmelCase = False		,__UpperCAmelCase = None		,)       ->					Mapping[str, Any]:
													if self.task in ["default", "seq2seq-lm"]:
																				A__						  =				self._generate_dummy_inputs_for_default_and_seqaseq_lm(
																				    __UpperCAmelCase		,batch_size=__UpperCAmelCase		,seq_length=__UpperCAmelCase		,is_pair=__UpperCAmelCase		,framework=__UpperCAmelCase  )

													elif self.task == "causal-lm":
																				A__						  =				self._generate_dummy_inputs_for_causal_lm(
																				    __UpperCAmelCase		,batch_size=__UpperCAmelCase		,seq_length=__UpperCAmelCase		,is_pair=__UpperCAmelCase		,framework=__UpperCAmelCase  )
													else:
																				A__						  =				self._generate_dummy_inputs_for_sequence_classification_and_question_answering(
																				    __UpperCAmelCase		,batch_size=__UpperCAmelCase		,seq_length=__UpperCAmelCase		,is_pair=__UpperCAmelCase		,framework=__UpperCAmelCase  )

													return common_inputs





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					List[Any]:
													if self.task in ["default", "seq2seq-lm"]:
																				A__						  =				super()._flatten_past_key_values_(__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )
													else:
																				A__						  =				super(__UpperCAmelCase		,self  )._flatten_past_key_values_(
																				    __UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )






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





