


'''simple docstring'''



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 							__SCREAMING_SNAKE_CASE       (							_lowerCAmelCase     ):




      def  __UpperCamelCase							(  self      )							->Optional[Any]:



         '''simple docstring'''
         __a				     =    tempfile.mkdtemp()
         __a				     =    8

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

         __a				     =    os.path.join(self.tmpdirname ,   'bart_tokenizer'      )
         os.makedirs(lowerCamelCase ,   exist_ok=lowerCamelCase      )
         __a				     =    os.path.join(lowerCamelCase ,   BART_VOCAB_FILES_NAMES['vocab_file']      )
         __a				     =    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  __UpperCamelCase							(  self      )							->DPRQuestionEncoderTokenizer:



         '''simple docstring'''
         return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname ,   'dpr_tokenizer'      )      )




      def  __UpperCamelCase							(  self      )							->DPRContextEncoderTokenizer:



         '''simple docstring'''
         return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname ,   'dpr_tokenizer'      )      )




      def  __UpperCamelCase							(  self      )							->BartTokenizer:



         '''simple docstring'''
         return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname ,   'bart_tokenizer'      )      )




      def  __UpperCamelCase							(  self      )							->Optional[Any]:



         '''simple docstring'''
         shutil.rmtree(self.tmpdirname      )




      def  __UpperCamelCase							(  self      )							->List[Any]:



         '''simple docstring'''
         __a				     =    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  __UpperCamelCase							(  self      )							->Dict:



         '''simple docstring'''
         __a				     =    self.get_dummy_dataset()
         __a				     =    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:
            __a				     =    dataset
            __a				     =    RagRetriever(
                lowerCamelCase ,   question_encoder_tokenizer=self.get_dpr_tokenizer() ,   generator_tokenizer=self.get_bart_tokenizer() ,   )
         return retriever




      def  __UpperCamelCase							(  self ,   lowerCamelCase      )							->int:



         '''simple docstring'''
         __a				     =    self.get_dummy_dataset()
         __a				     =    RagConfig(
             retrieval_vector_size=self.retrieval_vector_size ,   question_encoder=DPRConfig().to_dict() ,   generator=BartConfig().to_dict() ,   index_name='custom' ,   )
         if from_disk:
            __a				     =    os.path.join(self.tmpdirname ,   'dataset'      )
            __a				     =    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
            __a				     =    RagRetriever(
                lowerCamelCase ,   question_encoder_tokenizer=self.get_dpr_tokenizer() ,   generator_tokenizer=self.get_bart_tokenizer() ,   )
         else:
            __a				     =    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  __UpperCamelCase							(  self      )							->List[str]:



         '''simple docstring'''
         __a				     =    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      )

         __a				     =    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'      )      )

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

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




      def  __UpperCamelCase							(  self      )							->List[Any]:



         '''simple docstring'''
         __a				     =    1
         __a				     =    self.get_dummy_canonical_hf_index_retriever()
         __a				     =    np.array(
             [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )] ,   dtype=np.floataa      )
         __a     ,			__a     ,			__a				     =    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  __UpperCamelCase							(  self      )							->Optional[int]:



         '''simple docstring'''
         __a				     =    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:
               __a				     =    self.get_dummy_dataset()
               retriever.save_pretrained(lowerCamelCase      )
               __a				     =    RagRetriever.from_pretrained(lowerCamelCase      )
               self.assertIsInstance(lowerCamelCase ,   lowerCamelCase      )
               __a				     =    np.array(
                   [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )] ,   dtype=np.floataa      )
               __a				     =    retriever.retrieve(lowerCamelCase ,   n_docs=1      )
               self.assertTrue(out is not None      )




      def  __UpperCamelCase							(  self      )							->str:



         '''simple docstring'''
         __a				     =    1
         __a				     =    self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase      )
         __a				     =    np.array(
             [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )] ,   dtype=np.floataa      )
         __a     ,			__a     ,			__a				     =    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  __UpperCamelCase							(  self      )							->Tuple:



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




      def  __UpperCamelCase							(  self      )							->Any:



         '''simple docstring'''
         __a				     =    1
         __a				     =    self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase      )
         __a				     =    np.array(
             [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )] ,   dtype=np.floataa      )
         __a     ,			__a     ,			__a				     =    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  __UpperCamelCase							(  self      )							->Optional[int]:



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




      def  __UpperCamelCase							(  self      )							->Union[str, Any]:



         '''simple docstring'''
         __a				     =    1
         __a				     =    self.get_dummy_legacy_index_retriever()
         __a				     =    np.array(
             [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )] ,   dtype=np.floataa      )
         __a     ,			__a     ,			__a				     =    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  __UpperCamelCase							(  self      )							->Optional[int]:



         '''simple docstring'''
         __a				     =    self.get_dummy_legacy_index_retriever()
         with tempfile.TemporaryDirectory() as tmp_dirname:
            retriever.save_pretrained(lowerCamelCase      )
            __a				     =    RagRetriever.from_pretrained(lowerCamelCase      )
            self.assertIsInstance(lowerCamelCase ,   lowerCamelCase      )
            __a				     =    np.array(
                [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )] ,   dtype=np.floataa      )
            __a				     =    retriever.retrieve(lowerCamelCase ,   n_docs=1      )
            self.assertTrue(out is not None      )




      @require_torch
      @require_tokenizers
      @require_sentencepiece
      def  __UpperCamelCase							(  self      )							->Union[str, Any]:



         '''simple docstring'''
         import torch

         __a				     =    1
         __a				     =    self.get_dummy_canonical_hf_index_retriever()
         __a				     =    [[5, 7], [10, 11]]
         __a				     =    np.array(
             [np.ones(self.retrieval_vector_size      ), -np.ones(self.retrieval_vector_size      )] ,   dtype=np.floataa      )
         __a				     =    retriever(lowerCamelCase ,   lowerCamelCase ,   prefix=retriever.config.generator.prefix ,   n_docs=lowerCamelCase      )
         __a     ,			__a     ,			__a				     =    (
             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      )

         __a				     =    retriever(
             lowerCamelCase ,   lowerCamelCase ,   prefix=retriever.config.generator.prefix ,   n_docs=lowerCamelCase ,   return_tensors='pt' ,   )
         __a     ,			__a     ,			__a     ,			__a				     =    (  # 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  __UpperCamelCase							(  self      )							->str:



         '''simple docstring'''
         __a				     =    self.get_dpr_ctx_encoder_tokenizer()
         __a				     =    1
         __a				     =    self.get_dummy_custom_hf_index_retriever(from_disk=lowerCamelCase      )
         retriever.set_ctx_encoder_tokenizer(lowerCamelCase      )

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







'''simple docstring'''



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	int,					SCREAMING_SNAKE_CASE__:	bool = False    )  ->			bool:

   """simple docstring"""
   if n == 2:
      return True
   if not n % 2 or n < 2:
      return False
   if n > 5 and n % 10 not in (1, 3, 7, 9):  # can quickly check last digit
      return False
   if n > 3317044064679887385961981 and not allow_probable:
      raise ValueError(
          'Warning: upper bound of deterministic test is exceeded. '
          'Pass allow_probable=True to allow probabilistic test. '
          'A return value of True indicates a probable prime.'    )
   # array bounds provided by analysis
   __a				     =    [
       2047,
       1373653,
       25326001,
       3215031751,
       2152302898747,
       3474749660383,
       341550071728321,
       1,
       3825123056546413051,
       1,
       1,
       318665857834031151167461,
       3317044064679887385961981,
   ]

   __a				     =    [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41]
   for idx, _p in enumerate(SCREAMING_SNAKE_CASE__,					1    ):
      if n < _p:
         # then we have our last prime to check
         __a				     =    primes[:idx]
         break
   __a     ,			__a				     =    n - 1, 0
   # break up n -1 into a power of 2 (s) and
   # remaining odd component
   # essentially, solve for d * 2 ** s == n - 1
   while d % 2 == 0:
      d //= 2
      s += 1
   for prime in plist:
      __a				     =    False
      for r in range(SCREAMING_SNAKE_CASE__    ):
         __a				     =    pow(SCREAMING_SNAKE_CASE__,					d * 2**r,					SCREAMING_SNAKE_CASE__    )
         # see article for analysis explanation for m
         if (r == 0 and m == 1) or ((m + 1) % n == 0):
            __a				     =    True
            # this loop will not determine compositeness
            break
      if pr:
         continue
      # if pr is False, then the above loop never evaluated to true,
      # and the n MUST be composite
      return False
   return True



def    __UpperCAmelCase	(				)  ->			None:

   """simple docstring"""
   assert not miller_rabin(561    )
   assert miller_rabin(563    )
   # 2047

   assert not miller_rabin(838201    )
   assert miller_rabin(838207    )
   # 1_373_653

   assert not miller_rabin(17316001    )
   assert miller_rabin(17316017    )
   # 25_326_001

   assert not miller_rabin(3078386641    )
   assert miller_rabin(3078386653    )
   # 3_215_031_751

   assert not miller_rabin(1713045574801    )
   assert miller_rabin(1713045574819    )
   # 2_152_302_898_747

   assert not miller_rabin(2779799728307    )
   assert miller_rabin(2779799728327    )
   # 3_474_749_660_383

   assert not miller_rabin(113850023909441    )
   assert miller_rabin(113850023909527    )
   # 341_550_071_728_321

   assert not miller_rabin(1275041018848804351    )
   assert miller_rabin(1275041018848804391    )
   # 3_825_123_056_546_413_051

   assert not miller_rabin(79666464458507787791867    )
   assert miller_rabin(79666464458507787791951    )
   # 318_665_857_834_031_151_167_461

   assert not miller_rabin(552840677446647897660333    )
   assert miller_rabin(552840677446647897660359    )
   # 3_317_044_064_679_887_385_961_981
   # upper limit for probabilistic test


if __name__ == "__main__":
 test_miller_rabin()