




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    __lowerCAmelCase       (							lowerCAmelCase):






       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         lowercase    :List[Any]      	=						tempfile.mkdtemp()
         lowercase    :str      	=						8

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

         lowercase    :Dict      	=						os.path.join(self.tmpdirname       ,					"bart_tokenizer"				)
         os.makedirs(_lowerCAmelCase       ,					exist_ok=_lowerCAmelCase				)
         lowercase    :Union[str, Any]      	=						os.path.join(_lowerCAmelCase       ,					BART_VOCAB_FILES_NAMES["vocab_file"]				)
         lowercase    :Optional[int]      	=						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 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname       ,					"dpr_tokenizer"				)				)






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname       ,					"dpr_tokenizer"				)				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname       ,					"bart_tokenizer"				)				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Union[str, Any]				):
         shutil.rmtree(self.tmpdirname				)






       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         lowercase    :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 					SCREAMING_SNAKE_CASE    (   self:					int				):
         lowercase    :str      	=						self.get_dummy_dataset()
         lowercase    :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:
           lowercase    :List[str]      	=						dataset
           lowercase    :List[Any]      	=						RagRetriever(
               _lowerCAmelCase       ,					question_encoder_tokenizer=self.get_dpr_tokenizer()       ,					generator_tokenizer=self.get_bart_tokenizer()       ,					)
         return retriever






       def 					SCREAMING_SNAKE_CASE    (   self:					Dict       ,					_lowerCAmelCase:					bool				):
         lowercase    :List[str]      	=						self.get_dummy_dataset()
         lowercase    :Union[str, Any]      	=						RagConfig(
             retrieval_vector_size=self.retrieval_vector_size       ,					question_encoder=DPRConfig().to_dict()       ,					generator=BartConfig().to_dict()       ,					index_name="custom"       ,					)
         if from_disk:
           lowercase    :Dict      	=						os.path.join(self.tmpdirname       ,					"dataset"				)
           lowercase    :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
           lowercase    :Optional[Any]      	=						RagRetriever(
               _lowerCAmelCase       ,					question_encoder_tokenizer=self.get_dpr_tokenizer()       ,					generator_tokenizer=self.get_bart_tokenizer()       ,					)
         else:
           lowercase    :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 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
         lowercase    :List[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				)

         lowercase    :Optional[Any]      	=						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"				)				)

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

         lowercase    :Optional[Any]      	=						RagConfig(
             retrieval_vector_size=self.retrieval_vector_size       ,					question_encoder=DPRConfig().to_dict()       ,					generator=BartConfig().to_dict()       ,					index_name="legacy"       ,					index_path=self.tmpdirname       ,					)
         lowercase    :Union[str, Any]      	=						RagRetriever(
             _lowerCAmelCase       ,					question_encoder_tokenizer=self.get_dpr_tokenizer()       ,					generator_tokenizer=self.get_bart_tokenizer()				)
         return retriever






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






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[Any]				):
         lowercase    :Dict      	=						1
         lowercase    :Union[str, Any]      	=						self.get_dummy_custom_hf_index_retriever(from_disk=_lowerCAmelCase				)
         lowercase    :Tuple      	=						np.array(
             [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]       ,					dtype=np.floataa				)
         lowercase	,    lowercase	,    lowercase    :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]				)       ,					["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 					SCREAMING_SNAKE_CASE    (   self:					str				):
         lowercase    :Dict      	=						self.get_dummy_custom_hf_index_retriever(from_disk=_lowerCAmelCase				)
         with tempfile.TemporaryDirectory() as tmp_dirname:
           retriever.save_pretrained(_lowerCAmelCase				)
           lowercase    :Tuple      	=						RagRetriever.from_pretrained(_lowerCAmelCase				)
           self.assertIsInstance(_lowerCAmelCase       ,					_lowerCAmelCase				)
           lowercase    :Tuple      	=						np.array(
               [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]       ,					dtype=np.floataa				)
           lowercase    :int      	=						retriever.retrieve(_lowerCAmelCase       ,					n_docs=1				)
           self.assertTrue(out is not None				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Any				):
         lowercase    :Tuple      	=						1
         lowercase    :Any      	=						self.get_dummy_custom_hf_index_retriever(from_disk=_lowerCAmelCase				)
         lowercase    :Any      	=						np.array(
             [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]       ,					dtype=np.floataa				)
         lowercase	,    lowercase	,    lowercase    :Optional[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 					SCREAMING_SNAKE_CASE    (   self:					List[Any]				):
         lowercase    :Union[str, Any]      	=						self.get_dummy_custom_hf_index_retriever(from_disk=_lowerCAmelCase				)
         with tempfile.TemporaryDirectory() as tmp_dirname:
           retriever.save_pretrained(_lowerCAmelCase				)
           lowercase    :List[str]      	=						RagRetriever.from_pretrained(_lowerCAmelCase				)
           self.assertIsInstance(_lowerCAmelCase       ,					_lowerCAmelCase				)
           lowercase    :Tuple      	=						np.array(
               [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]       ,					dtype=np.floataa				)
           lowercase    :Optional[Any]      	=						retriever.retrieve(_lowerCAmelCase       ,					n_docs=1				)
           self.assertTrue(out is not None				)






       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         lowercase    :Optional[Any]      	=						1
         lowercase    :Optional[Any]      	=						self.get_dummy_legacy_index_retriever()
         lowercase    :Union[str, Any]      	=						np.array(
             [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]       ,					dtype=np.floataa				)
         lowercase	,    lowercase	,    lowercase    :Tuple      	=						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 					SCREAMING_SNAKE_CASE    (   self:					List[str]				):
         lowercase    :Dict      	=						self.get_dummy_legacy_index_retriever()
         with tempfile.TemporaryDirectory() as tmp_dirname:
           retriever.save_pretrained(_lowerCAmelCase				)
           lowercase    :str      	=						RagRetriever.from_pretrained(_lowerCAmelCase				)
           self.assertIsInstance(_lowerCAmelCase       ,					_lowerCAmelCase				)
           lowercase    :str      	=						np.array(
               [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]       ,					dtype=np.floataa				)
           lowercase    :Any      	=						retriever.retrieve(_lowerCAmelCase       ,					n_docs=1				)
           self.assertTrue(out is not None				)






       @require_torch
       @require_tokenizers
       @require_sentencepiece
       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
         import torch

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

         lowercase    :Optional[Any]      	=						retriever(
             _lowerCAmelCase       ,					_lowerCAmelCase       ,					prefix=retriever.config.generator.prefix       ,					n_docs=_lowerCAmelCase       ,					return_tensors="pt"       ,					)
         lowercase	,    lowercase	,    lowercase	,    lowercase    :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 					SCREAMING_SNAKE_CASE    (   self:					List[Any]				):
         lowercase    :List[str]      	=						self.get_dpr_ctx_encoder_tokenizer()
         lowercase    :Optional[int]      	=						1
         lowercase    :Any      	=						self.get_dummy_custom_hf_index_retriever(from_disk=_lowerCAmelCase				)
         retriever.set_ctx_encoder_tokenizer(_lowerCAmelCase				)

         lowercase    :List[str]      	=						[[5, 7], [10, 11]]
         lowercase    :List[str]      	=						np.array(
             [np.ones(self.retrieval_vector_size				), -np.ones(self.retrieval_vector_size				)]       ,					dtype=np.floataa				)
         lowercase    :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.









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

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


_UpperCAmelCase	:						Optional[int]			=     logging.get_logger(__name__)

_UpperCAmelCase	:						List[Any]			=     {
    "EleutherAI/gpt-j-6B": "https://huggingface.co/EleutherAI/gpt-j-6B/resolve/main/config.json",
    # See all GPT-J models at https://huggingface.co/models?filter=gpt_j
}




class    __lowerCAmelCase       (							lowerCAmelCase):
       _a 					=						'''gptj'''
       _a 					=						{
           '''max_position_embeddings''': '''n_positions''',
           '''hidden_size''': '''n_embd''',
           '''num_attention_heads''': '''n_head''',
           '''num_hidden_layers''': '''n_layer''',
       }






       def __init__(   self:					List[str]       ,					_lowerCAmelCase:					Union[str, Any]=5_04_00       ,					_lowerCAmelCase:					Optional[int]=20_48       ,					_lowerCAmelCase:					Dict=40_96       ,					_lowerCAmelCase:					str=28       ,					_lowerCAmelCase:					Dict=16       ,					_lowerCAmelCase:					str=64       ,					_lowerCAmelCase:					Dict=None       ,					_lowerCAmelCase:					Optional[int]="gelu_new"       ,					_lowerCAmelCase:					Optional[int]=0.0       ,					_lowerCAmelCase:					List[str]=0.0       ,					_lowerCAmelCase:					str=0.0       ,					_lowerCAmelCase:					Any=1e-5       ,					_lowerCAmelCase:					str=0.02       ,					_lowerCAmelCase:					str=True       ,					_lowerCAmelCase:					Any=5_02_56       ,					_lowerCAmelCase:					Union[str, Any]=5_02_56       ,					_lowerCAmelCase:					Optional[int]=False       ,					**_lowerCAmelCase:					int       ,					):
         lowercase    :str      	=						vocab_size
         lowercase    :Union[str, Any]      	=						n_positions
         lowercase    :Optional[int]      	=						n_embd
         lowercase    :int      	=						n_layer
         lowercase    :Tuple      	=						n_head
         lowercase    :Optional[Any]      	=						n_inner
         lowercase    :List[Any]      	=						rotary_dim
         lowercase    :List[Any]      	=						activation_function
         lowercase    :Tuple      	=						resid_pdrop
         lowercase    :Optional[int]      	=						embd_pdrop
         lowercase    :Any      	=						attn_pdrop
         lowercase    :Optional[int]      	=						layer_norm_epsilon
         lowercase    :str      	=						initializer_range
         lowercase    :Optional[int]      	=						use_cache

         lowercase    :Dict      	=						bos_token_id
         lowercase    :List[Any]      	=						eos_token_id

         super().__init__(
             bos_token_id=_lowerCAmelCase       ,					eos_token_id=_lowerCAmelCase       ,					tie_word_embeddings=_lowerCAmelCase       ,					**_lowerCAmelCase				)







class    __lowerCAmelCase       (							lowerCAmelCase):






       def __init__(   self:					Tuple       ,					_lowerCAmelCase:					PretrainedConfig       ,					_lowerCAmelCase:					str = "default"       ,					_lowerCAmelCase:					List[PatchingSpec] = None       ,					_lowerCAmelCase:					bool = False       ,					):
         super().__init__(_lowerCAmelCase       ,					task=_lowerCAmelCase       ,					patching_specs=_lowerCAmelCase       ,					use_past=_lowerCAmelCase				)
         if not getattr(self._config       ,					"pad_token_id"       ,					_lowerCAmelCase				):
           # TODO: how to do that better?
           lowercase    :int      	=						0






       @property
       def 					SCREAMING_SNAKE_CASE    (   self:					int				):
         lowercase    :Optional[int]      	=						OrderedDict({"input_ids": {0: "batch", 1: "sequence"}}				)
         if self.use_past:
           self.fill_with_past_key_values_(_lowerCAmelCase       ,					direction="inputs"				)
           lowercase    :Any      	=						{0: "batch", 1: "past_sequence + sequence"}
         else:
           lowercase    :Optional[int]      	=						{0: "batch", 1: "sequence"}

         return common_inputs






       @property
       def 					SCREAMING_SNAKE_CASE    (   self:					Optional[int]				):
         return self._config.n_layer






       @property
       def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]				):
         return self._config.n_head






       def 					SCREAMING_SNAKE_CASE    (   self:					List[Any]       ,					_lowerCAmelCase:					PreTrainedTokenizer       ,					_lowerCAmelCase:					int = -1       ,					_lowerCAmelCase:					int = -1       ,					_lowerCAmelCase:					bool = False       ,					_lowerCAmelCase:					Optional[TensorType] = None       ,					):
         lowercase    :Union[str, Any]      	=						super(_lowerCAmelCase       ,					self				).generate_dummy_inputs(
             _lowerCAmelCase       ,					batch_size=_lowerCAmelCase       ,					seq_length=_lowerCAmelCase       ,					is_pair=_lowerCAmelCase       ,					framework=_lowerCAmelCase				)

         # We need to order the input in the way they appears in the forward()
         lowercase    :Optional[Any]      	=						OrderedDict({"input_ids": common_inputs["input_ids"]}				)

         # Need to add the past_keys
         if self.use_past:
           if not is_torch_available():
             raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."				)
           else:
             import torch

             lowercase	,    lowercase    :Optional[int]      	=						common_inputs["input_ids"].shape
             # Not using the same length for past_key_values
             lowercase    :Union[str, Any]      	=						seqlen + 2
             lowercase    :Optional[Any]      	=						(
                 batch,
                 self.num_attention_heads,
                 past_key_values_length,
                 self._config.hidden_size // self.num_attention_heads,
             )
             lowercase    :Optional[int]      	=						[
                 (torch.zeros(_lowerCAmelCase				), torch.zeros(_lowerCAmelCase				)) for _ in range(self.num_layers				)
             ]

         lowercase    :List[Any]      	=						common_inputs["attention_mask"]
         if self.use_past:
           lowercase    :Tuple      	=						ordered_inputs["attention_mask"].dtype
           lowercase    :List[Any]      	=						torch.cat(
               [ordered_inputs["attention_mask"], torch.ones(_lowerCAmelCase       ,					_lowerCAmelCase       ,					dtype=_lowerCAmelCase				)]       ,					dim=1				)

         return ordered_inputs






       @property
       def 					SCREAMING_SNAKE_CASE    (   self:					Tuple				):
         return 13



