


'''simple docstring'''


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

from transformers import BigBirdTokenizerFast


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







def       snake_case_       (  SCREAMING_SNAKE_CASE__							,				SCREAMING_SNAKE_CASE__=0							,				SCREAMING_SNAKE_CASE__=None					):







				"""simple docstring"""





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







def       snake_case_       (  SCREAMING_SNAKE_CASE__							,				SCREAMING_SNAKE_CASE__=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			:	Any							=				(
								    given_answers + get_sub_answers(SCREAMING_SNAKE_CASE__							,				begin=1					) + get_sub_answers(SCREAMING_SNAKE_CASE__							,				end=-1					)
								)
				_SCREAMING_SNAKE_CASE			:	Optional[int]							=				[]
				for answer in given_answers:
								_SCREAMING_SNAKE_CASE			:	str							=				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(SCREAMING_SNAKE_CASE__					)







def       snake_case_       (  SCREAMING_SNAKE_CASE__							,				SCREAMING_SNAKE_CASE__							,				SCREAMING_SNAKE_CASE__=1							,				SCREAMING_SNAKE_CASE__=100					):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE			:	int							=				jax.lax.top_k(SCREAMING_SNAKE_CASE__							,				SCREAMING_SNAKE_CASE__					)
				_SCREAMING_SNAKE_CASE			:	List[Any]							=				jax.lax.top_k(SCREAMING_SNAKE_CASE__							,				SCREAMING_SNAKE_CASE__					)

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

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







def       snake_case_       (  SCREAMING_SNAKE_CASE__					):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE			:	int							=				sample["question"]["text"]
				_SCREAMING_SNAKE_CASE			:	List[Any]							=				sample["document"]["tokens"]["token"]
				_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				sample["document"]["tokens"]["is_html"]
				_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				sample["annotations"]["long_answer"]
				_SCREAMING_SNAKE_CASE			:	str							=				sample["annotations"]["short_answers"]

				_SCREAMING_SNAKE_CASE			:	Dict							=				" ".join([context[i] for i in range(len(SCREAMING_SNAKE_CASE__					)					) 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			:	Dict							=				[]
				for s in short_answers:
								short_targets.extend(s["""text"""]					)
				_SCREAMING_SNAKE_CASE			:	List[str]							=				list(set(SCREAMING_SNAKE_CASE__					)					)

				_SCREAMING_SNAKE_CASE			:	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			:	Tuple							=				is_html[s["start_token"] : s["end_token"]]
								_SCREAMING_SNAKE_CASE			:	Optional[int]							=				" ".join([answer[i] for i in range(len(SCREAMING_SNAKE_CASE__					)					) if not html[i]]					)
								if new_answer not in long_targets:
												long_targets.append(SCREAMING_SNAKE_CASE__					)

				_SCREAMING_SNAKE_CASE			:	Union[str, Any]							=				"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_       (  ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE			:	Dict							=				load_from_disk("""natural-questions-validation"""					)
				_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				dataset.map(SCREAMING_SNAKE_CASE__					).remove_columns(["""annotations""", """document""", """id"""]					)
				print(SCREAMING_SNAKE_CASE__					)

				_SCREAMING_SNAKE_CASE			:	Union[str, Any]							=				dataset.filter(lambda SCREAMING_SNAKE_CASE__					: (len(x["""question"""]					) + len(x["""context"""]					)) < 4 * 4096					)
				_SCREAMING_SNAKE_CASE			:	Any							=				short_validation_dataset.filter(lambda SCREAMING_SNAKE_CASE__					: x["category"] != "null"					)
				short_validation_dataset

				_SCREAMING_SNAKE_CASE			:	str							=				"vasudevgupta/flax-bigbird-natural-questions"
				_SCREAMING_SNAKE_CASE			:	List[str]							=				FlaxBigBirdForNaturalQuestions.from_pretrained(SCREAMING_SNAKE_CASE__					)
				_SCREAMING_SNAKE_CASE			:	Dict							=				BigBirdTokenizerFast.from_pretrained(SCREAMING_SNAKE_CASE__					)

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

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

								_SCREAMING_SNAKE_CASE			:	List[Any]							=				forward(**SCREAMING_SNAKE_CASE__					)

								_SCREAMING_SNAKE_CASE			:	str							=				CATEGORY_MAPPING[category.item()]

								_SCREAMING_SNAKE_CASE			:	Optional[int]							=				example["long"] + example["short"]
								if example["category"] in ["yes", "no", "null"]:
												_SCREAMING_SNAKE_CASE			:	Any							=				[example["category"]]
								_SCREAMING_SNAKE_CASE			:	Optional[int]							=				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			:	List[Any]							=				[predicted_category]
												_SCREAMING_SNAKE_CASE			:	Tuple							=				example["output"] == example["targets"]
												_SCREAMING_SNAKE_CASE			:	int							=				predicted_category != "null"
												return example

								_SCREAMING_SNAKE_CASE			:	List[Any]							=				38 if predicted_category == "short" else 1024
								_SCREAMING_SNAKE_CASE			:	Tuple							=				get_best_valid_start_end_idx(
								    start_scores[0]							,				end_scores[0]							,				top_k=8							,				max_size=SCREAMING_SNAKE_CASE__					)

								_SCREAMING_SNAKE_CASE			:	Dict							=				inputs["input_ids"][0].tolist()
								_SCREAMING_SNAKE_CASE			:	List[str]							=				[tokenizer.decode(input_ids[start_score : end_score + 1]					)]

								_SCREAMING_SNAKE_CASE			:	Optional[int]							=				expand_to_aliases(example["""targets"""]							,				make_sub_answers=SCREAMING_SNAKE_CASE__					)
								_SCREAMING_SNAKE_CASE			:	Union[str, Any]							=				expand_to_aliases(example["""output"""]					)

								# some preprocessing to both prediction and answer
								_SCREAMING_SNAKE_CASE			:	int							=				{"".join(a.split()					) for a in answers}
								_SCREAMING_SNAKE_CASE			:	List[str]							=				{"".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			:	Any							=				len(list(answers & predictions					)					) > 0
								_SCREAMING_SNAKE_CASE			:	str							=				predicted_category != "null" and len(SCREAMING_SNAKE_CASE__					) > 0

								return example

				_SCREAMING_SNAKE_CASE			:	List[str]							=				short_validation_dataset.map(SCREAMING_SNAKE_CASE__					)

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


if __name__ == "__main__":
				main()








'''simple docstring'''


from typing import Any, Dict, List, Union

import numpy as np

from ..utils import add_end_docstrings, is_torch_available, is_vision_available, logging, requires_backends
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
				from PIL import Image

				from ..image_utils import load_image

if is_torch_available():
				from ..models.auto.modeling_auto import (
				    MODEL_FOR_IMAGE_SEGMENTATION_MAPPING,
				    MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING,
				    MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING,
				    MODEL_FOR_UNIVERSAL_SEGMENTATION_MAPPING,
				)


UpperCAmelCase_		:       Optional[int]							       =		logging.get_logger(__name__)


UpperCAmelCase_		:       List[str]							       =		Dict[str, Any]
UpperCAmelCase_		:       List[str]							       =		List[Prediction]





@add_end_docstrings(_snake_case )
class    lowercase__  (						_snake_case ):





							'''simple docstring'''


							def __init__(   self   , *__snake_case   , **__snake_case  ):
											super().__init__(*__snake_case   , **__snake_case  )

											if self.framework == "tf":
															raise ValueError(f"""The {self.__class__} is only available in PyTorch."""  )

											requires_backends(self   , """vision"""  )
											self.check_model_type(
											    dict(
											        MODEL_FOR_IMAGE_SEGMENTATION_MAPPING.items()
											        + MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING.items()
											        + MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING.items()
											        + MODEL_FOR_UNIVERSAL_SEGMENTATION_MAPPING.items()  )  )






							def UpperCAmelCase_     (   self   , **__snake_case  ):
											_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				{}
											_SCREAMING_SNAKE_CASE			:	Any							=				{}
											if "subtask" in kwargs:
															_SCREAMING_SNAKE_CASE			:	Any							=				kwargs["""subtask"""]
															_SCREAMING_SNAKE_CASE			:	List[str]							=				kwargs["""subtask"""]
											if "threshold" in kwargs:
															_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				kwargs["""threshold"""]
											if "mask_threshold" in kwargs:
															_SCREAMING_SNAKE_CASE			:	Dict							=				kwargs["""mask_threshold"""]
											if "overlap_mask_area_threshold" in kwargs:
															_SCREAMING_SNAKE_CASE			:	Tuple							=				kwargs["""overlap_mask_area_threshold"""]

											return preprocess_kwargs, {}, postprocess_kwargs






							def __call__(   self   , __snake_case   , **__snake_case  ):
											return super().__call__(__snake_case   , **__snake_case  )






							def UpperCAmelCase_     (   self   , __snake_case   , __snake_case=None  ):
											_SCREAMING_SNAKE_CASE			:	List[Any]							=				load_image(__snake_case  )
											_SCREAMING_SNAKE_CASE			:	List[Any]							=				[(image.height, image.width)]
											if self.model.config.__class__.__name__ == "OneFormerConfig":
															if subtask is None:
																			_SCREAMING_SNAKE_CASE			:	int							=				{}
															else:
																			_SCREAMING_SNAKE_CASE			:	str							=				{"""task_inputs""": [subtask]}
															_SCREAMING_SNAKE_CASE			:	Any							=				self.image_processor(images=[image]   , return_tensors="""pt"""   , **__snake_case  )
															_SCREAMING_SNAKE_CASE			:	List[Any]							=				self.tokenizer(
															    inputs["""task_inputs"""]   , padding="""max_length"""   , max_length=self.model.config.task_seq_len   , return_tensors=self.framework   , )["""input_ids"""]
											else:
															_SCREAMING_SNAKE_CASE			:	int							=				self.image_processor(images=[image]   , return_tensors="""pt"""  )
											_SCREAMING_SNAKE_CASE			:	Dict							=				target_size
											return inputs






							def UpperCAmelCase_     (   self   , __snake_case  ):
											_SCREAMING_SNAKE_CASE			:	Dict							=				model_inputs.pop("""target_size"""  )
											_SCREAMING_SNAKE_CASE			:	int							=				self.model(**__snake_case  )
											_SCREAMING_SNAKE_CASE			:	Union[str, Any]							=				target_size
											return model_outputs






							def UpperCAmelCase_     (   self   , __snake_case   , __snake_case=None   , __snake_case=0.9   , __snake_case=0.5   , __snake_case=0.5  ):
											_SCREAMING_SNAKE_CASE			:	str							=				None
											if subtask in {"panoptic", None} and hasattr(self.image_processor   , """post_process_panoptic_segmentation"""  ):
															_SCREAMING_SNAKE_CASE			:	int							=				self.image_processor.post_process_panoptic_segmentation
											elif subtask in {"instance", None} and hasattr(self.image_processor   , """post_process_instance_segmentation"""  ):
															_SCREAMING_SNAKE_CASE			:	str							=				self.image_processor.post_process_instance_segmentation

											if fn is not None:
															_SCREAMING_SNAKE_CASE			:	Optional[Any]							=				fn(
															    __snake_case   , threshold=__snake_case   , mask_threshold=__snake_case   , overlap_mask_area_threshold=__snake_case   , target_sizes=model_outputs["""target_size"""]   , )[0]

															_SCREAMING_SNAKE_CASE			:	Dict							=				[]
															_SCREAMING_SNAKE_CASE			:	int							=				outputs["""segmentation"""]

															for segment in outputs["segments_info"]:
																			_SCREAMING_SNAKE_CASE			:	str							=				(segmentation == segment["""id"""]) * 255
																			_SCREAMING_SNAKE_CASE			:	int							=				Image.fromarray(mask.numpy().astype(np.uinta  )   , mode="""L"""  )
																			_SCREAMING_SNAKE_CASE			:	List[Any]							=				self.model.config.idalabel[segment["""label_id"""]]
																			_SCREAMING_SNAKE_CASE			:	List[Any]							=				segment["""score"""]
																			annotation.append({"""score""": score, """label""": label, """mask""": mask}  )

											elif subtask in {"semantic", None} and hasattr(self.image_processor   , """post_process_semantic_segmentation"""  ):
															_SCREAMING_SNAKE_CASE			:	Union[str, Any]							=				self.image_processor.post_process_semantic_segmentation(
															    __snake_case   , target_sizes=model_outputs["""target_size"""]  )[0]

															_SCREAMING_SNAKE_CASE			:	str							=				[]
															_SCREAMING_SNAKE_CASE			:	List[str]							=				outputs.numpy()
															_SCREAMING_SNAKE_CASE			:	Optional[int]							=				np.unique(__snake_case  )

															for label in labels:
																			_SCREAMING_SNAKE_CASE			:	Any							=				(segmentation == label) * 255
																			_SCREAMING_SNAKE_CASE			:	Tuple							=				Image.fromarray(mask.astype(np.uinta  )   , mode="""L"""  )
																			_SCREAMING_SNAKE_CASE			:	Tuple							=				self.model.config.idalabel[label]
																			annotation.append({"""score""": None, """label""": label, """mask""": mask}  )
											else:
															raise ValueError(f"""Subtask {subtask} is not supported for model {type(self.model  )}"""  )
											return annotation




