





import os

import jsonlines
import numpy as np
from tqdm import tqdm


SCREAMING_SNAKE_CASE_   :      Union[str, Any]			=   2048
SCREAMING_SNAKE_CASE_   :      str			=   4096
SCREAMING_SNAKE_CASE_   :      List[Any]			=   42
SCREAMING_SNAKE_CASE_   :      Optional[Any]			=   os.environ.pop('''PROCESS_TRAIN''', '''false''')
SCREAMING_SNAKE_CASE_   :      List[str]			=   {'''null''': 0, '''short''': 1, '''long''': 2, '''yes''': 3, '''no''': 4}



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

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

		return answer



def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case=False							)		->    int:
		__lowercase											=     _get_single_answer(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
				__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(snake_case							),
				    "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

		__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(snake_case							)							) if not is_html[i]]							)
				if new != old:
						print('ID:'						,   example['id']							)
						print('New:'						,   snake_case						,   end='\n'							)
						print('Old:'						,   snake_case						,   end='\n\n'							)

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

		__lowercase											=     get_context_and_ans(snake_case						,   assertion=snake_case							)
		__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(snake_case						,   len(snake_case							)						,   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": [-100] * len(snake_case							),
				        "end_token": [-100] * len(snake_case							),
				        "category": category,
				    },
				}

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

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

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

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

		__lowercase											=     input_ids[:q_len]
		__lowercase											=     range(snake_case						,   len(snake_case							)						,   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											=     -100
						__lowercase											=     -100
						answers_category.append('null'							)
				__lowercase											=     inputs[-1][start_token : end_token + 1]

				answers_start_token.append(snake_case							)
				answers_end_token.append(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(snake_case							)							)
								print('Old:'						,   tokenizer.decode(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  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case						,   snake_case=2_048						,   snake_case=4_096						,   snake_case=False							)		->    Tuple:
		__lowercase											=     get_strided_contexts_and_ans(
		    snake_case						,   snake_case						,   doc_stride=snake_case						,   max_length=snake_case						,   assertion=snake_case						,   )

		return example



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

		SCREAMING_SNAKE_CASE_   :      Tuple			=   load_dataset('''natural_questions''')
		SCREAMING_SNAKE_CASE_   :      List[Any]			=   BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base''')

		SCREAMING_SNAKE_CASE_   :      Any			=   data['''train''' if PROCESS_TRAIN == '''true''' else '''validation''']

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

		np.random.seed(SEED)
		SCREAMING_SNAKE_CASE_   :      Optional[int]			=   '''nq-training.jsonl''' if PROCESS_TRAIN == '''true''' else '''nq-validation.jsonl'''
		save_to_disk(data, file_name=cache_file_name)











from math import factorial



def  SCREAMING_SNAKE_CASE				(      snake_case						,   snake_case							)		->    int:

		# If either of the conditions are true, the function is being asked
		# to calculate a factorial of a negative number, which is not possible
		if n < k or k < 0:
				raise ValueError('Please enter positive integers for n and k where n >= k'							)
		return factorial(snake_case							) // (factorial(snake_case							) * factorial(n - k							))


if __name__ == "__main__":
		print(
		    '''The number of five-card hands possible from a standard''',
		    F"""fifty-two card deck is: {combinations(52, 5)}\n""",
		)

		print(
		    '''If a class of 40 students must be arranged into groups of''',
		    F"""4 for group projects, there are {combinations(40, 4)} ways""",
		    '''to arrange them.\n''',
		)

		print(
		    '''If 10 teams are competing in a Formula One race, there''',
		    F"""are {combinations(10, 3)} ways that first, second and""",
		    '''third place can be awarded.''',
		)




