



"""simple docstring"""





import os

import jsonlines
import numpy as np
from tqdm import tqdm


lowerCamelCase  			=		2_048
lowerCamelCase  			=		4_096
lowerCamelCase  			=		42
lowerCamelCase  			=		os.environ.pop("""PROCESS_TRAIN""", """false""")
lowerCamelCase  			=		{'''null''': 0, '''short''': 1, '''long''': 2, '''yes''': 3, '''no''': 4}
def   a__      ( lowerCAmelCase__ ):
      def choose_first(lowerCAmelCase__      ,     lowerCAmelCase__=False ):
            assert isinstance(lowerCAmelCase__      ,     lowerCAmelCase__ )
            if len(lowerCAmelCase__ ) == 1:
                  UpperCAmelCase_     					=				answer[0]
                  return {k: [answer[k]] for k in answer} if is_long_answer else answer
            for a in answer:
                  if is_long_answer:
                        UpperCAmelCase_     					=				{k: [a[k]] for k in a}
                  if len(a["start_token"] ) > 0:
                        break
            return a

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

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

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

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

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

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

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

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

      return {
          "context": " ".join(lowerCAmelCase__ ),
          "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__      ( lowerCAmelCase__      ,     lowerCAmelCase__      ,     lowerCAmelCase__=2048      ,     lowerCAmelCase__=4096      ,     lowerCAmelCase__=True ):
      # overlap will be of doc_stride - q_len

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

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

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

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

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

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

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

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

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

      UpperCAmelCase_     					=				input_ids[:q_len]
      UpperCAmelCase_     					=				range(lowerCAmelCase__      ,     len(lowerCAmelCase__ )      ,     max_length - doc_stride )

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

            answers_start_token.append(lowerCAmelCase__ )
            answers_end_token.append(lowerCAmelCase__ )
            if assertion:
                  if new != old and new != [tokenizer.cls_token_id]:
                        print("ISSUE in strided for ID:"      ,     example["id"] )
                        print("New:"      ,     tokenizer.decode(lowerCAmelCase__ ) )
                        print("Old:"      ,     tokenizer.decode(lowerCAmelCase__ )      ,     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__      ( lowerCAmelCase__      ,     lowerCAmelCase__      ,     lowerCAmelCase__=2048      ,     lowerCAmelCase__=4096      ,     lowerCAmelCase__=False ):
      UpperCAmelCase_     					=				get_strided_contexts_and_ans(
          lowerCAmelCase__      ,     lowerCAmelCase__      ,     doc_stride=lowerCAmelCase__      ,     max_length=lowerCAmelCase__      ,     assertion=lowerCAmelCase__      ,     )

      return example
def   a__      ( lowerCAmelCase__      ,     lowerCAmelCase__ ):
      with jsonlines.open(lowerCAmelCase__      ,     "a" ) as writer:
            for example in tqdm(lowerCAmelCase__      ,     total=len(lowerCAmelCase__ )      ,     desc="Saving samples ... " ):
                  UpperCAmelCase_     					=				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)











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.''',
		)




