




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

from transformers import BigBirdTokenizerFast


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


def 		_lowerCamelCase					(       UpperCamelCase       ,  UpperCamelCase=0       ,  UpperCamelCase=None				)      ->		Any:





  """simple docstring"""





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


def 		_lowerCamelCase					(       UpperCamelCase       ,  UpperCamelCase=False				)      ->		List[Any]:





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


def 		_lowerCamelCase					(       UpperCamelCase       ,  UpperCamelCase       ,  UpperCamelCase=1       ,  UpperCamelCase=100				)      ->		str:





  """simple docstring"""





  _lowercase  ,       _lowercase:			Dict      =     jax.lax.top_k(UpperCamelCase       ,  UpperCamelCase				)
  _lowercase  ,       _lowercase:			List[Any]      =     jax.lax.top_k(UpperCamelCase       ,  UpperCamelCase				)

  _lowercase:			int      =     best_end_idx[:, None] - best_start_idx[None, :]
  _lowercase:			Tuple      =     jnp.logical_or(widths < 0       ,  widths > max_size				)
  _lowercase:			Union[str, Any]      =     (best_end_scores[:, None] + best_start_scores[None, :]) - (1E8 * mask)
  _lowercase:			Tuple      =     jnp.argmax(UpperCamelCase				).item()

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


def 		_lowerCamelCase					(       UpperCamelCase				)      ->		List[Any]:





  """simple docstring"""





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

  _lowercase:			Tuple      =     ' '.join([context[i] for i in range(len(UpperCamelCase				)				) 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",
      }

  _lowercase:			Any      =     []
  for s in short_answers:
    short_targets.extend(s['text']				)
  _lowercase:			str      =     list(set(UpperCamelCase				)				)

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

  _lowercase:			Tuple      =     '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					(       )      ->		int:





  """simple docstring"""





  _lowercase:			Optional[int]      =     load_from_disk('natural-questions-validation'				)
  _lowercase:			Optional[Any]      =     dataset.map(UpperCamelCase				).remove_columns(['annotations', 'document', 'id']				)
  print(UpperCamelCase				)

  _lowercase:			Optional[Any]      =     dataset.filter(lambda UpperCamelCase				: (len(x['question']				) + len(x['context']				)) < 4 * 4_096				)
  _lowercase:			Union[str, Any]      =     short_validation_dataset.filter(lambda UpperCamelCase				: x["category"] != "null"				)
  short_validation_dataset

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

  @jax.jit
  def forward(*UpperCamelCase       ,  **UpperCamelCase				):
    _lowercase  ,       _lowercase  ,       _lowercase:			Optional[int]      =     model(*UpperCamelCase       ,  **UpperCamelCase				)
    return start_logits, end_logits, jnp.argmax(UpperCamelCase       ,  axis=-1				)

  def evaluate(UpperCamelCase				):
    # encode question and context so that they are separated by a tokenizer.sep_token and cut at max_length
    _lowercase:			Tuple      =     tokenizer(
        example['question']       ,  example['context']       ,  return_tensors='np'       ,  max_length=4_096       ,  padding='max_length'       ,  truncation=UpperCamelCase       ,  )

    _lowercase  ,       _lowercase  ,       _lowercase:			Union[str, Any]      =     forward(**UpperCamelCase				)

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

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

    if predicted_category in ["yes", "no", "null"]:
      _lowercase:			Union[str, Any]      =     [predicted_category]
      _lowercase:			List[str]      =     example['output'] == example['targets']
      _lowercase:			Union[str, Any]      =     predicted_category != 'null'
      return example

    _lowercase:			Tuple      =     38 if predicted_category == 'short' else 1_024
    _lowercase  ,       _lowercase:			List[Any]      =     get_best_valid_start_end_idx(
        start_scores[0]       ,  end_scores[0]       ,  top_k=8       ,  max_size=UpperCamelCase				)

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

    _lowercase:			int      =     expand_to_aliases(example['targets']       ,  make_sub_answers=UpperCamelCase				)
    _lowercase:			List[str]      =     expand_to_aliases(example['output']				)

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

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

    return example

  _lowercase:			Tuple      =     short_validation_dataset.map(UpperCamelCase				)

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


if __name__ == "__main__":
   main()






import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()

class   UpperCAmelCase    (      A ,       unittest.TestCase ):
   __SCREAMING_SNAKE_CASE		:				int      					=       KandinskyPipeline
   __SCREAMING_SNAKE_CASE		:				Union[str, Any]      					=       [
       "prompt",
       "image_embeds",
       "negative_image_embeds",
   ]
   __SCREAMING_SNAKE_CASE		:				Tuple      					=       ["prompt", "negative_prompt", "image_embeds", "negative_image_embeds"]
   __SCREAMING_SNAKE_CASE		:				Tuple      					=       [
       "generator",
       "height",
       "width",
       "latents",
       "guidance_scale",
       "negative_prompt",
       "num_inference_steps",
       "return_dict",
       "guidance_scale",
       "num_images_per_prompt",
       "output_type",
       "return_dict",
   ]
   __SCREAMING_SNAKE_CASE		:				Dict      					=       False






   @property
   def 			_a    (						self :							str  )->      Dict:
     return 32






   @property
   def 			_a    (						self :							Tuple  )->      Any:
     return 32






   @property
   def 			_a    (						self :							Union[str, Any]  )->      List[str]:
     return self.time_input_dim






   @property
   def 			_a    (						self :							Optional[Any]  )->      Tuple:
     return self.time_input_dim * 4






   @property
   def 			_a    (						self :							int  )->      Dict:
     return 1_00






   @property
   def 			_a    (						self :							Tuple  )->      Tuple:
     _lowercase:			List[Any]      =     XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base'  )
     return tokenizer






   @property
   def 			_a    (						self :							List[str]  )->      Optional[int]:
     torch.manual_seed(0  )
     _lowercase:			int      =     MCLIPConfig(
         numDims=self.cross_attention_dim    ,       transformerDimensions=self.text_embedder_hidden_size    ,       hidden_size=self.text_embedder_hidden_size    ,       intermediate_size=37    ,       num_attention_heads=4    ,       num_hidden_layers=5    ,       vocab_size=10_05    ,       )

     _lowercase:			Optional[Any]      =     MultilingualCLIP(UpperCamelCase_  )
     _lowercase:			int      =     text_encoder.eval()

     return text_encoder






   @property
   def 			_a    (						self :							Optional[int]  )->      Optional[Any]:
     torch.manual_seed(0  )

     _lowercase:			Optional[Any]      =     {
         'in_channels': 4,
         # Out channels is double in channels because predicts mean and variance
         'out_channels': 8,
         'addition_embed_type': 'text_image',
         'down_block_types': ('ResnetDownsampleBlock2D', 'SimpleCrossAttnDownBlock2D'),
         'up_block_types': ('SimpleCrossAttnUpBlock2D', 'ResnetUpsampleBlock2D'),
         'mid_block_type': 'UNetMidBlock2DSimpleCrossAttn',
         'block_out_channels': (self.block_out_channels_a, self.block_out_channels_a * 2),
         'layers_per_block': 1,
         'encoder_hid_dim': self.text_embedder_hidden_size,
         'encoder_hid_dim_type': 'text_image_proj',
         'cross_attention_dim': self.cross_attention_dim,
         'attention_head_dim': 4,
         'resnet_time_scale_shift': 'scale_shift',
         'class_embed_type': None,
     }

     _lowercase:			str      =     UNetaDConditionModel(**UpperCamelCase_  )
     return model






   @property
   def 			_a    (						self :							Dict  )->      Optional[Any]:
     return {
         "block_out_channels": [32, 64],
         "down_block_types": ["DownEncoderBlock2D", "AttnDownEncoderBlock2D"],
         "in_channels": 3,
         "latent_channels": 4,
         "layers_per_block": 1,
         "norm_num_groups": 8,
         "norm_type": "spatial",
         "num_vq_embeddings": 12,
         "out_channels": 3,
         "up_block_types": [
             "AttnUpDecoderBlock2D",
             "UpDecoderBlock2D",
         ],
         "vq_embed_dim": 4,
     }






   @property
   def 			_a    (						self :							List[str]  )->      int:
     torch.manual_seed(0  )
     _lowercase:			List[str]      =     VQModel(**self.dummy_movq_kwargs  )
     return model






   def 			_a    (						self :							int  )->      List[Any]:
     _lowercase:			List[str]      =     self.dummy_text_encoder
     _lowercase:			List[Any]      =     self.dummy_tokenizer
     _lowercase:			Optional[Any]      =     self.dummy_unet
     _lowercase:			List[str]      =     self.dummy_movq

     _lowercase:			Optional[Any]      =     DDIMScheduler(
         num_train_timesteps=10_00    ,       beta_schedule='linear'    ,       beta_start=0.00_085    ,       beta_end=0.012    ,       clip_sample=UpperCamelCase_    ,       set_alpha_to_one=UpperCamelCase_    ,       steps_offset=1    ,       prediction_type='epsilon'    ,       thresholding=UpperCamelCase_    ,       )

     _lowercase:			List[Any]      =     {
         'text_encoder': text_encoder,
         'tokenizer': tokenizer,
         'unet': unet,
         'scheduler': scheduler,
         'movq': movq,
     }
     return components






   def 			_a    (						self :							str    ,       UpperCamelCase_ :							int    ,       UpperCamelCase_ :							Optional[int]=0  )->      Union[str, Any]:
     _lowercase:			str      =     floats_tensor((1, self.cross_attention_dim)    ,       rng=random.Random(UpperCamelCase_  )  ).to(UpperCamelCase_  )
     _lowercase:			Optional[int]      =     floats_tensor((1, self.cross_attention_dim)    ,       rng=random.Random(seed + 1  )  ).to(UpperCamelCase_  )
     if str(UpperCamelCase_  ).startswith('mps'  ):
       _lowercase:			Any      =     torch.manual_seed(UpperCamelCase_  )
     else:
       _lowercase:			Tuple      =     torch.Generator(device=UpperCamelCase_  ).manual_seed(UpperCamelCase_  )
     _lowercase:			List[str]      =     {
         'prompt': 'horse',
         'image_embeds': image_embeds,
         'negative_image_embeds': negative_image_embeds,
         'generator': generator,
         'height': 64,
         'width': 64,
         'guidance_scale': 4.0,
         'num_inference_steps': 2,
         'output_type': 'np',
     }
     return inputs






   def 			_a    (						self :							List[str]  )->      List[str]:
     _lowercase:			str      =     'cpu'

     _lowercase:			Optional[Any]      =     self.get_dummy_components()

     _lowercase:			int      =     self.pipeline_class(**UpperCamelCase_  )
     _lowercase:			List[str]      =     pipe.to(UpperCamelCase_  )

     pipe.set_progress_bar_config(disable=UpperCamelCase_  )

     _lowercase:			Tuple      =     pipe(**self.get_dummy_inputs(UpperCamelCase_  )  )
     _lowercase:			Union[str, Any]      =     output.images

     _lowercase:			Union[str, Any]      =     pipe(
         **self.get_dummy_inputs(UpperCamelCase_  )    ,       return_dict=UpperCamelCase_    ,       )[0]

     _lowercase:			int      =     image[0, -3:, -3:, -1]
     _lowercase:			List[str]      =     image_from_tuple[0, -3:, -3:, -1]

     assert image.shape == (1, 64, 64, 3)

     _lowercase:			Union[str, Any]      =     np.array(
         [0.328_663, 1.0, 0.23_216_873, 1.0, 0.92_717_564, 0.4_639_046, 0.96_894_777, 0.31_713_378, 0.6_293_953]  )

     assert (
         np.abs(image_slice.flatten() - expected_slice  ).max() < 1E-2
     ), f' expected_slice {expected_slice}, but got {image_slice.flatten()}'
     assert (
         np.abs(image_from_tuple_slice.flatten() - expected_slice  ).max() < 1E-2
     ), f' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'

@slow
@require_torch_gpu
class   UpperCAmelCase    (      unittest.TestCase ):






   def 			_a    (						self :							Optional[int]  )->      Any:
     # clean up the VRAM after each test
     super().tearDown()
     gc.collect()
     torch.cuda.empty_cache()






   def 			_a    (						self :							List[Any]  )->      Optional[int]:
     _lowercase:			Optional[int]      =     load_numpy(
         'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
         '/kandinsky/kandinsky_text2img_cat_fp16.npy'  )

     _lowercase:			List[Any]      =     KandinskyPriorPipeline.from_pretrained(
         'kandinsky-community/kandinsky-2-1-prior'    ,       torch_dtype=torch.floataa  )
     pipe_prior.to(UpperCamelCase_  )

     _lowercase:			str      =     KandinskyPipeline.from_pretrained('kandinsky-community/kandinsky-2-1'    ,       torch_dtype=torch.floataa  )
     _lowercase:			List[str]      =     pipeline.to(UpperCamelCase_  )
     pipeline.set_progress_bar_config(disable=UpperCamelCase_  )

     _lowercase:			Dict      =     'red cat, 4k photo'

     _lowercase:			List[Any]      =     torch.Generator(device='cuda'  ).manual_seed(0  )
     _lowercase  ,       _lowercase:			Dict      =     pipe_prior(
         UpperCamelCase_    ,       generator=UpperCamelCase_    ,       num_inference_steps=5    ,       negative_prompt=''    ,       ).to_tuple()

     _lowercase:			str      =     torch.Generator(device='cuda'  ).manual_seed(0  )
     _lowercase:			Optional[int]      =     pipeline(
         UpperCamelCase_    ,       image_embeds=UpperCamelCase_    ,       negative_image_embeds=UpperCamelCase_    ,       generator=UpperCamelCase_    ,       num_inference_steps=1_00    ,       output_type='np'    ,       )

     _lowercase:			Optional[int]      =     output.images[0]

     assert image.shape == (5_12, 5_12, 3)

     assert_mean_pixel_difference(UpperCamelCase_    ,       UpperCamelCase_  )
