

import os

import jsonlines
import numpy as np
from tqdm import tqdm


A__						=							2048
A__						=							4096
A__						=							42
A__						=							os.environ.pop("PROCESS_TRAIN", "false")
A__						=							{"null": 0, "short": 1, "long": 2, "yes": 3, "no": 4}




def 					_lowercase     (				a_				:					List[Any]				)				->	str:




					'''simple docstring'''







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

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

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

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

					return answer




def 					_lowercase     (				a_				:					Optional[Any]					,a_				:					List[Any]=False				)				->	Union[str, Any]:




					'''simple docstring'''







					__magic_name__							=    _get_single_answer(SCREAMING_SNAKE_CASE_				)
					# 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
										__magic_name__							=    example['document']['tokens']
										__magic_name__							=    []
										for i in range(len(doc['token']				)				):
															if not doc["is_html"][i]:
																				context.append(doc['token'][i]				)
										return {
										    "context": " ".join(SCREAMING_SNAKE_CASE_				),
										    "answer": {
										        "start_token": -1_0_0,  # ignore index in cross-entropy
										        "end_token": -1_0_0,  # 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

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

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

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

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

					return {
					    "context": " ".join(SCREAMING_SNAKE_CASE_				),
					    "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 					_lowercase     (				a_				:					Any					,a_				:					List[Any]					,a_				:					Dict=2_0_4_8					,a_				:					Tuple=4_0_9_6					,a_				:					Union[str, Any]=True				)				->	Dict:




					'''simple docstring'''







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

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

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

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

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

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

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

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

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

					__magic_name__							=    input_ids[:q_len]
					__magic_name__							=    range(SCREAMING_SNAKE_CASE_					,len(SCREAMING_SNAKE_CASE_				)					,max_length - doc_stride				)

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

										answers_start_token.append(SCREAMING_SNAKE_CASE_				)
										answers_end_token.append(SCREAMING_SNAKE_CASE_				)
										if assertion:
															if new != old and new != [tokenizer.cls_token_id]:
																				print('ISSUE in strided for ID:'					,example['id']				)
																				print('New:'					,tokenizer.decode(SCREAMING_SNAKE_CASE_				)				)
																				print('Old:'					,tokenizer.decode(SCREAMING_SNAKE_CASE_				)					,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 					_lowercase     (				a_				:					int					,a_				:					List[str]					,a_				:					Union[str, Any]=2_0_4_8					,a_				:					Dict=4_0_9_6					,a_				:					Optional[int]=False				)				->	Union[str, Any]:




					'''simple docstring'''







					__magic_name__							=    get_strided_contexts_and_ans(
					    SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,doc_stride=SCREAMING_SNAKE_CASE_					,max_length=SCREAMING_SNAKE_CASE_					,assertion=SCREAMING_SNAKE_CASE_					,)

					return example




def 					_lowercase     (				a_				:					List[Any]					,a_				:					Tuple				)				->	Optional[int]:




					'''simple docstring'''







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

						A__						=							load_dataset("natural_questions")
						A__						=							BigBirdTokenizer.from_pretrained("google/bigbird-roberta-base")

						A__						=							data["train" if PROCESS_TRAIN == "true" else "validation"]

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

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







from typing import Dict, List, Optional, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_DEFAULT_MEAN,
    IMAGENET_DEFAULT_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_vision_available, logging


if is_vision_available():
						import PIL


A__						=							logging.get_logger(__name__)

class 				__UpperCamelCase				(					SCREAMING_SNAKE_CASE			):
	_lowercase  :   Tuple       				=  ["pixel_values"]







	def __init__(   self: Optional[Any]  , __UpperCamelCase: bool = True  , __UpperCamelCase: Dict[str, int] = None  , __UpperCamelCase: int = 0.9  , __UpperCamelCase: PILImageResampling = PILImageResampling.BICUBIC  , __UpperCamelCase: bool = True  , __UpperCamelCase: Dict[str, int] = None  , __UpperCamelCase: Union[int, float] = 1 / 2_55  , __UpperCamelCase: bool = True  , __UpperCamelCase: bool = True  , __UpperCamelCase: Optional[Union[float, List[float]]] = None  , __UpperCamelCase: Optional[Union[float, List[float]]] = None  , **__UpperCamelCase: List[str]  , ):

						'''simple docstring'''





						super().__init__(**__UpperCamelCase      )
						__magic_name__							=    size if size is not None else {'shortest_edge': 2_24}
						__magic_name__							=    get_size_dict(__UpperCamelCase  , default_to_square=__UpperCamelCase      )
						__magic_name__							=    crop_size if crop_size is not None else {'height': 2_24, 'width': 2_24}
						__magic_name__							=    get_size_dict(__UpperCamelCase  , param_name='crop_size'      )

						__magic_name__							=    do_resize
						__magic_name__							=    size
						__magic_name__							=    crop_pct
						__magic_name__							=    resample
						__magic_name__							=    do_center_crop
						__magic_name__							=    crop_size
						__magic_name__							=    do_rescale
						__magic_name__							=    rescale_factor
						__magic_name__							=    do_normalize
						__magic_name__							=    image_mean if image_mean is not None else IMAGENET_DEFAULT_MEAN
						__magic_name__							=    image_std if image_std is not None else IMAGENET_DEFAULT_STD







	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]  , __UpperCamelCase: np.ndarray  , __UpperCamelCase: Dict[str, int]  , __UpperCamelCase: Optional[float] = None  , __UpperCamelCase: PILImageResampling = PILImageResampling.BICUBIC  , __UpperCamelCase: Optional[Union[str, ChannelDimension]] = None  , **__UpperCamelCase: Optional[Any]  , ):

						'''simple docstring'''





						__magic_name__							=    get_size_dict(__UpperCamelCase  , default_to_square=__UpperCamelCase      )
						if "shortest_edge" not in size and ("height" not in size or "width" not in size):
											raise ValueError(F'size must contain \'height\' and \'width\' or \'shortest_edge\' as keys. Got {size.keys()}'      )
						if crop_pct is not None:
											if "shortest_edge" in size:
																__magic_name__							=    int(size['shortest_edge'] / crop_pct      )
											elif "height" in size and "width" in size:
																if size["height"] == size["width"]:
																					__magic_name__							=    int(size['height'] / crop_pct      )
																else:
																					__magic_name__							=    (int(size['height'] / crop_pct      ), int(size['width'] / crop_pct      ))
											else:
																raise ValueError('Invalid size for resize: {}'.format(__UpperCamelCase      )      )

											__magic_name__							=    get_resize_output_image_size(__UpperCamelCase  , size=__UpperCamelCase  , default_to_square=__UpperCamelCase      )
						else:
											if "shortest_edge" in size:
																__magic_name__							=    get_resize_output_image_size(__UpperCamelCase  , size=size['shortest_edge']  , default_to_square=__UpperCamelCase      )
											elif "height" in size and "width" in size:
																__magic_name__							=    (size['height'], size['width'])
											else:
																raise ValueError('Invalid size for resize: {}'.format(__UpperCamelCase      )      )

						return resize(__UpperCamelCase  , size=__UpperCamelCase  , resample=__UpperCamelCase  , data_format=__UpperCamelCase  , **__UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Any  , __UpperCamelCase: np.ndarray  , __UpperCamelCase: Dict[str, int]  , __UpperCamelCase: Optional[Union[str, ChannelDimension]] = None  , **__UpperCamelCase: List[str]  , ):

						'''simple docstring'''





						__magic_name__							=    get_size_dict(__UpperCamelCase      )
						if "height" not in size or "width" not in size:
											raise ValueError(F'size must contain \'height\' and \'width\' as keys. Got {size.keys()}'      )
						return center_crop(__UpperCamelCase  , size=(size['height'], size['width'])  , data_format=__UpperCamelCase  , **__UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: List[Any]  , __UpperCamelCase: np.ndarray  , __UpperCamelCase: Union[int, float]  , __UpperCamelCase: Optional[Union[str, ChannelDimension]] = None  , **__UpperCamelCase: Optional[Any]  , ):

						'''simple docstring'''





						return rescale(__UpperCamelCase  , scale=__UpperCamelCase  , data_format=__UpperCamelCase  , **__UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Dict  , __UpperCamelCase: np.ndarray  , __UpperCamelCase: Union[float, List[float]]  , __UpperCamelCase: Union[float, List[float]]  , __UpperCamelCase: Optional[Union[str, ChannelDimension]] = None  , **__UpperCamelCase: Optional[int]  , ):

						'''simple docstring'''





						return normalize(__UpperCamelCase  , mean=__UpperCamelCase  , std=__UpperCamelCase  , data_format=__UpperCamelCase  , **__UpperCamelCase      )







	def 	_SCREAMING_SNAKE_CASE							(   self: Union[str, Any]  , __UpperCamelCase: ImageInput  , __UpperCamelCase: bool = None  , __UpperCamelCase: Dict[str, int] = None  , __UpperCamelCase: int = None  , __UpperCamelCase: PILImageResampling = None  , __UpperCamelCase: bool = None  , __UpperCamelCase: Dict[str, int] = None  , __UpperCamelCase: bool = None  , __UpperCamelCase: float = None  , __UpperCamelCase: bool = None  , __UpperCamelCase: Optional[Union[float, List[float]]] = None  , __UpperCamelCase: Optional[Union[float, List[float]]] = None  , __UpperCamelCase: Optional[Union[str, TensorType]] = None  , __UpperCamelCase: ChannelDimension = ChannelDimension.FIRST  , **__UpperCamelCase: Any  , ):

						'''simple docstring'''





						__magic_name__							=    do_resize if do_resize is not None else self.do_resize
						__magic_name__							=    crop_pct if crop_pct is not None else self.crop_pct
						__magic_name__							=    resample if resample is not None else self.resample
						__magic_name__							=    do_center_crop if do_center_crop is not None else self.do_center_crop
						__magic_name__							=    do_rescale if do_rescale is not None else self.do_rescale
						__magic_name__							=    rescale_factor if rescale_factor is not None else self.rescale_factor
						__magic_name__							=    do_normalize if do_normalize is not None else self.do_normalize
						__magic_name__							=    image_mean if image_mean is not None else self.image_mean
						__magic_name__							=    image_std if image_std is not None else self.image_std

						__magic_name__							=    size if size is not None else self.size
						__magic_name__							=    get_size_dict(__UpperCamelCase  , default_to_square=__UpperCamelCase      )
						__magic_name__							=    crop_size if crop_size is not None else self.crop_size
						__magic_name__							=    get_size_dict(__UpperCamelCase  , param_name='crop_size'      )

						__magic_name__							=    make_list_of_images(__UpperCamelCase      )

						if not valid_images(__UpperCamelCase      ):
											raise ValueError(
											    'Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, '
											    'torch.Tensor, tf.Tensor or jax.ndarray.'      )

						if do_resize and size is None or resample is None:
											raise ValueError('Size and resample must be specified if do_resize is True.'      )

						if do_center_crop and crop_pct is None:
											raise ValueError('Crop_pct must be specified if do_center_crop is True.'      )

						if do_rescale and rescale_factor is None:
											raise ValueError('Rescale factor must be specified if do_rescale is True.'      )

						if do_normalize and (image_mean is None or image_std is None):
											raise ValueError('Image mean and std must be specified if do_normalize is True.'      )

						# All transformations expect numpy arrays.
						__magic_name__							=    [to_numpy_array(__UpperCamelCase      ) for image in images]

						if do_resize:
											__magic_name__							=    [self.resize(image=__UpperCamelCase  , size=__UpperCamelCase  , crop_pct=__UpperCamelCase  , resample=__UpperCamelCase      ) for image in images]

						if do_center_crop:
											__magic_name__							=    [self.center_crop(image=__UpperCamelCase  , size=__UpperCamelCase      ) for image in images]

						if do_rescale:
											__magic_name__							=    [self.rescale(image=__UpperCamelCase  , scale=__UpperCamelCase      ) for image in images]

						if do_normalize:
											__magic_name__							=    [self.normalize(image=__UpperCamelCase  , mean=__UpperCamelCase  , std=__UpperCamelCase      ) for image in images]

						__magic_name__							=    [to_channel_dimension_format(__UpperCamelCase  , __UpperCamelCase      ) for image in images]

						__magic_name__							=    {'pixel_values': images}
						return BatchFeature(data=__UpperCamelCase  , tensor_type=__UpperCamelCase      )




