
import json
import unittest

import numpy as np
from huggingface_hub import hf_hub_download

from transformers.testing_utils import require_torch, require_vision
from transformers.utils import is_torch_available, is_vision_available

from ...test_image_processing_common import ImageProcessingSavingTestMixin, prepare_image_inputs


if is_torch_available():
					import torch

					if is_vision_available():
										from transformers import OneFormerImageProcessor
										from transformers.models.oneformer.image_processing_oneformer import binary_mask_to_rle
										from transformers.models.oneformer.modeling_oneformer import OneFormerForUniversalSegmentationOutput

if is_vision_available():
					from PIL import Image


def 		UpperCAmelCase     (      a_						,  a_="shi-labs/oneformer_demo"		)						->      Optional[int]:







				"""simple docstring"""



				with open(hf_hub_download(a_						,  a_						,  repo_type="dataset"		)						,  "r"		) as f:
								__A       =					json.load(a_		)
				__A       =					{}
				__A       =					[]
				__A       =					[]
				for key, info in class_info.items():
								__A       =					info["name"]
								class_names.append(info["name"]		)
								if info["isthing"]:
												thing_ids.append(int(a_		)		)
				__A       =					thing_ids
				__A       =					class_names
				return metadata





class 				UpperCAmelCase (    unittest.TestCase    ):




	'''simple docstring'''





	def __init__(			self   :					str				,A   :					Optional[int]				,A   :					List[Any]=7				,A   :					List[str]=3				,A   :					Tuple=30				,A   :					List[Any]=4_00				,A   :					List[str]=None				,A   :					List[str]=True				,A   :					str=True				,A   :					Optional[Any]=[0.5, 0.5, 0.5]				,A   :					Optional[int]=[0.5, 0.5, 0.5]				,A   :					List[Any]=10				,A   :					Dict=False				,A   :					Any=2_55				,A   :					int="shi-labs/oneformer_demo"				,A   :					Tuple="ade20k_panoptic.json"				,A   :					List[Any]=10				,):
					__A       =					parent
					__A       =					batch_size
					__A       =					num_channels
					__A       =					min_resolution
					__A       =					max_resolution
					__A       =					do_resize
					__A       =					{"shortest_edge": 32, "longest_edge": 13_33} if size is None else size
					__A       =					do_normalize
					__A       =					image_mean
					__A       =					image_std
					__A       =					class_info_file
					__A       =					prepare_metadata(A				,A					)
					__A       =					num_text
					__A       =					repo_path

					# for the post_process_functions
					__A       =					2
					__A       =					10
					__A       =					10
					__A       =					3
					__A       =					4
					__A       =					num_labels
					__A       =					do_reduce_labels
					__A       =					ignore_index






	def 							UpperCamelCase_    (			self   :					Dict					):
					return {
					    "do_resize": self.do_resize,
					    "size": self.size,
					    "do_normalize": self.do_normalize,
					    "image_mean": self.image_mean,
					    "image_std": self.image_std,
					    "num_labels": self.num_labels,
					    "do_reduce_labels": self.do_reduce_labels,
					    "ignore_index": self.ignore_index,
					    "class_info_file": self.class_info_file,
					    "metadata": self.metadata,
					    "num_text": self.num_text,
					}






	def 							UpperCamelCase_    (			self   :					Dict				,A   :					Optional[int]				,A   :					Tuple=False					):
					if not batched:
									__A       =					image_inputs[0]
									if isinstance(A				,Image.Image					):
													__A     ,				__A       =					image.size
									else:
													__A     ,				__A       =					image.shape[1], image.shape[2]
									if w < h:
													__A       =					int(self.size["shortest_edge"] * h / w					)
													__A       =					self.size["shortest_edge"]
									elif w > h:
													__A       =					self.size["shortest_edge"]
													__A       =					int(self.size["shortest_edge"] * w / h					)
									else:
													__A       =					self.size["shortest_edge"]
													__A       =					self.size["shortest_edge"]

					else:
									__A       =					[]
									for image in image_inputs:
													__A     ,				__A       =					self.get_expected_values([image]					)
													expected_values.append((expected_height, expected_width)					)
									__A       =					max(A				,key=lambda A					: item[0]					)[0]
									__A       =					max(A				,key=lambda A					: item[1]					)[1]

					return expected_height, expected_width







	def 							UpperCamelCase_    (			self   :					Union[str, Any]					):
					return OneFormerForUniversalSegmentationOutput(
					    # +1 for null class
					    class_queries_logits=torch.randn((self.batch_size, self.num_queries, self.num_classes + 1)					)				,masks_queries_logits=torch.randn((self.batch_size, self.num_queries, self.height, self.width)					)				,)





@require_torch
@require_vision
class 				UpperCAmelCase (    __SCREAMING_SNAKE_CASE     ,     unittest.TestCase    ):




	'''simple docstring'''





	snake_case_          =	OneFormerImageProcessor if (is_vision_available() and is_torch_available()) else None
	# only for test_image_processing_common.test_image_proc_to_json_string
	snake_case_          =	image_processing_class






	def 							UpperCamelCase_    (			self   :					Union[str, Any]					):
					__A       =					OneFormerImageProcessorTester(self					)






	@property
	def 							UpperCamelCase_    (			self   :					Optional[Any]					):
					return self.image_processing_tester.prepare_image_processor_dict()






	def 							UpperCamelCase_    (			self   :					Union[str, Any]					):
					__A       =					self.image_processing_class(**self.image_processor_dict					)
					self.assertTrue(hasattr(A				,"image_mean"					)					)
					self.assertTrue(hasattr(A				,"image_std"					)					)
					self.assertTrue(hasattr(A				,"do_normalize"					)					)
					self.assertTrue(hasattr(A				,"do_resize"					)					)
					self.assertTrue(hasattr(A				,"size"					)					)
					self.assertTrue(hasattr(A				,"ignore_index"					)					)
					self.assertTrue(hasattr(A				,"class_info_file"					)					)
					self.assertTrue(hasattr(A				,"num_text"					)					)
					self.assertTrue(hasattr(A				,"repo_path"					)					)
					self.assertTrue(hasattr(A				,"metadata"					)					)
					self.assertTrue(hasattr(A				,"do_reduce_labels"					)					)






	def 							UpperCamelCase_    (			self   :					Optional[int]					):
					pass






	def 							UpperCamelCase_    (			self   :					Any					):
					# Initialize image_processor
					__A       =					self.image_processing_class(**self.image_processor_dict					)
					# create random PIL images
					__A       =					prepare_image_inputs(self.image_processing_tester				,equal_resolution=A					)
					for image in image_inputs:
									self.assertIsInstance(A				,Image.Image					)

					# Test not batched input
					__A       =					image_processor(image_inputs[0]				,["semantic"]				,return_tensors="pt"					).pixel_values

					__A     ,				__A       =					self.image_processing_tester.get_expected_values(A					)

					self.assertEqual(
					    encoded_images.shape				,(1, self.image_processing_tester.num_channels, expected_height, expected_width)				,)

					# Test batched
					__A     ,				__A       =					self.image_processing_tester.get_expected_values(A				,batched=A					)

					__A       =					image_processor(
					    A				,["semantic"] * len(A					)				,return_tensors="pt"					).pixel_values
					self.assertEqual(
					    encoded_images.shape				,(
					        self.image_processing_tester.batch_size,
					        self.image_processing_tester.num_channels,
					        expected_height,
					        expected_width,
					    )				,)






	def 							UpperCamelCase_    (			self   :					Union[str, Any]					):
					# Initialize image_processor
					__A       =					self.image_processing_class(**self.image_processor_dict					)
					# create random numpy tensors
					__A       =					prepare_image_inputs(self.image_processing_tester				,equal_resolution=A				,numpify=A					)
					for image in image_inputs:
									self.assertIsInstance(A				,np.ndarray					)

					# Test not batched input
					__A       =					image_processor(image_inputs[0]				,["semantic"]				,return_tensors="pt"					).pixel_values

					__A     ,				__A       =					self.image_processing_tester.get_expected_values(A					)

					self.assertEqual(
					    encoded_images.shape				,(1, self.image_processing_tester.num_channels, expected_height, expected_width)				,)

					# Test batched
					__A     ,				__A       =					self.image_processing_tester.get_expected_values(A				,batched=A					)

					__A       =					image_processor(
					    A				,["semantic"] * len(A					)				,return_tensors="pt"					).pixel_values
					self.assertEqual(
					    encoded_images.shape				,(
					        self.image_processing_tester.batch_size,
					        self.image_processing_tester.num_channels,
					        expected_height,
					        expected_width,
					    )				,)






	def 							UpperCamelCase_    (			self   :					Tuple					):
					# Initialize image_processor
					__A       =					self.image_processing_class(**self.image_processor_dict					)
					# create random PyTorch tensors
					__A       =					prepare_image_inputs(self.image_processing_tester				,equal_resolution=A				,torchify=A					)
					for image in image_inputs:
									self.assertIsInstance(A				,torch.Tensor					)

					# Test not batched input
					__A       =					image_processor(image_inputs[0]				,["semantic"]				,return_tensors="pt"					).pixel_values

					__A     ,				__A       =					self.image_processing_tester.get_expected_values(A					)

					self.assertEqual(
					    encoded_images.shape				,(1, self.image_processing_tester.num_channels, expected_height, expected_width)				,)

					# Test batched
					__A     ,				__A       =					self.image_processing_tester.get_expected_values(A				,batched=A					)

					__A       =					image_processor(
					    A				,["semantic"] * len(A					)				,return_tensors="pt"					).pixel_values
					self.assertEqual(
					    encoded_images.shape				,(
					        self.image_processing_tester.batch_size,
					        self.image_processing_tester.num_channels,
					        expected_height,
					        expected_width,
					    )				,)






	def 							UpperCamelCase_    (			self   :					Union[str, Any]				,A   :					Tuple=False				,A   :					List[str]=False				,A   :					Optional[Any]="np"					):
					__A       =					self.image_processing_class(**self.image_processor_dict					)
					# prepare image and target
					__A       =					self.image_processing_tester.num_labels
					__A       =					None
					__A       =					None
					__A       =					prepare_image_inputs(self.image_processing_tester				,equal_resolution=A					)
					if with_segmentation_maps:
									__A       =					num_labels
									if is_instance_map:
													__A       =					list(range(A					)					) * 2
													__A       =					dict(enumerate(A					)					)
									__A       =					[
									    np.random.randint(0				,high * 2				,(img.size[1], img.size[0])					).astype(np.uinta					) for img in image_inputs
									]
									if segmentation_type == "pil":
													__A       =					[Image.fromarray(A					) for annotation in annotations]

					__A       =					image_processor(
					    A				,["semantic"] * len(A					)				,A				,return_tensors="pt"				,instance_id_to_semantic_id=A				,pad_and_return_pixel_mask=A				,)

					return inputs






	def 							UpperCamelCase_    (			self   :					int					):
					pass






	def 							UpperCamelCase_    (			self   :					Tuple					):
					def common(A   :					Optional[Any]=False				,A   :					str=None					):
									__A       =					self.comm_get_image_processor_inputs(
									    with_segmentation_maps=A				,is_instance_map=A				,segmentation_type=A					)

									__A       =					inputs["mask_labels"]
									__A       =					inputs["class_labels"]
									__A       =					inputs["pixel_values"]
									__A       =					inputs["text_inputs"]

									# check the batch_size
									for mask_label, class_label, text_input in zip(A				,A				,A					):
													self.assertEqual(mask_label.shape[0]				,class_label.shape[0]					)
													# this ensure padding has happened
													self.assertEqual(mask_label.shape[1:]				,pixel_values.shape[2:]					)
													self.assertEqual(len(A					)				,self.image_processing_tester.num_text					)

					common()
					common(is_instance_map=A					)
					common(is_instance_map=A				,segmentation_type="pil"					)
					common(is_instance_map=A				,segmentation_type="pil"					)






	def 							UpperCamelCase_    (			self   :					List[Any]					):
					__A       =					np.zeros((20, 50)					)
					__A       =					1
					__A       =					1
					__A       =					1

					__A       =					binary_mask_to_rle(A					)
					self.assertEqual(len(A					)				,4					)
					self.assertEqual(rle[0]				,21					)
					self.assertEqual(rle[1]				,45					)






	def 							UpperCamelCase_    (			self   :					Dict					):
					__A       =					self.image_processing_class(
					    num_labels=self.image_processing_tester.num_classes				,max_seq_length=77				,task_seq_length=77				,class_info_file="ade20k_panoptic.json"				,num_text=self.image_processing_tester.num_text				,repo_path="shi-labs/oneformer_demo"				,)
					__A       =					self.image_processing_tester.get_fake_oneformer_outputs()

					__A       =					fature_extractor.post_process_semantic_segmentation(A					)

					self.assertEqual(len(A					)				,self.image_processing_tester.batch_size					)
					self.assertEqual(
					    segmentation[0].shape				,(
					        self.image_processing_tester.height,
					        self.image_processing_tester.width,
					    )				,)

					__A       =					[(1, 4) for i in range(self.image_processing_tester.batch_size					)]
					__A       =					fature_extractor.post_process_semantic_segmentation(A				,target_sizes=A					)

					self.assertEqual(segmentation[0].shape				,target_sizes[0]					)






	def 							UpperCamelCase_    (			self   :					int					):
					__A       =					self.image_processing_class(
					    num_labels=self.image_processing_tester.num_classes				,max_seq_length=77				,task_seq_length=77				,class_info_file="ade20k_panoptic.json"				,num_text=self.image_processing_tester.num_text				,repo_path="shi-labs/oneformer_demo"				,)
					__A       =					self.image_processing_tester.get_fake_oneformer_outputs()
					__A       =					image_processor.post_process_instance_segmentation(A				,threshold=0					)

					self.assertTrue(len(A					) == self.image_processing_tester.batch_size					)
					for el in segmentation:
									self.assertTrue("segmentation" in el					)
									self.assertTrue("segments_info" in el					)
									self.assertEqual(type(el["segments_info"]					)				,A					)
									self.assertEqual(
									    el["segmentation"].shape				,(self.image_processing_tester.height, self.image_processing_tester.width)					)







	def 							UpperCamelCase_    (			self   :					Tuple					):
					__A       =					self.image_processing_class(
					    num_labels=self.image_processing_tester.num_classes				,max_seq_length=77				,task_seq_length=77				,class_info_file="ade20k_panoptic.json"				,num_text=self.image_processing_tester.num_text				,repo_path="shi-labs/oneformer_demo"				,)
					__A       =					self.image_processing_tester.get_fake_oneformer_outputs()
					__A       =					image_processor.post_process_panoptic_segmentation(A				,threshold=0					)

					self.assertTrue(len(A					) == self.image_processing_tester.batch_size					)
					for el in segmentation:
									self.assertTrue("segmentation" in el					)
									self.assertTrue("segments_info" in el					)
									self.assertEqual(type(el["segments_info"]					)				,A					)
									self.assertEqual(
									    el["segmentation"].shape				,(self.image_processing_tester.height, self.image_processing_tester.width)					)


"""simple docstring"""






import os

import jsonlines
import numpy as np
from tqdm import tqdm


_lowercase							=				2_048
_lowercase							=				4_096
_lowercase							=				42
_lowercase							=				os.environ.pop('''PROCESS_TRAIN''', '''false''')
_lowercase							=				{'''null''': 0, '''short''': 1, '''long''': 2, '''yes''': 3, '''no''': 4}



def      lowerCAmelCase__			(		__magic_name__      )						->str:
			def choose_first(__magic_name__			,   __magic_name__=False      ):
						assert isinstance(__magic_name__			,   __magic_name__      )
						if len(__magic_name__      ) == 1:
									__lowercase       =       answer[0]
									return {k: [answer[k]] for k in answer} if is_long_answer else answer
						for a in answer:
									if is_long_answer:
												__lowercase       =       {k: [a[k]] for k in a}
									if len(a["start_token"]      ) > 0:
												break
						return a

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

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

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

			return answer



def      lowerCAmelCase__			(		__magic_name__			,   __magic_name__=False      )						->int:
			__lowercase       =       _get_single_answer(__magic_name__      )
			# 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
						__lowercase       =       example["document"]["tokens"]
						__lowercase       =       []
						for i in range(len(doc["token"]      )      ):
									if not doc["is_html"][i]:
												context.append(doc["token"][i]      )
						return {
						    "context": " ".join(__magic_name__      ),
						    "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

			__lowercase       =       ["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

			__lowercase       =       example["document"]["tokens"]
			__lowercase       =       answer["start_token"]
			__lowercase       =       answer["end_token"]

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

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

			return {
			    "context": " ".join(__magic_name__      ),
			    "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__			(		__magic_name__			,   __magic_name__			,   __magic_name__=2_0_4_8			,   __magic_name__=4_0_9_6			,   __magic_name__=True      )						->Optional[Any]:
			# overlap will be of doc_stride - q_len

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

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

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

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

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

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

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

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

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

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

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

						answers_start_token.append(__magic_name__      )
						answers_end_token.append(__magic_name__      )
						if assertion:
									if new != old and new != [tokenizer.cls_token_id]:
												print("ISSUE in strided for ID:"			,   example["id"]      )
												print("New:"			,   tokenizer.decode(__magic_name__      )      )
												print("Old:"			,   tokenizer.decode(__magic_name__      )			,   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__			(		__magic_name__			,   __magic_name__			,   __magic_name__=2_0_4_8			,   __magic_name__=4_0_9_6			,   __magic_name__=False      )						->List[str]:
			__lowercase       =       get_strided_contexts_and_ans(
			    __magic_name__			,   __magic_name__			,   doc_stride=__magic_name__			,   max_length=__magic_name__			,   assertion=__magic_name__			,   )

			return example



def      lowerCAmelCase__			(		__magic_name__			,   __magic_name__      )						->Any:
			with jsonlines.open(__magic_name__			,   "a"      ) as writer:
						for example in tqdm(__magic_name__			,   total=len(__magic_name__      )			,   desc="Saving samples ... "      ):
									__lowercase       =       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							=				load_dataset('''natural_questions''')
					_lowercase							=				BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base''')

					_lowercase							=				data['''train''' if PROCESS_TRAIN == '''true''' else '''validation''']

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

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





