"""simple docstring"""




from __future__ import annotations

import unittest

from transformers import AutoTokenizer, GPTJConfig, is_tf_available
from transformers.testing_utils import require_tf, slow, tooslow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin
from ...utils.test_modeling_tf_core import TFCoreModelTesterMixin


if is_tf_available():
				import tensorflow as tf

				from transformers.models.gptj.modeling_tf_gptj import (
				    TFGPTJForCausalLM,
				    TFGPTJForQuestionAnswering,
				    TFGPTJForSequenceClassification,
				    TFGPTJModel,
				    shape_list,
				)
class      UpperCamelCase__:





						def __init__(							self		,__UpperCAmelCase  )       ->					Union[str, Any]:
													A__						  =				parent
													A__						  =				13
													A__						  =				7
													A__						  =				True
													A__						  =				True
													A__						  =				True
													A__						  =				True
													A__						  =				True
													A__						  =				99
													A__						  =				32
													A__						  =				4
													A__						  =				2
													A__						  =				4
													A__						  =				37
													A__						  =				'gelu'
													A__						  =				0.1
													A__						  =				0.1
													A__						  =				5_12
													A__						  =				16
													A__						  =				2
													A__						  =				0.0_2
													A__						  =				3
													A__						  =				4
													A__						  =				None
													A__						  =				self.vocab_size - 1
													A__						  =				self.vocab_size - 1
													A__						  =				self.vocab_size - 1





						def 					snake_case__						(							self  )       ->					Dict:
													A__						  =				ids_tensor([self.batch_size, self.seq_length]		,self.vocab_size  )

													A__						  =				None
													if self.use_input_mask:
																				A__						  =				random_attention_mask([self.batch_size, self.seq_length]  )

													A__						  =				None
													if self.use_token_type_ids:
																				A__						  =				ids_tensor([self.batch_size, self.seq_length]		,self.type_vocab_size  )

													A__						  =				None
													if self.use_mc_token_ids:
																				A__						  =				ids_tensor([self.batch_size, self.num_choices]		,self.seq_length  )

													A__						  =				None
													A__						  =				None
													A__						  =				None
													if self.use_labels:
																				A__						  =				ids_tensor([self.batch_size]		,self.type_sequence_label_size  )
																				A__						  =				ids_tensor([self.batch_size, self.seq_length]		,self.num_labels  )
																				A__						  =				ids_tensor([self.batch_size]		,self.num_choices  )

													A__						  =				GPTJConfig(
													    vocab_size=self.vocab_size		,n_embd=self.hidden_size		,n_layer=self.num_hidden_layers		,n_head=self.num_attention_heads		,intermediate_size=self.intermediate_size		,hidden_act=self.hidden_act		,hidden_dropout_prob=self.hidden_dropout_prob		,attention_probs_dropout_prob=self.attention_probs_dropout_prob		,n_positions=self.max_position_embeddings		,type_vocab_size=self.type_vocab_size		,initializer_range=self.initializer_range		,bos_token_id=self.bos_token_id		,eos_token_id=self.eos_token_id		,pad_token_id=self.pad_token_id		,rotary_dim=self.rotary_dim		,return_dict=__UpperCAmelCase		,)

													A__						  =				ids_tensor([self.num_hidden_layers, self.num_attention_heads]		,2  )

													return (
													    config,
													    input_ids,
													    input_mask,
													    head_mask,
													    token_type_ids,
													    mc_token_ids,
													    sequence_labels,
													    token_labels,
													    choice_labels,
													)





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,*__UpperCAmelCase  )       ->					List[str]:
													A__						  =				TFGPTJModel(config=__UpperCAmelCase  )
													A__						  =				{
													    'input_ids': input_ids,
													    'attention_mask': input_mask,
													    'token_type_ids': token_type_ids,
													}
													A__						  =				model(__UpperCAmelCase  )

													A__						  =				[input_ids, None, input_mask]  # None is the input for 'past'
													A__						  =				model(__UpperCAmelCase  )

													A__						  =				model(__UpperCAmelCase  )

													self.parent.assertEqual(result.last_hidden_state.shape		,(self.batch_size, self.seq_length, self.hidden_size)  )





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,*__UpperCAmelCase  )       ->					Optional[int]:
													A__						  =				TFGPTJModel(config=__UpperCAmelCase  )

													# first forward pass
													A__						  =				model(__UpperCAmelCase		,token_type_ids=__UpperCAmelCase		,use_cache=__UpperCAmelCase  )
													A__						  =				model(__UpperCAmelCase		,token_type_ids=__UpperCAmelCase  )
													A__						  =				model(__UpperCAmelCase		,token_type_ids=__UpperCAmelCase		,use_cache=__UpperCAmelCase  )

													self.parent.assertTrue(len(__UpperCAmelCase  ) == len(__UpperCAmelCase  )  )
													self.parent.assertTrue(len(__UpperCAmelCase  ) == len(__UpperCAmelCase  ) + 1  )

													A__  ,	A__						  =				outputs.to_tuple()

													# create hypothetical next token and extent to next_input_ids
													A__						  =				ids_tensor((self.batch_size, 1)		,config.vocab_size  )
													A__						  =				ids_tensor([self.batch_size, 1]		,self.type_vocab_size  )

													# append to next input_ids and token_type_ids
													A__						  =				tf.concat([input_ids, next_tokens]		,axis=-1  )
													A__						  =				tf.concat([token_type_ids, next_token_types]		,axis=-1  )

													A__						  =				model(__UpperCAmelCase		,token_type_ids=__UpperCAmelCase  )['last_hidden_state']
													A__						  =				model(__UpperCAmelCase		,token_type_ids=__UpperCAmelCase		,past_key_values=__UpperCAmelCase  )[
													    'last_hidden_state'
													]

													# select random slice
													A__						  =				int(ids_tensor((1,)		,shape_list(__UpperCAmelCase  )[-1]  )  )
													A__						  =				output_from_no_past[:, -1, random_slice_idx]
													A__						  =				output_from_past[:, 0, random_slice_idx]

													# test that outputs are equal for slice
													tf.debugging.assert_near(__UpperCAmelCase		,__UpperCAmelCase		,rtol=1e-6  )





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

													# create attention mask
													A__						  =				self.seq_length // 2
													A__						  =				tf.ones((self.batch_size, half_seq_length)		,dtype=tf.intaa  )
													A__						  =				tf.zeros((self.batch_size, self.seq_length - half_seq_length)		,dtype=tf.intaa  )
													A__						  =				tf.concat([attn_mask_begin, attn_mask_end]		,axis=1  )

													# first forward pass
													A__  ,	A__						  =				model(__UpperCAmelCase		,attention_mask=__UpperCAmelCase  ).to_tuple()

													# create hypothetical next token and extent to next_input_ids
													A__						  =				ids_tensor((self.batch_size, 1)		,config.vocab_size  )

													# change a random masked slice from input_ids
													A__						  =				ids_tensor((1,)		,__UpperCAmelCase  ).numpy() + 1
													A__						  =				ids_tensor((self.batch_size, self.seq_length)		,config.vocab_size  )
													A__						  =				tf.range(self.seq_length  ) == (self.seq_length - random_seq_idx_to_change)
													A__						  =				tf.transpose(
													    tf.broadcast_to(tf.expand_dims(__UpperCAmelCase		,-1  )		,(self.seq_length, self.batch_size)  )  )
													A__						  =				tf.where(__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )

													# append to next input_ids and attn_mask
													A__						  =				tf.concat([input_ids, next_tokens]		,axis=-1  )
													A__						  =				tf.concat([attn_mask, tf.ones((shape_list(__UpperCAmelCase  )[0], 1)		,dtype=tf.intaa  )]		,axis=1  )

													# get two different outputs
													A__						  =				model(__UpperCAmelCase		,attention_mask=__UpperCAmelCase  )['last_hidden_state']
													A__						  =				model(__UpperCAmelCase		,past_key_values=__UpperCAmelCase		,attention_mask=__UpperCAmelCase  )[
													    'last_hidden_state'
													]

													# select random slice
													A__						  =				int(ids_tensor((1,)		,shape_list(__UpperCAmelCase  )[-1]  )  )
													A__						  =				output_from_no_past[:, -1, random_slice_idx]
													A__						  =				output_from_past[:, 0, random_slice_idx]

													# test that outputs are equal for slice
													tf.debugging.assert_near(__UpperCAmelCase		,__UpperCAmelCase		,rtol=1e-12  )





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,*__UpperCAmelCase  )       ->					str:
													A__						  =				TFGPTJModel(config=__UpperCAmelCase  )

													A__						  =				input_ids[:1, :]
													A__						  =				input_mask[:1, :]
													A__						  =				token_type_ids[:1, :]
													A__						  =				1

													# first forward pass
													A__						  =				model(__UpperCAmelCase		,attention_mask=__UpperCAmelCase		,token_type_ids=__UpperCAmelCase		,use_cache=__UpperCAmelCase  )

													A__  ,	A__						  =				outputs.to_tuple()

													# create hypothetical next token and extent to next_input_ids
													A__						  =				ids_tensor((self.batch_size, 3)		,config.vocab_size  )
													A__						  =				ids_tensor((self.batch_size, 3)		,2  )
													A__						  =				ids_tensor((self.batch_size, 3)		,self.type_vocab_size  )

													# append to next input_ids and token_type_ids
													A__						  =				tf.concat([input_ids, next_tokens]		,axis=-1  )
													A__						  =				tf.concat([input_mask, next_attn_mask]		,axis=-1  )
													A__						  =				tf.concat([token_type_ids, next_token_types]		,axis=-1  )

													A__						  =				model(
													    __UpperCAmelCase		,token_type_ids=__UpperCAmelCase		,attention_mask=__UpperCAmelCase  )['last_hidden_state']
													A__						  =				model(
													    __UpperCAmelCase		,token_type_ids=__UpperCAmelCase		,attention_mask=__UpperCAmelCase		,past_key_values=__UpperCAmelCase		,)['last_hidden_state']
													self.parent.assertTrue(output_from_past.shape[1] == next_tokens.shape[1]  )

													# select random slice
													A__						  =				int(ids_tensor((1,)		,shape_list(__UpperCAmelCase  )[-1]  )  )
													A__						  =				output_from_no_past[:, -3:, random_slice_idx]
													A__						  =				output_from_past[:, :, random_slice_idx]

													# test that outputs are equal for slice
													tf.debugging.assert_near(__UpperCAmelCase		,__UpperCAmelCase		,rtol=1e-3  )





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,*__UpperCAmelCase  )       ->					Optional[Any]:
													A__						  =				TFGPTJForCausalLM(config=__UpperCAmelCase  )
													A__						  =				{
													    'input_ids': input_ids,
													    'attention_mask': input_mask,
													    'token_type_ids': token_type_ids,
													}
													A__						  =				model(__UpperCAmelCase  )
													self.parent.assertEqual(result.logits.shape		,(self.batch_size, self.seq_length, self.vocab_size)  )





						def 					snake_case__						(							self  )       ->					Optional[int]:
													A__						  =				self.prepare_config_and_inputs()

													(
													    (
													    A__
													)  ,	(
													    A__
													)  ,	(
													    A__
													)  ,	(
													    A__
													)  ,	(
													    A__
													)  ,	(
													    A__
													)  ,	(
													    A__
													)  ,	(
													    A__
													)  ,	(
													    A__
													)  ,	
													)						  =				config_and_inputs

													A__						  =				{
													    'input_ids': input_ids,
													    'token_type_ids': token_type_ids,
													    'attention_mask': input_mask,
													}
													return config, inputs_dict






@require_tf
class      UpperCamelCase__(      __A		,    __A		,    __A		,    unittest.TestCase			):
						lowerCAmelCase__	:    Any    				=		(
						    (TFGPTJForCausalLM, TFGPTJForSequenceClassification, TFGPTJForQuestionAnswering, TFGPTJModel)
						    if is_tf_available()
						    else ()
						)

						lowerCAmelCase__	:    Optional[int]    				=		(TFGPTJForCausalLM,) if is_tf_available() else ()
						lowerCAmelCase__	:    Tuple    				=		(
						    {
						        'feature-extraction': TFGPTJModel,
						        'question-answering': TFGPTJForQuestionAnswering,
						        'text-classification': TFGPTJForSequenceClassification,
						        'text-generation': TFGPTJForCausalLM,
						        'zero-shot': TFGPTJForSequenceClassification,
						    }
						    if is_tf_available()
						    else {}
						)
						lowerCAmelCase__	:    List[Any]    				=		False
						lowerCAmelCase__	:    str    				=		False
						lowerCAmelCase__	:    List[str]    				=		False
						lowerCAmelCase__	:    str    				=		False





						def 					snake_case__						(							self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase  )       ->					List[Any]:
													if (
													    pipeline_test_casse_name == "QAPipelineTests"
													    and tokenizer_name is not None
													    and not tokenizer_name.endswith('Fast'  )
													):
																				# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
																				# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
																				# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
																				return True

													return False





						def 					snake_case__						(							self  )       ->					Optional[Any]:
													A__						  =				TFGPTJModelTester(self  )
													A__						  =				ConfigTester(self		,config_class=__UpperCAmelCase		,n_embd=37  )





						def 					snake_case__						(							self  )       ->					int:
													self.config_tester.run_common_tests()





						def 					snake_case__						(							self  )       ->					List[Any]:
													A__						  =				self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_gptj_model(*__UpperCAmelCase  )





						def 					snake_case__						(							self  )       ->					Dict:
													A__						  =				self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_gptj_model_past(*__UpperCAmelCase  )





						def 					snake_case__						(							self  )       ->					str:
													A__						  =				self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_gptj_model_attention_mask_past(*__UpperCAmelCase  )





						def 					snake_case__						(							self  )       ->					List[Any]:
													A__						  =				self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_gptj_model_past_large_inputs(*__UpperCAmelCase  )





						def 					snake_case__						(							self  )       ->					str:
													A__						  =				self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_gptj_lm_head_model(*__UpperCAmelCase  )





						@slow
						@unittest.skipIf(
						    not is_tf_available() or len(tf.config.list_physical_devices('GPU'  )  ) > 0		,'skip testing on GPU for now to avoid GPU OOM.'		,)
						def 					snake_case__						(							self  )       ->					List[str]:
													A__						  =				TFGPTJModel.from_pretrained('EleutherAI/gpt-j-6B'		,from_pt=__UpperCAmelCase  )
													self.assertIsNotNone(__UpperCAmelCase  )





						@unittest.skip(reason='Currently, model embeddings are going to undergo a major refactor.'  )
						def 					snake_case__						(							self  )       ->					Dict:
													super().test_resize_token_embeddings()






@require_tf
@tooslow
# Marked as @tooslow due to GPU OOM -- but still useful to run locally. Requires ~39GB of RAM.
class      UpperCamelCase__(      unittest.TestCase			):





						def 					snake_case__						(							self  )       ->					Any:
													A__						  =				TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'		,from_pt=__UpperCAmelCase  )
													A__						  =				tf.convert_to_tensor([[4_64, 32_90]]		,dtype=tf.intaa  )  # The dog
													# fmt: off
													# The dog is a man's best friend. It is a loyal companion, and it is a friend
													A__						  =				[4_64, 32_90, 3_18, 2_57, 5_82, 3_38, 12_66, 15_45, 13, 6_32, 3_18, 2_57, 91_12, 1_51_85, 11, 2_90, 3_40, 3_18, 2_57, 15_45]
													# fmt: on
													A__						  =				model.generate(__UpperCAmelCase		,do_sample=__UpperCAmelCase  )
													self.assertListEqual(output_ids[0].numpy().tolist()		,__UpperCAmelCase  )





						def 					snake_case__						(							self  )       ->					str:
													A__						  =				AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B'		,revision='float16'  )
													A__						  =				TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'		,revision='float16'		,from_pt=__UpperCAmelCase  )

													A__						  =				tokenizer('Today is a nice day and'		,return_tensors='tf'  )
													# forces the generation to happen on CPU, to avoid GPU-related quirks
													with tf.device(':/CPU:0'  ):
																				A__						  =				model.generate(**__UpperCAmelCase		,do_sample=__UpperCAmelCase		,seed=[42, 0]  )
													A__						  =				tokenizer.decode(output_ids[0]		,skip_special_tokens=__UpperCAmelCase  )

													A__						  =				'Today is a nice day and I’m going to go for a walk. I’'
													self.assertEqual(__UpperCAmelCase		,__UpperCAmelCase  )





						def 					snake_case__						(							self  )       ->					Union[str, Any]:
													A__						  =				TFGPTJForCausalLM.from_pretrained('EleutherAI/gpt-j-6B'		,revision='float16'		,from_pt=__UpperCAmelCase  )
													A__						  =				AutoTokenizer.from_pretrained('EleutherAI/gpt-j-6B'		,revision='float16'  )

													A__						  =				'left'

													# Define PAD Token = EOS Token = 50256
													A__						  =				tokenizer.eos_token
													A__						  =				model.config.eos_token_id

													# use different length sentences to test batching
													A__						  =				[
													    'Hello, my dog is a little',
													    'Today, I',
													]
													A__						  =				[
													    'Hello, my dog is a little over a year old and has been diagnosed with hip dysplasia',
													    'Today, I’m going to be talking about a topic that’',
													]
													return model, tokenizer, sentences, expected_output_sentences





						def 					snake_case__						(							self  )       ->					Union[str, Any]:
													# Confirms that we get the expected results with left-padded beam search
													A__  ,	A__  ,	A__  ,	A__						  =				self._get_beam_search_test_objects()

													A__						  =				tokenizer(__UpperCAmelCase		,return_tensors='tf'		,padding=__UpperCAmelCase  )
													A__						  =				model.generate(**__UpperCAmelCase		,do_sample=__UpperCAmelCase		,num_beams=2  )
													A__						  =				tokenizer.batch_decode(__UpperCAmelCase		,skip_special_tokens=__UpperCAmelCase  )
													self.assertListEqual(__UpperCAmelCase		,__UpperCAmelCase  )





						def 					snake_case__						(							self  )       ->					int:
													# Confirms that left-padding is working properly
													A__  ,	A__  ,	A__  ,	A__						  =				self._get_beam_search_test_objects()

													A__						  =				tokenizer(__UpperCAmelCase		,return_tensors='tf'		,padding=__UpperCAmelCase  )
													A__						  =				tokenizer(sentences[0]		,return_tensors='tf'  )
													A__						  =				model.generate(**__UpperCAmelCase		,do_sample=__UpperCAmelCase		,num_beams=2  )
													A__						  =				(
													    shape_list(inputs_non_padded['input_ids']  )[-1]
													    - tf.reduce_sum(tf.cast(inputs['attention_mask'][-1]		,tf.intaa  )  ).numpy()
													)
													A__						  =				tokenizer(sentences[1]		,return_tensors='tf'  )
													A__						  =				model.generate(
													    **__UpperCAmelCase		,do_sample=__UpperCAmelCase		,num_beams=2		,max_length=model.config.max_length - num_paddings  )
													A__						  =				tokenizer.decode(output_non_padded[0]		,skip_special_tokens=__UpperCAmelCase  )
													A__						  =				tokenizer.decode(output_padded[0]		,skip_special_tokens=__UpperCAmelCase  )
													self.assertListEqual(__UpperCAmelCase		,[non_padded_sentence, padded_sentence]  )





						def 					snake_case__						(							self  )       ->					int:
													# Confirms that XLA is working properly
													A__  ,	A__  ,	A__  ,	A__						  =				self._get_beam_search_test_objects()

													A__						  =				tokenizer(__UpperCAmelCase		,return_tensors='tf'		,padding=__UpperCAmelCase  )
													A__						  =				tf.function(model.generate		,jit_compile=__UpperCAmelCase  )
													A__						  =				xla_generate(**__UpperCAmelCase		,do_sample=__UpperCAmelCase		,num_beams=2  )
													A__						  =				tokenizer.batch_decode(__UpperCAmelCase		,skip_special_tokens=__UpperCAmelCase  )
													self.assertListEqual(__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





