





import os

import jsonlines
import numpy as np
from tqdm import tqdm


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





def   lowerCamelCase__       (       UpperCamelCase__     :    Optional[int]  )     ->	Dict:



					'''simple docstring'''




					def choose_first(UpperCamelCase__     :    Union[str, Any]    ,      UpperCamelCase__     :    Union[str, Any]=False  ):
										assert isinstance(UpperCamelCase__    ,      UpperCamelCase__  )
										if len(UpperCamelCase__  ) == 1:
															_snake_case													=			answer[0]
															return {k: [answer[k]] for k in answer} if is_long_answer else answer
										for a in answer:
															if is_long_answer:
																				_snake_case													=			{k: [a[k]] for k in a}
															if len(a['start_token']  ) > 0:
																				break
										return a

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

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

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

					return answer





def   lowerCamelCase__       (       UpperCamelCase__     :    Optional[int]    ,      UpperCamelCase__     :    Any=False  )     ->	Any:



					'''simple docstring'''




					_snake_case													=			_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
										_snake_case													=			example['document']['tokens']
										_snake_case													=			[]
										for i in range(len(doc['token']  )  ):
															if not doc["is_html"][i]:
																				context.append(doc['token'][i]  )
										return {
										    "context": " ".join(UpperCamelCase__  ),
										    "answer": {
										        "start_token": -100,  # ignore index in cross-entropy
										        "end_token": -100,  # ignore index in cross-entropy
										        "category": answer["category"],
										        "span": answer["category"],  # extra
										    },
										}

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

					# handling normal samples

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

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

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

					# checking above code
					if assertion:
										_snake_case													=			doc['is_html'][answer['start_token'] : answer['end_token']]
										_snake_case													=			doc['token'][answer['start_token'] : answer['end_token']]
										_snake_case													=			' '.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__     :    int    ,      UpperCamelCase__     :    Tuple    ,      UpperCamelCase__     :    List[str]=2_048    ,      UpperCamelCase__     :    Dict=4_096    ,      UpperCamelCase__     :    Dict=True  )     ->	Tuple:



					'''simple docstring'''




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

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

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

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

					_snake_case													=			out['context'].split()
					_snake_case													=			splitted_context[answer['end_token']]
					_snake_case													=			len(
					    tokenizer(
					        ' '.join(splitted_context[: answer['start_token']]  )    ,      add_special_tokens=UpperCamelCase__    ,      ).input_ids  )
					_snake_case													=			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
					_snake_case													=			len(tokenizer(UpperCamelCase__    ,      add_special_tokens=UpperCamelCase__  ).input_ids  )
					if num_sub_tokens > 1:
										answer["end_token"] += num_sub_tokens - 1

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

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

					_snake_case													=			input_ids[:q_len]
					_snake_case													=			range(UpperCamelCase__    ,      len(UpperCamelCase__  )    ,      max_length - doc_stride  )

					_snake_case													=			[]
					_snake_case													=			[]
					_snake_case													=			[]
					_snake_case													=			[]  # null, yes, no, long, short
					for i in doc_start_indices:
										_snake_case													=			i + max_length - q_len
										_snake_case													=			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:
															_snake_case													=			start_token - i + q_len
															_snake_case													=			end_token - i + q_len
															answers_category.append(answer['category'][0]  )  # ["short"] -> "short"
										else:
															_snake_case													=			-100
															_snake_case													=			-100
															answers_category.append('null'  )
										_snake_case													=			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__     :    Tuple    ,      UpperCamelCase__     :    Dict    ,      UpperCamelCase__     :    List[Any]=2_048    ,      UpperCamelCase__     :    int=4_096    ,      UpperCamelCase__     :    List[str]=False  )     ->	List[str]:



					'''simple docstring'''




					_snake_case													=			get_strided_contexts_and_ans(
					    UpperCamelCase__    ,      UpperCamelCase__    ,      doc_stride=UpperCamelCase__    ,      max_length=UpperCamelCase__    ,      assertion=UpperCamelCase__    ,      )

					return example





def   lowerCamelCase__       (       UpperCamelCase__     :    Optional[Any]    ,      UpperCamelCase__     :    Dict  )     ->	Tuple:



					'''simple docstring'''




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

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

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

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










import os
from tempfile import TemporaryDirectory
from unittest import TestCase

import pytest
from absl.testing import parameterized

from datasets import config
from datasets.arrow_reader import HF_GCP_BASE_URL
from datasets.builder import DatasetBuilder
from datasets.dataset_dict import IterableDatasetDict
from datasets.iterable_dataset import IterableDataset
from datasets.load import dataset_module_factory, import_main_class
from datasets.utils.file_utils import cached_path


UpperCAmelCase_ 							=       [
    {"""dataset""": """wikipedia""", """config_name""": """20220301.de"""},
    {"""dataset""": """wikipedia""", """config_name""": """20220301.en"""},
    {"""dataset""": """wikipedia""", """config_name""": """20220301.fr"""},
    {"""dataset""": """wikipedia""", """config_name""": """20220301.frr"""},
    {"""dataset""": """wikipedia""", """config_name""": """20220301.it"""},
    {"""dataset""": """wikipedia""", """config_name""": """20220301.simple"""},
    {"""dataset""": """snli""", """config_name""": """plain_text"""},
    {"""dataset""": """eli5""", """config_name""": """LFQA_reddit"""},
    {"""dataset""": """wiki40b""", """config_name""": """en"""},
    {"""dataset""": """wiki_dpr""", """config_name""": """psgs_w100.nq.compressed"""},
    {"""dataset""": """wiki_dpr""", """config_name""": """psgs_w100.nq.no_index"""},
    {"""dataset""": """wiki_dpr""", """config_name""": """psgs_w100.multiset.no_index"""},
    {"""dataset""": """natural_questions""", """config_name""": """default"""},
]





def   lowerCamelCase__       (       UpperCamelCase__     :    Dict=True  )     ->	Dict:



					'''simple docstring'''




					if with_config:
										return [
										    {
										        "testcase_name": d["dataset"] + "/" + d["config_name"],
										        "dataset": d["dataset"],
										        "config_name": d["config_name"],
										    }
										    for d in DATASETS_ON_HF_GCP
										]
					else:
										return [
										    {"testcase_name": dataset, "dataset": dataset} for dataset in {d["dataset"] for d in DATASETS_ON_HF_GCP}
										]
@parameterized.named_parameters(list_datasets_on_hf_gcp_parameters(with_config=_lowerCamelCase   )   )
class      UpperCamelCase_						(					_lowerCamelCase   ):
					lowerCAmelCase_							=		None
					lowerCAmelCase_							=		None




					def 		lowerCAmelCase     (						self			,  lowerCAmelCase_			,  lowerCAmelCase_					)						->							Union[str, Any]:
										with TemporaryDirectory() as tmp_dir:
															_snake_case													=			dataset_module_factory(lowerCAmelCase_			,  cache_dir=lowerCAmelCase_					)

															_snake_case													=			import_main_class(dataset_module.module_path			,  dataset=lowerCAmelCase_					)

															_snake_case													=			builder_cls(
															    cache_dir=lowerCAmelCase_			,  config_name=lowerCAmelCase_			,  hash=dataset_module.hash			,  )

															_snake_case													=			'/'.join(
															    [
															        HF_GCP_BASE_URL,
															        builder_instance._relative_data_dir(with_hash=lowerCAmelCase_					).replace(os.sep			,  '/'					),
															        config.DATASET_INFO_FILENAME,
															    ]					)
															_snake_case													=			cached_path(lowerCAmelCase_			,  cache_dir=lowerCAmelCase_					)
															self.assertTrue(os.path.exists(lowerCAmelCase_					)					)





@pytest.mark.integration
def   lowerCamelCase__       (       UpperCamelCase__     :    Any  )     ->	Tuple:



					'''simple docstring'''




					_snake_case													=			tmp_path_factory.mktemp('test_hf_gcp'  ) / 'test_wikipedia_simple'
					_snake_case													=			dataset_module_factory('wikipedia'    ,      cache_dir=UpperCamelCase__  )
					_snake_case													=			import_main_class(dataset_module.module_path  )
					_snake_case													=			builder_cls(
					    cache_dir=UpperCamelCase__    ,      config_name='20220301.frr'    ,      hash=dataset_module.hash    ,      )
					# use the HF cloud storage, not the original download_and_prepare that uses apache-beam
					_snake_case													=			None
					builder_instance.download_and_prepare()
					_snake_case													=			builder_instance.as_dataset()
					assert ds





@pytest.mark.integration
def   lowerCamelCase__       (       UpperCamelCase__     :    Union[str, Any]  )     ->	Union[str, Any]:



					'''simple docstring'''




					_snake_case													=			dataset_module_factory('wikipedia'    ,      cache_dir=UpperCamelCase__  )
					_snake_case													=			import_main_class(dataset_module.module_path    ,      dataset=UpperCamelCase__  )
					_snake_case													=			builder_cls(
					    cache_dir=UpperCamelCase__    ,      config_name='20220301.frr'    ,      hash=dataset_module.hash    ,      )
					_snake_case													=			builder_instance.as_streaming_dataset()
					assert ds
					assert isinstance(UpperCamelCase__    ,      UpperCamelCase__  )
					assert "train" in ds
					assert isinstance(ds['train']    ,      UpperCamelCase__  )
					assert next(iter(ds['train']  )  )



