





from __future__ import annotations

import inspect
import unittest
from math import floor

import numpy as np

from transformers import CvtConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import TFCvtForImageClassification, TFCvtModel
  from transformers.models.cvt.modeling_tf_cvt import TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import AutoImageProcessor






class 							__lowercase			(     SCREAMING_SNAKE_CASE_       ):






      def 			_a						(		self)     ->       Tuple:
       __snake_case         =			self.config_class(**self.inputs_dict)
       self.parent.assertTrue(hasattr(__SCREAMING_SNAKE_CASE			,							'embed_dim'))
       self.parent.assertTrue(hasattr(__SCREAMING_SNAKE_CASE			,							'num_heads'))






class 							__lowercase			:






      def __init__(		self			,							lowercase_			,							lowercase_=1_3			,							lowercase_=6_4			,							lowercase_=3			,							lowercase_=[1_6, 4_8, 9_6]			,							lowercase_=[1, 3, 6]			,							lowercase_=[1, 2, 1_0]			,							lowercase_=[7, 3, 3]			,							lowercase_=[4, 2, 2]			,							lowercase_=[2, 1, 1]			,							lowercase_=[2, 2, 2]			,							lowercase_=[False, False, True]			,							lowercase_=[0.0, 0.0, 0.0]			,							lowercase_=0.02			,							lowercase_=1e-12			,							lowercase_=True			,							lowercase_=True			,							lowercase_=2			,							)     ->       Optional[int]:
       __snake_case         =			parent
       __snake_case         =			batch_size
       __snake_case         =			image_size
       __snake_case         =			patch_sizes
       __snake_case         =			patch_stride
       __snake_case         =			patch_padding
       __snake_case         =			is_training
       __snake_case         =			use_labels
       __snake_case         =			num_labels
       __snake_case         =			num_channels
       __snake_case         =			embed_dim
       __snake_case         =			num_heads
       __snake_case         =			stride_kv
       __snake_case         =			depth
       __snake_case         =			cls_token
       __snake_case         =			attention_drop_rate
       __snake_case         =			initializer_range
       __snake_case         =			layer_norm_eps






      def 			_a						(		self)     ->       Tuple:
       __snake_case         =			floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

       __snake_case         =			None
       if self.use_labels:
        # create a random int32 tensor of given shape
        __snake_case         =			ids_tensor([self.batch_size]			,							self.num_labels)

       __snake_case         =			self.get_config()
       return config, pixel_values, labels






      def 			_a						(		self)     ->       Union[str, Any]:
       return CvtConfig(
           image_size=self.image_size			,							num_labels=self.num_labels			,							num_channels=self.num_channels			,							embed_dim=self.embed_dim			,							num_heads=self.num_heads			,							patch_sizes=self.patch_sizes			,							patch_padding=self.patch_padding			,							patch_stride=self.patch_stride			,							stride_kv=self.stride_kv			,							depth=self.depth			,							cls_token=self.cls_token			,							attention_drop_rate=self.attention_drop_rate			,							initializer_range=self.initializer_range			,							)






      def 			_a						(		self			,							lowercase_			,							lowercase_			,							lowercase_)     ->       Optional[Any]:
       __snake_case         =			TFCvtModel(config=__SCREAMING_SNAKE_CASE)
       __snake_case         =			model(__SCREAMING_SNAKE_CASE			,							training=__SCREAMING_SNAKE_CASE)
       __snake_case         =			(self.image_size, self.image_size)
       __snake_case         =			image_size[0], image_size[1]
       for i in range(len(self.depth)):
        __snake_case         =			floor(((height + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1)
        __snake_case         =			floor(((width + 2 * self.patch_padding[i] - self.patch_sizes[i]) / self.patch_stride[i]) + 1)
       self.parent.assertEqual(result.last_hidden_state.shape			,							(self.batch_size, self.embed_dim[-1], height, width))






      def 			_a						(		self			,							lowercase_			,							lowercase_			,							lowercase_)     ->       List[Any]:
       __snake_case         =			self.num_labels
       __snake_case         =			TFCvtForImageClassification(__SCREAMING_SNAKE_CASE)
       __snake_case         =			model(__SCREAMING_SNAKE_CASE			,							labels=__SCREAMING_SNAKE_CASE			,							training=__SCREAMING_SNAKE_CASE)
       self.parent.assertEqual(result.logits.shape			,							(self.batch_size, self.num_labels))







      def 			_a						(		self)     ->       int:
       __snake_case         =			self.prepare_config_and_inputs()
       __snake_case         =			config_and_inputs
       __snake_case         =			{'''pixel_values''': pixel_values}
       return config, inputs_dict






@require_tf
class 							__lowercase			(     SCREAMING_SNAKE_CASE_   ,							SCREAMING_SNAKE_CASE_   ,							unittest.TestCase       ):
      __UpperCAmelCase		      =						(TFCvtModel, TFCvtForImageClassification) if is_tf_available() else ()
      __UpperCAmelCase		      =						(
          {'''feature-extraction''': TFCvtModel, '''image-classification''': TFCvtForImageClassification}
          if is_tf_available()
          else {}
      )
      __UpperCAmelCase		      =						False
      __UpperCAmelCase		      =						False
      __UpperCAmelCase		      =						False
      __UpperCAmelCase		      =						False
      __UpperCAmelCase		      =						False






      def 			_a						(		self)     ->       Dict:
       __snake_case         =			TFCvtModelTester(self)
       __snake_case         =			TFCvtConfigTester(self			,							config_class=__SCREAMING_SNAKE_CASE			,							has_text_modality=__SCREAMING_SNAKE_CASE			,							hidden_size=3_7)






      def 			_a						(		self)     ->       List[str]:
       self.config_tester.create_and_test_config_common_properties()
       self.config_tester.create_and_test_config_to_json_string()
       self.config_tester.create_and_test_config_to_json_file()
       self.config_tester.create_and_test_config_from_and_save_pretrained()
       self.config_tester.create_and_test_config_with_num_labels()
       self.config_tester.check_config_can_be_init_without_params()
       self.config_tester.check_config_arguments_init()






      @unittest.skip(reason='Cvt does not output attentions')
      def 			_a						(		self)     ->       int:
       pass






      @unittest.skip(reason='Cvt does not use inputs_embeds')
      def 			_a						(		self)     ->       Any:
       pass






      @unittest.skip(reason='Cvt does not support input and output embeddings')
      def 			_a						(		self)     ->       Optional[int]:
       pass






      @unittest.skipIf(
          not is_tf_available() or len(tf.config.list_physical_devices('GPU')) == 0			,							reason='TF does not support backprop for grouped convolutions on CPU.'			,							)
      def 			_a						(		self)     ->       Optional[Any]:
       super().test_dataset_conversion()






      @unittest.skipIf(
          not is_tf_available() or len(tf.config.list_physical_devices('GPU')) == 0			,							reason='TF does not support backprop for grouped convolutions on CPU.'			,							)
      @slow
      def 			_a						(		self)     ->       str:
       super().test_keras_fit()






      @unittest.skip(reason='Get `Failed to determine best cudnn convolution algo.` error after using TF 2.12+cuda 11.8')
      def 			_a						(		self)     ->       Tuple:
       __snake_case         =			tf.keras.mixed_precision.Policy('mixed_float16')
       tf.keras.mixed_precision.set_global_policy(__SCREAMING_SNAKE_CASE)
       super().test_keras_fit()
       tf.keras.mixed_precision.set_global_policy('float32')






      def 			_a						(		self)     ->       Optional[int]:
       __snake_case         =			self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        __snake_case         =			model_class(__SCREAMING_SNAKE_CASE)
        __snake_case         =			inspect.signature(model.call)
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        __snake_case         =			[*signature.parameters.keys()]

        __snake_case         =			['''pixel_values''']
        self.assertListEqual(arg_names[:1]			,							__SCREAMING_SNAKE_CASE)






      def 			_a						(		self)     ->       str:
       def check_hidden_states_output(lowercase_			,							lowercase_			,							lowercase_):
        __snake_case         =			model_class(__SCREAMING_SNAKE_CASE)

        __snake_case         =			model(**self._prepare_for_class(__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE))
        __snake_case         =			outputs.hidden_states

        __snake_case         =			len(self.model_tester.depth)
        self.assertEqual(len(__SCREAMING_SNAKE_CASE)			,							__SCREAMING_SNAKE_CASE)

        # verify the first hidden states (first block)
        self.assertListEqual(
            list(hidden_states[0].shape[-3:])			,							[
                self.model_tester.embed_dim[0],
                self.model_tester.image_size // 4,
                self.model_tester.image_size // 4,
            ]			,							)

       __snake_case         =			self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        __snake_case         =			True
        check_hidden_states_output(__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE)

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        __snake_case         =			True

        check_hidden_states_output(__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE			,							__SCREAMING_SNAKE_CASE)






      def 			_a						(		self)     ->       List[str]:
       __snake_case         =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE)






      def 			_a						(		self)     ->       int:
       __snake_case         =			self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*__SCREAMING_SNAKE_CASE)







      @slow
      def 			_a						(		self)     ->       Optional[Any]:
       for model_name in TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __snake_case         =			TFCvtModel.from_pretrained(__SCREAMING_SNAKE_CASE)
        self.assertIsNotNone(__SCREAMING_SNAKE_CASE)






def  A							(			)					->				List[Any]:
 '''simple docstring'''
 __snake_case         =			Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'				)
 return image







@require_tf
@require_vision
class 							__lowercase			(     unittest.TestCase       ):






      @cached_property
      def 			_a						(		self)     ->       List[str]:
       return AutoImageProcessor.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0])







      @slow
      def 			_a						(		self)     ->       Any:
       __snake_case         =			TFCvtForImageClassification.from_pretrained(TF_CVT_PRETRAINED_MODEL_ARCHIVE_LIST[0])

       __snake_case         =			self.default_image_processor
       __snake_case         =			prepare_img()
       __snake_case         =			image_processor(images=__SCREAMING_SNAKE_CASE			,							return_tensors='tf')

       # forward pass
       __snake_case         =			model(**__SCREAMING_SNAKE_CASE)

       # verify the logits
       __snake_case         =			tf.TensorShape((1, 1_0_0_0))
       self.assertEqual(outputs.logits.shape			,							__SCREAMING_SNAKE_CASE)

       __snake_case         =			tf.constant([0.9285, 0.9015, -0.3150])
       self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()			,							__SCREAMING_SNAKE_CASE			,							atol=1e-4))









'''simple docstring'''







import logging
import os
from dataclasses import dataclass, field
from functools import partial
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import List, Optional

import faiss
import torch
from datasets import Features, Sequence, Value, load_dataset

from transformers import DPRContextEncoder, DPRContextEncoderTokenizerFast, HfArgumentParser


__lowerCAmelCase									=   logging.getLogger(__name__)
torch.set_grad_enabled(False)
__lowerCAmelCase									=   'cuda' if torch.cuda.is_available() else 'cpu'





def 		_UpperCAmelCase (				__A				: str    ,		__A				: Optional[Any]=1_00    ,		__A				: int=" "   ):
				a_		:	Optional[int]									=   text.split(__A   )
				return [character.join(text[i : i + n]   ).strip() for i in range(0    ,		len(__A   )    ,		__A   )]





def 		_UpperCAmelCase (				__A				: dict   ):
				a_					,      a_		:	List[Any]									=   [], []
				for title, text in zip(documents['''title''']    ,		documents['''text''']   ):
								if text is not None:
												for passage in split_text(__A   ):
																titles.append(title if title is not None else ''''''   )
																texts.append(__A   )
				return {"title": titles, "text": texts}





def 		_UpperCAmelCase (				__A				: dict    ,		__A				: DPRContextEncoder    ,		__A				: DPRContextEncoderTokenizerFast   ):
				a_		:	int									=   ctx_tokenizer(
				    documents['''title''']    ,		documents['''text''']    ,		truncation=__A    ,		padding='''longest'''    ,		return_tensors='''pt'''   )['''input_ids''']
				a_		:	List[Any]									=   ctx_encoder(input_ids.to(device=__A   )    ,		return_dict=__A   ).pooler_output
				return {"embeddings": embeddings.detach().cpu().numpy()}





def 		_UpperCAmelCase (				__A				: "RagExampleArguments"    ,		__A				: "ProcessingArguments"    ,		__A				: "IndexHnswArguments"    ,		):
				######################################
				logger.info('''Step 1 - Create the dataset'''   )
				######################################

				# The dataset needed for RAG must have three columns:
				# - title (string): title of the document
				# - text (string): text of a passage of the document
				# - embeddings (array of dimension d): DPR representation of the passage
				# Let's say you have documents in tab-separated csv files with columns "title" and "text"
				assert os.path.isfile(rag_example_args.csv_path   ), "Please provide a valid path to a csv file"

				# You can load a Dataset object this way
				a_		:	List[str]									=   load_dataset(
				    '''csv'''    ,		data_files=[rag_example_args.csv_path]    ,		split='''train'''    ,		delimiter='''\t'''    ,		column_names=['''title''', '''text''']   )

				# More info about loading csv files in the documentation: https://huggingface.co/docs/datasets/loading_datasets.html?highlight=csv#csv-files

				# Then split the documents into passages of 100 words
				a_		:	Dict									=   dataset.map(__A    ,		batched=__A    ,		num_proc=processing_args.num_proc   )

				# And compute the embeddings
				a_		:	str									=   DPRContextEncoder.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name   ).to(device=__A   )
				a_		:	Tuple									=   DPRContextEncoderTokenizerFast.from_pretrained(rag_example_args.dpr_ctx_encoder_model_name   )
				a_		:	Union[str, Any]									=   Features(
				    {'''text''': Value('''string'''   ), '''title''': Value('''string'''   ), '''embeddings''': Sequence(Value('''float32'''   )   )}   )  # optional, save as float32 instead of float64 to save space
				a_		:	Union[str, Any]									=   dataset.map(
				    partial(__A    ,		ctx_encoder=__A    ,		ctx_tokenizer=__A   )    ,		batched=__A    ,		batch_size=processing_args.batch_size    ,		features=__A    ,		)

				# And finally save your dataset
				a_		:	Tuple									=   os.path.join(rag_example_args.output_dir    ,		'''my_knowledge_dataset'''   )
				dataset.save_to_disk(__A   )
				# from datasets import load_from_disk
				# dataset = load_from_disk(passages_path)  # to reload the dataset

				######################################
				logger.info('''Step 2 - Index the dataset'''   )
				######################################

				# Let's use the Faiss implementation of HNSW for fast approximate nearest neighbor search
				a_		:	Dict									=   faiss.IndexHNSWFlat(index_hnsw_args.d    ,		index_hnsw_args.m    ,		faiss.METRIC_INNER_PRODUCT   )
				dataset.add_faiss_index('''embeddings'''    ,		custom_index=__A   )

				# And save the index
				a_		:	Optional[Any]									=   os.path.join(rag_example_args.output_dir    ,		'''my_knowledge_dataset_hnsw_index.faiss'''   )
				dataset.get_index('''embeddings'''   ).save(__A   )
				# dataset.load_faiss_index("embeddings", index_path)  # to reload the index





@dataclass
class 						SCREAMING_SNAKE_CASE       :
			snake_case__	=					field(
			    default=str(Path(SCREAMING_SNAKE_CASE_							).parent / "test_run" / "dummy-kb" / "my_knowledge_dataset.csv"							)						, metadata={"help": "Path to a tab-separated csv file with columns 'title' and 'text'"}						, )
			snake_case__	=					field(
			    default=SCREAMING_SNAKE_CASE_						, metadata={"help": "Question that is passed as input to RAG. Default is 'What does Moses' rod turn into ?'."}						, )
			snake_case__	=					field(
			    default="facebook/rag-sequence-nq"						, metadata={"help": "The RAG model to use. Either 'facebook/rag-sequence-nq' or 'facebook/rag-token-nq'"}						, )
			snake_case__	=					field(
			    default="facebook/dpr-ctx_encoder-multiset-base"						, metadata={
			        "help": (
			            "The DPR context encoder model to use. Either 'facebook/dpr-ctx_encoder-single-nq-base' or"
			            " 'facebook/dpr-ctx_encoder-multiset-base'"
			        )
			    }						, )
			snake_case__	=					field(
			    default=str(Path(SCREAMING_SNAKE_CASE_							).parent / "test_run" / "dummy-kb"							)						, metadata={"help": "Path to a directory where the dataset passages and the index will be saved"}						, )





@dataclass
class 						SCREAMING_SNAKE_CASE       :
			snake_case__	=					field(
			    default=SCREAMING_SNAKE_CASE_						, metadata={
			        "help": "The number of processes to use to split the documents into passages. Default is single process."
			    }						, )
			snake_case__	=					field(
			    default=16						, metadata={
			        "help": "The batch size to use when computing the passages embeddings using the DPR context encoder."
			    }						, )





@dataclass
class 						SCREAMING_SNAKE_CASE       :
			snake_case__	=					field(
			    default=768						, metadata={"help": "The dimension of the embeddings to pass to the HNSW Faiss index."}						, )
			snake_case__	=					field(
			    default=128						, metadata={
			        "help": (
			            "The number of bi-directional links created for every new element during the HNSW index construction."
			        )
			    }						, )


if __name__ == "__main__":
		logging.basicConfig(level=logging.WARNING)
		logger.setLevel(logging.INFO)

		__lowerCAmelCase									=   HfArgumentParser((RagExampleArguments, ProcessingArguments, IndexHnswArguments))
		__lowerCAmelCase		,				__lowerCAmelCase		,				__lowerCAmelCase									=   parser.parse_args_into_dataclasses()
		with TemporaryDirectory() as tmp_dir:
				__lowerCAmelCase									=   rag_example_args.output_dir or tmp_dir
				main(rag_example_args, processing_args, index_hnsw_args)




