




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:    List[str] ):
         lowercase			:Optional[int]							=  tempfile.mkdtemp()
         lowercase			:Optional[Any]							=  8

         # DPR tok
         lowercase			:Any							=  [
             "[UNK]",
             "[CLS]",
             "[SEP]",
             "[PAD]",
             "[MASK]",
             "want",
             "##want",
             "##ed",
             "wa",
             "un",
             "runn",
             "##ing",
             ",",
             "low",
             "lowest",
         ]
         lowercase			:Optional[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			:Tuple							=  [
             "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			:str							=  ["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
         lowercase			:List[Any]							=  {"unk_token": "<unk>"}

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






       def     SCREAMING_SNAKE_CASE			(      self:    List[str] ):
         return DPRContextEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname				,    "dpr_tokenizer" ) )






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






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






       def     SCREAMING_SNAKE_CASE			(      self:    Any				,    _lowerCAmelCase:    bool ):
         lowercase			:Tuple							=  self.get_dummy_dataset()
         lowercase			:str							=  RagConfig(
             retrieval_vector_size=self.retrieval_vector_size				,    question_encoder=DPRConfig().to_dict()				,    generator=BartConfig().to_dict()				,    index_name="custom"				,    )
         if from_disk:
           lowercase			:Any							=  os.path.join(self.tmpdirname				,    "dataset" )
           lowercase			:Optional[int]							=  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			:List[str]							=  RagRetriever(
               _lowerCAmelCase				,    question_encoder_tokenizer=self.get_dpr_tokenizer()				,    generator_tokenizer=self.get_bart_tokenizer()				,    )
         else:
           lowercase			: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     SCREAMING_SNAKE_CASE			(      self:    Optional[int] ):
         lowercase			:Union[str, 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			:List[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			:Optional[Any]							=  os.path.join(self.tmpdirname				,    "psgs_w100.tsv.pkl" )
         lowercase			:Tuple							=  {sample["id"]: [sample["text"], sample["title"]] for sample in dataset}
         pickle.dump(_lowerCAmelCase				,    open(_lowerCAmelCase				,    "wb" ) )

         lowercase			:List[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			:Any							=  RagRetriever(
             _lowerCAmelCase				,    question_encoder_tokenizer=self.get_dpr_tokenizer()				,    generator_tokenizer=self.get_bart_tokenizer() )
         return retriever






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






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






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






       @require_torch
       @require_tokenizers
       @require_sentencepiece
       def     SCREAMING_SNAKE_CASE			(      self:    Optional[int] ):
         import torch

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

         lowercase			:List[Any]							=  [[5, 7], [10, 11]]
         lowercase			:Union[str, Any]							=  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.









import math
import sys



def 					UpperCAmelCase__     (   lowerCamelCase      ):
  if number != int(lowerCamelCase      ):
    raise ValueError("the value of input must be a natural number"      )
  if number < 0:
    raise ValueError("the value of input must not be a negative number"      )
  if number == 0:
    return 1
  lowercase			:Dict							=  [-1] * (number + 1)
  lowercase			:Optional[Any]							=  0
  for i in range(1,	number + 1      ):
    lowercase			:Optional[int]							=  sys.maxsize
    lowercase			:Any							=  int(math.sqrt(lowerCamelCase      )      )
    for j in range(1,	root + 1      ):
      lowercase			:Any							=  1 + answers[i - (j**2)]
      lowercase			:Dict							=  min(lowerCamelCase,	lowerCamelCase      )
    lowercase			:Union[str, Any]							=  answer
  return answers[number]


if __name__ == "__main__":
     import doctest

     doctest.testmod()



