





import json
import os
import pickle
import shutil
import tempfile
from unittest import TestCase
from unittest.mock import patch

import numpy as np
from datasets import Dataset

from transformers import is_faiss_available
from transformers.models.bart.configuration_bart import BartConfig
from transformers.models.bart.tokenization_bart import BartTokenizer
from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
from transformers.models.dpr.configuration_dpr import DPRConfig
from transformers.models.dpr.tokenization_dpr import DPRContextEncoderTokenizer, DPRQuestionEncoderTokenizer
from transformers.models.rag.configuration_rag import RagConfig
from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
from transformers.testing_utils import require_faiss, require_sentencepiece, require_tokenizers, require_torch


if is_faiss_available():
       import faiss







@require_faiss
class       __A							(   lowerCamelCase__	):




      """simple docstring"""







      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Optional[Any]						=		tempfile.mkdtemp()
             _lowerCamelCase				:  Tuple						=		8

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

             _lowerCamelCase				:  Tuple						=		os.path.join(self.tmpdirname   ,				'''bart_tokenizer''')
             os.makedirs(a__   ,				exist_ok=a__)
             _lowerCamelCase				:  Tuple						=		os.path.join(a__   ,				BART_VOCAB_FILES_NAMES['''vocab_file'''])
             _lowerCamelCase				:  Optional[Any]						=		os.path.join(a__   ,				BART_VOCAB_FILES_NAMES['''merges_file'''])
             with open(self.vocab_file   ,				'''w'''   ,				encoding='''utf-8''') as fp:
                    fp.write(json.dumps(a__) + '''\n''')
             with open(self.merges_file   ,				'''w'''   ,				encoding='''utf-8''') as fp:
                    fp.write('''\n'''.join(a__))


      def 					__snake_case      (		self):



             """simple docstring"""





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


      def 					__snake_case      (		self):



             """simple docstring"""





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


      def 					__snake_case      (		self):



             """simple docstring"""





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


      def 					__snake_case      (		self):



             """simple docstring"""





             shutil.rmtree(self.tmpdirname)


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  List[str]						=		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 					__snake_case      (		self):



             """simple docstring"""





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


      def 					__snake_case      (		self   ,				a__):



             """simple docstring"""





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


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  List[str]						=		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)

             _lowerCamelCase				:  int						=		os.path.join(self.tmpdirname   ,				'''hf_bert_base.hnswSQ8_correct_phi_128.c_index''')
             dataset.save_faiss_index('''embeddings'''   ,				index_file_name + '''.index.dpr''')
             pickle.dump(dataset['''id''']   ,				open(index_file_name + '''.index_meta.dpr'''   ,				'''wb'''))

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

             _lowerCamelCase				:  int						=		RagConfig(
                 retrieval_vector_size=self.retrieval_vector_size   ,				question_encoder=DPRConfig().to_dict()   ,				generator=BartConfig().to_dict()   ,				index_name='''legacy'''   ,				index_path=self.tmpdirname   ,				)
             _lowerCamelCase				:  Optional[Any]						=		RagRetriever(
                 a__   ,				question_encoder_tokenizer=self.get_dpr_tokenizer()   ,				generator_tokenizer=self.get_bart_tokenizer())
             return retriever


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  int						=		1
             _lowerCamelCase				:  str						=		self.get_dummy_canonical_hf_index_retriever()
             _lowerCamelCase				:  int						=		np.array(
                 [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]   ,				dtype=np.floataa)
             _lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase				:  str						=		retriever.retrieve(a__   ,				n_docs=a__)
             self.assertEqual(retrieved_doc_embeds.shape   ,				(2, n_docs, self.retrieval_vector_size))
             self.assertEqual(len(a__)   ,				2)
             self.assertEqual(sorted(doc_dicts[0])   ,				['''embeddings''', '''id''', '''text''', '''title'''])
             self.assertEqual(len(doc_dicts[0]['''id'''])   ,				a__)
             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 					__snake_case      (		self):



             """simple docstring"""





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


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Union[str, Any]						=		1
             _lowerCamelCase				:  Tuple						=		self.get_dummy_custom_hf_index_retriever(from_disk=a__)
             _lowerCamelCase				:  Tuple						=		np.array(
                 [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]   ,				dtype=np.floataa)
             _lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase				:  Tuple						=		retriever.retrieve(a__   ,				n_docs=a__)
             self.assertEqual(retrieved_doc_embeds.shape   ,				(2, n_docs, self.retrieval_vector_size))
             self.assertEqual(len(a__)   ,				2)
             self.assertEqual(sorted(doc_dicts[0])   ,				['''embeddings''', '''id''', '''text''', '''title'''])
             self.assertEqual(len(doc_dicts[0]['''id'''])   ,				a__)
             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 					__snake_case      (		self):



             """simple docstring"""





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


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Dict						=		1
             _lowerCamelCase				:  Dict						=		self.get_dummy_custom_hf_index_retriever(from_disk=a__)
             _lowerCamelCase				:  str						=		np.array(
                 [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]   ,				dtype=np.floataa)
             _lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase				:  Dict						=		retriever.retrieve(a__   ,				n_docs=a__)
             self.assertEqual(retrieved_doc_embeds.shape   ,				(2, n_docs, self.retrieval_vector_size))
             self.assertEqual(len(a__)   ,				2)
             self.assertEqual(sorted(doc_dicts[0])   ,				['''embeddings''', '''id''', '''text''', '''title'''])
             self.assertEqual(len(doc_dicts[0]['''id'''])   ,				a__)
             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 					__snake_case      (		self):



             """simple docstring"""





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


      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  int						=		1
             _lowerCamelCase				:  Optional[Any]						=		self.get_dummy_legacy_index_retriever()
             _lowerCamelCase				:  int						=		np.array(
                 [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)]   ,				dtype=np.floataa)
             _lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase				:  Optional[int]						=		retriever.retrieve(a__   ,				n_docs=a__)
             self.assertEqual(retrieved_doc_embeds.shape   ,				(2, n_docs, self.retrieval_vector_size))
             self.assertEqual(len(a__)   ,				2)
             self.assertEqual(sorted(doc_dicts[0])   ,				['''text''', '''title'''])
             self.assertEqual(len(doc_dicts[0]['''text'''])   ,				a__)
             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 					__snake_case      (		self):



             """simple docstring"""





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


      @require_torch
      @require_tokenizers
      @require_sentencepiece
      def 					__snake_case      (		self):



             """simple docstring"""





             import torch

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

             _lowerCamelCase				:  Optional[Any]						=		retriever(
                 a__   ,				a__   ,				prefix=retriever.config.generator.prefix   ,				n_docs=a__   ,				return_tensors='''pt'''   ,				)
             _lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase,	_lowerCamelCase				:  int						=		(  # 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(a__   ,				torch.Tensor)
             self.assertIsInstance(a__   ,				torch.Tensor)
             self.assertIsInstance(a__   ,				torch.Tensor)



      @require_torch
      @require_tokenizers
      @require_sentencepiece
      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  str						=		self.get_dpr_ctx_encoder_tokenizer()
             _lowerCamelCase				:  Dict						=		1
             _lowerCamelCase				:  Dict						=		self.get_dummy_custom_hf_index_retriever(from_disk=a__)
             retriever.set_ctx_encoder_tokenizer(a__)

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

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








import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
       from sagemaker import Session, TrainingJobAnalytics
       from sagemaker.huggingface import HuggingFace







@pytest.mark.skipif(
    literal_eval(os.getenv("""TEST_SAGEMAKER"""			,"""False"""	)	) is not True			,reason="""Skipping test because should only be run when releasing minor transformers version"""			,)
@pytest.mark.usefixtures("""sm_env"""	)
@parameterized_class(
    [
        {
            """framework""": """pytorch""",
            """script""": """run_glue.py""",
            """model_name_or_path""": """distilbert-base-cased""",
            """instance_type""": """ml.g4dn.xlarge""",
            """results""": {"""train_runtime""": 650, """eval_accuracy""": 0.6, """eval_loss""": 0.9},
        },
        {
            """framework""": """tensorflow""",
            """script""": """run_tf.py""",
            """model_name_or_path""": """distilbert-base-cased""",
            """instance_type""": """ml.g4dn.xlarge""",
            """results""": {"""train_runtime""": 600, """eval_accuracy""": 0.3, """eval_loss""": 0.9},
        },
    ]	)
class       __A							(   unittest.TestCase	):




      """simple docstring"""







      def 					__snake_case      (		self):



             """simple docstring"""





             if self.framework == "pytorch":
                    subprocess.run(
                        F"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()   ,				encoding='''utf-8'''   ,				check=a__   ,				)
             assert hasattr(self   ,				'''env''')


      def 					__snake_case      (		self   ,				a__=1):



             """simple docstring"""





             return HuggingFace(
                 entry_point=self.script   ,				source_dir=self.env.test_path   ,				role=self.env.role   ,				image_uri=self.env.image_uri   ,				base_job_name=F"""{self.env.base_job_name}-single"""   ,				instance_count=a__   ,				instance_type=self.instance_type   ,				debugger_hook_config=a__   ,				hyperparameters={**self.env.hyperparameters, '''model_name_or_path''': self.model_name_or_path}   ,				metric_definitions=self.env.metric_definitions   ,				py_version='''py36'''   ,				)


      def 					__snake_case      (		self   ,				a__):



             """simple docstring"""





             TrainingJobAnalytics(a__).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv""")



      def 					__snake_case      (		self):



             """simple docstring"""





             _lowerCamelCase				:  Dict						=		self.create_estimator()

             # run training
             estimator.fit()

             # result dataframe
             _lowerCamelCase				:  Tuple						=		TrainingJobAnalytics(estimator.latest_training_job.name).dataframe()

             # extract kpis
             _lowerCamelCase				:  Tuple						=		list(result_metrics_df[result_metrics_df.metric_name == '''eval_accuracy''']['''value'''])
             _lowerCamelCase				:  Dict						=		list(result_metrics_df[result_metrics_df.metric_name == '''eval_loss''']['''value'''])
             # get train time from SageMaker job, this includes starting, preprocessing, stopping
             _lowerCamelCase				:  Optional[Any]						=		(
                 Session().describe_training_job(estimator.latest_training_job.name).get('''TrainingTimeInSeconds'''   ,				99_9999)
             )

             # assert kpis
             assert train_runtime <= self.results["train_runtime"]
             assert all(t >= self.results['''eval_accuracy'''] for t in eval_accuracy)
             assert all(t <= self.results['''eval_loss'''] for t in eval_loss)

             # dump tests result into json file to share in PR
             with open(F"""{estimator.latest_training_job.name}.json"""   ,				'''w''') as outfile:
                    json.dump({'''train_time''': train_runtime, '''eval_accuracy''': eval_accuracy, '''eval_loss''': eval_loss}   ,				a__)

