
"""simple docstring"""




import argparse
import json

import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from timm.models import create_model

from transformers import (
    BeitImageProcessor,
    DataaVecVisionConfig,
    DataaVecVisionForImageClassification,
    DataaVecVisionModel,
)

def   UpperCamelCase_  (		SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_=False,					SCREAMING_SNAKE_CASE_=False,					SCREAMING_SNAKE_CASE_="data2vec."      ):
					SCREAMING_SNAKE_CASE           =     'backbone.' if is_semantic else ''

					SCREAMING_SNAKE_CASE           =     []
					for i in range(config.num_hidden_layers      ):
										# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
										rename_keys.append(
										    (f'''{prefix}blocks.{i}.norm1.weight''', f'''{hf_prefix}encoder.layer.{i}.layernorm_before.weight''')      )
										rename_keys.append((f'''{prefix}blocks.{i}.norm1.bias''', f'''{hf_prefix}encoder.layer.{i}.layernorm_before.bias''')      )
										rename_keys.append(
										    (f'''{prefix}blocks.{i}.attn.proj.weight''', f'''{hf_prefix}encoder.layer.{i}.attention.output.dense.weight''')      )
										rename_keys.append(
										    (f'''{prefix}blocks.{i}.attn.proj.bias''', f'''{hf_prefix}encoder.layer.{i}.attention.output.dense.bias''')      )
										rename_keys.append(
										    (f'''{prefix}blocks.{i}.norm2.weight''', f'''{hf_prefix}encoder.layer.{i}.layernorm_after.weight''')      )
										rename_keys.append((f'''{prefix}blocks.{i}.norm2.bias''', f'''{hf_prefix}encoder.layer.{i}.layernorm_after.bias''')      )
										rename_keys.append(
										    (f'''{prefix}blocks.{i}.mlp.fc1.weight''', f'''{hf_prefix}encoder.layer.{i}.intermediate.dense.weight''')      )
										rename_keys.append(
										    (f'''{prefix}blocks.{i}.mlp.fc1.bias''', f'''{hf_prefix}encoder.layer.{i}.intermediate.dense.bias''')      )
										rename_keys.append((f'''{prefix}blocks.{i}.mlp.fc2.weight''', f'''{hf_prefix}encoder.layer.{i}.output.dense.weight''')      )
										rename_keys.append((f'''{prefix}blocks.{i}.mlp.fc2.bias''', f'''{hf_prefix}encoder.layer.{i}.output.dense.bias''')      )

					# projection layer + position embeddings
					rename_keys.extend(
					    [
					        (f'''{prefix}cls_token''', f'''{hf_prefix}embeddings.cls_token'''),
					        (f'''{prefix}patch_embed.proj.weight''', f'''{hf_prefix}embeddings.patch_embeddings.projection.weight'''),
					        (f'''{prefix}patch_embed.proj.bias''', f'''{hf_prefix}embeddings.patch_embeddings.projection.bias'''),
					    ]      )

					if has_lm_head:
										# mask token + shared relative position bias + layernorm
										rename_keys.extend(
										    [
										        ('mask_token', f'''{hf_prefix}embeddings.mask_token'''),
										        (
										            'rel_pos_bias.relative_position_bias_table',
										            f'''{hf_prefix}encoder.relative_position_bias.relative_position_bias_table''',
										        ),
										        (
										            'rel_pos_bias.relative_position_index',
										            f'''{hf_prefix}encoder.relative_position_bias.relative_position_index''',
										        ),
										        ('norm.weight', 'layernorm.weight'),
										        ('norm.bias', 'layernorm.bias'),
										    ]      )
					elif is_semantic:
										# semantic segmentation classification heads
										rename_keys.extend(
										    [
										        ('decode_head.conv_seg.weight', 'decode_head.classifier.weight'),
										        ('decode_head.conv_seg.bias', 'decode_head.classifier.bias'),
										        ('auxiliary_head.conv_seg.weight', 'auxiliary_head.classifier.weight'),
										        ('auxiliary_head.conv_seg.bias', 'auxiliary_head.classifier.bias'),
										    ]      )
					else:
										# layernorm + classification head
										rename_keys.extend(
										    [
										        ('fc_norm.weight', f'''{hf_prefix}pooler.layernorm.weight'''),
										        ('fc_norm.bias', f'''{hf_prefix}pooler.layernorm.bias'''),
										        ('head.weight', 'classifier.weight'),
										        ('head.bias', 'classifier.bias'),
										    ]      )

					return rename_keys

def   UpperCamelCase_  (		SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_=False,					SCREAMING_SNAKE_CASE_=False,					SCREAMING_SNAKE_CASE_="data2vec_vision."      ):
					for i in range(config.num_hidden_layers      ):
										SCREAMING_SNAKE_CASE           =     'backbone.' if is_semantic else ''
										# queries, keys and values
										SCREAMING_SNAKE_CASE           =     state_dict.pop(f'''{prefix}blocks.{i}.attn.qkv.weight'''      )
										SCREAMING_SNAKE_CASE           =     state_dict.pop(f'''{prefix}blocks.{i}.attn.q_bias'''      )
										SCREAMING_SNAKE_CASE           =     state_dict.pop(f'''{prefix}blocks.{i}.attn.v_bias'''      )

										SCREAMING_SNAKE_CASE           =     in_proj_weight[
										    : config.hidden_size, :
										]
										SCREAMING_SNAKE_CASE           =     q_bias
										SCREAMING_SNAKE_CASE           =     in_proj_weight[
										    config.hidden_size : config.hidden_size * 2, :
										]
										SCREAMING_SNAKE_CASE           =     in_proj_weight[
										    -config.hidden_size :, :
										]
										SCREAMING_SNAKE_CASE           =     v_bias

										# gamma_1 and gamma_2
										# we call them lambda because otherwise they are renamed when using .from_pretrained
										SCREAMING_SNAKE_CASE           =     state_dict.pop(f'''{prefix}blocks.{i}.gamma_1'''      )
										SCREAMING_SNAKE_CASE           =     state_dict.pop(f'''{prefix}blocks.{i}.gamma_2'''      )

										SCREAMING_SNAKE_CASE           =     gamma_a
										SCREAMING_SNAKE_CASE           =     gamma_a

										# relative_position bias table + index
										if not has_lm_head:
															# each layer has its own relative position bias
															SCREAMING_SNAKE_CASE           =     state_dict.pop(f'''{prefix}blocks.{i}.attn.relative_position_bias_table'''      )
															SCREAMING_SNAKE_CASE           =     state_dict.pop(f'''{prefix}blocks.{i}.attn.relative_position_index'''      )

															SCREAMING_SNAKE_CASE           =     table
															SCREAMING_SNAKE_CASE           =     index

def   UpperCamelCase_  (		):
					SCREAMING_SNAKE_CASE           =     argparse.ArgumentParser(
					    'Convert Data2VecVision to HF for image classification and pretraining',					add_help=SCREAMING_SNAKE_CASE_      )
					parser.add_argument('--hf_checkpoint_name',					type=SCREAMING_SNAKE_CASE_      )
					parser.add_argument('--input_size',					default=2_2_4,					type=SCREAMING_SNAKE_CASE_,					help='images input size'      )
					parser.add_argument('--beit_checkpoint',					default='',					help='beit checkpoint'      )

					return parser.parse_args()

def   UpperCamelCase_  (		SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_      ):
					def load_state_dict(SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_="",					SCREAMING_SNAKE_CASE_="relative_position_index"      ):
										SCREAMING_SNAKE_CASE           =     []
										SCREAMING_SNAKE_CASE           =     []
										SCREAMING_SNAKE_CASE           =     []
										# copy state_dict so _load_from_state_dict can modify it
										SCREAMING_SNAKE_CASE           =     getattr(SCREAMING_SNAKE_CASE_,					'_metadata',					SCREAMING_SNAKE_CASE_      )
										SCREAMING_SNAKE_CASE           =     state_dict.copy()
										if metadata is not None:
															SCREAMING_SNAKE_CASE           =     metadata

										def load(SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_=""      ):
															SCREAMING_SNAKE_CASE           =     {} if metadata is None else metadata.get(prefix[:-1],					{}      )
															module._load_from_state_dict(
															    SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_      )
															for name, child in module._modules.items():
																				if child is not None:
																									load(SCREAMING_SNAKE_CASE_,					prefix + name + '.'      )

										load(SCREAMING_SNAKE_CASE_,					prefix=SCREAMING_SNAKE_CASE_      )

										SCREAMING_SNAKE_CASE           =     []
										SCREAMING_SNAKE_CASE           =     []
										for key in missing_keys:
															SCREAMING_SNAKE_CASE           =     True
															for ignore_key in ignore_missing.split('|'      ):
																				if ignore_key in key:
																									SCREAMING_SNAKE_CASE           =     False
																									break
															if keep_flag:
																				warn_missing_keys.append(SCREAMING_SNAKE_CASE_      )
															else:
																				ignore_missing_keys.append(SCREAMING_SNAKE_CASE_      )

										SCREAMING_SNAKE_CASE           =     warn_missing_keys

										if len(SCREAMING_SNAKE_CASE_      ) > 0:
															print(
															    'Weights of {} not initialized from pretrained model: {}'.format(
															        model.__class__.__name__,					SCREAMING_SNAKE_CASE_      )      )
										if len(SCREAMING_SNAKE_CASE_      ) > 0:
															print('Weights from pretrained model not used in {}: {}'.format(model.__class__.__name__,					SCREAMING_SNAKE_CASE_      )      )
										if len(SCREAMING_SNAKE_CASE_      ) > 0:
															print(
															    'Ignored weights of {} not initialized from pretrained model: {}'.format(
															        model.__class__.__name__,					SCREAMING_SNAKE_CASE_      )      )
										if len(SCREAMING_SNAKE_CASE_      ) > 0:
															print('\n'.join(SCREAMING_SNAKE_CASE_      )      )

					SCREAMING_SNAKE_CASE           =     {
					    'pretrained': False,
					    'use_shared_rel_pos_bias': True,
					    'use_abs_pos_emb': False,
					    'init_values': 0.1,
					}

					if is_finetuned:
										model_kwargs.update(
										    {
										        'num_classes': 1_0_0_0,
										        'use_mean_pooling': True,
										        'init_scale': 0.001,
										        'use_rel_pos_bias': True,
										    }      )

					SCREAMING_SNAKE_CASE           =     create_model(
					    'beit_large_patch16_224' if is_large else 'beit_base_patch16_224',					**SCREAMING_SNAKE_CASE_,					)
					SCREAMING_SNAKE_CASE           =     model.patch_embed.patch_size
					SCREAMING_SNAKE_CASE           =     (args.input_size // patch_size[0], args.input_size // patch_size[1])
					SCREAMING_SNAKE_CASE           =     torch.load(args.beit_checkpoint,					map_location='cpu'      )

					print(f'''Load ckpt from {args.beit_checkpoint}'''      )
					SCREAMING_SNAKE_CASE           =     None
					for model_key in ("model", "module"):
										if model_key in checkpoint:
															SCREAMING_SNAKE_CASE           =     checkpoint[model_key]
															print(f'''Load state_dict by model_key = {model_key}'''      )
															break

					SCREAMING_SNAKE_CASE           =     list(checkpoint_model.keys()      )
					for key in all_keys:
										if "relative_position_index" in key:
															checkpoint_model.pop(SCREAMING_SNAKE_CASE_      )

										if "relative_position_bias_table" in key:
															SCREAMING_SNAKE_CASE           =     checkpoint_model[key]
															SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE           =     rel_pos_bias.size()
															SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE           =     model.state_dict()[key].size()
															SCREAMING_SNAKE_CASE           =     model.patch_embed.patch_shape
															if dst_patch_shape[0] != dst_patch_shape[1]:
																				raise NotImplementedError()

					load_state_dict(SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					prefix=''      )

					return model

def   UpperCamelCase_  (		):
					SCREAMING_SNAKE_CASE           =     get_args()

					SCREAMING_SNAKE_CASE           =     'ft1k' in args.hf_checkpoint_name
					SCREAMING_SNAKE_CASE           =     'large' in args.hf_checkpoint_name

					if is_finetuned:
										# To convert Beit's data2vec_vision to HF you need to copy
										# https://github.com/facebookresearch/data2vec_vision/blob/main/beit/modeling_finetune.py
										# into this folder.
										import modeling_finetune  # noqa: F401
					else:
										# To convert Beit's data2vec_vision to HF you need to copy
										# https://github.com/facebookresearch/data2vec_vision/blob/main/beit/modeling_cyclical.py
										# into this folder
										# IMPORTANT: Note that for now we've only converted the down-stream
										# model and not the full pretrained model. This means for the integration
										# test you need to add a `return x` after the following line:
										# https://github.com/facebookresearch/data2vec_vision/blob/af9a36349aaed59ae66e69b5dabeef2d62fdc5da/beit/modeling_cyclical.py#L197
										# to make the integration test pass.
										import modeling_cyclical  # noqa: F401

					# 1. Create model config
					SCREAMING_SNAKE_CASE           =     DataaVecVisionConfig()
					if is_finetuned:
										SCREAMING_SNAKE_CASE           =     True
										SCREAMING_SNAKE_CASE           =     False
										SCREAMING_SNAKE_CASE           =     True
										SCREAMING_SNAKE_CASE           =     1_0_0_0

										SCREAMING_SNAKE_CASE           =     'huggingface/label-files'
										SCREAMING_SNAKE_CASE           =     'imagenet-1k-id2label.json'
										SCREAMING_SNAKE_CASE           =     json.load(open(hf_hub_download(SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					repo_type='dataset'      ),					'r'      )      )
										SCREAMING_SNAKE_CASE           =     {int(SCREAMING_SNAKE_CASE_      ): v for k, v in idalabel.items()}
										SCREAMING_SNAKE_CASE           =     idalabel
										SCREAMING_SNAKE_CASE           =     {v: k for k, v in idalabel.items()}
					else:
										SCREAMING_SNAKE_CASE           =     False
										SCREAMING_SNAKE_CASE           =     True
										SCREAMING_SNAKE_CASE           =     False

					if is_large:
										SCREAMING_SNAKE_CASE           =     1_0_2_4
										SCREAMING_SNAKE_CASE           =     4_0_9_6
										SCREAMING_SNAKE_CASE           =     2_4
										SCREAMING_SNAKE_CASE           =     1_6

					# 2. Load Beit model
					SCREAMING_SNAKE_CASE           =     load_beit_model(SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_      )
					orig_model.eval()

					# 3. Forward Beit model
					SCREAMING_SNAKE_CASE           =     BeitImageProcessor(size=config.image_size,					do_center_crop=SCREAMING_SNAKE_CASE_      )
					SCREAMING_SNAKE_CASE           =     Image.open('../../../../tests/fixtures/tests_samples/COCO/000000039769.png'      )
					SCREAMING_SNAKE_CASE           =     image_processor(images=SCREAMING_SNAKE_CASE_,					return_tensors='pt'      )
					SCREAMING_SNAKE_CASE           =     encoding['pixel_values']

					SCREAMING_SNAKE_CASE           =     (pixel_values,) if is_finetuned else (pixel_values, None)
					with torch.no_grad():
										SCREAMING_SNAKE_CASE           =     orig_model(*SCREAMING_SNAKE_CASE_      )

					# 4. Load HF Data2VecVision model
					if is_finetuned:
										SCREAMING_SNAKE_CASE           =     DataaVecVisionForImageClassification(SCREAMING_SNAKE_CASE_      )
										hf_model.eval()
										SCREAMING_SNAKE_CASE           =     False
										SCREAMING_SNAKE_CASE           =     'data2vec_vision.'
					else:
										SCREAMING_SNAKE_CASE           =     DataaVecVisionModel(SCREAMING_SNAKE_CASE_      )
										hf_model.eval()
										SCREAMING_SNAKE_CASE           =     True
										SCREAMING_SNAKE_CASE           =     ''

					SCREAMING_SNAKE_CASE           =     create_rename_keys(SCREAMING_SNAKE_CASE_,					hf_prefix=SCREAMING_SNAKE_CASE_,					has_lm_head=SCREAMING_SNAKE_CASE_      )
					SCREAMING_SNAKE_CASE           =     orig_model.state_dict()
					for src, dest in rename_keys:
										SCREAMING_SNAKE_CASE           =     state_dict.pop(SCREAMING_SNAKE_CASE_      )
										SCREAMING_SNAKE_CASE           =     val

					read_in_q_k_v(SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					hf_prefix=SCREAMING_SNAKE_CASE_,					has_lm_head=SCREAMING_SNAKE_CASE_      )
					SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE           =     hf_model.load_state_dict(SCREAMING_SNAKE_CASE_,					strict=SCREAMING_SNAKE_CASE_      )
					print('HF missing',					SCREAMING_SNAKE_CASE_      )
					print('HF unexpected_keys',					SCREAMING_SNAKE_CASE_      )

					# 5. Forward HF Data2VecVision model
					with torch.no_grad():
										SCREAMING_SNAKE_CASE           =     hf_model(SCREAMING_SNAKE_CASE_      )

					SCREAMING_SNAKE_CASE           =     hf_model_output.logits if is_finetuned else hf_model_output.last_hidden_state

					# 6. Compare
					SCREAMING_SNAKE_CASE           =     torch.max(torch.abs(hf_output - orig_model_output      )      ).item()

					print(f'''max_absolute_diff = {max_absolute_diff}'''      )
					SCREAMING_SNAKE_CASE           =     torch.allclose(SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_,					atol=1E-3      )
					print('Do both models output the same tensors?',					'🔥' if success else '💩'      )
					if not success:
										raise Exception('Something went wRoNg'      )

					# 7. Save
					print(f'''Saving to {args.hf_checkpoint_name}'''      )
					hf_model.save_pretrained(args.hf_checkpoint_name      )
					image_processor.save_pretrained(args.hf_checkpoint_name      )


if __name__ == "__main__":
	main()
	# Run the following to convert checkpoints
	#  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
	#          --beit_checkpoint ./pretrained_base.pt \
	#          --hf_checkpoint_name "./data2vec-vision-base"
	#  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
	#          --beit_checkpoint ./finetuned_base.pt \
	#          --hf_checkpoint_name "./data2vec-vision-base-ft1k"
	#  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
	#          --beit_checkpoint ./pretrained_large.pt \
	#          --hf_checkpoint_name "./data2vec-vision-large"
	#  python ./convert_data2vec_vision_original_pytorch_checkpoint_to_pytorch.py \
	#          --beit_checkpoint ./finetuned_large.pt \
	#          --hf_checkpoint_name "./data2vec-vision-large-ft1k"


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


snake_case			       =		'block_records.npy'


snake_case			       =		logging.get_logger(__name__)

def   UpperCamelCase_  (		SCREAMING_SNAKE_CASE_,					SCREAMING_SNAKE_CASE_      ):
					import tensorflow.compat.va as tf

					SCREAMING_SNAKE_CASE           =     tf.data.TFRecordDataset(SCREAMING_SNAKE_CASE_,					buffer_size=5_1_2 * 1_0_2_4 * 1_0_2_4      )
					SCREAMING_SNAKE_CASE           =     blocks_dataset.batch(SCREAMING_SNAKE_CASE_,					drop_remainder=SCREAMING_SNAKE_CASE_      )
					SCREAMING_SNAKE_CASE           =     next(blocks_dataset.take(1      ).as_numpy_iterator()      )

					return np_record




class 					UpperCamelCase     :

				"""simple docstring"""



				def __init__(    self		,     lowercase__		,     lowercase__		,     lowercase__=2		,     lowercase__=1000		,     lowercase__=100		,     lowercase__=100000		,     )   ->   Dict:
									"""simple docstring"""







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

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

									SCREAMING_SNAKE_CASE           =     builder.build()




				def 							A (    self		,     lowercase__      )   ->   str:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE     ,   SCREAMING_SNAKE_CASE           =     self.searcher.search_batched(question_projection.detach().cpu()      )
									return retrieved_block_ids.astype('int64'      )






class 					UpperCamelCase     :

				"""simple docstring"""



				def __init__(    self		,     lowercase__		,     lowercase__      )   ->   Dict:
									"""simple docstring"""







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


				def __call__(    self		,     lowercase__		,     lowercase__		,     lowercase__		,     lowercase__=None		,     lowercase__="pt"      )   ->   Optional[int]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     np.take(self.block_records		,     indices=lowercase__		,     axis=0      )

									SCREAMING_SNAKE_CASE           =     self.tokenizer.decode(question_input_ids[0]		,     skip_special_tokens=lowercase__      )

									SCREAMING_SNAKE_CASE           =     []
									SCREAMING_SNAKE_CASE           =     []
									for retrieved_block in retrieved_blocks:
														text.append(lowercase__      )
														text_pair.append(retrieved_block.decode()      )

									SCREAMING_SNAKE_CASE           =     self.tokenizer(
									    lowercase__		,     lowercase__		,     padding=lowercase__		,     truncation=lowercase__		,     return_special_tokens_mask=lowercase__		,     max_length=lowercase__      )
									SCREAMING_SNAKE_CASE           =     concat_inputs.convert_to_tensors(lowercase__      )

									if answer_ids is not None:
														return self.block_has_answer(lowercase__		,     lowercase__      ) + (concat_inputs_tensors,)
									else:
														return (None, None, None, concat_inputs_tensors)


				@classmethod
				def 							A (    cls		,     lowercase__		,     *lowercase__		,     **lowercase__      )   ->   Optional[Any]:
									"""simple docstring"""







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

									SCREAMING_SNAKE_CASE           =     AutoTokenizer.from_pretrained(lowercase__		,     *lowercase__		,     **lowercase__      )

									return cls(lowercase__		,     lowercase__      )


				def 							A (    self		,     lowercase__      )   ->   List[str]:
									"""simple docstring"""







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




				def 							A (    self		,     lowercase__		,     lowercase__      )   ->   Optional[Any]:
									"""simple docstring"""







									SCREAMING_SNAKE_CASE           =     []
									SCREAMING_SNAKE_CASE           =     []
									SCREAMING_SNAKE_CASE           =     []
									SCREAMING_SNAKE_CASE           =     0

									for input_id in concat_inputs.input_ids:
														SCREAMING_SNAKE_CASE           =     input_id.tolist()
														# Check answers between two [SEP] tokens
														SCREAMING_SNAKE_CASE           =     input_id_list.index(self.tokenizer.sep_token_id      )
														SCREAMING_SNAKE_CASE           =     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		,     lowercase__      ):
																								if answer[0] == input_id_list[idx]:
																													if input_id_list[idx : idx + len(lowercase__      )] == answer:
																																		start_pos[-1].append(lowercase__      )
																																		end_pos[-1].append(idx + len(lowercase__      ) - 1      )

														if len(start_pos[-1]      ) == 0:
																			has_answers.append(lowercase__      )
														else:
																			has_answers.append(lowercase__      )
																			if len(start_pos[-1]      ) > max_answers:
																								SCREAMING_SNAKE_CASE           =     len(start_pos[-1]      )

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