


import argparse
import json
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from typing import Callable, Dict, List, Tuple

import timm
import torch
import torch.nn as nn
from classy_vision.models.regnet import RegNet, RegNetParams, RegNetYaagf, RegNetYaagf, RegNetYaaagf
from huggingface_hub import cached_download, hf_hub_url
from torch import Tensor
from vissl.models.model_helpers import get_trunk_forward_outputs

from transformers import AutoImageProcessor, RegNetConfig, RegNetForImageClassification, RegNetModel
from transformers.utils import logging


logging.set_verbosity_info()
a       	=  logging.get_logger()







@dataclass
class 							UpperCamelCase__ :
					__SCREAMING_SNAKE_CASE							:	nn.Module
					__SCREAMING_SNAKE_CASE							:	List[nn.Module]            = field(default_factory=lowercase_ )
					__SCREAMING_SNAKE_CASE							:	list            = field(default_factory=lowercase_ )

					def 	UpperCAmelCase__							(  self							: str				,       UpperCamelCase__							: Any				,       UpperCamelCase__							: Dict				,       UpperCamelCase__							: Optional[int]							):






										'''simple docstring'''






										lowercase_					=      len(list(m.modules()							)							) == 1 or isinstance(UpperCamelCase__				,       nn.Convad							) or isinstance(UpperCamelCase__				,       nn.BatchNormad							)
										if has_not_submodules:
															self.traced.append(UpperCamelCase__							)

					def __call__(  self							: str				,       UpperCamelCase__							: Tuple							):






										'''simple docstring'''






										for m in self.module.modules():
															self.handles.append(m.register_forward_hook(self._forward_hook							)							)
										self.module(UpperCamelCase__							)
										[x.remove() for x in self.handles]
										return self







					@property
					def 	UpperCAmelCase__							(  self							: Any							):






										'''simple docstring'''






										return list(filter(lambda UpperCamelCase__							: len(list(x.state_dict().keys()							)							) > 0				,       self.traced							)							)







@dataclass
class 							UpperCamelCase__ :
					__SCREAMING_SNAKE_CASE							:	nn.Module
					__SCREAMING_SNAKE_CASE							:	nn.Module
					__SCREAMING_SNAKE_CASE							:	int            = 1
					__SCREAMING_SNAKE_CASE							:	List            = field(default_factory=lowercase_ )
					__SCREAMING_SNAKE_CASE							:	List            = field(default_factory=lowercase_ )
					__SCREAMING_SNAKE_CASE							:	bool            = True

					def __call__(  self							: int				,       UpperCamelCase__							: Tuple							):






										'''simple docstring'''






										lowercase_					=      Tracker(self.dest							)(UpperCamelCase__							).parametrized
										lowercase_					=      Tracker(self.src							)(UpperCamelCase__							).parametrized

										lowercase_					=      list(filter(lambda UpperCamelCase__							: type(UpperCamelCase__							) not in self.src_skip				,       UpperCamelCase__							)							)
										lowercase_					=      list(filter(lambda UpperCamelCase__							: type(UpperCamelCase__							) not in self.dest_skip				,       UpperCamelCase__							)							)

										if len(UpperCamelCase__							) != len(UpperCamelCase__							) and self.raise_if_mismatch:
															raise Exception(
															    F'''Numbers of operations are different. Source module has {len(UpperCamelCase__							)} operations while'''
															    F''' destination module has {len(UpperCamelCase__							)}.'''							)

										for dest_m, src_m in zip(UpperCamelCase__				,       UpperCamelCase__							):
															dest_m.load_state_dict(src_m.state_dict()							)
															if self.verbose == 1:
																				print(F'''Transfered from={src_m} to={dest_m}'''							)







class 							UpperCamelCase__ (						nn.Module ):

					def __init__(  self							: List[str]				,       UpperCamelCase__							: List[str]							):






										'''simple docstring'''






										super().__init__()

										lowercase_					=      []
										# - get the stem
										feature_blocks.append(("""conv1""", model.stem)							)
										# - get all the feature blocks
										for k, v in model.trunk_output.named_children():
															assert k.startswith("""block"""							), F'''Unexpected layer name {k}'''
															lowercase_					=      len(UpperCamelCase__							) + 1
															feature_blocks.append((F'''res{block_index}''', v)							)

										lowercase_					=      nn.ModuleDict(UpperCamelCase__							)







					def 	UpperCAmelCase__							(  self							: int				,       UpperCamelCase__							: Optional[int]							):






										'''simple docstring'''






										return get_trunk_forward_outputs(
										    UpperCamelCase__				,       out_feat_keys=UpperCamelCase__				,       feature_blocks=self._feature_blocks				,       )







class 							UpperCamelCase__ (						lowercase_ ):

					def 	UpperCAmelCase__							(  self							: Tuple				,       UpperCamelCase__							: str							):






										'''simple docstring'''






										lowercase_					=      x.split("""-"""							)
										return x_split[0] + x_split[1] + "_" + "".join(x_split[2:]							)







					def __getitem__(  self							: Optional[Any]				,       UpperCamelCase__							: Optional[int]							):






										'''simple docstring'''






										if x not in self:
															lowercase_					=      self.convert_name_to_timm(UpperCamelCase__							)
															lowercase_					=      partial(lambda: (timm.create_model(UpperCamelCase__				,       pretrained=UpperCamelCase__							).eval(), None)							)

										else:
															lowercase_					=      super().__getitem__(UpperCamelCase__							)

										return val







class 							UpperCamelCase__ (						lowercase_ ):

					def __getitem__(  self							: int				,       UpperCamelCase__							: Dict							):






										'''simple docstring'''






										if "seer" in x and "in1k" not in x:
															lowercase_					=      RegNetModel
										else:
															lowercase_					=      RegNetForImageClassification
										return val



def       UpperCAmelCase_					(   UpperCAmelCase__    ,				UpperCAmelCase__    ,				UpperCAmelCase__	):






					for from_key, to_key in keys:
										lowercase_					=      from_state_dict[from_key].clone()
										print(F'''Copied key={from_key} to={to_key}'''	)
					return to_state_dict


def       UpperCAmelCase_					(   UpperCAmelCase__    ,				UpperCAmelCase__    ,				UpperCAmelCase__    ,				UpperCAmelCase__    ,				UpperCAmelCase__    ,				UpperCAmelCase__ = True    ,				):
					print(F'''Converting {name}...'''	)
					with torch.no_grad():
										lowercase_				,       lowercase_					=      from_model_func()
										lowercase_					=      our_model_func(UpperCamelCase__	).eval()
										lowercase_					=      ModuleTransfer(src=UpperCamelCase__    ,				dest=UpperCamelCase__    ,				raise_if_mismatch=UpperCamelCase__	)
										lowercase_					=      torch.randn((1, 3, 2_2_4, 2_2_4)	)
										module_transfer(UpperCamelCase__	)

					if from_state_dict is not None:
										lowercase_					=      []
										# for seer - in1k finetuned we have to manually copy the head
										if "seer" in name and "in1k" in name:
															lowercase_					=      [("""0.clf.0.weight""", """classifier.1.weight"""), ("""0.clf.0.bias""", """classifier.1.bias""")]
										lowercase_					=      manually_copy_vissl_head(UpperCamelCase__    ,				our_model.state_dict()    ,				UpperCamelCase__	)
										our_model.load_state_dict(UpperCamelCase__	)

					lowercase_					=      our_model(UpperCamelCase__    ,				output_hidden_states=UpperCamelCase__	)
					lowercase_					=      (
					    our_outputs.logits if isinstance(UpperCamelCase__    ,				UpperCamelCase__	) else our_outputs.last_hidden_state
					)

					lowercase_					=      from_model(UpperCamelCase__	)
					lowercase_					=      from_output[-1] if type(UpperCamelCase__	) is list else from_output

					# now since I don't want to use any config files, vissl seer model doesn't actually have an head, so let's just check the last hidden state
					if "seer" in name and "in1k" in name:
										lowercase_					=      our_outputs.hidden_states[-1]

					assert torch.allclose(UpperCamelCase__    ,				UpperCamelCase__	), "The model logits don't match the original one."

					if push_to_hub:
										our_model.push_to_hub(
										    repo_path_or_name=save_directory / name    ,				commit_message="""Add model"""    ,				use_temp_dir=UpperCamelCase__    ,				)

										lowercase_					=      2_2_4 if """seer""" not in name else 3_8_4
										# we can use the convnext one
										lowercase_					=      AutoImageProcessor.from_pretrained("""facebook/convnext-base-224-22k-1k"""    ,				size=UpperCamelCase__	)
										image_processor.push_to_hub(
										    repo_path_or_name=save_directory / name    ,				commit_message="""Add image processor"""    ,				use_temp_dir=UpperCamelCase__    ,				)

										print(F'''Pushed {name}'''	)


def       UpperCAmelCase_					(   UpperCAmelCase__    ,				UpperCAmelCase__ = None    ,				UpperCAmelCase__ = True	):
					lowercase_					=      """imagenet-1k-id2label.json"""
					lowercase_					=      1_0_0_0
					lowercase_					=      (1, num_labels)

					lowercase_					=      """huggingface/label-files"""
					lowercase_					=      num_labels
					lowercase_					=      json.load(open(cached_download(hf_hub_url(UpperCamelCase__    ,				UpperCamelCase__    ,				repo_type="""dataset"""	)	)    ,				"""r"""	)	)
					lowercase_					=      {int(UpperCamelCase__	): v for k, v in idalabel.items()}

					lowercase_					=      idalabel
					lowercase_					=      {v: k for k, v in idalabel.items()}

					lowercase_					=      partial(UpperCamelCase__    ,				num_labels=UpperCamelCase__    ,				idalabel=UpperCamelCase__    ,				labelaid=UpperCamelCase__	)

					lowercase_					=      {
					    """regnet-x-002""": ImageNetPreTrainedConfig(
					        depths=[1, 1, 4, 7]    ,				hidden_sizes=[2_4, 5_6, 1_5_2, 3_6_8]    ,				groups_width=8    ,				layer_type="""x"""	),
					    """regnet-x-004""": ImageNetPreTrainedConfig(
					        depths=[1, 2, 7, 1_2]    ,				hidden_sizes=[3_2, 6_4, 1_6_0, 3_8_4]    ,				groups_width=1_6    ,				layer_type="""x"""	),
					    """regnet-x-006""": ImageNetPreTrainedConfig(
					        depths=[1, 3, 5, 7]    ,				hidden_sizes=[4_8, 9_6, 2_4_0, 5_2_8]    ,				groups_width=2_4    ,				layer_type="""x"""	),
					    """regnet-x-008""": ImageNetPreTrainedConfig(
					        depths=[1, 3, 7, 5]    ,				hidden_sizes=[6_4, 1_2_8, 2_8_8, 6_7_2]    ,				groups_width=1_6    ,				layer_type="""x"""	),
					    """regnet-x-016""": ImageNetPreTrainedConfig(
					        depths=[2, 4, 1_0, 2]    ,				hidden_sizes=[7_2, 1_6_8, 4_0_8, 9_1_2]    ,				groups_width=2_4    ,				layer_type="""x"""	),
					    """regnet-x-032""": ImageNetPreTrainedConfig(
					        depths=[2, 6, 1_5, 2]    ,				hidden_sizes=[9_6, 1_9_2, 4_3_2, 1_0_0_8]    ,				groups_width=4_8    ,				layer_type="""x"""	),
					    """regnet-x-040""": ImageNetPreTrainedConfig(
					        depths=[2, 5, 1_4, 2]    ,				hidden_sizes=[8_0, 2_4_0, 5_6_0, 1_3_6_0]    ,				groups_width=4_0    ,				layer_type="""x"""	),
					    """regnet-x-064""": ImageNetPreTrainedConfig(
					        depths=[2, 4, 1_0, 1]    ,				hidden_sizes=[1_6_8, 3_9_2, 7_8_4, 1_6_2_4]    ,				groups_width=5_6    ,				layer_type="""x"""	),
					    """regnet-x-080""": ImageNetPreTrainedConfig(
					        depths=[2, 5, 1_5, 1]    ,				hidden_sizes=[8_0, 2_4_0, 7_2_0, 1_9_2_0]    ,				groups_width=1_2_0    ,				layer_type="""x"""	),
					    """regnet-x-120""": ImageNetPreTrainedConfig(
					        depths=[2, 5, 1_1, 1]    ,				hidden_sizes=[2_2_4, 4_4_8, 8_9_6, 2_2_4_0]    ,				groups_width=1_1_2    ,				layer_type="""x"""	),
					    """regnet-x-160""": ImageNetPreTrainedConfig(
					        depths=[2, 6, 1_3, 1]    ,				hidden_sizes=[2_5_6, 5_1_2, 8_9_6, 2_0_4_8]    ,				groups_width=1_2_8    ,				layer_type="""x"""	),
					    """regnet-x-320""": ImageNetPreTrainedConfig(
					        depths=[2, 7, 1_3, 1]    ,				hidden_sizes=[3_3_6, 6_7_2, 1_3_4_4, 2_5_2_0]    ,				groups_width=1_6_8    ,				layer_type="""x"""	),
					    # y variant
					    """regnet-y-002""": ImageNetPreTrainedConfig(depths=[1, 1, 4, 7]    ,				hidden_sizes=[2_4, 5_6, 1_5_2, 3_6_8]    ,				groups_width=8	),
					    """regnet-y-004""": ImageNetPreTrainedConfig(
					        depths=[1, 3, 6, 6]    ,				hidden_sizes=[4_8, 1_0_4, 2_0_8, 4_4_0]    ,				groups_width=8	),
					    """regnet-y-006""": ImageNetPreTrainedConfig(
					        depths=[1, 3, 7, 4]    ,				hidden_sizes=[4_8, 1_1_2, 2_5_6, 6_0_8]    ,				groups_width=1_6	),
					    """regnet-y-008""": ImageNetPreTrainedConfig(
					        depths=[1, 3, 8, 2]    ,				hidden_sizes=[6_4, 1_2_8, 3_2_0, 7_6_8]    ,				groups_width=1_6	),
					    """regnet-y-016""": ImageNetPreTrainedConfig(
					        depths=[2, 6, 1_7, 2]    ,				hidden_sizes=[4_8, 1_2_0, 3_3_6, 8_8_8]    ,				groups_width=2_4	),
					    """regnet-y-032""": ImageNetPreTrainedConfig(
					        depths=[2, 5, 1_3, 1]    ,				hidden_sizes=[7_2, 2_1_6, 5_7_6, 1_5_1_2]    ,				groups_width=2_4	),
					    """regnet-y-040""": ImageNetPreTrainedConfig(
					        depths=[2, 6, 1_2, 2]    ,				hidden_sizes=[1_2_8, 1_9_2, 5_1_2, 1_0_8_8]    ,				groups_width=6_4	),
					    """regnet-y-064""": ImageNetPreTrainedConfig(
					        depths=[2, 7, 1_4, 2]    ,				hidden_sizes=[1_4_4, 2_8_8, 5_7_6, 1_2_9_6]    ,				groups_width=7_2	),
					    """regnet-y-080""": ImageNetPreTrainedConfig(
					        depths=[2, 4, 1_0, 1]    ,				hidden_sizes=[1_6_8, 4_4_8, 8_9_6, 2_0_1_6]    ,				groups_width=5_6	),
					    """regnet-y-120""": ImageNetPreTrainedConfig(
					        depths=[2, 5, 1_1, 1]    ,				hidden_sizes=[2_2_4, 4_4_8, 8_9_6, 2_2_4_0]    ,				groups_width=1_1_2	),
					    """regnet-y-160""": ImageNetPreTrainedConfig(
					        depths=[2, 4, 1_1, 1]    ,				hidden_sizes=[2_2_4, 4_4_8, 1_2_3_2, 3_0_2_4]    ,				groups_width=1_1_2	),
					    """regnet-y-320""": ImageNetPreTrainedConfig(
					        depths=[2, 5, 1_2, 1]    ,				hidden_sizes=[2_3_2, 6_9_6, 1_3_9_2, 3_7_1_2]    ,				groups_width=2_3_2	),
					    # models created by SEER -> https://arxiv.org/abs/2202.08360
					    """regnet-y-320-seer""": RegNetConfig(depths=[2, 5, 1_2, 1]    ,				hidden_sizes=[2_3_2, 6_9_6, 1_3_9_2, 3_7_1_2]    ,				groups_width=2_3_2	),
					    """regnet-y-640-seer""": RegNetConfig(depths=[2, 5, 1_2, 1]    ,				hidden_sizes=[3_2_8, 9_8_4, 1_9_6_8, 4_9_2_0]    ,				groups_width=3_2_8	),
					    """regnet-y-1280-seer""": RegNetConfig(
					        depths=[2, 7, 1_7, 1]    ,				hidden_sizes=[5_2_8, 1_0_5_6, 2_9_0_4, 7_3_9_2]    ,				groups_width=2_6_4	),
					    """regnet-y-2560-seer""": RegNetConfig(
					        depths=[3, 7, 1_6, 1]    ,				hidden_sizes=[6_4_0, 1_6_9_6, 2_5_4_4, 5_0_8_8]    ,				groups_width=6_4_0	),
					    """regnet-y-10b-seer""": ImageNetPreTrainedConfig(
					        depths=[2, 7, 1_7, 1]    ,				hidden_sizes=[2_0_2_0, 4_0_4_0, 1_1_1_1_0, 2_8_2_8_0]    ,				groups_width=1_0_1_0	),
					    # finetuned on imagenet
					    """regnet-y-320-seer-in1k""": ImageNetPreTrainedConfig(
					        depths=[2, 5, 1_2, 1]    ,				hidden_sizes=[2_3_2, 6_9_6, 1_3_9_2, 3_7_1_2]    ,				groups_width=2_3_2	),
					    """regnet-y-640-seer-in1k""": ImageNetPreTrainedConfig(
					        depths=[2, 5, 1_2, 1]    ,				hidden_sizes=[3_2_8, 9_8_4, 1_9_6_8, 4_9_2_0]    ,				groups_width=3_2_8	),
					    """regnet-y-1280-seer-in1k""": ImageNetPreTrainedConfig(
					        depths=[2, 7, 1_7, 1]    ,				hidden_sizes=[5_2_8, 1_0_5_6, 2_9_0_4, 7_3_9_2]    ,				groups_width=2_6_4	),
					    """regnet-y-2560-seer-in1k""": ImageNetPreTrainedConfig(
					        depths=[3, 7, 1_6, 1]    ,				hidden_sizes=[6_4_0, 1_6_9_6, 2_5_4_4, 5_0_8_8]    ,				groups_width=6_4_0	),
					    """regnet-y-10b-seer-in1k""": ImageNetPreTrainedConfig(
					        depths=[2, 7, 1_7, 1]    ,				hidden_sizes=[2_0_2_0, 4_0_4_0, 1_1_1_1_0, 2_8_2_8_0]    ,				groups_width=1_0_1_0	),
					}

					lowercase_					=      NameToOurModelFuncMap()
					lowercase_					=      NameToFromModelFuncMap()
					# add seer weights logic

					def load_using_classy_vision(UpperCAmelCase__    ,				UpperCAmelCase__	) -> Tuple[nn.Module, Dict]:
										lowercase_					=      torch.hub.load_state_dict_from_url(UpperCamelCase__    ,				model_dir=str(UpperCamelCase__	)    ,				map_location="""cpu"""	)
										lowercase_					=      model_func()
										# check if we have a head, if yes add it
										lowercase_					=      files["""classy_state_dict"""]["""base_model"""]["""model"""]
										lowercase_					=      model_state_dict["""trunk"""]
										model.load_state_dict(UpperCamelCase__	)
										return model.eval(), model_state_dict["heads"]

					# pretrained
					lowercase_					=      partial(
					    UpperCamelCase__    ,				"""https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet32d/seer_regnet32gf_model_iteration244000.torch"""    ,				lambda: FakeRegNetVisslWrapper(RegNetYaagf()	)    ,				)

					lowercase_					=      partial(
					    UpperCamelCase__    ,				"""https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet64/seer_regnet64gf_model_final_checkpoint_phase0.torch"""    ,				lambda: FakeRegNetVisslWrapper(RegNetYaagf()	)    ,				)

					lowercase_					=      partial(
					    UpperCamelCase__    ,				"""https://dl.fbaipublicfiles.com/vissl/model_zoo/swav_ig1b_regnet128Gf_cnstant_bs32_node16_sinkhorn10_proto16k_syncBN64_warmup8k/model_final_checkpoint_phase0.torch"""    ,				lambda: FakeRegNetVisslWrapper(RegNetYaaagf()	)    ,				)

					lowercase_					=      partial(
					    UpperCamelCase__    ,				"""https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_regnet10B/model_iteration124500_conso.torch"""    ,				lambda: FakeRegNetVisslWrapper(
					        RegNet(RegNetParams(depth=2_7    ,				group_width=1_0_1_0    ,				w_a=1_7_4_4    ,				w_a=6_2_0.8_3    ,				w_m=2.52	)	)	)    ,				)

					# IN1K finetuned
					lowercase_					=      partial(
					    UpperCamelCase__    ,				"""https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_regnet32_finetuned_in1k_model_final_checkpoint_phase78.torch"""    ,				lambda: FakeRegNetVisslWrapper(RegNetYaagf()	)    ,				)

					lowercase_					=      partial(
					    UpperCamelCase__    ,				"""https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_regnet64_finetuned_in1k_model_final_checkpoint_phase78.torch"""    ,				lambda: FakeRegNetVisslWrapper(RegNetYaagf()	)    ,				)

					lowercase_					=      partial(
					    UpperCamelCase__    ,				"""https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_regnet128_finetuned_in1k_model_final_checkpoint_phase78.torch"""    ,				lambda: FakeRegNetVisslWrapper(RegNetYaaagf()	)    ,				)

					lowercase_					=      partial(
					    UpperCamelCase__    ,				"""https://dl.fbaipublicfiles.com/vissl/model_zoo/seer_finetuned/seer_10b_finetuned_in1k_model_phase28_conso.torch"""    ,				lambda: FakeRegNetVisslWrapper(
					        RegNet(RegNetParams(depth=2_7    ,				group_width=1_0_1_0    ,				w_a=1_7_4_4    ,				w_a=6_2_0.8_3    ,				w_m=2.52	)	)	)    ,				)

					if model_name:
										convert_weight_and_push(
										    UpperCamelCase__    ,				names_to_from_model_map[model_name]    ,				names_to_ours_model_map[model_name]    ,				names_to_config[model_name]    ,				UpperCamelCase__    ,				UpperCamelCase__    ,				)
					else:
										for model_name, config in names_to_config.items():
															convert_weight_and_push(
															    UpperCamelCase__    ,				names_to_from_model_map[model_name]    ,				names_to_ours_model_map[model_name]    ,				UpperCamelCase__    ,				UpperCamelCase__    ,				UpperCamelCase__    ,				)
					return config, expected_shape


if __name__ == "__main__":
							a       	=  argparse.ArgumentParser()
							# Required parameters
							parser.add_argument(
							    '--model_name',
							    default=None,
							    type=str,
							    help=(
							        'The name of the model you wish to convert, it must be one of the supported regnet* architecture,'
							        ' currently: regnetx-*, regnety-*. If `None`, all of them will the converted.'
							    ),
							)
							parser.add_argument(
							    '--pytorch_dump_folder_path',
							    default=None,
							    type=Path,
							    required=True,
							    help='Path to the output PyTorch model directory.',
							)
							parser.add_argument(
							    '--push_to_hub',
							    default=True,
							    type=bool,
							    required=False,
							    help='If True, push model and image processor to the hub.',
							)

							a       	=  parser.parse_args()

							a       	=  args.pytorch_dump_folder_path
							pytorch_dump_folder_path.mkdir(exist_ok=True, parents=True)
							convert_weights_and_push(pytorch_dump_folder_path, args.model_name, args.push_to_hub)





import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
					import faiss







@require_faiss
class 					lowercase							(			lowercase_     ):




				def        a				(   self			):
											snake_case_     							=       tempfile.mkdtemp()
											snake_case_     							=       8

											# DPR tok
											snake_case_     							=       [
											    '[UNK]',
											    '[CLS]',
											    '[SEP]',
											    '[PAD]',
											    '[MASK]',
											    'want',
											    '##want',
											    '##ed',
											    'wa',
											    'un',
											    'runn',
											    '##ing',
											    ',',
											    'low',
											    'lowest',
											]
											snake_case_     							=       os.path.join(self.tmpdirname ,  'dpr_tokenizer'			)
											os.makedirs(snake_case ,  exist_ok=snake_case			)
											snake_case_     							=       os.path.join(snake_case ,  DPR_VOCAB_FILES_NAMES['vocab_file']			)
											with open(self.vocab_file ,  'w' ,  encoding='utf-8'			) as vocab_writer:
																		vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]			)			)

											# BART tok
											snake_case_     							=       [
											    'l',
											    'o',
											    'w',
											    'e',
											    'r',
											    's',
											    't',
											    'i',
											    'd',
											    'n',
											    '\u0120',
											    '\u0120l',
											    '\u0120n',
											    '\u0120lo',
											    '\u0120low',
											    'er',
											    '\u0120lowest',
											    '\u0120newer',
											    '\u0120wider',
											    '<unk>',
											]
											snake_case_     							=       dict(zip(snake_case ,  range(len(snake_case			)			)			)			)
											snake_case_     							=       ['#version: 0.2', '\u0120 l', '\u0120l o', '\u0120lo w', 'e r', '']
											snake_case_     							=       {'unk_token': '<unk>'}

											snake_case_     							=       os.path.join(self.tmpdirname ,  'bart_tokenizer'			)
											os.makedirs(snake_case ,  exist_ok=snake_case			)
											snake_case_     							=       os.path.join(snake_case ,  BART_VOCAB_FILES_NAMES['vocab_file']			)
											snake_case_     							=       os.path.join(snake_case ,  BART_VOCAB_FILES_NAMES['merges_file']			)
											with open(self.vocab_file ,  'w' ,  encoding='utf-8'			) as fp:
																		fp.write(json.dumps(snake_case			) + '\n'			)
											with open(self.merges_file ,  'w' ,  encoding='utf-8'			) as fp:
																		fp.write('\n'.join(snake_case			)			)




				def        a				(   self			):
											return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname ,  'dpr_tokenizer'			)			)




				def        a				(   self			):
											return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname ,  'dpr_tokenizer'			)			)




				def        a				(   self			):
											return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname ,  'bart_tokenizer'			)			)




				def        a				(   self			):
											shutil.rmtree(self.tmpdirname			)




				def        a				(   self			):
											snake_case_     							=       Dataset.from_dict(
											    {
											        'id': ['0', '1'],
											        'text': ['foo', 'bar'],
											        'title': ['Foo', 'Bar'],
											        'embeddings': [np.ones(self.retrieval_vector_size			), 2 * np.ones(self.retrieval_vector_size			)],
											    }			)
											dataset.add_faiss_index('embeddings' ,  string_factory='Flat' ,  metric_type=faiss.METRIC_INNER_PRODUCT			)
											return dataset




				def        a				(   self			):
											snake_case_     							=       self.get_dummy_dataset()
											snake_case_     							=       RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size ,  question_encoder=DPRConfig().to_dict() ,  generator=BartConfig().to_dict() ,  )
											with patch('transformers.models.rag.retrieval_rag.load_dataset'			) as mock_load_dataset:
																		snake_case_     							=       dataset
																		snake_case_     							=       RagRetriever(
																		    snake_case ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer() ,  )
											return retriever




				def        a				(   self ,  snake_case			):
											snake_case_     							=       self.get_dummy_dataset()
											snake_case_     							=       RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size ,  question_encoder=DPRConfig().to_dict() ,  generator=BartConfig().to_dict() ,  index_name='custom' ,  )
											if from_disk:
																		snake_case_     							=       os.path.join(self.tmpdirname ,  'dataset'			)
																		snake_case_     							=       os.path.join(self.tmpdirname ,  'index.faiss'			)
																		dataset.get_index('embeddings'			).save(os.path.join(self.tmpdirname ,  'index.faiss'			)			)
																		dataset.drop_index('embeddings'			)
																		dataset.save_to_disk(os.path.join(self.tmpdirname ,  'dataset'			)			)
																		del dataset
																		snake_case_     							=       RagRetriever(
																		    snake_case ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer() ,  )
											else:
																		snake_case_     							=       RagRetriever(
																		    snake_case ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer() ,  index=CustomHFIndex(config.retrieval_vector_size ,  snake_case			) ,  )
											return retriever




				def        a				(   self			):
											snake_case_     							=       Dataset.from_dict(
											    {
											        'id': ['0', '1'],
											        'text': ['foo', 'bar'],
											        'title': ['Foo', 'Bar'],
											        'embeddings': [np.ones(self.retrieval_vector_size + 1			), 2 * np.ones(self.retrieval_vector_size + 1			)],
											    }			)
											dataset.add_faiss_index('embeddings' ,  string_factory='Flat' ,  metric_type=faiss.METRIC_INNER_PRODUCT			)

											snake_case_     							=       os.path.join(self.tmpdirname ,  'hf_bert_base.hnswSQ8_correct_phi_128.c_index'			)
											dataset.save_faiss_index('embeddings' ,  index_file_name + '.index.dpr'			)
											pickle.dump(dataset['id'] ,  open(index_file_name + '.index_meta.dpr' ,  'wb'			)			)

											snake_case_     							=       os.path.join(self.tmpdirname ,  'psgs_w100.tsv.pkl'			)
											snake_case_     							=       {sample['id']: [sample['text'], sample['title']] for sample in dataset}
											pickle.dump(snake_case ,  open(snake_case ,  'wb'			)			)

											snake_case_     							=       RagConfig(
											    retrieval_vector_size=self.retrieval_vector_size ,  question_encoder=DPRConfig().to_dict() ,  generator=BartConfig().to_dict() ,  index_name='legacy' ,  index_path=self.tmpdirname ,  )
											snake_case_     							=       RagRetriever(
											    snake_case ,  question_encoder_tokenizer=self.get_dpr_tokenizer() ,  generator_tokenizer=self.get_bart_tokenizer()			)
											return retriever




				def        a				(   self			):
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_canonical_hf_index_retriever()
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_							,       snake_case_							,       snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=snake_case			)
											self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)			)
											self.assertEqual(len(snake_case			) ,  2			)
											self.assertEqual(sorted(doc_dicts[0]			) ,  ['embeddings', 'id', 'text', 'title']			)
											self.assertEqual(len(doc_dicts[0]['id']			) ,  snake_case			)
											self.assertEqual(doc_dicts[0]['id'][0] ,  '1'			)  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['id'][0] ,  '0'			)  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist() ,  [[1], [0]]			)




				def        a				(   self			):
											snake_case_     							=       self.get_dummy_canonical_hf_index_retriever()
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		with patch('transformers.models.rag.retrieval_rag.load_dataset'			) as mock_load_dataset:
																									snake_case_     							=       self.get_dummy_dataset()
																									retriever.save_pretrained(snake_case			)
																									snake_case_     							=       RagRetriever.from_pretrained(snake_case			)
																									self.assertIsInstance(snake_case ,  snake_case			)
																									snake_case_     							=       np.array(
																									    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
																									snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=1			)
																									self.assertTrue(out is not None			)




				def        a				(   self			):
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_							,       snake_case_							,       snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=snake_case			)
											self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)			)
											self.assertEqual(len(snake_case			) ,  2			)
											self.assertEqual(sorted(doc_dicts[0]			) ,  ['embeddings', 'id', 'text', 'title']			)
											self.assertEqual(len(doc_dicts[0]['id']			) ,  snake_case			)
											self.assertEqual(doc_dicts[0]['id'][0] ,  '1'			)  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['id'][0] ,  '0'			)  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist() ,  [[1], [0]]			)




				def        a				(   self			):
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		retriever.save_pretrained(snake_case			)
																		snake_case_     							=       RagRetriever.from_pretrained(snake_case			)
																		self.assertIsInstance(snake_case ,  snake_case			)
																		snake_case_     							=       np.array(
																		    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
																		snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=1			)
																		self.assertTrue(out is not None			)




				def        a				(   self			):
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_							,       snake_case_							,       snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=snake_case			)
											self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)			)
											self.assertEqual(len(snake_case			) ,  2			)
											self.assertEqual(sorted(doc_dicts[0]			) ,  ['embeddings', 'id', 'text', 'title']			)
											self.assertEqual(len(doc_dicts[0]['id']			) ,  snake_case			)
											self.assertEqual(doc_dicts[0]['id'][0] ,  '1'			)  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['id'][0] ,  '0'			)  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist() ,  [[1], [0]]			)




				def        a				(   self			):
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		retriever.save_pretrained(snake_case			)
																		snake_case_     							=       RagRetriever.from_pretrained(snake_case			)
																		self.assertIsInstance(snake_case ,  snake_case			)
																		snake_case_     							=       np.array(
																		    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
																		snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=1			)
																		self.assertTrue(out is not None			)




				def        a				(   self			):
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_legacy_index_retriever()
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_							,       snake_case_							,       snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=snake_case			)
											self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)			)
											self.assertEqual(len(snake_case			) ,  2			)
											self.assertEqual(sorted(doc_dicts[0]			) ,  ['text', 'title']			)
											self.assertEqual(len(doc_dicts[0]['text']			) ,  snake_case			)
											self.assertEqual(doc_dicts[0]['text'][0] ,  'bar'			)  # max inner product is reached with second doc
											self.assertEqual(doc_dicts[1]['text'][0] ,  'foo'			)  # max inner product is reached with first doc
											self.assertListEqual(doc_ids.tolist() ,  [[1], [0]]			)




				def        a				(   self			):
											snake_case_     							=       self.get_dummy_legacy_index_retriever()
											with tempfile.TemporaryDirectory() as tmp_dirname:
																		retriever.save_pretrained(snake_case			)
																		snake_case_     							=       RagRetriever.from_pretrained(snake_case			)
																		self.assertIsInstance(snake_case ,  snake_case			)
																		snake_case_     							=       np.array(
																		    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
																		snake_case_     							=       retriever.retrieve(snake_case ,  n_docs=1			)
																		self.assertTrue(out is not None			)




				@require_torch
				@require_tokenizers
				@require_sentencepiece
				def        a				(   self			):
											import torch

											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_canonical_hf_index_retriever()
											snake_case_     							=       [[5, 7], [10, 11]]
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_     							=       retriever(snake_case ,  snake_case ,  prefix=retriever.config.generator.prefix ,  n_docs=snake_case			)
											snake_case_							,       snake_case_							,       snake_case_     							=       (
											    out['context_input_ids'],
											    out['context_attention_mask'],
											    out['retrieved_doc_embeds'],
											)
											self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)			)
											self.assertIsInstance(snake_case ,  snake_case			)
											self.assertIsInstance(snake_case ,  snake_case			)
											self.assertIsInstance(snake_case ,  np.ndarray			)

											snake_case_     							=       retriever(
											    snake_case ,  snake_case ,  prefix=retriever.config.generator.prefix ,  n_docs=snake_case ,  return_tensors='pt' ,  )
											snake_case_							,       snake_case_							,       snake_case_							,       snake_case_     							=       (  # noqa: F841
											    out['context_input_ids'],
											    out['context_attention_mask'],
											    out['retrieved_doc_embeds'],
											    out['doc_ids'],
											)
											self.assertEqual(retrieved_doc_embeds.shape ,  (2, n_docs, self.retrieval_vector_size)			)
											self.assertIsInstance(snake_case ,  torch.Tensor			)
											self.assertIsInstance(snake_case ,  torch.Tensor			)
											self.assertIsInstance(snake_case ,  torch.Tensor			)




				@require_torch
				@require_tokenizers
				@require_sentencepiece
				def        a				(   self			):
											snake_case_     							=       self.get_dpr_ctx_encoder_tokenizer()
											snake_case_     							=       1
											snake_case_     							=       self.get_dummy_custom_hf_index_retriever(from_disk=snake_case			)
											retriever.set_ctx_encoder_tokenizer(snake_case			)

											snake_case_     							=       [[5, 7], [10, 11]]
											snake_case_     							=       np.array(
											    [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)] ,  dtype=np.floataa			)
											snake_case_     							=       retriever(snake_case ,  snake_case ,  prefix=retriever.config.generator.prefix ,  n_docs=snake_case			)

											self.assertEqual(
											    len(snake_case			) ,  6			)  # check whether the retriever output consist of 6 attributes including tokenized docs
											self.assertEqual(
											    all(k in out for k in ('tokenized_doc_ids', 'tokenized_doc_attention_mask')			) ,  snake_case			)  # check for doc token related keys in dictionary.




