




'''simple docstring'''




import tempfile
import unittest

from diffusers import (
    DDIMScheduler,
    DDPMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    PNDMScheduler,
    logging,
)
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.utils.testing_utils import CaptureLogger






class      UpperCamelCase_  (a__  ):







						"""simple docstring"""


						_lowerCAmelCase         =       'config.json'







						@register_to_config
						def __init__(  self   :							Optional[int]				,			_lowerCamelCase   :							Tuple=2				,			_lowerCamelCase   :							List[str]=5				,			_lowerCamelCase   :							Optional[int]=(2, 5)				,			_lowerCamelCase   :							Union[str, Any]="for diffusion"				,			_lowerCamelCase   :							Union[str, Any]=[1, 3]				,			):





											"""simple docstring"""







											pass






class      UpperCamelCase_  (a__  ):







						"""simple docstring"""


						_lowerCAmelCase         =       'config.json'







						@register_to_config
						def __init__(  self   :							int				,			_lowerCamelCase   :							List[str]=2				,			_lowerCamelCase   :							Any=5				,			_lowerCamelCase   :							Tuple=(2, 5)				,			_lowerCamelCase   :							Optional[int]="for diffusion"				,			_lowerCamelCase   :							Any=[1, 3]				,			):





											"""simple docstring"""







											pass






class      UpperCamelCase_  (a__  ):







						"""simple docstring"""


						_lowerCAmelCase         =       'config.json'







						@register_to_config
						def __init__(  self   :							Tuple				,			_lowerCamelCase   :							Dict=2				,			_lowerCamelCase   :							Optional[int]=5				,			_lowerCamelCase   :							Tuple=(2, 5)				,			_lowerCamelCase   :							Optional[Any]="for diffusion"				,			_lowerCamelCase   :							int=[1, 3]				,			_lowerCamelCase   :							int=[1, 3]				,			):





											"""simple docstring"""







											pass






class      UpperCamelCase_  (a__  ):







						"""simple docstring"""


						_lowerCAmelCase         =       'config.json'







						@register_to_config
						def __init__(  self   :							List[str]				,			_lowerCamelCase   :							List[str]=2				,			_lowerCamelCase   :							Tuple=5				,			_lowerCamelCase   :							int=(2, 5)				,			_lowerCamelCase   :							Union[str, Any]="for diffusion"				,			_lowerCamelCase   :							Optional[int]=[1, 5]				,			_lowerCamelCase   :							List[Any]=[5, 4]				,			):





											"""simple docstring"""







											pass






class      UpperCamelCase_  (unittest.TestCase  ):







						"""simple docstring"""


						def    _a     (  self   :							List[Any]     ):





											"""simple docstring"""







											with self.assertRaises(_lowerCamelCase     ):
																ConfigMixin.load_config('''dummy_path'''     )







						def    _a     (  self   :							Dict     ):





											"""simple docstring"""







											A_			:  Optional[int]			=				SampleObject()
											A_			:  List[Any]			=				obj.config
											assert config["a"] == 2
											assert config["b"] == 5
											assert config["c"] == (2, 5)
											assert config["d"] == "for diffusion"
											assert config["e"] == [1, 3]

											# init ignore private arguments
											A_			:  List[Any]			=				SampleObject(_name_or_path='''lalala'''     )
											A_			:  List[str]			=				obj.config
											assert config["a"] == 2
											assert config["b"] == 5
											assert config["c"] == (2, 5)
											assert config["d"] == "for diffusion"
											assert config["e"] == [1, 3]

											# can override default
											A_			:  Union[str, Any]			=				SampleObject(c=6     )
											A_			:  Any			=				obj.config
											assert config["a"] == 2
											assert config["b"] == 5
											assert config["c"] == 6
											assert config["d"] == "for diffusion"
											assert config["e"] == [1, 3]

											# can use positional arguments.
											A_			:  Optional[Any]			=				SampleObject(1				,			c=6     )
											A_			:  List[Any]			=				obj.config
											assert config["a"] == 1
											assert config["b"] == 5
											assert config["c"] == 6
											assert config["d"] == "for diffusion"
											assert config["e"] == [1, 3]







						def    _a     (  self   :							str     ):





											"""simple docstring"""







											A_			:  List[Any]			=				SampleObject()
											A_			:  List[str]			=				obj.config

											assert config["a"] == 2
											assert config["b"] == 5
											assert config["c"] == (2, 5)
											assert config["d"] == "for diffusion"
											assert config["e"] == [1, 3]

											with tempfile.TemporaryDirectory() as tmpdirname:
																obj.save_config(_lowerCamelCase     )
																A_			:  Union[str, Any]			=				SampleObject.from_config(SampleObject.load_config(_lowerCamelCase     )     )
																A_			:  Any			=				new_obj.config

											# unfreeze configs
											A_			:  Dict			=				dict(_lowerCamelCase     )
											A_			:  Optional[Any]			=				dict(_lowerCamelCase     )

											assert config.pop('''c'''     ) == (2, 5)  # instantiated as tuple
											assert new_config.pop('''c'''     ) == [2, 5]  # saved & loaded as list because of json
											config.pop('''_use_default_values'''     )
											assert config == new_config







						def    _a     (  self   :							List[str]     ):





											"""simple docstring"""







											A_			:  Any			=				logging.get_logger('''diffusers.configuration_utils'''     )
											# 30 for warning
											logger.setLevel(30     )

											with CaptureLogger(_lowerCamelCase     ) as cap_logger:
																A_			:  List[Any]			=				DDIMScheduler.from_pretrained(
																    '''hf-internal-testing/tiny-stable-diffusion-torch'''				,			subfolder='''scheduler'''     )

											assert ddim.__class__ == DDIMScheduler
											# no warning should be thrown
											assert cap_logger.out == ""







						def    _a     (  self   :							List[str]     ):





											"""simple docstring"""







											A_			:  Any			=				logging.get_logger('''diffusers.configuration_utils'''     )
											# 30 for warning
											logger.setLevel(30     )

											with CaptureLogger(_lowerCamelCase     ) as cap_logger:
																A_			:  Optional[Any]			=				EulerDiscreteScheduler.from_pretrained(
																    '''hf-internal-testing/tiny-stable-diffusion-torch'''				,			subfolder='''scheduler'''     )

											assert euler.__class__ == EulerDiscreteScheduler
											# no warning should be thrown
											assert cap_logger.out == ""







						def    _a     (  self   :							Any     ):





											"""simple docstring"""







											A_			:  Tuple			=				logging.get_logger('''diffusers.configuration_utils'''     )
											# 30 for warning
											logger.setLevel(30     )

											with CaptureLogger(_lowerCamelCase     ) as cap_logger:
																A_			:  Optional[int]			=				EulerAncestralDiscreteScheduler.from_pretrained(
																    '''hf-internal-testing/tiny-stable-diffusion-torch'''				,			subfolder='''scheduler'''     )

											assert euler.__class__ == EulerAncestralDiscreteScheduler
											# no warning should be thrown
											assert cap_logger.out == ""







						def    _a     (  self   :							str     ):





											"""simple docstring"""







											A_			:  Dict			=				logging.get_logger('''diffusers.configuration_utils'''     )
											# 30 for warning
											logger.setLevel(30     )

											with CaptureLogger(_lowerCamelCase     ) as cap_logger:
																A_			:  Dict			=				PNDMScheduler.from_pretrained(
																    '''hf-internal-testing/tiny-stable-diffusion-torch'''				,			subfolder='''scheduler'''     )

											assert pndm.__class__ == PNDMScheduler
											# no warning should be thrown
											assert cap_logger.out == ""







						def    _a     (  self   :							Dict     ):





											"""simple docstring"""







											A_			:  List[str]			=				logging.get_logger('''diffusers.configuration_utils'''     )
											# 30 for warning
											logger.setLevel(30     )

											with CaptureLogger(_lowerCamelCase     ) as cap_logger:
																A_			:  List[str]			=				DDPMScheduler.from_pretrained(
																    '''hf-internal-testing/tiny-stable-diffusion-torch'''				,			subfolder='''scheduler'''				,			prediction_type='''sample'''				,			beta_end=8				,			)

											with CaptureLogger(_lowerCamelCase     ) as cap_logger_a:
																A_			:  Any			=				DDPMScheduler.from_pretrained('''google/ddpm-celebahq-256'''				,			beta_start=88     )

											assert ddpm.__class__ == DDPMScheduler
											assert ddpm.config.prediction_type == "sample"
											assert ddpm.config.beta_end == 8
											assert ddpm_a.config.beta_start == 88

											# no warning should be thrown
											assert cap_logger.out == ""
											assert cap_logger_a.out == ""







						def    _a     (  self   :							List[str]     ):





											"""simple docstring"""







											A_			:  List[str]			=				logging.get_logger('''diffusers.configuration_utils'''     )
											# 30 for warning
											logger.setLevel(30     )

											with CaptureLogger(_lowerCamelCase     ) as cap_logger:
																A_			:  Optional[Any]			=				DPMSolverMultistepScheduler.from_pretrained(
																    '''hf-internal-testing/tiny-stable-diffusion-torch'''				,			subfolder='''scheduler'''     )

											assert dpm.__class__ == DPMSolverMultistepScheduler
											# no warning should be thrown
											assert cap_logger.out == ""







						def    _a     (  self   :							str     ):





											"""simple docstring"""







											A_			:  str			=				SampleObject()

											A_			:  str			=				{k: v for k, v in config.config.items() if not k.startswith('''_'''     )}

											# make sure that default config has all keys in `_use_default_values`
											assert set(config_dict.keys()     ) == set(config.config._use_default_values     )

											with tempfile.TemporaryDirectory() as tmpdirname:
																config.save_config(_lowerCamelCase     )

																# now loading it with SampleObject2 should put f into `_use_default_values`
																A_			:  int			=				SampleObjecta.from_config(_lowerCamelCase     )

																assert "f" in config._use_default_values
																assert config.f == [1, 3]

											# now loading the config, should **NOT** use [1, 3] for `f`, but the default [1, 4] value
											# **BECAUSE** it is part of `config._use_default_values`
											A_			:  str			=				SampleObjecta.from_config(config.config     )
											assert new_config.f == [5, 4]

											config.config._use_default_values.pop()
											A_			:  List[str]			=				SampleObjecta.from_config(config.config     )
											assert new_config_a.f == [1, 3]

											# Nevertheless "e" should still be correctly loaded to [1, 3] from SampleObject2 instead of defaulting to [1, 5]
											assert new_config_a.e == [1, 3]









'''simple docstring'''




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

from transformers import BigBirdTokenizerFast


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




def     snake_case__     (			lowerCamelCase__					: List[str]      ,		lowerCamelCase__					: Tuple=0      ,		lowerCamelCase__					: Tuple=None      ) -> Optional[Any]:
					return [" ".join(x.split(''' '''      )[begin:end]      ) for x in answers if len(x.split(''' '''      )      ) > 1]




def     snake_case__     (			lowerCamelCase__					: Any      ,		lowerCamelCase__					: Any=False      ) -> List[Any]:
					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"
										A_			:  Union[str, Any]			=				(
										    given_answers + get_sub_answers(lowerCamelCase__      ,		begin=1      ) + get_sub_answers(lowerCamelCase__      ,		end=-1      )
										)
					A_			:  List[Any]			=				[]
					for answer in given_answers:
										A_			:  Dict			=				answer.replace('''_'''      ,		''' '''      ).lower()
										A_			:  str			=				''''''.join(c if c not in PUNCTUATION_SET_TO_EXCLUDE else ''' ''' for c in alias      )
										answers.append(''' '''.join(alias.split()      ).strip()      )
					return set(lowerCamelCase__      )




def     snake_case__     (			lowerCamelCase__					: Optional[Any]      ,		lowerCamelCase__					: Dict      ,		lowerCamelCase__					: Union[str, Any]=1      ,		lowerCamelCase__					: Any=1_0_0      ) -> Any:
					A_			,A_			:  str			=				jax.lax.top_k(lowerCamelCase__      ,		lowerCamelCase__      )
					A_			,A_			:  str			=				jax.lax.top_k(lowerCamelCase__      ,		lowerCamelCase__      )

					A_			:  List[Any]			=				best_end_idx[:, None] - best_start_idx[None, :]
					A_			:  Optional[Any]			=				jnp.logical_or(widths < 0      ,		widths > max_size      )
					A_			:  Optional[int]			=				(best_end_scores[:, None] + best_start_scores[None, :]) - (1e8 * mask)
					A_			:  str			=				jnp.argmax(lowerCamelCase__      ).item()

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




def     snake_case__     (			lowerCamelCase__					: List[str]      ) -> Union[str, Any]:
					A_			:  str			=				sample['''question''']['''text''']
					A_			:  Tuple			=				sample['''document''']['''tokens''']['''token''']
					A_			:  List[Any]			=				sample['''document''']['''tokens''']['''is_html''']
					A_			:  Tuple			=				sample['''annotations''']['''long_answer''']
					A_			:  List[Any]			=				sample['''annotations''']['''short_answers''']

					A_			:  Optional[Any]			=				''' '''.join([context[i] for i in range(len(lowerCamelCase__      )      ) 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",
															}

					A_			:  List[Any]			=				[]
					for s in short_answers:
										short_targets.extend(s['''text''']      )
					A_			:  Tuple			=				list(set(lowerCamelCase__      )      )

					A_			:  Tuple			=				[]
					for s in long_answers:
										if s["start_token"] == -1:
															continue
										A_			:  Dict			=				context[s['''start_token'''] : s['''end_token''']]
										A_			:  Tuple			=				is_html[s['''start_token'''] : s['''end_token''']]
										A_			:  Optional[Any]			=				''' '''.join([answer[i] for i in range(len(lowerCamelCase__      )      ) if not html[i]]      )
										if new_answer not in long_targets:
															long_targets.append(lowerCamelCase__      )

					A_			:  str			=				'''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__     (			) -> List[Any]:
					A_			:  int			=				load_from_disk('''natural-questions-validation'''      )
					A_			:  Optional[int]			=				dataset.map(lowerCamelCase__      ).remove_columns(['''annotations''', '''document''', '''id''']      )
					print(lowerCamelCase__      )

					A_			:  Tuple			=				dataset.filter(lambda lowerCamelCase__      : (len(x['''question''']      ) + len(x['''context''']      )) < 4 * 4_0_9_6      )
					A_			:  Optional[Any]			=				short_validation_dataset.filter(lambda lowerCamelCase__      : x["category"] != "null"      )
					short_validation_dataset

					A_			:  Union[str, Any]			=				'''vasudevgupta/flax-bigbird-natural-questions'''
					A_			:  List[str]			=				FlaxBigBirdForNaturalQuestions.from_pretrained(lowerCamelCase__      )
					A_			:  int			=				BigBirdTokenizerFast.from_pretrained(lowerCamelCase__      )

					@jax.jit
					def forward(*lowerCamelCase__					: Any      ,		**lowerCamelCase__					: Optional[Any]      ):
										A_			,A_			,A_			:  Tuple			=				model(*lowerCamelCase__      ,		**lowerCamelCase__      )
										return start_logits, end_logits, jnp.argmax(lowerCamelCase__      ,		axis=-1      )

					def evaluate(lowerCamelCase__					: List[str]      ):
										# encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
										A_			:  List[Any]			=				tokenizer(
										    example['''question''']      ,		example['''context''']      ,		return_tensors='''np'''      ,		max_length=4_0_9_6      ,		padding='''max_length'''      ,		truncation=lowerCamelCase__      ,		)

										A_			,A_			,A_			:  Optional[Any]			=				forward(**lowerCamelCase__      )

										A_			:  Optional[Any]			=				CATEGORY_MAPPING[category.item()]

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

										if predicted_category in ["yes", "no", "null"]:
															A_			:  int			=				[predicted_category]
															A_			:  Optional[int]			=				example['''output'''] == example['''targets''']
															A_			:  Dict			=				predicted_category != '''null'''
															return example

										A_			:  List[Any]			=				3_8 if predicted_category == '''short''' else 1_0_2_4
										A_			,A_			:  str			=				get_best_valid_start_end_idx(
										    start_scores[0]      ,		end_scores[0]      ,		top_k=8      ,		max_size=lowerCamelCase__      )

										A_			:  List[Any]			=				inputs['''input_ids'''][0].tolist()
										A_			:  Dict			=				[tokenizer.decode(input_ids[start_score : end_score + 1]      )]

										A_			:  Optional[int]			=				expand_to_aliases(example['''targets''']      ,		make_sub_answers=lowerCamelCase__      )
										A_			:  str			=				expand_to_aliases(example['''output''']      )

										# some preprocessing to both prediction and answer
										A_			:  Tuple			=				{''''''.join(a.split()      ) for a in answers}
										A_			:  Optional[int]			=				{''''''.join(p.split()      ) for p in predictions}
										A_			:  List[str]			=				{s for s in predictions if s not in ['''``''', '''\'\'''', '''`''', '''\'''']}

										# if there is a common element, it's a exact match
										A_			:  List[Any]			=				len(list(answers & predictions      )      ) > 0
										A_			:  Dict			=				predicted_category != '''null''' and len(lowerCamelCase__      ) > 0

										return example

					A_			:  Dict			=				short_validation_dataset.map(lowerCamelCase__      )

					A_			:  List[str]			=				len(lowerCamelCase__      )
					A_			:  Any			=				len(short_validation_dataset.filter(lambda lowerCamelCase__      : x["match"] == 1      )      )
					print('''EM score:'''      ,		(matched / total) * 1_0_0      ,		'''%'''      )


if __name__ == "__main__":
							main()



