




import datasets

from .evaluate import (
    apply_no_ans_threshold,
    find_all_best_thresh,
    get_raw_scores,
    make_eval_dict,
    make_qid_to_has_ans,
    merge_eval,
)


lowercase__   :    str 					=      "\\n@inproceedings{Rajpurkar2016SQuAD10,\n  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},\n  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},\n  booktitle={EMNLP},\n  year={2016}\n}\n"

lowercase__   :    Optional[Any] 					=      "\nThis metric wrap the official scoring script for version 2 of the Stanford Question\nAnswering Dataset (SQuAD).\n\nStanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by\ncrowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,\nfrom the corresponding reading passage, or the question might be unanswerable.\n\nSQuAD2.0 combines the 100,000 questions in SQuAD1.1 with over 50,000 unanswerable questions\nwritten adversarially by crowdworkers to look similar to answerable ones.\nTo do well on SQuAD2.0, systems must not only answer questions when possible, but also\ndetermine when no answer is supported by the paragraph and abstain from answering.\n"

lowercase__   :    Union[str, Any] 					=      "\nComputes SQuAD v2 scores (F1 and EM).\nArgs:\n    predictions: List of triple for question-answers to score with the following elements:\n        - the question-answer 'id' field as given in the references (see below)\n        - the text of the answer\n        - the probability that the question has no answer\n    references: List of question-answers dictionaries with the following key-values:\n            - 'id': id of the question-answer pair (see above),\n            - 'answers': a list of Dict {'text': text of the answer as a string}\n    no_answer_threshold: float\n        Probability threshold to decide that a question has no answer.\nReturns:\n    'exact': Exact match (the normalized answer exactly match the gold answer)\n    'f1': The F-score of predicted tokens versus the gold answer\n    'total': Number of score considered\n    'HasAns_exact': Exact match (the normalized answer exactly match the gold answer)\n    'HasAns_f1': The F-score of predicted tokens versus the gold answer\n    'HasAns_total': Number of score considered\n    'NoAns_exact': Exact match (the normalized answer exactly match the gold answer)\n    'NoAns_f1': The F-score of predicted tokens versus the gold answer\n    'NoAns_total': Number of score considered\n    'best_exact': Best exact match (with varying threshold)\n    'best_exact_thresh': No-answer probability threshold associated to the best exact match\n    'best_f1': Best F1 (with varying threshold)\n    'best_f1_thresh': No-answer probability threshold associated to the best F1\nExamples:\n\n    >>> predictions = [{'prediction_text': '1976', 'id': '56e10a3be3433e1400422b22', 'no_answer_probability': 0.}]\n    >>> references = [{'answers': {'answer_start': [97], 'text': ['1976']}, 'id': '56e10a3be3433e1400422b22'}]\n    >>> squad_v2_metric = datasets.load_metric(\"squad_v2\")\n    >>> results = squad_v2_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {'exact': 100.0, 'f1': 100.0, 'total': 1, 'HasAns_exact': 100.0, 'HasAns_f1': 100.0, 'HasAns_total': 1, 'best_exact': 100.0, 'best_exact_thresh': 0.0, 'best_f1': 100.0, 'best_f1_thresh': 0.0}\n"




@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION     ,      _KWARGS_DESCRIPTION   )
class     a__ (      datasets.Metric   ):







     def     lowerCAmelCase_       (			self				)						->					str:



          '''simple docstring'''



          return datasets.MetricInfo(
              description=_DESCRIPTION						,	citation=_CITATION						,	inputs_description=_KWARGS_DESCRIPTION						,	features=datasets.Features(
                  {
                      "predictions": {
                          "id": datasets.Value("string"				),
                          "prediction_text": datasets.Value("string"				),
                          "no_answer_probability": datasets.Value("float32"				),
                      },
                      "references": {
                          "id": datasets.Value("string"				),
                          "answers": datasets.features.Sequence(
                              {"text": datasets.Value("string"				), "answer_start": datasets.Value("int32"				)}				),
                      },
                  }				)						,	codebase_urls=["https://rajpurkar.github.io/SQuAD-explorer/"]						,	reference_urls=["https://rajpurkar.github.io/SQuAD-explorer/"]						,	)







     def     lowerCAmelCase_       (			self						,	A						,	A						,	A=1.0				)						->					Dict:



          '''simple docstring'''



          a											=   {p["id"]: p["no_answer_probability"] for p in predictions}
          a											=   [{"paragraphs": [{"qas": references}]}]
          a											=   {p["id"]: p["prediction_text"] for p in predictions}

          a											=   make_qid_to_has_ans(A				)  # maps qid to True/False
          a											=   [k for k, v in qid_to_has_ans.items() if v]
          a											=   [k for k, v in qid_to_has_ans.items() if not v]

          a						,   a											=   get_raw_scores(A						,	A				)
          a											=   apply_no_ans_threshold(A						,	A						,	A						,	A				)
          a											=   apply_no_ans_threshold(A						,	A						,	A						,	A				)
          a											=   make_eval_dict(A						,	A				)

          if has_ans_qids:
               a											=   make_eval_dict(A						,	A						,	qid_list=A				)
               merge_eval(A						,	A						,	"HasAns"				)
          if no_ans_qids:
               a											=   make_eval_dict(A						,	A						,	qid_list=A				)
               merge_eval(A						,	A						,	"NoAns"				)
          find_all_best_thresh(A						,	A						,	A						,	A						,	A						,	A				)
          return dict(A				)







import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device

def       SCREAMING_SNAKE_CASE							(     __UpperCamelCase)    ->  Tuple:
     a											=   {}
     for name in unet.attn_processors.keys():
          a											=   None if name.endswith("attn1.processor") else unet.config.cross_attention_dim
          if name.startswith("mid_block"):
               a											=   unet.config.block_out_channels[-1]
          elif name.startswith("up_blocks"):
               a											=   int(name[len("up_blocks.")])
               a											=   list(reversed(unet.config.block_out_channels))[block_id]
          elif name.startswith("down_blocks"):
               a											=   int(name[len("down_blocks.")])
               a											=   unet.config.block_out_channels[block_id]
          a											=   (
              LoRAAttnProcessor2_0 if hasattr(__UpperCamelCase	,      "scaled_dot_product_attention") else LoRAAttnProcessor
          )
          a											=   lora_attn_processor_class(
              hidden_size=__UpperCamelCase	,      cross_attention_dim=__UpperCamelCase)
     a											=   AttnProcsLayers(__UpperCamelCase)
     return lora_attn_procs, unet_lora_layers

def       SCREAMING_SNAKE_CASE							(     __UpperCamelCase)    ->  Optional[Any]:
     a											=   {}
     a											=   (
         LoRAAttnProcessor2_0 if hasattr(__UpperCamelCase	,      "scaled_dot_product_attention") else LoRAAttnProcessor
     )
     for name, module in text_encoder.named_modules():
          if name.endswith(__UpperCamelCase):
               a											=   lora_attn_processor_class(
                   hidden_size=module.out_proj.out_features	,      cross_attention_dim=__UpperCamelCase)
     return text_lora_attn_procs

def       SCREAMING_SNAKE_CASE							(     __UpperCamelCase)    ->  str:
     a											=   create_text_encoder_lora_attn_procs(__UpperCamelCase)
     a											=   AttnProcsLayers(__UpperCamelCase)
     return text_encoder_lora_layers

def       SCREAMING_SNAKE_CASE							(     __UpperCamelCase	,      __UpperCamelCase=False)    ->  Any:
     for _, attn_proc in text_lora_attn_procs.items():
          # set up.weights
          for layer_name, layer_module in attn_proc.named_modules():
               if layer_name.endswith("_lora"):
                    a											=   (
                        torch.randn_like(layer_module.up.weight)
                        if randn_weight
                        else torch.zeros_like(layer_module.up.weight)
                    )
                    a											=   torch.nn.Parameter(__UpperCamelCase)





class     a__ (      unittest.TestCase   ):







     def     lowerCAmelCase_       (			self				)						->					Union[str, Any]:



          '''simple docstring'''



          torch.manual_seed(0				)
          a											=   UNetaDConditionModel(
              block_out_channels=(32, 64)						,	layers_per_block=2						,	sample_size=32						,	in_channels=4						,	out_channels=4						,	down_block_types=("DownBlock2D", "CrossAttnDownBlock2D")						,	up_block_types=("CrossAttnUpBlock2D", "UpBlock2D")						,	cross_attention_dim=32						,	)
          a											=   DDIMScheduler(
              beta_start=0.0_0_0_8_5						,	beta_end=0.0_1_2						,	beta_schedule="scaled_linear"						,	clip_sample=A						,	set_alpha_to_one=A						,	steps_offset=1						,	)
          torch.manual_seed(0				)
          a											=   AutoencoderKL(
              block_out_channels=[32, 64]						,	in_channels=3						,	out_channels=3						,	down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"]						,	up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"]						,	latent_channels=4						,	)
          a											=   CLIPTextConfig(
              bos_token_id=0						,	eos_token_id=2						,	hidden_size=32						,	intermediate_size=37						,	layer_norm_eps=1e-05						,	num_attention_heads=4						,	num_hidden_layers=5						,	pad_token_id=1						,	vocab_size=1000						,	)
          a											=   CLIPTextModel(A				)
          a											=   CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"				)

          a						,   a											=   create_unet_lora_layers(A				)
          a											=   create_text_encoder_lora_layers(A				)

          a											=   {
              "unet": unet,
              "scheduler": scheduler,
              "vae": vae,
              "text_encoder": text_encoder,
              "tokenizer": tokenizer,
              "safety_checker": None,
              "feature_extractor": None,
          }
          a											=   {
              "unet_lora_layers": unet_lora_layers,
              "text_encoder_lora_layers": text_encoder_lora_layers,
              "unet_lora_attn_procs": unet_lora_attn_procs,
          }
          return pipeline_components, lora_components







     def     lowerCAmelCase_       (			self				)						->					Optional[Any]:



          '''simple docstring'''



          a											=   1
          a											=   10
          a											=   4
          a											=   (32, 32)

          a											=   torch.manual_seed(0				)
          a											=   floats_tensor((batch_size, num_channels) + sizes				)
          a											=   torch.randint(1						,	A						,	size=(batch_size, sequence_length)						,	generator=A				)

          a											=   {
              "prompt": "A painting of a squirrel eating a burger",
              "generator": generator,
              "num_inference_steps": 2,
              "guidance_scale": 6.0,
              "output_type": "numpy",
          }

          return noise, input_ids, pipeline_inputs

          # copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb







     def     lowerCAmelCase_       (			self				)						->					Optional[int]:



          '''simple docstring'''



          a											=   77

          a											=   torch.randint(2						,	56						,	size=(1, max_seq_length)						,	generator=torch.manual_seed(0				)				)

          a											=   {}
          a											=   inputs
          return prepared_inputs







     def     lowerCAmelCase_       (			self						,	A				)						->					Optional[int]:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          LoraLoaderMixin.save_lora_weights(
              save_directory=A						,	unet_lora_layers=lora_components["unet_lora_layers"]						,	text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]						,	)
          self.assertTrue(os.path.isfile(os.path.join(A						,	"pytorch_lora_weights.bin"				)				)				)







     def     lowerCAmelCase_       (			self				)						->					Any:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   StableDiffusionPipeline(**A				)
          a											=   sd_pipe.to(A				)
          sd_pipe.set_progress_bar_config(disable=A				)

          a						,   a						,   a											=   self.get_dummy_inputs()

          a											=   sd_pipe(**A				).images
          a											=   original_images[0, -3:, -3:, -1]

          with tempfile.TemporaryDirectory() as tmpdirname:
               LoraLoaderMixin.save_lora_weights(
                   save_directory=A						,	unet_lora_layers=lora_components["unet_lora_layers"]						,	text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]						,	)
               self.assertTrue(os.path.isfile(os.path.join(A						,	"pytorch_lora_weights.bin"				)				)				)
               sd_pipe.load_lora_weights(A				)

          a											=   sd_pipe(**A				).images
          a											=   lora_images[0, -3:, -3:, -1]

          # Outputs shouldn't match.
          self.assertFalse(torch.allclose(torch.from_numpy(A				)						,	torch.from_numpy(A				)				)				)







     def     lowerCAmelCase_       (			self				)						->					Optional[int]:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   StableDiffusionPipeline(**A				)
          a											=   sd_pipe.to(A				)
          sd_pipe.set_progress_bar_config(disable=A				)

          a						,   a						,   a											=   self.get_dummy_inputs()

          a											=   sd_pipe(**A				).images
          a											=   original_images[0, -3:, -3:, -1]

          with tempfile.TemporaryDirectory() as tmpdirname:
               LoraLoaderMixin.save_lora_weights(
                   save_directory=A						,	unet_lora_layers=lora_components["unet_lora_layers"]						,	text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]						,	safe_serialization=A						,	)
               self.assertTrue(os.path.isfile(os.path.join(A						,	"pytorch_lora_weights.safetensors"				)				)				)
               sd_pipe.load_lora_weights(A				)

          a											=   sd_pipe(**A				).images
          a											=   lora_images[0, -3:, -3:, -1]

          # Outputs shouldn't match.
          self.assertFalse(torch.allclose(torch.from_numpy(A				)						,	torch.from_numpy(A				)				)				)







     def     lowerCAmelCase_       (			self				)						->					Union[str, Any]:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   lora_components["unet_lora_attn_procs"]
          a											=   StableDiffusionPipeline(**A				)
          a											=   sd_pipe.to(A				)
          sd_pipe.set_progress_bar_config(disable=A				)

          a						,   a						,   a											=   self.get_dummy_inputs()

          a											=   sd_pipe(**A				).images
          a											=   original_images[0, -3:, -3:, -1]

          with tempfile.TemporaryDirectory() as tmpdirname:
               a											=   sd_pipe.unet
               unet.set_attn_processor(A				)
               unet.save_attn_procs(A				)
               self.assertTrue(os.path.isfile(os.path.join(A						,	"pytorch_lora_weights.bin"				)				)				)
               sd_pipe.load_lora_weights(A				)

          a											=   sd_pipe(**A				).images
          a											=   lora_images[0, -3:, -3:, -1]

          # Outputs shouldn't match.
          self.assertFalse(torch.allclose(torch.from_numpy(A				)						,	torch.from_numpy(A				)				)				)







     def     lowerCAmelCase_       (			self				)						->					Optional[Any]:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   StableDiffusionPipeline(**A				)

          a											=   self.get_dummy_tokens()

          # inference without lora
          a											=   pipe.text_encoder(**A				)[0]
          assert outputs_without_lora.shape == (1, 77, 32)

          # create lora_attn_procs with zeroed out up.weights
          a											=   create_text_encoder_lora_attn_procs(pipe.text_encoder				)
          set_lora_up_weights(A						,	randn_weight=A				)

          # monkey patch
          pipe._modify_text_encoder(A				)

          # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
          del text_attn_procs
          gc.collect()

          # inference with lora
          a											=   pipe.text_encoder(**A				)[0]
          assert outputs_with_lora.shape == (1, 77, 32)

          assert torch.allclose(
              A						,	A				), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

          # create lora_attn_procs with randn up.weights
          a											=   create_text_encoder_lora_attn_procs(pipe.text_encoder				)
          set_lora_up_weights(A						,	randn_weight=A				)

          # monkey patch
          pipe._modify_text_encoder(A				)

          # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
          del text_attn_procs
          gc.collect()

          # inference with lora
          a											=   pipe.text_encoder(**A				)[0]
          assert outputs_with_lora.shape == (1, 77, 32)

          assert not torch.allclose(
              A						,	A				), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"







     def     lowerCAmelCase_       (			self				)						->					Any:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   StableDiffusionPipeline(**A				)

          a											=   self.get_dummy_tokens()

          # inference without lora
          a											=   pipe.text_encoder(**A				)[0]
          assert outputs_without_lora.shape == (1, 77, 32)

          # create lora_attn_procs with randn up.weights
          a											=   create_text_encoder_lora_attn_procs(pipe.text_encoder				)
          set_lora_up_weights(A						,	randn_weight=A				)

          # monkey patch
          pipe._modify_text_encoder(A				)

          # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
          del text_attn_procs
          gc.collect()

          # inference with lora
          a											=   pipe.text_encoder(**A				)[0]
          assert outputs_with_lora.shape == (1, 77, 32)

          assert not torch.allclose(
              A						,	A				), "lora outputs should be different to without lora outputs"

          # remove monkey patch
          pipe._remove_text_encoder_monkey_patch()

          # inference with removed lora
          a											=   pipe.text_encoder(**A				)[0]
          assert outputs_without_lora_removed.shape == (1, 77, 32)

          assert torch.allclose(
              A						,	A				), "remove lora monkey patch should restore the original outputs"







     def     lowerCAmelCase_       (			self				)						->					str:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   StableDiffusionPipeline(**A				)
          a											=   sd_pipe.to(A				)
          sd_pipe.set_progress_bar_config(disable=A				)

          a						,   a						,   a											=   self.get_dummy_inputs()

          with tempfile.TemporaryDirectory() as tmpdirname:
               LoraLoaderMixin.save_lora_weights(
                   save_directory=A						,	unet_lora_layers=lora_components["unet_lora_layers"]						,	text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]						,	)
               self.assertTrue(os.path.isfile(os.path.join(A						,	"pytorch_lora_weights.bin"				)				)				)
               sd_pipe.load_lora_weights(A				)

          a											=   sd_pipe(**A				).images
          a											=   lora_images[0, -3:, -3:, -1]

          a											=   sd_pipe(**A						,	cross_attention_kwargs={"scale": 0.5}				).images
          a											=   lora_images_with_scale[0, -3:, -3:, -1]

          # Outputs shouldn't match.
          self.assertFalse(
              torch.allclose(torch.from_numpy(A				)						,	torch.from_numpy(A				)				)				)







     def     lowerCAmelCase_       (			self				)						->					Optional[Any]:



          '''simple docstring'''



          with tempfile.TemporaryDirectory() as tmpdirname:
               self.create_lora_weight_file(A				)

               a						,   a											=   self.get_dummy_components()
               a											=   StableDiffusionPipeline(**A				)
               a											=   sd_pipe.to(A				)
               sd_pipe.set_progress_bar_config(disable=A				)

               # check if vanilla attention processors are used
               for _, module in sd_pipe.unet.named_modules():
                    if isinstance(A						,	A				):
                         self.assertIsInstance(module.processor						,	(AttnProcessor, AttnProcessor2_0)				)

            # load LoRA weight file
               sd_pipe.load_lora_weights(A				)

               # check if lora attention processors are used
               for _, module in sd_pipe.unet.named_modules():
                    if isinstance(A						,	A				):
                         a											=   (
                             LoRAAttnProcessor2_0 if hasattr(A						,	"scaled_dot_product_attention"				) else LoRAAttnProcessor
                         )
                         self.assertIsInstance(module.processor						,	A				)







     @unittest.skipIf(torch_device != "cuda"						,	"This test is supposed to run on GPU"				)
     def     lowerCAmelCase_       (			self				)						->					Union[str, Any]:



          '''simple docstring'''



          with tempfile.TemporaryDirectory() as tmpdirname:
               self.create_lora_weight_file(A				)

               a						,   a											=   self.get_dummy_components()
               a											=   StableDiffusionPipeline(**A				)
               a											=   sd_pipe.to(A				)
               sd_pipe.set_progress_bar_config(disable=A				)

               # enable XFormers
               sd_pipe.enable_xformers_memory_efficient_attention()

               # check if xFormers attention processors are used
               for _, module in sd_pipe.unet.named_modules():
                    if isinstance(A						,	A				):
                         self.assertIsInstance(module.processor						,	A				)

            # load LoRA weight file
               sd_pipe.load_lora_weights(A				)

               # check if lora attention processors are used
               for _, module in sd_pipe.unet.named_modules():
                    if isinstance(A						,	A				):
                         self.assertIsInstance(module.processor						,	A				)







     @unittest.skipIf(torch_device != "cuda"						,	"This test is supposed to run on GPU"				)
     def     lowerCAmelCase_       (			self				)						->					int:



          '''simple docstring'''



          a						,   a											=   self.get_dummy_components()
          a											=   StableDiffusionPipeline(**A				)
          a											=   sd_pipe.to(A				)
          sd_pipe.set_progress_bar_config(disable=A				)

          a						,   a						,   a											=   self.get_dummy_inputs()

          # enable XFormers
          sd_pipe.enable_xformers_memory_efficient_attention()

          a											=   sd_pipe(**A				).images
          a											=   original_images[0, -3:, -3:, -1]

          with tempfile.TemporaryDirectory() as tmpdirname:
               LoraLoaderMixin.save_lora_weights(
                   save_directory=A						,	unet_lora_layers=lora_components["unet_lora_layers"]						,	text_encoder_lora_layers=lora_components["text_encoder_lora_layers"]						,	)
               self.assertTrue(os.path.isfile(os.path.join(A						,	"pytorch_lora_weights.bin"				)				)				)
               sd_pipe.load_lora_weights(A				)

          a											=   sd_pipe(**A				).images
          a											=   lora_images[0, -3:, -3:, -1]

          # Outputs shouldn't match.
          self.assertFalse(torch.allclose(torch.from_numpy(A				)						,	torch.from_numpy(A				)				)				)

