






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 		_A						( UpperCAmelCase_ ):

  def    a					(    self						:	str			):
     """simple docstring"""

     __UpperCamelCase	:			str					=					tempfile.mkdtemp()
     __UpperCamelCase	:			Union[str, Any]					=					8

     # DPR tok
     __UpperCamelCase	:			Tuple					=					[
         """[UNK]""",
         """[CLS]""",
         """[SEP]""",
         """[PAD]""",
         """[MASK]""",
         """want""",
         """##want""",
         """##ed""",
         """wa""",
         """un""",
         """runn""",
         """##ing""",
         """,""",
         """low""",
         """lowest""",
     ]
     __UpperCamelCase	:			Any					=					os.path.join(self.tmpdirname   ,				"""dpr_tokenizer"""			)
     os.makedirs(lowerCamelCase__   ,				exist_ok=lowerCamelCase__			)
     __UpperCamelCase	:			str					=					os.path.join(lowerCamelCase__   ,				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
     __UpperCamelCase	:			Tuple					=					[
         """l""",
         """o""",
         """w""",
         """e""",
         """r""",
         """s""",
         """t""",
         """i""",
         """d""",
         """n""",
         """\u0120""",
         """\u0120l""",
         """\u0120n""",
         """\u0120lo""",
         """\u0120low""",
         """er""",
         """\u0120lowest""",
         """\u0120newer""",
         """\u0120wider""",
         """<unk>""",
     ]
     __UpperCamelCase	:			Optional[Any]					=					dict(zip(lowerCamelCase__   ,				range(len(lowerCamelCase__			)			)			)			)
     __UpperCamelCase	:			List[str]					=					["""#version: 0.2""", """\u0120 l""", """\u0120l o""", """\u0120lo w""", """e r""", """"""]
     __UpperCamelCase	:			str					=					{"""unk_token""": """<unk>"""}

     __UpperCamelCase	:			str					=					os.path.join(self.tmpdirname   ,				"""bart_tokenizer"""			)
     os.makedirs(lowerCamelCase__   ,				exist_ok=lowerCamelCase__			)
     __UpperCamelCase	:			Any					=					os.path.join(lowerCamelCase__   ,				BART_VOCAB_FILES_NAMES["""vocab_file"""]			)
     __UpperCamelCase	:			Union[str, Any]					=					os.path.join(lowerCamelCase__   ,				BART_VOCAB_FILES_NAMES["""merges_file"""]			)
     with open(self.vocab_file   ,				"""w"""   ,				encoding="""utf-8"""			) as fp:
        fp.write(json.dumps(lowerCamelCase__			) + """\n"""			)
     with open(self.merges_file   ,				"""w"""   ,				encoding="""utf-8"""			) as fp:
        fp.write("""\n""".join(lowerCamelCase__			)			)

  def    a					(    self						:	List[str]			):
     """simple docstring"""

     return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname   ,				"""dpr_tokenizer"""			)			)

  def    a					(    self						:	Optional[int]			):
     """simple docstring"""

     return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname   ,				"""dpr_tokenizer"""			)			)

  def    a					(    self						:	Union[str, Any]			):
     """simple docstring"""

     return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname   ,				"""bart_tokenizer"""			)			)

  def    a					(    self						:	Tuple			):
     """simple docstring"""

     shutil.rmtree(self.tmpdirname			)

  def    a					(    self						:	Any			):
     """simple docstring"""

     __UpperCamelCase	:			Optional[Any]					=					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						:	List[str]			):
     """simple docstring"""

     __UpperCamelCase	:			Optional[int]					=					self.get_dummy_dataset()
     __UpperCamelCase	:			Optional[Any]					=					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:
        __UpperCamelCase	:			int					=					dataset
        __UpperCamelCase	:			Dict					=					RagRetriever(
            lowerCamelCase__   ,				question_encoder_tokenizer=self.get_dpr_tokenizer()   ,				generator_tokenizer=self.get_bart_tokenizer()   ,				)
     return retriever

  def    a					(    self						:	List[Any]   ,				lowerCamelCase__						:	bool			):
     """simple docstring"""

     __UpperCamelCase	:			List[str]					=					self.get_dummy_dataset()
     __UpperCamelCase	:			Optional[int]					=					RagConfig(
         retrieval_vector_size=self.retrieval_vector_size   ,				question_encoder=DPRConfig().to_dict()   ,				generator=BartConfig().to_dict()   ,				index_name="""custom"""   ,				)
     if from_disk:
        __UpperCamelCase	:			Tuple					=					os.path.join(self.tmpdirname   ,				"""dataset"""			)
        __UpperCamelCase	:			Any					=					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
        __UpperCamelCase	:			int					=					RagRetriever(
            lowerCamelCase__   ,				question_encoder_tokenizer=self.get_dpr_tokenizer()   ,				generator_tokenizer=self.get_bart_tokenizer()   ,				)
     else:
        __UpperCamelCase	:			List[Any]					=					RagRetriever(
            lowerCamelCase__   ,				question_encoder_tokenizer=self.get_dpr_tokenizer()   ,				generator_tokenizer=self.get_bart_tokenizer()   ,				index=CustomHFIndex(config.retrieval_vector_size   ,				lowerCamelCase__			)   ,				)
     return retriever

  def    a					(    self						:	int			):
     """simple docstring"""

     __UpperCamelCase	:			Any					=					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			)

     __UpperCamelCase	:			Optional[int]					=					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"""			)			)

     __UpperCamelCase	:			str					=					os.path.join(self.tmpdirname   ,				"""psgs_w100.tsv.pkl"""			)
     __UpperCamelCase	:			List[str]					=					{sample["""id"""]: [sample["""text"""], sample["""title"""]] for sample in dataset}
     pickle.dump(lowerCamelCase__   ,				open(lowerCamelCase__   ,				"""wb"""			)			)

     __UpperCamelCase	:			Optional[int]					=					RagConfig(
         retrieval_vector_size=self.retrieval_vector_size   ,				question_encoder=DPRConfig().to_dict()   ,				generator=BartConfig().to_dict()   ,				index_name="""legacy"""   ,				index_path=self.tmpdirname   ,				)
     __UpperCamelCase	:			Dict					=					RagRetriever(
         lowerCamelCase__   ,				question_encoder_tokenizer=self.get_dpr_tokenizer()   ,				generator_tokenizer=self.get_bart_tokenizer()			)
     return retriever

  def    a					(    self						:	Tuple			):
     """simple docstring"""

     __UpperCamelCase	:			List[Any]					=					1
     __UpperCamelCase	:			Tuple					=					self.get_dummy_canonical_hf_index_retriever()
     __UpperCamelCase	:			List[str]					=					np.array(
         [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]   ,				dtype=np.floataa			)
     __UpperCamelCase	,     __UpperCamelCase	,     __UpperCamelCase	:			int					=					retriever.retrieve(lowerCamelCase__   ,				n_docs=lowerCamelCase__			)
     self.assertEqual(retrieved_doc_embeds.shape   ,				(2, n_docs, self.retrieval_vector_size)			)
     self.assertEqual(len(lowerCamelCase__			)   ,				2			)
     self.assertEqual(sorted(doc_dicts[0]			)   ,				["""embeddings""", """id""", """text""", """title"""]			)
     self.assertEqual(len(doc_dicts[0]["""id"""]			)   ,				lowerCamelCase__			)
     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						:	Tuple			):
     """simple docstring"""

     __UpperCamelCase	:			Union[str, Any]					=					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:
           __UpperCamelCase	:			Any					=					self.get_dummy_dataset()
           retriever.save_pretrained(lowerCamelCase__			)
           __UpperCamelCase	:			List[Any]					=					RagRetriever.from_pretrained(lowerCamelCase__			)
           self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)
           __UpperCamelCase	:			Tuple					=					np.array(
               [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]   ,				dtype=np.floataa			)
           __UpperCamelCase	:			Optional[int]					=					retriever.retrieve(lowerCamelCase__   ,				n_docs=1			)
           self.assertTrue(out is not None			)

  def    a					(    self						:	List[Any]			):
     """simple docstring"""

     __UpperCamelCase	:			str					=					1
     __UpperCamelCase	:			Dict					=					self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase__			)
     __UpperCamelCase	:			Optional[int]					=					np.array(
         [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]   ,				dtype=np.floataa			)
     __UpperCamelCase	,     __UpperCamelCase	,     __UpperCamelCase	:			List[Any]					=					retriever.retrieve(lowerCamelCase__   ,				n_docs=lowerCamelCase__			)
     self.assertEqual(retrieved_doc_embeds.shape   ,				(2, n_docs, self.retrieval_vector_size)			)
     self.assertEqual(len(lowerCamelCase__			)   ,				2			)
     self.assertEqual(sorted(doc_dicts[0]			)   ,				["""embeddings""", """id""", """text""", """title"""]			)
     self.assertEqual(len(doc_dicts[0]["""id"""]			)   ,				lowerCamelCase__			)
     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						:	List[str]			):
     """simple docstring"""

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

  def    a					(    self						:	Optional[int]			):
     """simple docstring"""

     __UpperCamelCase	:			Tuple					=					1
     __UpperCamelCase	:			int					=					self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase__			)
     __UpperCamelCase	:			Dict					=					np.array(
         [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]   ,				dtype=np.floataa			)
     __UpperCamelCase	,     __UpperCamelCase	,     __UpperCamelCase	:			List[Any]					=					retriever.retrieve(lowerCamelCase__   ,				n_docs=lowerCamelCase__			)
     self.assertEqual(retrieved_doc_embeds.shape   ,				(2, n_docs, self.retrieval_vector_size)			)
     self.assertEqual(len(lowerCamelCase__			)   ,				2			)
     self.assertEqual(sorted(doc_dicts[0]			)   ,				["""embeddings""", """id""", """text""", """title"""]			)
     self.assertEqual(len(doc_dicts[0]["""id"""]			)   ,				lowerCamelCase__			)
     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						:	int			):
     """simple docstring"""

     __UpperCamelCase	:			Optional[Any]					=					self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase__			)
     with tempfile.TemporaryDirectory() as tmp_dirname:
        retriever.save_pretrained(lowerCamelCase__			)
        __UpperCamelCase	:			Union[str, Any]					=					RagRetriever.from_pretrained(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)
        __UpperCamelCase	:			Union[str, Any]					=					np.array(
            [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]   ,				dtype=np.floataa			)
        __UpperCamelCase	:			Optional[int]					=					retriever.retrieve(lowerCamelCase__   ,				n_docs=1			)
        self.assertTrue(out is not None			)

  def    a					(    self						:	Tuple			):
     """simple docstring"""

     __UpperCamelCase	:			List[Any]					=					1
     __UpperCamelCase	:			int					=					self.get_dummy_legacy_index_retriever()
     __UpperCamelCase	:			int					=					np.array(
         [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]   ,				dtype=np.floataa			)
     __UpperCamelCase	,     __UpperCamelCase	,     __UpperCamelCase	:			str					=					retriever.retrieve(lowerCamelCase__   ,				n_docs=lowerCamelCase__			)
     self.assertEqual(retrieved_doc_embeds.shape   ,				(2, n_docs, self.retrieval_vector_size)			)
     self.assertEqual(len(lowerCamelCase__			)   ,				2			)
     self.assertEqual(sorted(doc_dicts[0]			)   ,				["""text""", """title"""]			)
     self.assertEqual(len(doc_dicts[0]["""text"""]			)   ,				lowerCamelCase__			)
     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						:	Dict			):
     """simple docstring"""

     __UpperCamelCase	:			Optional[int]					=					self.get_dummy_legacy_index_retriever()
     with tempfile.TemporaryDirectory() as tmp_dirname:
        retriever.save_pretrained(lowerCamelCase__			)
        __UpperCamelCase	:			List[str]					=					RagRetriever.from_pretrained(lowerCamelCase__			)
        self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)
        __UpperCamelCase	:			List[Any]					=					np.array(
            [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]   ,				dtype=np.floataa			)
        __UpperCamelCase	:			Optional[Any]					=					retriever.retrieve(lowerCamelCase__   ,				n_docs=1			)
        self.assertTrue(out is not None			)

  @require_torch
  @require_tokenizers
  @require_sentencepiece
  def    a					(    self						:	List[Any]			):
     """simple docstring"""

     import torch

     __UpperCamelCase	:			Tuple					=					1
     __UpperCamelCase	:			Tuple					=					self.get_dummy_canonical_hf_index_retriever()
     __UpperCamelCase	:			Dict					=					[[5, 7], [10, 11]]
     __UpperCamelCase	:			List[Any]					=					np.array(
         [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]   ,				dtype=np.floataa			)
     __UpperCamelCase	:			Optional[int]					=					retriever(lowerCamelCase__   ,				lowerCamelCase__   ,				prefix=retriever.config.generator.prefix   ,				n_docs=lowerCamelCase__			)
     __UpperCamelCase	,     __UpperCamelCase	,     __UpperCamelCase	:			Optional[int]					=					(
         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(lowerCamelCase__   ,				lowerCamelCase__			)
     self.assertIsInstance(lowerCamelCase__   ,				lowerCamelCase__			)
     self.assertIsInstance(lowerCamelCase__   ,				np.ndarray			)

     __UpperCamelCase	:			List[str]					=					retriever(
         lowerCamelCase__   ,				lowerCamelCase__   ,				prefix=retriever.config.generator.prefix   ,				n_docs=lowerCamelCase__   ,				return_tensors="""pt"""   ,				)
     __UpperCamelCase	,     __UpperCamelCase	,     __UpperCamelCase	,     __UpperCamelCase	:			List[Any]					=					(  # 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(lowerCamelCase__   ,				torch.Tensor			)
     self.assertIsInstance(lowerCamelCase__   ,				torch.Tensor			)
     self.assertIsInstance(lowerCamelCase__   ,				torch.Tensor			)

  @require_torch
  @require_tokenizers
  @require_sentencepiece
  def    a					(    self						:	List[str]			):
     """simple docstring"""

     __UpperCamelCase	:			Union[str, Any]					=					self.get_dpr_ctx_encoder_tokenizer()
     __UpperCamelCase	:			str					=					1
     __UpperCamelCase	:			Optional[Any]					=					self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase__			)
     retriever.set_ctx_encoder_tokenizer(lowerCamelCase__			)

     __UpperCamelCase	:			Any					=					[[5, 7], [10, 11]]
     __UpperCamelCase	:			Optional[int]					=					np.array(
         [np.ones(self.retrieval_vector_size			), -np.ones(self.retrieval_vector_size			)]   ,				dtype=np.floataa			)
     __UpperCamelCase	:			List[str]					=					retriever(lowerCamelCase__   ,				lowerCamelCase__   ,				prefix=retriever.config.generator.prefix   ,				n_docs=lowerCamelCase__			)

     self.assertEqual(
         len(lowerCamelCase__			)   ,				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""")			)   ,				lowerCamelCase__			)  # check for doc token related keys in dictionary.











import os
from pickle import UnpicklingError
from typing import Dict, Tuple

import jax
import jax.numpy as jnp
import numpy as np
from flax.serialization import from_bytes
from flax.traverse_util import flatten_dict, unflatten_dict

import transformers

from .utils import logging


UpperCamelCase    							=      logging.get_logger(__name__)


def 				__lowerCamelCase     (							__lowerCAmelCase							:					Any					,							__lowerCAmelCase							:					Any					,							__lowerCAmelCase							:					Any					,							__lowerCAmelCase							:					List[Any]=False							)     ->  Union[str, Any]:
   try:
      import torch  # noqa: F401
   except ImportError:
      logger.error(
          """Loading a PyTorch model in Flax, requires both PyTorch and Flax to be installed. Please see"""
          """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
          """ instructions."""							)
      raise

   if not is_sharded:
      __UpperCamelCase	:			Tuple					=					os.path.abspath(__lowerCAmelCase							)
      logger.info(f'Loading PyTorch weights from {pt_path}'							)

      __UpperCamelCase	:			int					=					torch.load(__lowerCAmelCase					,							map_location="""cpu"""							)
      logger.info(f'PyTorch checkpoint contains {sum(t.numel() for t in pt_state_dict.values()							):,} parameters.'							)

      __UpperCamelCase	:			List[Any]					=					convert_pytorch_state_dict_to_flax(__lowerCAmelCase					,							__lowerCAmelCase							)
   else:
      # model is sharded and pytorch_checkpoint_path already contains the list of .pt shard files
      __UpperCamelCase	:			List[str]					=					convert_pytorch_sharded_state_dict_to_flax(__lowerCAmelCase					,							__lowerCAmelCase							)
   return flax_state_dict


def 				__lowerCamelCase     (							__lowerCAmelCase							:					Tuple[str]					,							__lowerCAmelCase							:					np.ndarray					,							__lowerCAmelCase							:					Dict[str, jnp.ndarray]					,							__lowerCAmelCase							:					str					,							)     ->  (Tuple[str], np.ndarray):

   def is_key_or_prefix_key_in_dict(__lowerCAmelCase							:					Tuple[str]							) -> bool:
      return len(set(__lowerCAmelCase							) & {key, (model_prefix,) + key}							) > 0

   # layer norm
   __UpperCamelCase	:			Optional[int]					=					pt_tuple_key[:-1] + ("""scale""",)
   if pt_tuple_key[-1] in ["weight", "gamma"] and is_key_or_prefix_key_in_dict(__lowerCAmelCase							):
      return renamed_pt_tuple_key, pt_tensor

   # batch norm layer mean
   __UpperCamelCase	:			Optional[Any]					=					pt_tuple_key[:-1] + ("""mean""",)
   if pt_tuple_key[-1] == "running_mean" and not is_key_or_prefix_key_in_dict(__lowerCAmelCase							):
      return renamed_pt_tuple_key, pt_tensor

   # batch norm layer var
   __UpperCamelCase	:			Optional[Any]					=					pt_tuple_key[:-1] + ("""var""",)
   if pt_tuple_key[-1] == "running_var" and not is_key_or_prefix_key_in_dict(__lowerCAmelCase							):
      return renamed_pt_tuple_key, pt_tensor

   # embedding
   __UpperCamelCase	:			Union[str, Any]					=					pt_tuple_key[:-1] + ("""embedding""",)
   if pt_tuple_key[-1] == "weight" and is_key_or_prefix_key_in_dict(__lowerCAmelCase							):
      return renamed_pt_tuple_key, pt_tensor

   # conv layer
   __UpperCamelCase	:			Tuple					=					pt_tuple_key[:-1] + ("""kernel""",)
   if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4 and not is_key_or_prefix_key_in_dict(__lowerCAmelCase							):
      __UpperCamelCase	:			Optional[int]					=					pt_tensor.transpose(2					,							3					,							1					,							0							)
      return renamed_pt_tuple_key, pt_tensor

   # linear layer
   __UpperCamelCase	:			int					=					pt_tuple_key[:-1] + ("""kernel""",)
   if pt_tuple_key[-1] == "weight" and not is_key_or_prefix_key_in_dict(__lowerCAmelCase							):
      __UpperCamelCase	:			Optional[int]					=					pt_tensor.T
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm weight
   __UpperCamelCase	:			Tuple					=					pt_tuple_key[:-1] + ("""weight""",)
   if pt_tuple_key[-1] == "gamma":
      return renamed_pt_tuple_key, pt_tensor

   # old PyTorch layer norm bias
   __UpperCamelCase	:			str					=					pt_tuple_key[:-1] + ("""bias""",)
   if pt_tuple_key[-1] == "beta":
      return renamed_pt_tuple_key, pt_tensor

   # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
   __UpperCamelCase	:			Optional[int]					=					None
   if pt_tuple_key[-3::2] == ("parametrizations", "original0"):
      __UpperCamelCase	:			Tuple					=					pt_tuple_key[-2] + """_g"""
   elif pt_tuple_key[-3::2] == ("parametrizations", "original1"):
      __UpperCamelCase	:			Dict					=					pt_tuple_key[-2] + """_v"""
   if name is not None:
      __UpperCamelCase	:			Optional[Any]					=					pt_tuple_key[:-3] + (name,)
      return renamed_pt_tuple_key, pt_tensor

   return pt_tuple_key, pt_tensor


def 				__lowerCamelCase     (							__lowerCAmelCase							:					int					,							__lowerCAmelCase							:					int							)     ->  Union[str, Any]:
   # convert pytorch tensor to numpy
   __UpperCamelCase	:			Optional[Any]					=					{k: v.numpy() for k, v in pt_state_dict.items()}

   __UpperCamelCase	:			List[Any]					=					flax_model.base_model_prefix

   # use params dict if the model contains batch norm layers
   if "params" in flax_model.params:
      __UpperCamelCase	:			Union[str, Any]					=					flax_model.params["""params"""]
   else:
      __UpperCamelCase	:			List[Any]					=					flax_model.params
   __UpperCamelCase	:			List[Any]					=					flatten_dict(__lowerCAmelCase							)

   # add batch_stats keys,values to dict
   if "batch_stats" in flax_model.params:
      __UpperCamelCase	:			List[str]					=					flatten_dict(flax_model.params["""batch_stats"""]							)
      random_flax_state_dict.update(__lowerCAmelCase							)

   __UpperCamelCase	:			List[Any]					=					{}

   __UpperCamelCase	:			int					=					(model_prefix not in flax_model_params) and (
       model_prefix in {k.split("""."""							)[0] for k in pt_state_dict.keys()}
   )
   __UpperCamelCase	:			int					=					(model_prefix in flax_model_params) and (
       model_prefix not in {k.split("""."""							)[0] for k in pt_state_dict.keys()}
   )

   # Need to change some parameters name to match Flax names
   for pt_key, pt_tensor in pt_state_dict.items():
      __UpperCamelCase	:			List[str]					=					tuple(pt_key.split("""."""							)							)

      # remove base model prefix if necessary
      __UpperCamelCase	:			Optional[Any]					=					pt_tuple_key[0] == model_prefix
      if load_model_with_head_into_base_model and has_base_model_prefix:
         __UpperCamelCase	:			List[str]					=					pt_tuple_key[1:]

      # Correctly rename weight parameters
      __UpperCamelCase	,     __UpperCamelCase	:			Optional[int]					=					rename_key_and_reshape_tensor(
          __lowerCAmelCase					,							__lowerCAmelCase					,							__lowerCAmelCase					,							__lowerCAmelCase							)

      # add model prefix if necessary
      __UpperCamelCase	:			Optional[Any]					=					(model_prefix,) + flax_key in random_flax_state_dict
      if load_base_model_into_model_with_head and require_base_model_prefix:
         __UpperCamelCase	:			Dict					=					(model_prefix,) + flax_key

      if flax_key in random_flax_state_dict:
         if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
            raise ValueError(
                f'PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '
                f'{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'							)

        # add batch stats if the model contains batchnorm layers
      if "batch_stats" in flax_model.params:
         if "mean" in flax_key[-1] or "var" in flax_key[-1]:
            __UpperCamelCase	:			int					=					jnp.asarray(__lowerCAmelCase							)
            continue
         # remove num_batches_tracked key
         if "num_batches_tracked" in flax_key[-1]:
            flax_state_dict.pop(__lowerCAmelCase					,							__lowerCAmelCase							)
            continue

         # also add unexpected weight so that warning is thrown
         __UpperCamelCase	:			Optional[Any]					=					jnp.asarray(__lowerCAmelCase							)

      else:
         # also add unexpected weight so that warning is thrown
         __UpperCamelCase	:			Optional[int]					=					jnp.asarray(__lowerCAmelCase							)

   return unflatten_dict(__lowerCAmelCase							)


def 				__lowerCamelCase     (							__lowerCAmelCase							:					List[str]					,							__lowerCAmelCase							:					int							)     ->  Dict:
   import torch

   # Load the index
   __UpperCamelCase	:			List[str]					=					{}
   for shard_file in shard_filenames:
      # load using msgpack utils
      __UpperCamelCase	:			List[str]					=					torch.load(__lowerCAmelCase							)
      __UpperCamelCase	:			int					=					{k: v.numpy() for k, v in pt_state_dict.items()}

      __UpperCamelCase	:			List[str]					=					flax_model.base_model_prefix

      # use params dict if the model contains batch norm layers and then add batch_stats keys,values to dict
      if "batch_stats" in flax_model.params:
         __UpperCamelCase	:			Optional[Any]					=					flax_model.params["""params"""]

         __UpperCamelCase	:			int					=					flatten_dict(__lowerCAmelCase							)
         random_flax_state_dict.update(flatten_dict(flax_model.params["""batch_stats"""]							)							)
      else:
         __UpperCamelCase	:			Dict					=					flax_model.params
         __UpperCamelCase	:			int					=					flatten_dict(__lowerCAmelCase							)

      __UpperCamelCase	:			Union[str, Any]					=					(model_prefix not in flax_model_params) and (
          model_prefix in {k.split("""."""							)[0] for k in pt_state_dict.keys()}
      )
      __UpperCamelCase	:			Any					=					(model_prefix in flax_model_params) and (
          model_prefix not in {k.split("""."""							)[0] for k in pt_state_dict.keys()}
      )
      # Need to change some parameters name to match Flax names
      for pt_key, pt_tensor in pt_state_dict.items():
         __UpperCamelCase	:			Any					=					tuple(pt_key.split("""."""							)							)

         # remove base model prefix if necessary
         __UpperCamelCase	:			Dict					=					pt_tuple_key[0] == model_prefix
         if load_model_with_head_into_base_model and has_base_model_prefix:
            __UpperCamelCase	:			str					=					pt_tuple_key[1:]

         # Correctly rename weight parameters
         __UpperCamelCase	,     __UpperCamelCase	:			Optional[Any]					=					rename_key_and_reshape_tensor(
             __lowerCAmelCase					,							__lowerCAmelCase					,							__lowerCAmelCase					,							__lowerCAmelCase							)
         # add model prefix if necessary
         __UpperCamelCase	:			Tuple					=					(model_prefix,) + flax_key in random_flax_state_dict
         if load_base_model_into_model_with_head and require_base_model_prefix:
            __UpperCamelCase	:			Optional[Any]					=					(model_prefix,) + flax_key

         if flax_key in random_flax_state_dict:
            if flax_tensor.shape != random_flax_state_dict[flax_key].shape:
               raise ValueError(
                   f'PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape '
                   f'{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}.'							)

            # add batch stats if the model contains batchnorm layers
         if "batch_stats" in flax_model.params:
            if "mean" in flax_key[-1]:
               __UpperCamelCase	:			List[Any]					=					jnp.asarray(__lowerCAmelCase							)
               continue
            if "var" in flax_key[-1]:
               __UpperCamelCase	:			Optional[Any]					=					jnp.asarray(__lowerCAmelCase							)
               continue
            # remove num_batches_tracked key
            if "num_batches_tracked" in flax_key[-1]:
               flax_state_dict.pop(__lowerCAmelCase					,							__lowerCAmelCase							)
               continue

            # also add unexpected weight so that warning is thrown
            __UpperCamelCase	:			int					=					jnp.asarray(__lowerCAmelCase							)

         else:
            # also add unexpected weight so that warning is thrown
            __UpperCamelCase	:			int					=					jnp.asarray(__lowerCAmelCase							)
   return unflatten_dict(__lowerCAmelCase							)


def 				__lowerCamelCase     (							__lowerCAmelCase							:					Any					,							__lowerCAmelCase							:					Union[str, Any]							)     ->  Union[str, Any]:
   __UpperCamelCase	:			List[Any]					=					os.path.abspath(__lowerCAmelCase							)
   logger.info(f'Loading Flax weights from {flax_checkpoint_path}'							)

   # import correct flax class
   __UpperCamelCase	:			str					=					getattr(__lowerCAmelCase					,							"""Flax""" + model.__class__.__name__							)

   # load flax weight dict
   with open(__lowerCAmelCase					,							"""rb"""							) as state_f:
      try:
         __UpperCamelCase	:			Optional[int]					=					from_bytes(__lowerCAmelCase					,							state_f.read()							)
      except UnpicklingError:
         raise EnvironmentError(f'Unable to convert {flax_checkpoint_path} to Flax deserializable object. '							)

   return load_flax_weights_in_pytorch_model(__lowerCAmelCase					,							__lowerCAmelCase							)



def 				__lowerCamelCase     (							__lowerCAmelCase							:					Optional[int]					,							__lowerCAmelCase							:					Optional[int]							)     ->  List[Any]:

   try:
      import torch  # noqa: F401
   except ImportError:
      logger.error(
          """Loading a Flax weights in PyTorch, requires both PyTorch and Flax to be installed. Please see"""
          """ https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation"""
          """ instructions."""							)
      raise

   # check if we have bf16 weights
   __UpperCamelCase	:			Tuple					=					flatten_dict(jax.tree_util.tree_map(lambda __lowerCAmelCase							: x.dtype == jnp.bfloataa					,							__lowerCAmelCase							)							).values()
   if any(__lowerCAmelCase							):
      # convert all weights to fp32 if the are bf16 since torch.from_numpy can-not handle bf16
      # and bf16 is not fully supported in PT yet.
      logger.warning(
          """Found ``bfloat16`` weights in Flax model. Casting all ``bfloat16`` weights to ``float32`` """
          """before loading those in PyTorch model."""							)
      __UpperCamelCase	:			List[Any]					=					jax.tree_util.tree_map(
          lambda __lowerCAmelCase							: params.astype(np.floataa							) if params.dtype == jnp.bfloataa else params					,							__lowerCAmelCase							)

   __UpperCamelCase	:			Tuple					=					flatten_dict(__lowerCAmelCase							)
   __UpperCamelCase	:			int					=					pt_model.state_dict()

   __UpperCamelCase	:			Tuple					=					(pt_model.base_model_prefix in flax_state) and (
       pt_model.base_model_prefix not in {k.split("""."""							)[0] for k in pt_model_dict.keys()}
   )
   __UpperCamelCase	:			Any					=					(pt_model.base_model_prefix not in flax_state) and (
       pt_model.base_model_prefix in {k.split("""."""							)[0] for k in pt_model_dict.keys()}
   )

   # keep track of unexpected & missing keys
   __UpperCamelCase	:			Dict					=					[]
   __UpperCamelCase	:			List[str]					=					set(pt_model_dict.keys()							)

   for flax_key_tuple, flax_tensor in flax_state_dict.items():
      __UpperCamelCase	:			Union[str, Any]					=					flax_key_tuple[0] == pt_model.base_model_prefix
      __UpperCamelCase	:			Optional[Any]					=					""".""".join((pt_model.base_model_prefix,) + flax_key_tuple							) in pt_model_dict

      # adapt flax_key to prepare for loading from/to base model only
      if load_model_with_head_into_base_model and has_base_model_prefix:
         __UpperCamelCase	:			List[Any]					=					flax_key_tuple[1:]
      elif load_base_model_into_model_with_head and require_base_model_prefix:
         __UpperCamelCase	:			Tuple					=					(pt_model.base_model_prefix,) + flax_key_tuple

      # rename flax weights to PyTorch format
      if flax_key_tuple[-1] == "kernel" and flax_tensor.ndim == 4 and ".".join(__lowerCAmelCase							) not in pt_model_dict:
         # conv layer
         __UpperCamelCase	:			Optional[int]					=					flax_key_tuple[:-1] + ("""weight""",)
         __UpperCamelCase	:			Dict					=					jnp.transpose(__lowerCAmelCase					,							(3, 2, 0, 1)							)
      elif flax_key_tuple[-1] == "kernel" and ".".join(__lowerCAmelCase							) not in pt_model_dict:
         # linear layer
         __UpperCamelCase	:			Optional[int]					=					flax_key_tuple[:-1] + ("""weight""",)
         __UpperCamelCase	:			List[Any]					=					flax_tensor.T
      elif flax_key_tuple[-1] in ["scale", "embedding"]:
         __UpperCamelCase	:			Union[str, Any]					=					flax_key_tuple[:-1] + ("""weight""",)

      # adding batch stats from flax batch norm to pt
      elif "mean" in flax_key_tuple[-1]:
         __UpperCamelCase	:			List[str]					=					flax_key_tuple[:-1] + ("""running_mean""",)
      elif "var" in flax_key_tuple[-1]:
         __UpperCamelCase	:			Tuple					=					flax_key_tuple[:-1] + ("""running_var""",)

      if "batch_stats" in flax_state:
         __UpperCamelCase	:			Optional[int]					=					""".""".join(flax_key_tuple[1:]							)  # Remove the params/batch_stats header
      else:
         __UpperCamelCase	:			Optional[int]					=					""".""".join(__lowerCAmelCase							)

      # We also need to look at `pt_model_dict` and see if there are keys requiring further transformation.
      __UpperCamelCase	:			Union[str, Any]					=					{}
      # New `weight_norm` from https://github.com/huggingface/transformers/pull/24030
      for key in pt_model_dict:
         __UpperCamelCase	:			int					=					key.split("""."""							)
         __UpperCamelCase	:			str					=					None
         if key_components[-3::2] == ["parametrizations", "original0"]:
            __UpperCamelCase	:			Union[str, Any]					=					key_components[-2] + """_g"""
         elif key_components[-3::2] == ["parametrizations", "original1"]:
            __UpperCamelCase	:			Tuple					=					key_components[-2] + """_v"""
         if name is not None:
            __UpperCamelCase	:			Optional[Any]					=					key_components[:-3] + [name]
            __UpperCamelCase	:			int					=					""".""".join(__lowerCAmelCase							)
            __UpperCamelCase	:			List[str]					=					key

      if flax_key in special_pt_names:
         __UpperCamelCase	:			Any					=					special_pt_names[flax_key]

      if flax_key in pt_model_dict:
         if flax_tensor.shape != pt_model_dict[flax_key].shape:
            raise ValueError(
                f'Flax checkpoint seems to be incorrect. Weight {flax_key_tuple} was expected '
                f'to be of shape {pt_model_dict[flax_key].shape}, but is {flax_tensor.shape}.'							)
         else:
            # add weight to pytorch dict
            __UpperCamelCase	:			str					=					np.asarray(__lowerCAmelCase							) if not isinstance(__lowerCAmelCase					,							np.ndarray							) else flax_tensor
            __UpperCamelCase	:			Optional[int]					=					torch.from_numpy(__lowerCAmelCase							)
            # remove from missing keys
            missing_keys.remove(__lowerCAmelCase							)
      else:
         # weight is not expected by PyTorch model
         unexpected_keys.append(__lowerCAmelCase							)

   pt_model.load_state_dict(__lowerCAmelCase							)

   # re-transform missing_keys to list
   __UpperCamelCase	:			Union[str, Any]					=					list(__lowerCAmelCase							)

   if len(__lowerCAmelCase							) > 0:
      logger.warning(
          """Some weights of the Flax model were not used when initializing the PyTorch model"""
          f' {pt_model.__class__.__name__}: {unexpected_keys}\n- This IS expected if you are initializing'
          f' {pt_model.__class__.__name__} from a Flax model trained on another task or with another architecture'
          """ (e.g. initializing a BertForSequenceClassification model from a FlaxBertForPreTraining model).\n- This"""
          f' IS NOT expected if you are initializing {pt_model.__class__.__name__} from a Flax model that you expect'
          """ to be exactly identical (e.g. initializing a BertForSequenceClassification model from a"""
          """ FlaxBertForSequenceClassification model)."""							)
   else:
      logger.warning(f'All Flax model weights were used when initializing {pt_model.__class__.__name__}.\n'							)
   if len(__lowerCAmelCase							) > 0:
      logger.warning(
          f'Some weights of {pt_model.__class__.__name__} were not initialized from the Flax model and are newly'
          f' initialized: {missing_keys}\nYou should probably TRAIN this model on a down-stream task to be able to'
          """ use it for predictions and inference."""							)
   else:
      logger.warning(
          f'All the weights of {pt_model.__class__.__name__} were initialized from the Flax model.\n'
          """If your task is similar to the task the model of the checkpoint was trained on, """
          f'you can already use {pt_model.__class__.__name__} for predictions without further training.'							)

   return pt_model



