






"""simple docstring"""



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      _lowerCAmelCase		(							UpperCAmelCase__       :					Any      )		->List[str]:
						def choose_first(UpperCAmelCase__       :					Optional[int],						UpperCAmelCase__       :					Tuple=False      ):
												assert isinstance(UpperCAmelCase__,						UpperCAmelCase__      )
												if len(UpperCAmelCase__      ) == 1:
																		A__						:  str		  =  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__						:  int		  =  {k: [a[k]] for k in a}
																		if len(a["""start_token"""]      ) > 0:
																								break
												return a

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

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

						A__						:  List[str]		  =  ["""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__       :					int,						UpperCAmelCase__       :					Optional[int]=False      )		->Dict:
						A__						:  Optional[int]		  =  _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
												A__						:  Dict		  =  example["""document"""]["""tokens"""]
												A__						:  Any		  =  []
												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": -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

						A__						:  Dict		  =  ["""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__						:  Tuple		  =  example["""document"""]["""tokens"""]
						A__						:  str		  =  answer["""start_token"""]
						A__						:  str		  =  answer["""end_token"""]

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

						# checking above code
						if assertion:
												A__						:  List[str]		  =  doc["""is_html"""][answer["""start_token"""] : answer["""end_token"""]]
												A__						:  List[str]		  =  doc["""token"""][answer["""start_token"""] : answer["""end_token"""]]
												A__						:  Optional[int]		  =  """ """.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__       :					List[Any],						UpperCAmelCase__       :					Dict,						UpperCAmelCase__       :					Tuple=2_0_4_8,						UpperCAmelCase__       :					Union[str, Any]=4_0_9_6,						UpperCAmelCase__       :					Any=True      )		->List[str]:
						# overlap will be of doc_stride - q_len

						A__						:  List[str]		  =  get_context_and_ans(UpperCAmelCase__,						assertion=UpperCAmelCase__      )
						A__						:  Union[str, Any]		  =  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__						:  Optional[Any]		  =  tokenizer(example["""question"""]["""text"""],						out["""context"""]      ).input_ids
						A__						:  Any		  =  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__						:  Union[str, Any]		  =  []
												A__						:  List[Any]		  =  []
												A__						:  List[str]		  =  input_ids[:q_len]
												A__						:  Optional[Any]		  =  range(UpperCAmelCase__,						len(UpperCAmelCase__      ),						max_length - doc_stride      )
												for i in doc_start_indices:
																		A__						:  List[str]		  =  i + max_length - q_len
																		A__						:  List[Any]		  =  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(UpperCAmelCase__      ),
												        "end_token": [-1_0_0] * len(UpperCAmelCase__      ),
												        "category": category,
												    },
												}

						A__						:  Any		  =  out["""context"""].split()
						A__						:  Optional[int]		  =  splitted_context[answer["""end_token"""]]
						A__						:  int		  =  len(
						    tokenizer(
						        """ """.join(splitted_context[: answer["""start_token"""]]      ),						add_special_tokens=UpperCAmelCase__,						).input_ids      )
						A__						:  Tuple		  =  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
						A__						:  Dict		  =  len(tokenizer(UpperCAmelCase__,						add_special_tokens=UpperCAmelCase__      ).input_ids      )
						if num_sub_tokens > 1:
												answer["end_token"] += num_sub_tokens - 1

						A__						:  Optional[int]		  =  input_ids[answer["""start_token"""] : answer["""end_token"""] + 1]  # right & left are inclusive
						A__						:  int		  =  answer["""start_token"""]
						A__						:  List[Any]		  =  answer["""end_token"""]

						if assertion:
												A__						:  Any		  =  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"],
												    },
												}

						A__						:  Union[str, Any]		  =  input_ids[:q_len]
						A__						:  Any		  =  range(UpperCAmelCase__,						len(UpperCAmelCase__      ),						max_length - doc_stride      )

						A__						:  List[str]		  =  []
						A__						:  Any		  =  []
						A__						:  Tuple		  =  []
						A__						:  List[str]		  =  []  # null, yes, no, long, short
						for i in doc_start_indices:
												A__						:  List[str]		  =  i + max_length - q_len
												A__						:  List[str]		  =  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__						:  Tuple		  =  start_token - i + q_len
																		A__						:  int		  =  end_token - i + q_len
																		answers_category.append(answer["""category"""][0]      )  # ["short"] -> "short"
												else:
																		A__						:  Optional[int]		  =  -1_0_0
																		A__						:  List[str]		  =  -1_0_0
																		answers_category.append("""null"""      )
												A__						:  Optional[Any]		  =  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__       :					int,						UpperCAmelCase__       :					Tuple=2_0_4_8,						UpperCAmelCase__       :					int=4_0_9_6,						UpperCAmelCase__       :					str=False      )		->str:
						A__						:  Tuple		  =  get_strided_contexts_and_ans(
						    UpperCAmelCase__,						UpperCAmelCase__,						doc_stride=UpperCAmelCase__,						max_length=UpperCAmelCase__,						assertion=UpperCAmelCase__,						)

						return example




def      _lowerCAmelCase		(							UpperCAmelCase__       :					int,						UpperCAmelCase__       :					List[Any]      )		->str:
						with jsonlines.open(UpperCAmelCase__,						"""a"""      ) as writer:
												for example in tqdm(UpperCAmelCase__,						total=len(UpperCAmelCase__      ),						desc="""Saving samples ... """      ):
																		A__						:  List[Any]		  =  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)














"""simple docstring"""



import multiprocessing
import os
from typing import BinaryIO, Optional, Union

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader






class 							__SCREAMING_SNAKE_CASE						(				UpperCamelCase		):


							def __init__(			self      :						List[Any]			,    snake_case      :						NestedDataStructureLike[PathLike]			,    snake_case      :						Optional[NamedSplit] = None			,    snake_case      :						Optional[Features] = None			,    snake_case      :						str = None			,    snake_case      :						bool = False			,    snake_case      :						bool = False			,    snake_case      :						Optional[str] = None			,    snake_case      :						Optional[int] = None			,    **snake_case      :						Optional[int]			,    ):
													'''simple docstring'''





													super().__init__(
													    snake_case			,    split=snake_case			,    features=snake_case			,    cache_dir=snake_case			,    keep_in_memory=snake_case			,    streaming=snake_case			,    num_proc=snake_case			,    **snake_case			,    )
													A__						:  Optional[int]		  =  field
													A__						:  Tuple		  =  path_or_paths if isinstance(snake_case			,    snake_case							) else {self.split: path_or_paths}
													A__						:  List[Any]		  =  Json(
													    cache_dir=snake_case			,    data_files=snake_case			,    features=snake_case			,    field=snake_case			,    **snake_case			,    )


							def 						_UpperCamelCase       (			self      :						Tuple							):
													'''simple docstring'''





													if self.streaming:
																			A__						:  List[Any]		  =  self.builder.as_streaming_dataset(split=self.split							)
													# Build regular (map-style) dataset
													else:
																			A__						:  int		  =  None
																			A__						:  Dict		  =  None
																			A__						:  Tuple		  =  None
																			A__						:  Optional[Any]		  =  None

																			self.builder.download_and_prepare(
																			    download_config=snake_case			,    download_mode=snake_case			,    verification_mode=snake_case			,    base_path=snake_case			,    num_proc=self.num_proc			,    )
																			A__						:  str		  =  self.builder.as_dataset(
																			    split=self.split			,    verification_mode=snake_case			,    in_memory=self.keep_in_memory							)
													return dataset






class 							__SCREAMING_SNAKE_CASE						:


							def __init__(			self      :						Any			,    snake_case      :						Dataset			,    snake_case      :						Union[PathLike, BinaryIO]			,    snake_case      :						Optional[int] = None			,    snake_case      :						Optional[int] = None			,    **snake_case      :						Union[str, Any]			,    ):
													'''simple docstring'''





													if num_proc is not None and num_proc <= 0:
																			raise ValueError(F'num_proc {num_proc} must be an integer > 0.'							)

													A__						:  Union[str, Any]		  =  dataset
													A__						:  Any		  =  path_or_buf
													A__						:  str		  =  batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
													A__						:  str		  =  num_proc
													A__						:  Union[str, Any]		  =  """utf-8"""
													A__						:  Optional[int]		  =  to_json_kwargs


							def 						_UpperCamelCase       (			self      :						List[str]							):
													'''simple docstring'''





													A__						:  Any		  =  self.to_json_kwargs.pop("""path_or_buf"""			,    snake_case							)
													A__						:  List[Any]		  =  self.to_json_kwargs.pop("""orient"""			,    """records"""							)
													A__						:  List[str]		  =  self.to_json_kwargs.pop("""lines"""			,    True if orient == """records""" else False							)
													A__						:  str		  =  self.to_json_kwargs.pop("""index"""			,    False if orient in ["""split""", """table"""] else True							)
													A__						:  Optional[int]		  =  self.to_json_kwargs.pop("""compression"""			,    snake_case							)

													if compression not in [None, "infer", "gzip", "bz2", "xz"]:
																			raise NotImplementedError(F'`datasets` currently does not support {compression} compression'							)

													if isinstance(self.path_or_buf			,    (str, bytes, os.PathLike)							):
																			with fsspec.open(self.path_or_buf			,    """wb"""			,    compression=snake_case							) as buffer:
																									A__						:  Tuple		  =  self._write(file_obj=snake_case			,    orient=snake_case			,    lines=snake_case			,    index=snake_case			,    **self.to_json_kwargs							)
													else:
																			if compression:
																									raise NotImplementedError(
																									    F'The compression parameter is not supported when writing to a buffer, but compression={compression}'
																									    """ was passed. Please provide a local path instead."""							)
																			A__						:  List[Any]		  =  self._write(
																			    file_obj=self.path_or_buf			,    orient=snake_case			,    lines=snake_case			,    index=snake_case			,    **self.to_json_kwargs							)
													return written


							def 						_UpperCamelCase       (			self      :						List[Any]			,    snake_case      :						List[str]							):
													'''simple docstring'''





													A__  ,  A__  ,  A__  ,  A__  ,  A__						:  int		  =  args

													A__						:  Any		  =  query_table(
													    table=self.dataset.data			,    key=slice(snake_case			,    offset + self.batch_size							)			,    indices=self.dataset._indices			,    )
													A__						:  List[Any]		  =  batch.to_pandas().to_json(
													    path_or_buf=snake_case			,    orient=snake_case			,    lines=snake_case			,    index=snake_case			,    **snake_case							)
													if not json_str.endswith("""\n"""							):
																			json_str += "\n"
													return json_str.encode(self.encoding							)


							def 						_UpperCamelCase       (			self      :						int			,    snake_case      :						BinaryIO			,    snake_case      :						List[Any]			,    snake_case      :						Tuple			,    snake_case      :						List[str]			,    **snake_case      :						str			,    ):
													'''simple docstring'''





													A__						:  List[Any]		  =  0

													if self.num_proc is None or self.num_proc == 1:
																			for offset in logging.tqdm(
																			    range(0			,    len(self.dataset							)			,    self.batch_size							)			,    unit="""ba"""			,    disable=not logging.is_progress_bar_enabled()			,    desc="""Creating json from Arrow format"""			,    ):
																									A__						:  List[str]		  =  self._batch_json((offset, orient, lines, index, to_json_kwargs)							)
																									written += file_obj.write(snake_case							)
													else:
																			A__  ,  A__						:  Optional[int]		  =  len(self.dataset							), self.batch_size
																			with multiprocessing.Pool(self.num_proc							) as pool:
																									for json_str in logging.tqdm(
																									    pool.imap(
																									        self._batch_json			,    [(offset, orient, lines, index, to_json_kwargs) for offset in range(0			,    snake_case			,    snake_case							)]			,    )			,    total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size			,    unit="""ba"""			,    disable=not logging.is_progress_bar_enabled()			,    desc="""Creating json from Arrow format"""			,    ):
																															written += file_obj.write(snake_case							)

													return written






