


import os

import jsonlines
import numpy as np
from tqdm import tqdm


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


		def choose_first(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE=False					):
				assert isinstance(__A						,    __A					)
				if len(__A					) == 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=__A					)
						lowercase__						=				[]
				answer.update(__A					)

		# 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]						,    __A					) for k in cols					):
				raise ValueError('''Issue in ID'''						,    example['''id''']					)

		return answer
def       _a			(				SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE=False					):
		lowercase__						=				_get_single_answer(__A					)
		# 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(__A					),
				    "answer": {
				        "start_token": -1_00,  # ignore index in cross-entropy
				        "end_token": -1_00,  # 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(__A					)					) if not is_html[i]]					)
				if new != old:
						print('''ID:'''						,    example['''id''']					)
						print('''New:'''						,    __A						,    end='''\n'''					)
						print('''Old:'''						,    __A						,    end='''\n\n'''					)

		return {
		    "context": " ".join(__A					),
		    "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       _a			(				SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE=20_48						,    SCREAMING_SNAKE_CASE=40_96						,    SCREAMING_SNAKE_CASE=True					):
		lowercase__						=				get_context_and_ans(__A						,    assertion=__A					)
		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(__A						,    len(__A					)						,    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_00] * len(__A					),
				        "end_token": [-1_00] * len(__A					),
				        "category": category,
				    },
				}

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

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

		# fixing end token
		lowercase__						=				len(tokenizer(__A						,    add_special_tokens=__A					).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(__A					)
				if answer["span"] != new:
						print('''ISSUE IN TOKENIZATION'''					)
						print('''OLD:'''						,    answer['''span''']					)
						print('''NEW:'''						,    __A						,    end='''\n\n'''					)

		if len(__A					) <= 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(__A						,    len(__A					)						,    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_00
						lowercase__						=				-1_00
						answers_category.append('''null'''					)
				lowercase__						=				inputs[-1][start_token : end_token + 1]

				answers_start_token.append(__A					)
				answers_end_token.append(__A					)
				if assertion:
						if new != old and new != [tokenizer.cls_token_id]:
								print('''ISSUE in strided for ID:'''						,    example['''id''']					)
								print('''New:'''						,    tokenizer.decode(__A					)					)
								print('''Old:'''						,    tokenizer.decode(__A					)						,    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       _a			(				SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE=20_48						,    SCREAMING_SNAKE_CASE=40_96						,    SCREAMING_SNAKE_CASE=False					):
		lowercase__						=				get_strided_contexts_and_ans(
		    __A						,    __A						,    doc_stride=__A						,    max_length=__A						,    assertion=__A						,    )

		return example
def       _a			(				SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE					):


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

			lowerCAmelCase	=							load_dataset('natural_questions')
			lowerCAmelCase	=							BigBirdTokenizer.from_pretrained('google/bigbird-roberta-base')

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

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

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






import json
import logging
import os
import socket

import git
import numpy as np
import torch


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - PID: %(process)d -  %(message)s',
    datefmt='%m/%d/%Y %H:%M:%S',
    level=logging.INFO,
)
lowerCAmelCase	=							logging.getLogger(__name__)
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		lowercase__						=				git.Repo(search_parent_directories=SCREAMING_SNAKE_CASE					)
		lowercase__						=				{
		    '''repo_id''': str(SCREAMING_SNAKE_CASE					),
		    '''repo_sha''': str(repo.head.object.hexsha					),
		    '''repo_branch''': str(repo.active_branch					),
		}

		with open(os.path.join(SCREAMING_SNAKE_CASE						,    '''git_log.json'''					)						,    '''w'''					) as f:
				json.dump(SCREAMING_SNAKE_CASE						,    SCREAMING_SNAKE_CASE						,    indent=4					)
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		if params.n_gpu <= 0:
				lowercase__						=				0
				lowercase__						=				-1
				lowercase__						=				True
				lowercase__						=				False
				return

		assert torch.cuda.is_available()

		logger.info('''Initializing GPUs'''					)
		if params.n_gpu > 1:
				assert params.local_rank != -1

				lowercase__						=				int(os.environ['''WORLD_SIZE''']					)
				lowercase__						=				int(os.environ['''N_GPU_NODE''']					)
				lowercase__						=				int(os.environ['''RANK''']					)

				# number of nodes / node ID
				lowercase__						=				params.world_size // params.n_gpu_per_node
				lowercase__						=				params.global_rank // params.n_gpu_per_node
				lowercase__						=				True

				assert params.n_nodes == int(os.environ['''N_NODES''']					)
				assert params.node_id == int(os.environ['''NODE_RANK''']					)

		# local job (single GPU)
		else:
				assert params.local_rank == -1

				lowercase__						=				1
				lowercase__						=				0
				lowercase__						=				0
				lowercase__						=				0
				lowercase__						=				1
				lowercase__						=				1
				lowercase__						=				False

		# sanity checks
		assert params.n_nodes >= 1
		assert 0 <= params.node_id < params.n_nodes
		assert 0 <= params.local_rank <= params.global_rank < params.world_size
		assert params.world_size == params.n_nodes * params.n_gpu_per_node

		# define whether this is the master process / if we are in multi-node distributed mode
		lowercase__						=				params.node_id == 0 and params.local_rank == 0
		lowercase__						=				params.n_nodes > 1

		# summary
		lowercase__						=				f'--- Global rank: {params.global_rank} - '
		logger.info(PREFIX + '''Number of nodes: %i''' % params.n_nodes					)
		logger.info(PREFIX + '''Node ID        : %i''' % params.node_id					)
		logger.info(PREFIX + '''Local rank     : %i''' % params.local_rank					)
		logger.info(PREFIX + '''World size     : %i''' % params.world_size					)
		logger.info(PREFIX + '''GPUs per node  : %i''' % params.n_gpu_per_node					)
		logger.info(PREFIX + '''Master         : %s''' % str(params.is_master					)					)
		logger.info(PREFIX + '''Multi-node     : %s''' % str(params.multi_node					)					)
		logger.info(PREFIX + '''Multi-GPU      : %s''' % str(params.multi_gpu					)					)
		logger.info(PREFIX + '''Hostname       : %s''' % socket.gethostname()					)

		# set GPU device
		torch.cuda.set_device(params.local_rank					)

		# initialize multi-GPU
		if params.multi_gpu:
				logger.info('''Initializing PyTorch distributed'''					)
				torch.distributed.init_process_group(
				    init_method='''env://'''						,    backend='''nccl'''						,    )
def       _a			(				SCREAMING_SNAKE_CASE					):


		"""simple docstring"""




		np.random.seed(args.seed					)
		torch.manual_seed(args.seed					)
		if args.n_gpu > 0:
				torch.cuda.manual_seed_all(args.seed					)


