'''simple docstring'''
import os

import jsonlines
import numpy as np
from tqdm import tqdm


lowercase__  :      Dict        =					2048
lowercase__  :      Dict        =					4096
lowercase__  :      List[Any]        =					42
lowercase__  :      List[str]        =					os.environ.pop("PROCESS_TRAIN", "false")
lowercase__  :      Optional[int]        =					{"null": 0, "short": 1, "long": 2, "yes": 3, "no": 4}



def  __lowerCamelCase       (      _UpperCamelCase  :  Optional[int]					):


      '''simple docstring'''


      def choose_first(_UpperCamelCase  :  Any							,    _UpperCamelCase  :  Optional[int]=False					):
            assert isinstance(_UpperCamelCase							,    _UpperCamelCase					)
            if len(_UpperCamelCase					) == 1:
                  UpperCAmelCase_            =  answer[0]
                  return {k: [answer[k]] for k in answer} if is_long_answer else answer
            for a in answer:
                  if is_long_answer:
                        UpperCAmelCase_            =  {k: [a[k]] for k in a}
                  if len(a['''start_token''']					) > 0:
                        break
            return a

      UpperCAmelCase_            =  {'''id''': example['''id''']}
      UpperCAmelCase_            =  example['''annotations''']
      UpperCAmelCase_            =  annotation['''yes_no_answer''']
      if 0 in yes_no_answer or 1 in yes_no_answer:
            UpperCAmelCase_            =  ['''yes'''] if 1 in yes_no_answer else ['''no''']
            UpperCAmelCase_            =  UpperCAmelCase_            =  []
            UpperCAmelCase_            =  UpperCAmelCase_            =  []
            UpperCAmelCase_            =  ['''<cls>''']
      else:
            UpperCAmelCase_            =  ['''short''']
            UpperCAmelCase_            =  choose_first(annotation['''short_answers''']					)
            if len(out['''start_token''']					) == 0:
                  # answer will be long if short is not available
                  UpperCAmelCase_            =  ['''long''']
                  UpperCAmelCase_            =  choose_first(annotation['''long_answer''']							,    is_long_answer=_UpperCamelCase					)
                  UpperCAmelCase_            =  []
            answer.update(_UpperCamelCase					)

      # disregard some samples
      if len(answer['''start_token''']					) > 1 or answer["start_token"] == answer["end_token"]:
            UpperCAmelCase_            =  True
      else:
            UpperCAmelCase_            =  False

      UpperCAmelCase_            =  ['''start_token''', '''end_token''', '''start_byte''', '''end_byte''', '''text''']
      if not all(isinstance(answer[k]							,    _UpperCamelCase					) for k in cols					):
            raise ValueError('''Issue in ID'''							,    example['''id''']					)

      return answer



def  __lowerCamelCase       (      _UpperCamelCase  :  Optional[int]							,    _UpperCamelCase  :  List[Any]=False					):


      '''simple docstring'''


      UpperCAmelCase_            =  _get_single_answer(_UpperCamelCase					)
      # bytes are of no use
      del answer["start_byte"]
      del answer["end_byte"]

      # handle yes_no answers explicitly
      if answer["category"][0] in ["yes", "no"]:  # category is list with one element
            UpperCAmelCase_            =  example['''document''']['''tokens''']
            UpperCAmelCase_            =  []
            for i in range(len(doc['''token''']					)					):
                  if not doc["is_html"][i]:
                        context.append(doc['''token'''][i]					)
            return {
                "context": " ".join(_UpperCamelCase					),
                "answer": {
                    "start_token": -100,  # ignore index in cross-entropy
                    "end_token": -100,  # ignore index in cross-entropy
                    "category": answer["category"],
                    "span": answer["category"],  # extra
                },
            }

      # later, help in removing all no answers
      if answer["start_token"] == [-1]:
            return {
                "context": "None",
                "answer": {
                    "start_token": -1,
                    "end_token": -1,
                    "category": "null",
                    "span": "None",  # extra
                },
            }

      # handling normal samples

      UpperCAmelCase_            =  ['''start_token''', '''end_token''']
      answer.update({k: answer[k][0] if len(answer[k]					) > 0 else answer[k] for k in cols}					)  # e.g. [10] == 10

      UpperCAmelCase_            =  example['''document''']['''tokens''']
      UpperCAmelCase_            =  answer['''start_token''']
      UpperCAmelCase_            =  answer['''end_token''']

      UpperCAmelCase_            =  []
      for i in range(len(doc['''token''']					)					):
            if not doc["is_html"][i]:
                  context.append(doc['''token'''][i]					)
            else:
                  if answer["start_token"] > i:
                        start_token -= 1
                  if answer["end_token"] > i:
                        end_token -= 1
      UpperCAmelCase_            =  ''' '''.join(context[start_token:end_token]					)

      # checking above code
      if assertion:
            UpperCAmelCase_            =  doc['''is_html'''][answer['''start_token'''] : answer['''end_token''']]
            UpperCAmelCase_            =  doc['''token'''][answer['''start_token'''] : answer['''end_token''']]
            UpperCAmelCase_            =  ''' '''.join([old[i] for i in range(len(_UpperCamelCase					)					) if not is_html[i]]					)
            if new != old:
                  print('''ID:'''							,    example['''id''']					)
                  print('''New:'''							,    _UpperCamelCase							,    end='''\n'''					)
                  print('''Old:'''							,    _UpperCamelCase							,    end='''\n\n'''					)

      return {
          "context": " ".join(_UpperCamelCase					),
          "answer": {
              "start_token": start_token,
              "end_token": end_token - 1,  # this makes it inclusive
              "category": answer["category"],  # either long or short
              "span": new,  # extra
          },
      }



def  __lowerCamelCase       (      _UpperCamelCase  :  str							,    _UpperCamelCase  :  List[Any]							,    _UpperCamelCase  :  Any=2048							,    _UpperCamelCase  :  Dict=4096							,    _UpperCamelCase  :  Union[str, Any]=True					):


      '''simple docstring'''


      UpperCAmelCase_            =  get_context_and_ans(_UpperCamelCase							,    assertion=_UpperCamelCase					)
      UpperCAmelCase_            =  out['''answer''']

      # later, removing these samples
      if answer["start_token"] == -1:
            return {
                "example_id": example["id"],
                "input_ids": [[-1]],
                "labels": {
                    "start_token": [-1],
                    "end_token": [-1],
                    "category": ["null"],
                },
            }

      UpperCAmelCase_            =  tokenizer(example['''question''']['''text''']							,    out['''context''']					).input_ids
      UpperCAmelCase_            =  input_ids.index(tokenizer.sep_token_id					) + 1

      # return yes/no
      if answer["category"][0] in ["yes", "no"]:  # category is list with one element
            UpperCAmelCase_            =  []
            UpperCAmelCase_            =  []
            UpperCAmelCase_            =  input_ids[:q_len]
            UpperCAmelCase_            =  range(_UpperCamelCase							,    len(_UpperCamelCase					)							,    max_length - doc_stride					)
            for i in doc_start_indices:
                  UpperCAmelCase_            =  i + max_length - q_len
                  UpperCAmelCase_            =  input_ids[i:end_index]
                  inputs.append(q_indices + slice					)
                  category.append(answer['''category'''][0]					)
                  if slice[-1] == tokenizer.sep_token_id:
                        break

            return {
                "example_id": example["id"],
                "input_ids": inputs,
                "labels": {
                    "start_token": [-100] * len(_UpperCamelCase					),
                    "end_token": [-100] * len(_UpperCamelCase					),
                    "category": category,
                },
            }

      UpperCAmelCase_            =  out['''context'''].split()
      UpperCAmelCase_            =  splitted_context[answer['''end_token''']]
      UpperCAmelCase_            =  len(
          tokenizer(
              ''' '''.join(splitted_context[: answer['''start_token''']]					)							,    add_special_tokens=_UpperCamelCase							,    ).input_ids					)
      UpperCAmelCase_            =  len(
          tokenizer(''' '''.join(splitted_context[: answer['''end_token''']]					)							,    add_special_tokens=_UpperCamelCase					).input_ids					)

      answer["start_token"] += q_len
      answer["end_token"] += q_len

      # fixing end token
      UpperCAmelCase_            =  len(tokenizer(_UpperCamelCase							,    add_special_tokens=_UpperCamelCase					).input_ids					)
      if num_sub_tokens > 1:
            answer["end_token"] += num_sub_tokens - 1

      UpperCAmelCase_            =  input_ids[answer['''start_token'''] : answer['''end_token'''] + 1]  # right & left are inclusive
      UpperCAmelCase_            =  answer['''start_token''']
      UpperCAmelCase_            =  answer['''end_token''']

      if assertion:
            UpperCAmelCase_            =  tokenizer.decode(_UpperCamelCase					)
            if answer["span"] != new:
                  print('''ISSUE IN TOKENIZATION'''					)
                  print('''OLD:'''							,    answer['''span''']					)
                  print('''NEW:'''							,    _UpperCamelCase							,    end='''\n\n'''					)

      if len(_UpperCamelCase					) <= max_length:
            return {
                "example_id": example["id"],
                "input_ids": [input_ids],
                "labels": {
                    "start_token": [answer["start_token"]],
                    "end_token": [answer["end_token"]],
                    "category": answer["category"],
                },
            }

      UpperCAmelCase_            =  input_ids[:q_len]
      UpperCAmelCase_            =  range(_UpperCamelCase							,    len(_UpperCamelCase					)							,    max_length - doc_stride					)

      UpperCAmelCase_            =  []
      UpperCAmelCase_            =  []
      UpperCAmelCase_            =  []
      UpperCAmelCase_            =  []  # null, yes, no, long, short
      for i in doc_start_indices:
            UpperCAmelCase_            =  i + max_length - q_len
            UpperCAmelCase_            =  input_ids[i:end_index]
            inputs.append(q_indices + slice					)
            assert len(inputs[-1]					) <= max_length, "Issue in truncating length"

            if start_token >= i and end_token <= end_index - 1:
                  UpperCAmelCase_            =  start_token - i + q_len
                  UpperCAmelCase_            =  end_token - i + q_len
                  answers_category.append(answer['''category'''][0]					)  # ["short"] -> "short"
            else:
                  UpperCAmelCase_            =  -100
                  UpperCAmelCase_            =  -100
                  answers_category.append('''null'''					)
            UpperCAmelCase_            =  inputs[-1][start_token : end_token + 1]

            answers_start_token.append(_UpperCamelCase					)
            answers_end_token.append(_UpperCamelCase					)
            if assertion:
                  if new != old and new != [tokenizer.cls_token_id]:
                        print('''ISSUE in strided for ID:'''							,    example['''id''']					)
                        print('''New:'''							,    tokenizer.decode(_UpperCamelCase					)					)
                        print('''Old:'''							,    tokenizer.decode(_UpperCamelCase					)							,    end='''\n\n'''					)
            if slice[-1] == tokenizer.sep_token_id:
                  break

      return {
          "example_id": example["id"],
          "input_ids": inputs,
          "labels": {
              "start_token": answers_start_token,
              "end_token": answers_end_token,
              "category": answers_category,
          },
      }



def  __lowerCamelCase       (      _UpperCamelCase  :  Dict							,    _UpperCamelCase  :  Union[str, Any]							,    _UpperCamelCase  :  Any=2048							,    _UpperCamelCase  :  str=4096							,    _UpperCamelCase  :  int=False					):


      '''simple docstring'''


      UpperCAmelCase_            =  get_strided_contexts_and_ans(
          _UpperCamelCase							,    _UpperCamelCase							,    doc_stride=_UpperCamelCase							,    max_length=_UpperCamelCase							,    assertion=_UpperCamelCase							,    )

      return example



def  __lowerCamelCase       (      _UpperCamelCase  :  Optional[int]							,    _UpperCamelCase  :  str					):


      '''simple docstring'''


      with jsonlines.open(_UpperCamelCase							,    '''a'''					) as writer:
            for example in tqdm(_UpperCamelCase							,    total=len(_UpperCamelCase					)							,    desc='''Saving samples ... '''					):
                  UpperCAmelCase_            =  example['''labels''']
                  for ids, start, end, cat in zip(
                      example['''input_ids''']							,    labels['''start_token''']							,    labels['''end_token''']							,    labels['''category''']							,    ):
                        if start == -1 and end == -1:
                              continue  # leave waste samples with no answer
                        if cat == "null" and np.random.rand() < 0.6:
                              continue  # removing 50 % samples
                        writer.write(
                            {
                                '''input_ids''': ids,
                                '''start_token''': start,
                                '''end_token''': end,
                                '''category''': CATEGORY_MAPPING[cat],
                            }					)


if __name__ == "__main__":
      from datasets import load_dataset

      from transformers import BigBirdTokenizer

      lowercase__  :      List[Any]        =					load_dataset("natural_questions")
      lowercase__  :      Any        =					BigBirdTokenizer.from_pretrained("google/bigbird-roberta-base")

      lowercase__  :      Dict        =					data["train" if PROCESS_TRAIN == "true" else "validation"]

      lowercase__  :      List[Any]        =					{
          "tokenizer": tokenizer,
          "doc_stride": DOC_STRIDE,
          "max_length": MAX_LENGTH,
          "assertion": False,
      }
      lowercase__  :      List[str]        =					data.map(prepare_inputs, fn_kwargs=fn_kwargs)
      lowercase__  :      Optional[Any]        =					data.remove_columns(["annotations", "document", "id", "question"])
      print(data)

      np.random.seed(SEED)
      lowercase__  :      Optional[int]        =					"nq-training.jsonl" if PROCESS_TRAIN == "true" else "nq-validation.jsonl"
      save_to_disk(data, file_name=cache_file_name)






'''simple docstring'''
import gc
import unittest

import numpy as np
import torch

from diffusers import StableDiffusionKDiffusionPipeline
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu


enable_full_determinism()







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




  '''simple docstring'''



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






  def  lowerCAmelCase__ (      self			:       str  )				->List[str]:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''CompVis/stable-diffusion-v1-4'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_euler'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe([prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=9.0  ,							num_inference_steps=20  ,							output_type='''np'''  )

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array([0.0447, 0.0492, 0.0468, 0.0408, 0.0383, 0.0408, 0.0354, 0.0380, 0.0339]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2






  def  lowerCAmelCase__ (      self			:       List[str]  )				->int:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_euler'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe([prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=9.0  ,							num_inference_steps=20  ,							output_type='''np'''  )

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array([0.1237, 0.1320, 0.1438, 0.1359, 0.1390, 0.1132, 0.1277, 0.1175, 0.1112]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 5e-1






  def  lowerCAmelCase__ (      self			:       List[Any]  )				->Optional[int]:
        UpperCAmelCase_            =  StableDiffusionKDiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-1-base'''  )
        UpperCAmelCase_            =  sd_pipe.to(UpperCAmelCase__  )
        sd_pipe.set_progress_bar_config(disable=UpperCAmelCase__  )

        sd_pipe.set_scheduler('''sample_dpmpp_2m'''  )

        UpperCAmelCase_            =  '''A painting of a squirrel eating a burger'''
        UpperCAmelCase_            =  torch.manual_seed(0  )
        UpperCAmelCase_            =  sd_pipe(
            [prompt]  ,							generator=UpperCAmelCase__  ,							guidance_scale=7.5  ,							num_inference_steps=15  ,							output_type='''np'''  ,							use_karras_sigmas=UpperCAmelCase__  ,							)

        UpperCAmelCase_            =  output.images

        UpperCAmelCase_            =  image[0, -3:, -3:, -1]

        assert image.shape == (1, 512, 512, 3)
        UpperCAmelCase_            =  np.array(
            [0.1138_1689, 0.1211_2921, 0.138_9457, 0.1254_9606, 0.124_4964, 0.1083_1517, 0.1156_2866, 0.1086_7816, 0.1049_9048]  )

        assert np.abs(image_slice.flatten() - expected_slice  ).max() < 1e-2





