




'''simple docstring'''

import jax
import jax.numpy as jnp
from bigbird_flax import FlaxBigBirdForNaturalQuestions
from datasets import load_from_disk

from transformers import BigBirdTokenizerFast


a__         =     {0: '''null''', 1: '''short''', 2: '''long''', 3: '''yes''', 4: '''no'''}
a__         =     set(''''''.join(['''‘''', '''’''', '''´''', '''`''', '''.''', ''',''', '''-''', '''"''']))




def    snake_case__     (  a       ,      a=0       ,      a=None	)			->       Any:





			'''simple docstring'''






			return [" ".join(x.split(""" """	)[begin:end]	) for x in answers if len(x.split(""" """	)	) > 1]




def    snake_case__     (  a       ,      a=False	)			->       Dict:





			'''simple docstring'''






			if make_sub_answers:
						# if answers are longer than one word, make sure a predictions is correct if it coresponds to the complete 1: or :-1 sub word
						# *e.g.* if the correct answer contains a prefix such as "the", or "a"
						snake_case__          =						(
						    given_answers + get_sub_answers(a       ,      begin=1	) + get_sub_answers(a       ,      end=-1	)
						)
			snake_case__          =						[]
			for answer in given_answers:
						snake_case__          =						answer.replace("""_"""       ,      """ """	).lower()
						snake_case__          =						"""""".join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else """ """ for c in alias	)
						answers.append(""" """.join(alias.split()	).strip()	)
			return set(a	)




def    snake_case__     (  a       ,      a       ,      a=1       ,      a=100	)			->       str:





			'''simple docstring'''






			snake_case__    ,			snake_case__          =						jax.lax.top_k(a       ,      a	)
			snake_case__    ,			snake_case__          =						jax.lax.top_k(a       ,      a	)

			snake_case__          =						best_end_idx[:, None] - best_start_idx[None, :]
			snake_case__          =						jnp.logical_or(widths < 0       ,      widths > max_size	)
			snake_case__          =						(best_end_scores[:, None] + best_start_scores[None, :]) - (1e8 * mask)
			snake_case__          =						jnp.argmax(a	).item()

			return best_start_idx[best_score % top_k], best_end_idx[best_score // top_k]




def    snake_case__     (  a	)			->       Any:





			'''simple docstring'''






			snake_case__          =						sample["""question"""]["""text"""]
			snake_case__          =						sample["""document"""]["""tokens"""]["""token"""]
			snake_case__          =						sample["""document"""]["""tokens"""]["""is_html"""]
			snake_case__          =						sample["""annotations"""]["""long_answer"""]
			snake_case__          =						sample["""annotations"""]["""short_answers"""]

			snake_case__          =						""" """.join([context[i] for i in range(len(a	)	) if not is_html[i]]	)

			# 0 - No ; 1 - Yes
			for answer in sample["annotations"]["yes_no_answer"]:
						if answer == 0 or answer == 1:
									return {
									    "question": question,
									    "context": context_string,
									    "short": [],
									    "long": [],
									    "category": "no" if answer == 0 else "yes",
									}

			snake_case__          =						[]
			for s in short_answers:
						short_targets.extend(s["""text"""]	)
			snake_case__          =						list(set(a	)	)

			snake_case__          =						[]
			for s in long_answers:
						if s["start_token"] == -1:
									continue
						snake_case__          =						context[s["""start_token"""] : s["""end_token"""]]
						snake_case__          =						is_html[s["""start_token"""] : s["""end_token"""]]
						snake_case__          =						""" """.join([answer[i] for i in range(len(a	)	) if not html[i]]	)
						if new_answer not in long_targets:
									long_targets.append(a	)

			snake_case__          =						"""long_short""" if len(short_targets + long_targets	) > 0 else """null"""

			return {
			    "question": question,
			    "context": context_string,
			    "short": short_targets,
			    "long": long_targets,
			    "category": category,
			}




def    snake_case__     (  )			->       Dict:





			'''simple docstring'''






			snake_case__          =						load_from_disk("""natural-questions-validation"""	)
			snake_case__          =						dataset.map(a	).remove_columns(["""annotations""", """document""", """id"""]	)
			print(a	)

			snake_case__          =						dataset.filter(lambda a	: (len(x["""question"""]	) + len(x["""context"""]	)) < 4 * 4096	)
			snake_case__          =						short_validation_dataset.filter(lambda a	: x["category"] != "null"	)
			short_validation_dataset

			snake_case__          =						"""vasudevgupta/flax-bigbird-natural-questions"""
			snake_case__          =						FlaxBigBirdForNaturalQuestions.from_pretrained(a	)
			snake_case__          =						BigBirdTokenizerFast.from_pretrained(a	)

			@jax.jit
			def forward(*a       ,      **a	):
						snake_case__    ,			snake_case__    ,			snake_case__          =						model(*a       ,      **a	)
						return start_logits, end_logits, jnp.argmax(a       ,      axis=-1	)

			def evaluate(a	):
						# encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
						snake_case__          =						tokenizer(
						    example["""question"""]       ,      example["""context"""]       ,      return_tensors="""np"""       ,      max_length=4096       ,      padding="""max_length"""       ,      truncation=a       ,      )

						snake_case__    ,			snake_case__    ,			snake_case__          =						forward(**a	)

						snake_case__          =						CATEGORY_MAPPING[category.item()]

						snake_case__          =						example["""long"""] + example["""short"""]
						if example["category"] in ["yes", "no", "null"]:
									snake_case__          =						[example["""category"""]]
						snake_case__          =						example["""category"""] != """null"""
						# Now target can be: "yes", "no", "null", "list of long & short answers"

						if predicted_category in ["yes", "no", "null"]:
									snake_case__          =						[predicted_category]
									snake_case__          =						example["""output"""] == example["""targets"""]
									snake_case__          =						predicted_category != """null"""
									return example

						snake_case__          =						38 if predicted_category == """short""" else 1024
						snake_case__    ,			snake_case__          =						get_best_valid_start_end_idx(
						    start_scores[0]       ,      end_scores[0]       ,      top_k=8       ,      max_size=a	)

						snake_case__          =						inputs["""input_ids"""][0].tolist()
						snake_case__          =						[tokenizer.decode(input_ids[start_score : end_score + 1]	)]

						snake_case__          =						expand_to_aliases(example["""targets"""]       ,      make_sub_answers=a	)
						snake_case__          =						expand_to_aliases(example["""output"""]	)

						# some preprocessing to both prediction and answer
						snake_case__          =						{"""""".join(a.split()	) for a in answers}
						snake_case__          =						{"""""".join(p.split()	) for p in predictions}
						snake_case__          =						{s for s in predictions if s not in ["""``""", """''""", """`""", """'"""]}

						# if there is a common element, it's a exact match
						snake_case__          =						len(list(answers & predictions	)	) > 0
						snake_case__          =						predicted_category != """null""" and len(a	) > 0

						return example

			snake_case__          =						short_validation_dataset.map(a	)

			snake_case__          =						len(a	)
			snake_case__          =						len(short_validation_dataset.filter(lambda a	: x["match"] == 1	)	)
			print("""EM score:"""       ,      (matched / total) * 100       ,      """%"""	)


if __name__ == "__main__":
		main()



import inspect
import unittest

from transformers import ResNetConfig, is_flax_available
from transformers.testing_utils import require_flax, slow
from transformers.utils import cached_property, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
    import jax
    import jax.numpy as jnp

    from transformers.models.resnet.modeling_flax_resnet import FlaxResNetForImageClassification, FlaxResNetModel

if is_vision_available():
    from PIL import Image

    from transformers import AutoImageProcessor







class      _lowercase					(				unittest.TestCase):


       """simple docstring"""




       def __init__(  self				:     Dict   , __lowerCamelCase				:     int   , __lowerCamelCase				:     List[Any]=3   , __lowerCamelCase				:     Dict=32   , __lowerCamelCase				:     Optional[Any]=3   , __lowerCamelCase				:     List[str]=10   , __lowerCamelCase				:     List[str]=[10, 20, 30, 40]   , __lowerCamelCase				:     List[Any]=[1, 1, 2, 1]   , __lowerCamelCase				:     Any=True   , __lowerCamelCase				:     Optional[int]=True   , __lowerCamelCase				:     int="relu"   , __lowerCamelCase				:     List[str]=3   , __lowerCamelCase				:     Any=None   , ):


             '''simple docstring'''



             lowerCamelCase__						:							int       =    parent
             lowerCamelCase__						:							Optional[int]       =    batch_size
             lowerCamelCase__						:							List[Any]       =    image_size
             lowerCamelCase__						:							Union[str, Any]       =    num_channels
             lowerCamelCase__						:							Any       =    embeddings_size
             lowerCamelCase__						:							Tuple       =    hidden_sizes
             lowerCamelCase__						:							Tuple       =    depths
             lowerCamelCase__						:							List[str]       =    is_training
             lowerCamelCase__						:							Optional[Any]       =    use_labels
             lowerCamelCase__						:							str       =    hidden_act
             lowerCamelCase__						:							Any       =    num_labels
             lowerCamelCase__						:							List[str]       =    scope
             lowerCamelCase__						:							Any       =    len(__lowerCamelCase			)





       def 		lowerCAmelCase	(  self				:     List[str]			):


             '''simple docstring'''



             lowerCamelCase__						:							Any       =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

             lowerCamelCase__						:							Any       =    self.get_config()

             return config, pixel_values





       def 		lowerCAmelCase	(  self				:     Optional[int]			):


             '''simple docstring'''



             return ResNetConfig(
                 num_channels=self.num_channels   , embeddings_size=self.embeddings_size   , hidden_sizes=self.hidden_sizes   , depths=self.depths   , hidden_act=self.hidden_act   , num_labels=self.num_labels   , image_size=self.image_size   , )





       def 		lowerCAmelCase	(  self				:     str   , __lowerCamelCase				:     List[Any]   , __lowerCamelCase				:     Tuple			):


             '''simple docstring'''



             lowerCamelCase__						:							Any       =    FlaxResNetModel(config=__lowerCamelCase			)
             lowerCamelCase__						:							str       =    model(__lowerCamelCase			)

             # Output shape (b, c, h, w)
             self.parent.assertEqual(
                 result.last_hidden_state.shape   , (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)   , )





       def 		lowerCAmelCase	(  self				:     int   , __lowerCamelCase				:     Optional[Any]   , __lowerCamelCase				:     List[str]			):


             '''simple docstring'''



             lowerCamelCase__						:							Optional[Any]       =    self.num_labels
             lowerCamelCase__						:							int       =    FlaxResNetForImageClassification(config=__lowerCamelCase			)
             lowerCamelCase__						:							Optional[int]       =    model(__lowerCamelCase			)
             self.parent.assertEqual(result.logits.shape   , (self.batch_size, self.num_labels)			)





       def 		lowerCAmelCase	(  self				:     Dict			):


             '''simple docstring'''



             lowerCamelCase__						:							Tuple       =    self.prepare_config_and_inputs()
             lowerCamelCase__       ,  lowerCamelCase__						:							Tuple       =    config_and_inputs
             lowerCamelCase__						:							Optional[int]       =    {"pixel_values": pixel_values}
             return config, inputs_dict







@require_flax
class      _lowercase					(				lowercase__				,    unittest.TestCase):


       """simple docstring"""




       A__   		=     (FlaxResNetModel, FlaxResNetForImageClassification) if is_flax_available() else ()

       A__   		=     False
       A__   		=     False
       A__   		=     False





       def 		lowerCAmelCase	(  self				:     List[str]			):


             '''simple docstring'''



             lowerCamelCase__						:							List[str]       =    FlaxResNetModelTester(self			)
             lowerCamelCase__						:							Dict       =    ConfigTester(self   , config_class=__lowerCamelCase   , has_text_modality=__lowerCamelCase			)





       def 		lowerCAmelCase	(  self				:     Optional[Any]			):


             '''simple docstring'''



             self.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()





       def 		lowerCAmelCase	(  self				:     Dict			):


             '''simple docstring'''



             return





       def 		lowerCAmelCase	(  self				:     Dict			):


             '''simple docstring'''



             lowerCamelCase__						:							str       =    self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_model(*__lowerCamelCase			)





       def 		lowerCAmelCase	(  self				:     List[Any]			):


             '''simple docstring'''



             lowerCamelCase__						:							str       =    self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_for_image_classification(*__lowerCamelCase			)





       @unittest.skip(reason="ResNet does not use inputs_embeds"			)
       def 		lowerCAmelCase	(  self				:     Dict			):


             '''simple docstring'''



             pass





       @unittest.skip(reason="ResNet does not support input and output embeddings"			)
       def 		lowerCAmelCase	(  self				:     int			):


             '''simple docstring'''



             pass





       def 		lowerCAmelCase	(  self				:     Any			):


             '''simple docstring'''



             lowerCamelCase__       ,  lowerCamelCase__						:							Tuple       =    self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   lowerCamelCase__						:							Tuple       =    model_class(__lowerCamelCase			)
                   lowerCamelCase__						:							Tuple       =    inspect.signature(model.__call__			)
                   # signature.parameters is an OrderedDict => so arg_names order is deterministic
                   lowerCamelCase__						:							Optional[Any]       =    [*signature.parameters.keys()]

                   lowerCamelCase__						:							List[Any]       =    ["pixel_values"]
                   self.assertListEqual(arg_names[:1]   , __lowerCamelCase			)





       def 		lowerCAmelCase	(  self				:     Any			):


             '''simple docstring'''



             def check_hidden_states_output(__lowerCamelCase				:     List[str]   , __lowerCamelCase				:     Optional[int]   , __lowerCamelCase				:     str			):
                   lowerCamelCase__						:							Any       =    model_class(__lowerCamelCase			)

                   lowerCamelCase__						:							Dict       =    model(**self._prepare_for_class(__lowerCamelCase   , __lowerCamelCase			)			)

                   lowerCamelCase__						:							str       =    outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

                   lowerCamelCase__						:							Tuple       =    self.model_tester.num_stages
                   self.assertEqual(len(__lowerCamelCase			)   , expected_num_stages + 1			)





       @unittest.skip(reason="ResNet does not use feedforward chunking"			)
       def 		lowerCAmelCase	(  self				:     Any			):


             '''simple docstring'''



             pass





       def 		lowerCAmelCase	(  self				:     int			):


             '''simple docstring'''



             lowerCamelCase__       ,  lowerCamelCase__						:							Dict       =    self.model_tester.prepare_config_and_inputs_for_common()

             for model_class in self.all_model_classes:
                   with self.subTest(model_class.__name__			):
                         lowerCamelCase__						:							Any       =    self._prepare_for_class(__lowerCamelCase   , __lowerCamelCase			)
                         lowerCamelCase__						:							Optional[int]       =    model_class(__lowerCamelCase			)

                         @jax.jit
                         def model_jitted(__lowerCamelCase				:     Optional[Any]   , **__lowerCamelCase				:     List[Any]			):
                               return model(pixel_values=__lowerCamelCase   , **__lowerCamelCase			)

                         with self.subTest("JIT Enabled"			):
                               lowerCamelCase__						:							Any       =    model_jitted(**__lowerCamelCase			).to_tuple()

                         with self.subTest("JIT Disabled"			):
                               with jax.disable_jit():
                                     lowerCamelCase__						:							Union[str, Any]       =    model_jitted(**__lowerCamelCase			).to_tuple()

                         self.assertEqual(len(__lowerCamelCase			)   , len(__lowerCamelCase			)			)
                         for jitted_output, output in zip(__lowerCamelCase   , __lowerCamelCase			):
                               self.assertEqual(jitted_output.shape   , output.shape			)



def  lowercase_						(						):

      """simple docstring"""




      lowerCamelCase__						:							Optional[Any]       =    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"		)
      return image







@require_flax
class      _lowercase					(				unittest.TestCase):


       """simple docstring"""




       @cached_property
       def 		lowerCAmelCase	(  self				:     Tuple			):


             '''simple docstring'''



             return AutoImageProcessor.from_pretrained("microsoft/resnet-50"			) if is_vision_available() else None





       @slow
       def 		lowerCAmelCase	(  self				:     List[str]			):


             '''simple docstring'''



             lowerCamelCase__						:							List[Any]       =    FlaxResNetForImageClassification.from_pretrained("microsoft/resnet-50"			)

             lowerCamelCase__						:							List[str]       =    self.default_image_processor
             lowerCamelCase__						:							Optional[Any]       =    prepare_img()
             lowerCamelCase__						:							str       =    image_processor(images=__lowerCamelCase   , return_tensors="np"			)

             lowerCamelCase__						:							str       =    model(**__lowerCamelCase			)

             # verify the logits
             lowerCamelCase__						:							int       =    (1, 1000)
             self.assertEqual(outputs.logits.shape   , __lowerCamelCase			)

             lowerCamelCase__						:							Optional[int]       =    jnp.array([-1_1.1_0_6_9, -9.7_8_7_7, -8.3_7_7_7]			)

             self.assertTrue(jnp.allclose(outputs.logits[0, :3]   , __lowerCamelCase   , atol=1E-4			)			)


