


import itertools
import string
from collections.abc import Generator, Iterable


def       UpperCAmelCase_					(   UpperCAmelCase__    ,				UpperCAmelCase__	):
					lowercase_					=      iter(UpperCamelCase__	)
					while True:
										lowercase_					=      tuple(itertools.islice(UpperCamelCase__    ,				UpperCamelCase__	)	)
										if not chunk:
															return
										yield chunk


def       UpperCAmelCase_					(   UpperCAmelCase__	):
					lowercase_					=      """""".join([c.upper() for c in dirty if c in string.ascii_letters]	)
					lowercase_					=      """"""

					if len(UpperCamelCase__	) < 2:
										return dirty

					for i in range(len(UpperCamelCase__	) - 1	):
										clean += dirty[i]

										if dirty[i] == dirty[i + 1]:
															clean += "X"

					clean += dirty[-1]

					if len(UpperCamelCase__	) & 1:
										clean += "X"

					return clean


def       UpperCAmelCase_					(   UpperCAmelCase__	):
					lowercase_					=      """ABCDEFGHIKLMNOPQRSTUVWXYZ"""
					# we're using a list instead of a '2d' array because it makes the math
					# for setting up the table and doing the actual encoding/decoding simpler
					lowercase_					=      []

					# copy key chars into the table if they are in `alphabet` ignoring duplicates
					for char in key.upper():
										if char not in table and char in alphabet:
															table.append(UpperCamelCase__	)

    # fill the rest of the table in with the remaining alphabet chars
					for char in alphabet:
										if char not in table:
															table.append(UpperCamelCase__	)

					return table


def       UpperCAmelCase_					(   UpperCAmelCase__    ,				UpperCAmelCase__	):
					lowercase_					=      generate_table(UpperCamelCase__	)
					lowercase_					=      prepare_input(UpperCamelCase__	)
					lowercase_					=      """"""

					# https://en.wikipedia.org/wiki/Playfair_cipher#Description
					for chara, chara in chunker(UpperCamelCase__    ,				2	):
										lowercase_				,       lowercase_					=      divmod(table.index(UpperCamelCase__	)    ,				5	)
										lowercase_				,       lowercase_					=      divmod(table.index(UpperCamelCase__	)    ,				5	)

										if rowa == rowa:
															ciphertext += table[rowa * 5 + (cola + 1) % 5]
															ciphertext += table[rowa * 5 + (cola + 1) % 5]
										elif cola == cola:
															ciphertext += table[((rowa + 1) % 5) * 5 + cola]
															ciphertext += table[((rowa + 1) % 5) * 5 + cola]
										else:  # rectangle
															ciphertext += table[rowa * 5 + cola]
															ciphertext += table[rowa * 5 + cola]

					return ciphertext



def       UpperCAmelCase_					(   UpperCAmelCase__    ,				UpperCAmelCase__	):
					lowercase_					=      generate_table(UpperCamelCase__	)
					lowercase_					=      """"""

					# https://en.wikipedia.org/wiki/Playfair_cipher#Description
					for chara, chara in chunker(UpperCamelCase__    ,				2	):
										lowercase_				,       lowercase_					=      divmod(table.index(UpperCamelCase__	)    ,				5	)
										lowercase_				,       lowercase_					=      divmod(table.index(UpperCamelCase__	)    ,				5	)

										if rowa == rowa:
															plaintext += table[rowa * 5 + (cola - 1) % 5]
															plaintext += table[rowa * 5 + (cola - 1) % 5]
										elif cola == cola:
															plaintext += table[((rowa - 1) % 5) * 5 + cola]
															plaintext += table[((rowa - 1) % 5) * 5 + cola]
										else:  # rectangle
															plaintext += table[rowa * 5 + cola]
															plaintext += table[rowa * 5 + cola]

					return plaintext





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.




