






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        _lowerCamelCase(      lowerCAmelCase__					:		List[str]     ,     lowerCAmelCase__					:		List[str]=0     ,     lowerCAmelCase__					:		Any=None   ):





				'''simple docstring'''







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



def        _lowerCamelCase(      lowerCAmelCase__					:		str     ,     lowerCAmelCase__					:		Optional[int]=False   ):





				'''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"
								SCREAMING_SNAKE_CASE_ :						Optional[int]    				= (
								    given_answers + get_sub_answers(lowerCAmelCase__     ,     begin=1   ) + get_sub_answers(lowerCAmelCase__     ,     end=-1   )
								)
				SCREAMING_SNAKE_CASE_ :						List[Any]    				= []
				for answer in given_answers:
								SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= answer.replace('_'     ,     ' '   ).lower()
								SCREAMING_SNAKE_CASE_ :						Optional[int]    				= ''.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        _lowerCamelCase(      lowerCAmelCase__					:		Tuple     ,     lowerCAmelCase__					:		Union[str, Any]     ,     lowerCAmelCase__					:		List[str]=1     ,     lowerCAmelCase__					:		List[str]=100   ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						List[str]    				= jax.lax.top_k(lowerCAmelCase__     ,     lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= jax.lax.top_k(lowerCAmelCase__     ,     lowerCAmelCase__   )

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

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



def        _lowerCamelCase(      lowerCAmelCase__					:		List[str]   ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= sample['question']['text']
				SCREAMING_SNAKE_CASE_ :						List[Any]    				= sample['document']['tokens']['token']
				SCREAMING_SNAKE_CASE_ :						int    				= sample['document']['tokens']['is_html']
				SCREAMING_SNAKE_CASE_ :						str    				= sample['annotations']['long_answer']
				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= sample['annotations']['short_answers']

				SCREAMING_SNAKE_CASE_ :						List[str]    				= ' '.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",
												}

				SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= []
				for s in short_answers:
								short_targets.extend(s['text']   )
				SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= list(set(lowerCAmelCase__   )   )

				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= []
				for s in long_answers:
								if s["start_token"] == -1:
												continue
								SCREAMING_SNAKE_CASE_ :						Tuple    				= context[s['start_token'] : s['end_token']]
								SCREAMING_SNAKE_CASE_ :						Any    				= is_html[s['start_token'] : s['end_token']]
								SCREAMING_SNAKE_CASE_ :						List[str]    				= ' '.join([answer[i] for i in range(len(lowerCAmelCase__   )   ) if not html[i]]   )
								if new_answer not in long_targets:
												long_targets.append(lowerCAmelCase__   )

				SCREAMING_SNAKE_CASE_ :						Optional[int]    				= '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        _lowerCamelCase(      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= load_from_disk('natural-questions-validation'   )
				SCREAMING_SNAKE_CASE_ :						int    				= dataset.map(lowerCAmelCase__   ).remove_columns(['annotations', 'document', 'id']   )
				print(lowerCAmelCase__   )

				SCREAMING_SNAKE_CASE_ :						str    				= dataset.filter(lambda lowerCAmelCase__   : (len(x['question']   ) + len(x['context']   )) < 4 * 4096   )
				SCREAMING_SNAKE_CASE_ :						Any    				= short_validation_dataset.filter(lambda lowerCAmelCase__   : x["category"] != "null"   )
				short_validation_dataset

				SCREAMING_SNAKE_CASE_ :						Optional[int]    				= 'vasudevgupta/flax-bigbird-natural-questions'
				SCREAMING_SNAKE_CASE_ :						List[Any]    				= FlaxBigBirdForNaturalQuestions.from_pretrained(lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_ :						List[Any]    				= BigBirdTokenizerFast.from_pretrained(lowerCAmelCase__   )

				@jax.jit
				def forward(*lowerCAmelCase__					:		int     ,     **lowerCAmelCase__					:		Union[str, Any]   ):
								SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						Any    				= model(*lowerCAmelCase__     ,     **lowerCAmelCase__   )
								return start_logits, end_logits, jnp.argmax(lowerCAmelCase__     ,     axis=-1   )

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

								SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						Dict    				= forward(**lowerCAmelCase__   )

								SCREAMING_SNAKE_CASE_ :						Any    				= CATEGORY_MAPPING[category.item()]

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

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

								SCREAMING_SNAKE_CASE_ :						List[str]    				= 38 if predicted_category == 'short' else 1024
								SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						str    				= get_best_valid_start_end_idx(
								    start_scores[0]     ,     end_scores[0]     ,     top_k=8     ,     max_size=lowerCAmelCase__   )

								SCREAMING_SNAKE_CASE_ :						Dict    				= inputs['input_ids'][0].tolist()
								SCREAMING_SNAKE_CASE_ :						Tuple    				= [tokenizer.decode(input_ids[start_score : end_score + 1]   )]

								SCREAMING_SNAKE_CASE_ :						Optional[int]    				= expand_to_aliases(example['targets']     ,     make_sub_answers=lowerCAmelCase__   )
								SCREAMING_SNAKE_CASE_ :						Dict    				= expand_to_aliases(example['output']   )

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

								# if there is a common element, it's a exact match
								SCREAMING_SNAKE_CASE_ :						str    				= len(list(answers & predictions   )   ) > 0
								SCREAMING_SNAKE_CASE_ :						Any    				= predicted_category != 'null' and len(lowerCAmelCase__   ) > 0

								return example

				SCREAMING_SNAKE_CASE_ :						Dict    				= short_validation_dataset.map(lowerCAmelCase__   )

				SCREAMING_SNAKE_CASE_ :						Any    				= len(lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_ :						Dict    				= len(short_validation_dataset.filter(lambda lowerCAmelCase__   : x["match"] == 1   )   )
				print('EM score:'     ,     (matched / total) * 100     ,     '%'   )


if __name__ == "__main__":
				main()







import argparse
import collections
import importlib.util
import os
import re


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_table.py
A										=      'src/diffusers'
A										=      'docs/source/en'
A										=      '.'



def        _lowerCamelCase(      lowerCAmelCase__					:		List[Any]     ,     lowerCAmelCase__					:		int     ,     lowerCAmelCase__					:		Tuple   ):





				'''simple docstring'''







				with open(lowerCAmelCase__     ,     'r'     ,     encoding='utf-8'     ,     newline='\n'   ) as f:
								SCREAMING_SNAKE_CASE_ :						Dict    				= f.readlines()
				# Find the start prompt.
				SCREAMING_SNAKE_CASE_ :						List[Any]    				= 0
				while not lines[start_index].startswith(lowerCAmelCase__   ):
								start_index += 1
				start_index += 1

				SCREAMING_SNAKE_CASE_ :						List[str]    				= start_index
				while not lines[end_index].startswith(lowerCAmelCase__   ):
								end_index += 1
				end_index -= 1

				while len(lines[start_index]   ) <= 1:
								start_index += 1
				while len(lines[end_index]   ) <= 1:
								end_index -= 1
				end_index += 1
				return "".join(lines[start_index:end_index]   ), start_index, end_index, lines


# Add here suffixes that are used to identify models, separated by |
A										=      'Model|Encoder|Decoder|ForConditionalGeneration'
# Regexes that match TF/Flax/PT model names.
A										=      re.compile(R'TF(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')
A										=      re.compile(R'Flax(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')
# Will match any TF or Flax model too so need to be in an else branch afterthe two previous regexes.
A										=      re.compile(R'(.*)(?:Model|Encoder|Decoder|ForConditionalGeneration)')


# This is to make sure the diffusers module imported is the one in the repo.
A										=      importlib.util.spec_from_file_location(
    'diffusers',
    os.path.join(TRANSFORMERS_PATH, '__init__.py'),
    submodule_search_locations=[TRANSFORMERS_PATH],
)
A										=      spec.loader.load_module()



def        _lowerCamelCase(      lowerCAmelCase__					:		Any   ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= re.finditer('.+?(?:(?<=[a-z])(?=[A-Z])|(?<=[A-Z])(?=[A-Z][a-z])|$)'     ,     lowerCAmelCase__   )
				return [m.group(0   ) for m in matches]



def        _lowerCamelCase(      lowerCAmelCase__					:		Optional[Any]     ,     lowerCAmelCase__					:		Tuple   ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_ :						Any    				= 2 if text == '✅' or text == '❌' else len(lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_ :						int    				= (width - text_length) // 2
				SCREAMING_SNAKE_CASE_ :						str    				= width - text_length - left_indent
				return " " * left_indent + text + " " * right_indent



def        _lowerCamelCase(      ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_ :						Optional[Any]    				= diffusers_module.models.auto.configuration_auto.CONFIG_MAPPING_NAMES
				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= {
				    name: config_mapping_names[code]
				    for code, name in diffusers_module.MODEL_NAMES_MAPPING.items()
				    if code in config_mapping_names
				}
				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= {name: config.replace('ConfigMixin'     ,     ''   ) for name, config in model_name_to_config.items()}

				# Dictionaries flagging if each model prefix has a slow/fast tokenizer, backend in PT/TF/Flax.
				SCREAMING_SNAKE_CASE_ :						int    				= collections.defaultdict(lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_ :						List[Any]    				= collections.defaultdict(lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= collections.defaultdict(lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_ :						str    				= collections.defaultdict(lowerCAmelCase__   )
				SCREAMING_SNAKE_CASE_ :						Dict    				= collections.defaultdict(lowerCAmelCase__   )

				# Let's lookup through all diffusers object (once).
				for attr_name in dir(lowerCAmelCase__   ):
								SCREAMING_SNAKE_CASE_ :						List[str]    				= None
								if attr_name.endswith('Tokenizer'   ):
												SCREAMING_SNAKE_CASE_ :						List[Any]    				= slow_tokenizers
												SCREAMING_SNAKE_CASE_ :						Dict    				= attr_name[:-9]
								elif attr_name.endswith('TokenizerFast'   ):
												SCREAMING_SNAKE_CASE_ :						Any    				= fast_tokenizers
												SCREAMING_SNAKE_CASE_ :						int    				= attr_name[:-13]
								elif _re_tf_models.match(lowerCAmelCase__   ) is not None:
												SCREAMING_SNAKE_CASE_ :						List[Any]    				= tf_models
												SCREAMING_SNAKE_CASE_ :						Optional[int]    				= _re_tf_models.match(lowerCAmelCase__   ).groups()[0]
								elif _re_flax_models.match(lowerCAmelCase__   ) is not None:
												SCREAMING_SNAKE_CASE_ :						List[str]    				= flax_models
												SCREAMING_SNAKE_CASE_ :						int    				= _re_flax_models.match(lowerCAmelCase__   ).groups()[0]
								elif _re_pt_models.match(lowerCAmelCase__   ) is not None:
												SCREAMING_SNAKE_CASE_ :						Optional[int]    				= pt_models
												SCREAMING_SNAKE_CASE_ :						Any    				= _re_pt_models.match(lowerCAmelCase__   ).groups()[0]

								if lookup_dict is not None:
												while len(lowerCAmelCase__   ) > 0:
																if attr_name in model_name_to_prefix.values():
																				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= True
																				break
																# Try again after removing the last word in the name
																SCREAMING_SNAKE_CASE_ :						str    				= ''.join(camel_case_split(lowerCAmelCase__   )[:-1]   )

    # Let's build that table!
				SCREAMING_SNAKE_CASE_ :						Union[str, Any]    				= list(model_name_to_config.keys()   )
				model_names.sort(key=str.lower   )
				SCREAMING_SNAKE_CASE_ :						List[Any]    				= ['Model', 'Tokenizer slow', 'Tokenizer fast', 'PyTorch support', 'TensorFlow support', 'Flax Support']
				# We'll need widths to properly display everything in the center (+2 is to leave one extra space on each side).
				SCREAMING_SNAKE_CASE_ :						str    				= [len(lowerCAmelCase__   ) + 2 for c in columns]
				SCREAMING_SNAKE_CASE_ :						List[Any]    				= max([len(lowerCAmelCase__   ) for name in model_names]   ) + 2

				# Build the table per se
				SCREAMING_SNAKE_CASE_ :						Dict    				= '|' + '|'.join([_center_text(lowerCAmelCase__     ,     lowerCAmelCase__   ) for c, w in zip(lowerCAmelCase__     ,     lowerCAmelCase__   )]   ) + '|\n'
				# Use ":-----:" format to center-aligned table cell texts
				table += "|" + "|".join([':' + '-' * (w - 2) + ':' for w in widths]   ) + "|\n"

				SCREAMING_SNAKE_CASE_ :						int    				= {True: '✅', False: '❌'}
				for name in model_names:
								SCREAMING_SNAKE_CASE_ :						int    				= model_name_to_prefix[name]
								SCREAMING_SNAKE_CASE_ :						int    				= [
								    name,
								    check[slow_tokenizers[prefix]],
								    check[fast_tokenizers[prefix]],
								    check[pt_models[prefix]],
								    check[tf_models[prefix]],
								    check[flax_models[prefix]],
								]
								table += "|" + "|".join([_center_text(lowerCAmelCase__     ,     lowerCAmelCase__   ) for l, w in zip(lowerCAmelCase__     ,     lowerCAmelCase__   )]   ) + "|\n"
				return table



def        _lowerCamelCase(      lowerCAmelCase__					:		Optional[int]=False   ):





				'''simple docstring'''







				SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_				,							SCREAMING_SNAKE_CASE_ :						str    				= _find_text_in_file(
				    filename=os.path.join(lowerCAmelCase__     ,     'index.mdx'   )     ,     start_prompt='<!--This table is updated automatically from the auto modules'     ,     end_prompt='<!-- End table-->'     ,     )
				SCREAMING_SNAKE_CASE_ :						Tuple    				= get_model_table_from_auto_modules()

				if current_table != new_table:
								if overwrite:
												with open(os.path.join(lowerCAmelCase__     ,     'index.mdx'   )     ,     'w'     ,     encoding='utf-8'     ,     newline='\n'   ) as f:
																f.writelines(lines[:start_index] + [new_table] + lines[end_index:]   )
								else:
												raise ValueError(
												    'The model table in the `index.mdx` has not been updated. Run `make fix-copies` to fix this.'   )


if __name__ == "__main__":
				A										=      argparse.ArgumentParser()
				parser.add_argument('--fix_and_overwrite', action='store_true', help='Whether to fix inconsistencies.')
				A										=      parser.parse_args()

				check_model_table(args.fix_and_overwrite)