

from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
	import numpy
	import tensorflow as tf

	from transformers import (
	    TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
	    TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
	    TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
	    BertConfig,
	    DPRConfig,
	    TFDPRContextEncoder,
	    TFDPRQuestionEncoder,
	    TFDPRReader,
	)



class     _a						:


					'''simple docstring'''






					def __init__( self    ,       A__    ,       A__=13    ,       A__=7    ,       A__=True    ,       A__=True    ,       A__=True    ,       A__=True    ,       A__=99    ,       A__=32    ,       A__=2    ,       A__=4    ,       A__=37    ,       A__="gelu"    ,       A__=0.1    ,       A__=0.1    ,       A__=512    ,       A__=16    ,       A__=2    ,       A__=0.0_2    ,       A__=3    ,       A__=4    ,       A__=None    ,       A__=0    ,       ):
						A__							:    Tuple        =		parent
						A__							:    Any        =		batch_size
						A__							:    List[Any]        =		seq_length
						A__							:    List[str]        =		is_training
						A__							:    Union[str, Any]        =		use_input_mask
						A__							:    Dict        =		use_token_type_ids
						A__							:    str        =		use_labels
						A__							:    Union[str, Any]        =		vocab_size
						A__							:    Dict        =		hidden_size
						A__							:    Tuple        =		num_hidden_layers
						A__							:    int        =		num_attention_heads
						A__							:    Union[str, Any]        =		intermediate_size
						A__							:    Dict        =		hidden_act
						A__							:    Any        =		hidden_dropout_prob
						A__							:    List[str]        =		attention_probs_dropout_prob
						A__							:    Tuple        =		max_position_embeddings
						A__							:    List[str]        =		type_vocab_size
						A__							:    List[str]        =		type_sequence_label_size
						A__							:    List[Any]        =		initializer_range
						A__							:    Optional[int]        =		num_labels
						A__							:    str        =		num_choices
						A__							:    Optional[int]        =		scope
						A__							:    str        =		projection_dim






					def   __A			( self			):
						A__							:    List[str]        =		ids_tensor([self.batch_size, self.seq_length]    ,       self.vocab_size			)

						A__							:    int        =		None
						if self.use_input_mask:
							# follow test_modeling_tf_ctrl.py
							A__							:    int        =		random_attention_mask([self.batch_size, self.seq_length]			)

						A__							:    List[Any]        =		None
						if self.use_token_type_ids:
							A__							:    List[str]        =		ids_tensor([self.batch_size, self.seq_length]    ,       self.type_vocab_size			)

						A__							:    Any        =		None
						A__							:    Tuple        =		None
						A__							:    Dict        =		None
						if self.use_labels:
							A__							:    Any        =		ids_tensor([self.batch_size]    ,       self.type_sequence_label_size			)
							A__							:    Tuple        =		ids_tensor([self.batch_size, self.seq_length]    ,       self.num_labels			)
							A__							:    Optional[int]        =		ids_tensor([self.batch_size]    ,       self.num_choices			)

						A__							:    str        =		BertConfig(
						    vocab_size=self.vocab_size    ,       hidden_size=self.hidden_size    ,       num_hidden_layers=self.num_hidden_layers    ,       num_attention_heads=self.num_attention_heads    ,       intermediate_size=self.intermediate_size    ,       hidden_act=self.hidden_act    ,       hidden_dropout_prob=self.hidden_dropout_prob    ,       attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,       max_position_embeddings=self.max_position_embeddings    ,       type_vocab_size=self.type_vocab_size    ,       is_decoder=A__    ,       initializer_range=self.initializer_range    ,       )
						A__							:    Optional[int]        =		DPRConfig(projection_dim=self.projection_dim    ,       **config.to_dict()			)

						return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       A__    ,       A__    ,       A__    ,       A__			):
						A__							:    Any        =		TFDPRContextEncoder(config=A__			)
						A__							:    Dict        =		model(A__    ,       attention_mask=A__    ,       token_type_ids=A__			)
						A__							:    Tuple        =		model(A__    ,       token_type_ids=A__			)
						A__							:    Tuple        =		model(A__			)
						self.parent.assertEqual(result.pooler_output.shape    ,       (self.batch_size, self.projection_dim or self.hidden_size)			)






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       A__    ,       A__    ,       A__    ,       A__			):
						A__							:    List[str]        =		TFDPRQuestionEncoder(config=A__			)
						A__							:    Dict        =		model(A__    ,       attention_mask=A__    ,       token_type_ids=A__			)
						A__							:    int        =		model(A__    ,       token_type_ids=A__			)
						A__							:    Any        =		model(A__			)
						self.parent.assertEqual(result.pooler_output.shape    ,       (self.batch_size, self.projection_dim or self.hidden_size)			)






					def   __A			( self    ,       A__    ,       A__    ,       A__    ,       A__    ,       A__    ,       A__    ,       A__			):
						A__							:    Optional[Any]        =		TFDPRReader(config=A__			)
						A__							:    Dict        =		model(A__    ,       attention_mask=A__			)

						self.parent.assertEqual(result.start_logits.shape    ,       (self.batch_size, self.seq_length)			)
						self.parent.assertEqual(result.end_logits.shape    ,       (self.batch_size, self.seq_length)			)
						self.parent.assertEqual(result.relevance_logits.shape    ,       (self.batch_size,)			)






					def   __A			( self			):
						A__							:    Dict        =		self.prepare_config_and_inputs()
						(
						    (
						    A__
						)   ,   (
						    A__
						)   ,   (
						    A__
						)   ,   (
						    A__
						)   ,   (
						    A__
						)   ,   (
						    A__
						)   ,   (
						    A__
						)   ,   
						)							:    Union[str, Any]        =		config_and_inputs
						A__							:    Optional[int]        =		{"""input_ids""": input_ids}
						return config, inputs_dict



@require_tf
class     _a						(__magic_name__		,   __magic_name__		,   unittest.TestCase				):


					'''simple docstring'''






					UpperCAmelCase__:   int  							=       (
					    (
					        TFDPRContextEncoder,
					        TFDPRQuestionEncoder,
					        TFDPRReader,
					    )
					    if is_tf_available()
					    else ()
					)
					UpperCAmelCase__:   Optional[Any]  							=       {'''feature-extraction''': TFDPRQuestionEncoder} if is_tf_available() else {}

					UpperCAmelCase__:   Optional[Any]  							=       False
					UpperCAmelCase__:   Tuple  							=       False
					UpperCAmelCase__:   Dict  							=       False
					UpperCAmelCase__:   List[Any]  							=       False
					UpperCAmelCase__:   str  							=       False






					def   __A			( self			):
						A__							:    Any        =		TFDPRModelTester(self			)
						A__							:    List[str]        =		ConfigTester(self    ,       config_class=A__    ,       hidden_size=37			)






					def   __A			( self			):
						self.config_tester.run_common_tests()






					def   __A			( self			):
						A__							:    List[str]        =		self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_dpr_context_encoder(*A__			)






					def   __A			( self			):
						A__							:    Union[str, Any]        =		self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_dpr_question_encoder(*A__			)






					def   __A			( self			):
						A__							:    List[Any]        =		self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_dpr_reader(*A__			)






					@slow
					def   __A			( self			):
						for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							A__							:    int        =		TFDPRContextEncoder.from_pretrained(A__			)
							self.assertIsNotNone(A__			)

						for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							A__							:    str        =		TFDPRContextEncoder.from_pretrained(A__			)
							self.assertIsNotNone(A__			)

						for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							A__							:    str        =		TFDPRQuestionEncoder.from_pretrained(A__			)
							self.assertIsNotNone(A__			)

						for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							A__							:    Any        =		TFDPRReader.from_pretrained(A__			)
							self.assertIsNotNone(A__			)



@require_tf
class     _a						(unittest.TestCase				):


					'''simple docstring'''






					@slow
					def   __A			( self			):
						A__							:    int        =		TFDPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base"""			)

						A__							:    Any        =		tf.constant(
						    [[101, 7592, 1010, 2003, 2026, 3899, 1_0140, 1029, 102]]			)  # [CLS] hello, is my dog cute? [SEP]
						A__							:    str        =		model(A__			)[0]  # embedding shape = (1, 768)
						# compare the actual values for a slice.
						A__							:    Optional[Any]        =		tf.constant(
						    [
						        [
						            0.0_3_2_3_6_2_5_3,
						            0.1_2_7_5_3_3_3_5,
						            0.1_6_8_1_8_5_0_9,
						            0.0_0_2_7_9_7_8_6,
						            0.3_8_9_6_9_3_3,
						            0.2_4_2_6_4_9_4_5,
						            0.2_1_7_8_9_7_1,
						            -0.0_2_3_3_5_2_2_7,
						            -0.0_8_4_8_1_9_5_9,
						            -0.1_4_3_2_4_1_1_7,
						        ]
						    ]			)
						self.assertTrue(numpy.allclose(output[:, :10].numpy()    ,       expected_slice.numpy()    ,       atol=1e-4			)			)






import os

import pytest
import yaml

from datasets.features.features import Features, Value
from datasets.info import DatasetInfo, DatasetInfosDict

@pytest.mark.parametrize(
    """files"""    ,      [
        ["""full:README.md""", """dataset_infos.json"""],
        ["""empty:README.md""", """dataset_infos.json"""],
        ["""dataset_infos.json"""],
        ["""full:README.md"""],
    ]    ,      )
def 						UpperCamelCase					(lowercase_: List[str]    ,      lowercase_: Optional[int]					)     ->      List[str]:
	A__							:    Tuple        =		tmp_path_factory.mktemp("""dset_infos_dir"""					)
	if "full:README.md" in files:
		with open(dataset_infos_dir / """README.md"""    ,      """w"""					) as f:
			f.write("""---\ndataset_info:\n  dataset_size: 42\n---"""					)
	if "empty:README.md" in files:
		with open(dataset_infos_dir / """README.md"""    ,      """w"""					) as f:
			f.write(""""""					)
    # we want to support dataset_infos.json for backward compatibility
	if "dataset_infos.json" in files:
		with open(dataset_infos_dir / """dataset_infos.json"""    ,      """w"""					) as f:
			f.write("""{\"default\": {\"dataset_size\": 42}}"""					)
	A__							:    List[Any]        =		DatasetInfosDict.from_directory(lowercase_					)
	assert dataset_infos
	assert dataset_infos["default"].dataset_size == 42

@pytest.mark.parametrize(
    """dataset_info"""    ,      [
        DatasetInfo(),
        DatasetInfo(
            description="""foo"""    ,      features=Features({"""a""": Value("""int32"""					)}					)    ,      builder_name="""builder"""    ,      config_name="""config"""    ,      version="""1.0.0"""    ,      splits=[{"""name""": """train"""}]    ,      download_size=42    ,      ),
    ]    ,      )
def 						UpperCamelCase					(lowercase_: str    ,      lowercase_: DatasetInfo					)     ->      List[Any]:
	A__							:    Union[str, Any]        =		str(lowercase_					)
	dataset_info.write_to_directory(lowercase_					)
	A__							:    List[Any]        =		DatasetInfo.from_directory(lowercase_					)
	assert dataset_info == reloaded
	assert os.path.exists(os.path.join(lowercase_    ,      """dataset_info.json"""					)					)

def 						UpperCamelCase					()     ->      List[Any]:
	A__							:    Union[str, Any]        =		DatasetInfo(
	    description="""foo"""    ,      citation="""bar"""    ,      homepage="""https://foo.bar"""    ,      license="""CC0"""    ,      features=Features({"""a""": Value("""int32"""					)}					)    ,      post_processed={}    ,      supervised_keys=()    ,      task_templates=[]    ,      builder_name="""builder"""    ,      config_name="""config"""    ,      version="""1.0.0"""    ,      splits=[{"""name""": """train""", """num_examples""": 42}]    ,      download_checksums={}    ,      download_size=1337    ,      post_processing_size=442    ,      dataset_size=1234    ,      size_in_bytes=1337 + 442 + 1234    ,      )
	A__							:    Dict        =		dataset_info._to_yaml_dict()
	assert sorted(lowercase_					) == sorted(DatasetInfo._INCLUDED_INFO_IN_YAML					)
	for key in DatasetInfo._INCLUDED_INFO_IN_YAML:
		assert key in dataset_info_yaml_dict
		assert isinstance(dataset_info_yaml_dict[key]    ,      (list, dict, int, str)					)
	A__							:    Union[str, Any]        =		yaml.safe_dump(lowercase_					)
	A__							:    List[Any]        =		yaml.safe_load(lowercase_					)
	assert dataset_info_yaml_dict == reloaded

def 						UpperCamelCase					()     ->      List[str]:
	A__							:    Optional[int]        =		DatasetInfo()
	A__							:    List[Any]        =		dataset_info._to_yaml_dict()
	assert dataset_info_yaml_dict == {}


@pytest.mark.parametrize(
    """dataset_infos_dict"""    ,      [
        DatasetInfosDict(),
        DatasetInfosDict({"""default""": DatasetInfo()}					),
        DatasetInfosDict({"""my_config_name""": DatasetInfo()}					),
        DatasetInfosDict(
            {
                """default""": DatasetInfo(
                    description="""foo"""    ,      features=Features({"""a""": Value("""int32"""					)}					)    ,      builder_name="""builder"""    ,      config_name="""config"""    ,      version="""1.0.0"""    ,      splits=[{"""name""": """train"""}]    ,      download_size=42    ,      )
            }					),
        DatasetInfosDict(
            {
                """v1""": DatasetInfo(dataset_size=42					),
                """v2""": DatasetInfo(dataset_size=1337					),
            }					),
    ]    ,      )
def 						UpperCamelCase					(lowercase_: Tuple    ,      lowercase_: DatasetInfosDict					)     ->      Optional[Any]:
	A__							:    List[Any]        =		str(lowercase_					)
	dataset_infos_dict.write_to_directory(lowercase_					)
	A__							:    Dict        =		DatasetInfosDict.from_directory(lowercase_					)

	# the config_name of the dataset_infos_dict take over the attribute
	for config_name, dataset_info in dataset_infos_dict.items():
		A__							:    Optional[int]        =		config_name
		# the yaml representation doesn't include fields like description or citation
		# so we just test that we can recover what we can from the yaml
		A__							:    List[str]        =		DatasetInfo._from_yaml_dict(dataset_info._to_yaml_dict()					)
	assert dataset_infos_dict == reloaded

	if dataset_infos_dict:
		assert os.path.exists(os.path.join(lowercase_    ,      """README.md"""					)					)



