






'''simple docstring'''



import os

import jsonlines
import numpy as np
from tqdm import tqdm


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



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	Optional[Any]    )  ->			List[str]:

			"""simple docstring"""
			def choose_first(SCREAMING_SNAKE_CASE__:	Dict,					SCREAMING_SNAKE_CASE__:	str=False    ):
						assert isinstance(lowerCamelCase_,					lowerCamelCase_    )
						if len(lowerCamelCase_    ) == 1:
									__a				     =    answer[0]
									return {k: [answer[k]] for k in answer} if is_long_answer else answer
						for a in answer:
									if is_long_answer:
												__a				     =    {k: [a[k]] for k in a}
									if len(a['start_token']    ) > 0:
												break
						return a

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

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

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

			return answer



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	Union[str, Any],					SCREAMING_SNAKE_CASE__:	Optional[int]=False    )  ->			List[Any]:

			"""simple docstring"""
			__a				     =    _get_single_answer(lowerCamelCase_    )
			# 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
						__a				     =    example['document']['tokens']
						__a				     =    []
						for i in range(len(doc['token']    )    ):
									if not doc["is_html"][i]:
												context.append(doc['token'][i]    )
						return {
						    "context": " ".join(lowerCamelCase_    ),
						    "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

			__a				     =    ['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

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

			__a				     =    []
			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
			__a				     =    ' '.join(context[start_token:end_token]    )

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

			return {
			    "context": " ".join(lowerCamelCase_    ),
			    "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    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	Tuple,					SCREAMING_SNAKE_CASE__:	List[Any],					SCREAMING_SNAKE_CASE__:	Union[str, Any]=2048,					SCREAMING_SNAKE_CASE__:	str=4096,					SCREAMING_SNAKE_CASE__:	Union[str, Any]=True    )  ->			Optional[Any]:

			"""simple docstring"""
			# overlap will be of doc_stride - q_len

			__a				     =    get_context_and_ans(lowerCamelCase_,					assertion=lowerCamelCase_    )
			__a				     =    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"],
						    },
						}

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

			# return yes/no
			if answer["category"][0] in ["yes", "no"]:  # category is list with one element
						__a				     =    []
						__a				     =    []
						__a				     =    input_ids[:q_len]
						__a				     =    range(lowerCamelCase_,					len(lowerCamelCase_    ),					max_length - doc_stride    )
						for i in doc_start_indices:
									__a				     =    i + max_length - q_len
									__a				     =    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(lowerCamelCase_    ),
						        "end_token": [-100] * len(lowerCamelCase_    ),
						        "category": category,
						    },
						}

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

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

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

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

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

			if len(lowerCamelCase_    ) <= 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"],
						    },
						}

			__a				     =    input_ids[:q_len]
			__a				     =    range(lowerCamelCase_,					len(lowerCamelCase_    ),					max_length - doc_stride    )

			__a				     =    []
			__a				     =    []
			__a				     =    []
			__a				     =    []  # null, yes, no, long, short
			for i in doc_start_indices:
						__a				     =    i + max_length - q_len
						__a				     =    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:
									__a				     =    start_token - i + q_len
									__a				     =    end_token - i + q_len
									answers_category.append(answer['category'][0]    )  # ["short"] -> "short"
						else:
									__a				     =    -100
									__a				     =    -100
									answers_category.append('null'    )
						__a				     =    inputs[-1][start_token : end_token + 1]

						answers_start_token.append(lowerCamelCase_    )
						answers_end_token.append(lowerCamelCase_    )
						if assertion:
									if new != old and new != [tokenizer.cls_token_id]:
												print('ISSUE in strided for ID:',					example['id']    )
												print('New:',					tokenizer.decode(lowerCamelCase_    )    )
												print('Old:',					tokenizer.decode(lowerCamelCase_    ),					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    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	Tuple,					SCREAMING_SNAKE_CASE__:	str,					SCREAMING_SNAKE_CASE__:	List[str]=2048,					SCREAMING_SNAKE_CASE__:	Optional[int]=4096,					SCREAMING_SNAKE_CASE__:	Tuple=False    )  ->			Dict:

			"""simple docstring"""
			__a				     =    get_strided_contexts_and_ans(
			    lowerCamelCase_,					lowerCamelCase_,					doc_stride=lowerCamelCase_,					max_length=lowerCamelCase_,					assertion=lowerCamelCase_,					)

			return example



def    __UpperCAmelCase	(				SCREAMING_SNAKE_CASE__:	List[str],					SCREAMING_SNAKE_CASE__:	Dict    )  ->			Dict:

			"""simple docstring"""
			with jsonlines.open(lowerCamelCase_,					'a'    ) as writer:
						for example in tqdm(lowerCamelCase_,					total=len(lowerCamelCase_    ),					desc='Saving samples ... '    ):
									__a				     =    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

	__UpperCamelCase      :       Any      			=      load_dataset("""natural_questions""")
	__UpperCamelCase      :       str      			=      BigBirdTokenizer.from_pretrained("""google/bigbird-roberta-base""")

	__UpperCamelCase      :       Optional[Any]      			=      data["""train""" if PROCESS_TRAIN == """true""" else """validation"""]

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

	np.random.seed(SEED)
	__UpperCamelCase      :       Tuple      			=      """nq-training.jsonl""" if PROCESS_TRAIN == """true""" else """nq-validation.jsonl"""
	save_to_disk(data, file_name=cache_file_name)







'''simple docstring'''



__UpperCamelCase      :       Tuple      			=      """0.18.2"""

from .configuration_utils import ConfigMixin
from .utils import (
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_inflect_available,
    is_invisible_watermark_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_librosa_available,
    is_note_seq_available,
    is_onnx_available,
    is_scipy_available,
    is_torch_available,
    is_torchsde_available,
    is_transformers_available,
    is_transformers_version,
    is_unidecode_available,
    logging,
)


try:
 if not is_onnx_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_onnx_objects import *  # noqa F403
else:
 from .pipelines import OnnxRuntimeModel

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_pt_objects import *  # noqa F403
else:
 from .models import (
     AutoencoderKL,
     ControlNetModel,
     ModelMixin,
     PriorTransformer,
     TaFilmDecoder,
     TransformeraDModel,
     UNetaDModel,
     UNetaDConditionModel,
     UNetaDModel,
     UNetaDConditionModel,
     VQModel,
 )
 from .optimization import (
     get_constant_schedule,
     get_constant_schedule_with_warmup,
     get_cosine_schedule_with_warmup,
     get_cosine_with_hard_restarts_schedule_with_warmup,
     get_linear_schedule_with_warmup,
     get_polynomial_decay_schedule_with_warmup,
     get_scheduler,
 )
 from .pipelines import (
     AudioPipelineOutput,
     ConsistencyModelPipeline,
     DanceDiffusionPipeline,
     DDIMPipeline,
     DDPMPipeline,
     DiffusionPipeline,
     DiTPipeline,
     ImagePipelineOutput,
     KarrasVePipeline,
     LDMPipeline,
     LDMSuperResolutionPipeline,
     PNDMPipeline,
     RePaintPipeline,
     ScoreSdeVePipeline,
 )
 from .schedulers import (
     CMStochasticIterativeScheduler,
     DDIMInverseScheduler,
     DDIMParallelScheduler,
     DDIMScheduler,
     DDPMParallelScheduler,
     DDPMScheduler,
     DEISMultistepScheduler,
     DPMSolverMultistepInverseScheduler,
     DPMSolverMultistepScheduler,
     DPMSolverSinglestepScheduler,
     EulerAncestralDiscreteScheduler,
     EulerDiscreteScheduler,
     HeunDiscreteScheduler,
     IPNDMScheduler,
     KarrasVeScheduler,
     KDPMaAncestralDiscreteScheduler,
     KDPMaDiscreteScheduler,
     PNDMScheduler,
     RePaintScheduler,
     SchedulerMixin,
     ScoreSdeVeScheduler,
     UnCLIPScheduler,
     UniPCMultistepScheduler,
     VQDiffusionScheduler,
 )
 from .training_utils import EMAModel

try:
 if not (is_torch_available() and is_scipy_available()):
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_torch_and_scipy_objects import *  # noqa F403
else:
 from .schedulers import LMSDiscreteScheduler

try:
 if not (is_torch_available() and is_torchsde_available()):
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_torch_and_torchsde_objects import *  # noqa F403
else:
 from .schedulers import DPMSolverSDEScheduler

try:
 if not (is_torch_available() and is_transformers_available()):
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
 from .pipelines import (
     AltDiffusionImgaImgPipeline,
     AltDiffusionPipeline,
     AudioLDMPipeline,
     CycleDiffusionPipeline,
     IFImgaImgPipeline,
     IFImgaImgSuperResolutionPipeline,
     IFInpaintingPipeline,
     IFInpaintingSuperResolutionPipeline,
     IFPipeline,
     IFSuperResolutionPipeline,
     ImageTextPipelineOutput,
     KandinskyImgaImgPipeline,
     KandinskyInpaintPipeline,
     KandinskyPipeline,
     KandinskyPriorPipeline,
     KandinskyVaaControlnetImgaImgPipeline,
     KandinskyVaaControlnetPipeline,
     KandinskyVaaImgaImgPipeline,
     KandinskyVaaInpaintPipeline,
     KandinskyVaaPipeline,
     KandinskyVaaPriorEmbaEmbPipeline,
     KandinskyVaaPriorPipeline,
     LDMTextToImagePipeline,
     PaintByExamplePipeline,
     SemanticStableDiffusionPipeline,
     ShapEImgaImgPipeline,
     ShapEPipeline,
     StableDiffusionAttendAndExcitePipeline,
     StableDiffusionControlNetImgaImgPipeline,
     StableDiffusionControlNetInpaintPipeline,
     StableDiffusionControlNetPipeline,
     StableDiffusionDepthaImgPipeline,
     StableDiffusionDiffEditPipeline,
     StableDiffusionImageVariationPipeline,
     StableDiffusionImgaImgPipeline,
     StableDiffusionInpaintPipeline,
     StableDiffusionInpaintPipelineLegacy,
     StableDiffusionInstructPixaPixPipeline,
     StableDiffusionLatentUpscalePipeline,
     StableDiffusionLDMaDPipeline,
     StableDiffusionModelEditingPipeline,
     StableDiffusionPanoramaPipeline,
     StableDiffusionParadigmsPipeline,
     StableDiffusionPipeline,
     StableDiffusionPipelineSafe,
     StableDiffusionPixaPixZeroPipeline,
     StableDiffusionSAGPipeline,
     StableDiffusionUpscalePipeline,
     StableUnCLIPImgaImgPipeline,
     StableUnCLIPPipeline,
     TextToVideoSDPipeline,
     TextToVideoZeroPipeline,
     UnCLIPImageVariationPipeline,
     UnCLIPPipeline,
     UniDiffuserModel,
     UniDiffuserPipeline,
     UniDiffuserTextDecoder,
     VersatileDiffusionDualGuidedPipeline,
     VersatileDiffusionImageVariationPipeline,
     VersatileDiffusionPipeline,
     VersatileDiffusionTextToImagePipeline,
     VideoToVideoSDPipeline,
     VQDiffusionPipeline,
 )

try:
 if not (is_torch_available() and is_transformers_available() and is_invisible_watermark_available()):
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_torch_and_transformers_and_invisible_watermark_objects import *  # noqa F403
else:
 from .pipelines import StableDiffusionXLImgaImgPipeline, StableDiffusionXLPipeline

try:
 if not (is_torch_available() and is_transformers_available() and is_k_diffusion_available()):
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
else:
 from .pipelines import StableDiffusionKDiffusionPipeline

try:
 if not (is_torch_available() and is_transformers_available() and is_onnx_available()):
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_torch_and_transformers_and_onnx_objects import *  # noqa F403
else:
 from .pipelines import (
     OnnxStableDiffusionImgaImgPipeline,
     OnnxStableDiffusionInpaintPipeline,
     OnnxStableDiffusionInpaintPipelineLegacy,
     OnnxStableDiffusionPipeline,
     OnnxStableDiffusionUpscalePipeline,
     StableDiffusionOnnxPipeline,
 )

try:
 if not (is_torch_available() and is_librosa_available()):
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_torch_and_librosa_objects import *  # noqa F403
else:
 from .pipelines import AudioDiffusionPipeline, Mel

try:
 if not (is_transformers_available() and is_torch_available() and is_note_seq_available()):
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_transformers_and_torch_and_note_seq_objects import *  # noqa F403
else:
 from .pipelines import SpectrogramDiffusionPipeline

try:
 if not is_flax_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_flax_objects import *  # noqa F403
else:
 from .models.controlnet_flax import FlaxControlNetModel
 from .models.modeling_flax_utils import FlaxModelMixin
 from .models.unet_ad_condition_flax import FlaxUNetaDConditionModel
 from .models.vae_flax import FlaxAutoencoderKL
 from .pipelines import FlaxDiffusionPipeline
 from .schedulers import (
     FlaxDDIMScheduler,
     FlaxDDPMScheduler,
     FlaxDPMSolverMultistepScheduler,
     FlaxKarrasVeScheduler,
     FlaxLMSDiscreteScheduler,
     FlaxPNDMScheduler,
     FlaxSchedulerMixin,
     FlaxScoreSdeVeScheduler,
 )


try:
 if not (is_flax_available() and is_transformers_available()):
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_flax_and_transformers_objects import *  # noqa F403
else:
 from .pipelines import (
     FlaxStableDiffusionControlNetPipeline,
     FlaxStableDiffusionImgaImgPipeline,
     FlaxStableDiffusionInpaintPipeline,
     FlaxStableDiffusionPipeline,
 )

try:
 if not (is_note_seq_available()):
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 from .utils.dummy_note_seq_objects import *  # noqa F403
else:
 from .pipelines import MidiProcessor