





"""simple docstring"""





import os

import jsonlines
import numpy as np
from tqdm import tqdm


lowercase__				:Tuple				   =       2_0_4_8
lowercase__				:Any				   =       4_0_9_6
lowercase__				:Optional[Any]				   =       4_2
lowercase__				:int				   =       os.environ.pop('PROCESS_TRAIN', 'false')
lowercase__				:List[Any]				   =       {'null': 0, 'short': 1, 'long': 2, 'yes': 3, 'no': 4}


def       lowerCamelCase_       (   UpperCAmelCase_  )						->Optional[int]:





  """simple docstring"""


  def choose_first(UpperCAmelCase_  ,			UpperCAmelCase_=False  ):
    assert isinstance(UpperCAmelCase_  ,			UpperCAmelCase_  )
    if len(UpperCAmelCase_  ) == 1:
      __UpperCAmelCase    :						List[Any]       =						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    :						List[Any]       =						{k: [a[k]] for k in a}
      if len(a['''start_token''']  ) > 0:
        break
    return a

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

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

  __UpperCAmelCase    :						Tuple       =						['''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_  ,			UpperCAmelCase_=False  )						->List[str]:





  """simple docstring"""


  __UpperCAmelCase    :						List[Any]       =						_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
    __UpperCAmelCase    :						int       =						example['''document''']['''tokens''']
    __UpperCAmelCase    :						int       =						[]
    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_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

  __UpperCAmelCase    :						List[Any]       =						['''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    :						Optional[int]       =						example['''document''']['''tokens''']
  __UpperCAmelCase    :						Optional[int]       =						answer['''start_token''']
  __UpperCAmelCase    :						Union[str, Any]       =						answer['''end_token''']

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

  # checking above code
  if assertion:
    __UpperCAmelCase    :						Optional[int]       =						doc['''is_html'''][answer['''start_token'''] : answer['''end_token''']]
    __UpperCAmelCase    :						List[str]       =						doc['''token'''][answer['''start_token'''] : answer['''end_token''']]
    __UpperCAmelCase    :						Optional[Any]       =						''' '''.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_  ,			UpperCAmelCase_  ,			UpperCAmelCase_=20_48  ,			UpperCAmelCase_=40_96  ,			UpperCAmelCase_=True  )						->Any:





  """simple docstring"""


  __UpperCAmelCase    :						Tuple       =						get_context_and_ans(UpperCAmelCase_  ,			assertion=UpperCAmelCase_  )
  __UpperCAmelCase    :						int       =						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    :						str       =						tokenizer(example['''question''']['''text''']  ,			out['''context''']  ).input_ids
  __UpperCAmelCase    :						Tuple       =						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    :						str       =						[]
    __UpperCAmelCase    :						str       =						[]
    __UpperCAmelCase    :						List[str]       =						input_ids[:q_len]
    __UpperCAmelCase    :						Any       =						range(UpperCAmelCase_  ,			len(UpperCAmelCase_  )  ,			max_length - doc_stride  )
    for i in doc_start_indices:
      __UpperCAmelCase    :						Dict       =						i + max_length - q_len
      __UpperCAmelCase    :						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_00] * len(UpperCAmelCase_  ),
            "end_token": [-1_00] * len(UpperCAmelCase_  ),
            "category": category,
        },
    }

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

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

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

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

  __UpperCAmelCase    :						List[str]       =						[]
  __UpperCAmelCase    :						int       =						[]
  __UpperCAmelCase    :						Optional[int]       =						[]
  __UpperCAmelCase    :						List[str]       =						[]  # null, yes, no, long, short
  for i in doc_start_indices:
    __UpperCAmelCase    :						Union[str, Any]       =						i + max_length - q_len
    __UpperCAmelCase    :						Any       =						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    :						Any       =						start_token - i + q_len
      __UpperCAmelCase    :						List[Any]       =						end_token - i + q_len
      answers_category.append(answer['''category'''][0]  )  # ["short"] -> "short"
    else:
      __UpperCAmelCase    :						Union[str, Any]       =						-1_00
      __UpperCAmelCase    :						List[str]       =						-1_00
      answers_category.append('''null'''  )
    __UpperCAmelCase    :						str       =						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_  ,			UpperCAmelCase_  ,			UpperCAmelCase_=20_48  ,			UpperCAmelCase_=40_96  ,			UpperCAmelCase_=False  )						->List[Any]:





  """simple docstring"""


  __UpperCAmelCase    :						Tuple       =						get_strided_contexts_and_ans(
      UpperCAmelCase_  ,			UpperCAmelCase_  ,			doc_stride=UpperCAmelCase_  ,			max_length=UpperCAmelCase_  ,			assertion=UpperCAmelCase_  ,			)

  return example


def       lowerCamelCase_       (   UpperCAmelCase_  ,			UpperCAmelCase_  )						->Any:





  """simple docstring"""


  with jsonlines.open(UpperCAmelCase_  ,			'''a'''  ) as writer:
    for example in tqdm(UpperCAmelCase_  ,			total=len(UpperCAmelCase_  )  ,			desc='''Saving samples ... '''  ):
      __UpperCAmelCase    :						List[str]       =						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

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

  lowercase__				:Union[str, Any]				   =       data['train' if PROCESS_TRAIN == 'true' else 'validation']

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

  np.random.seed(SEED)
  lowercase__				:Tuple				   =       'nq-training.jsonl' if PROCESS_TRAIN == 'true' else 'nq-validation.jsonl'
  save_to_disk(data, file_name=cache_file_name)






"""simple docstring"""





from __future__ import annotations

import string
from itertools import cycle, product
from pathlib import Path

lowercase__				:str				   =       (
    string.ascii_letters + string.digits + string.punctuation + string.whitespace
)
lowercase__				:list[int]				   =       [ord(letter) for letter in string.ascii_lowercase]
lowercase__				:set[int]				   =       {ord(char) for char in VALID_CHARS}

lowercase__				:list[str]				   =       ["the", "be", "to", "of", "and", "in", "that", "have"]


def       lowerCamelCase_       (   UpperCAmelCase_  ,			UpperCAmelCase_  )						->str | None:





  """simple docstring"""


  __UpperCAmelCase    :						str       =						""
  __UpperCAmelCase    :						int
  __UpperCAmelCase    :						int
  __UpperCAmelCase    :						int

  for keychar, cipherchar in zip(cycle(UpperCAmelCase_  )  ,			UpperCAmelCase_  ):
    __UpperCAmelCase    :						Optional[Any]       =						cipherchar ^ keychar
    if decodedchar not in VALID_INTS:
      return None
    decoded += chr(UpperCAmelCase_  )

  return decoded


def       lowerCamelCase_       (   UpperCAmelCase_  )						->list[str]:





  """simple docstring"""


  __UpperCAmelCase    :						list[str]       =						[]
  for key in product(UpperCAmelCase_  ,			repeat=3  ):
    __UpperCAmelCase    :						List[Any]       =						try_key(UpperCAmelCase_  ,			UpperCAmelCase_  )
    if encoded is not None:
      possibles.append(UpperCAmelCase_  )
  return possibles


def       lowerCamelCase_       (   UpperCAmelCase_  ,			UpperCAmelCase_  )						->list[str]:





  """simple docstring"""


  return [possible for possible in possibles if common_word in possible.lower()]


def       lowerCamelCase_       (   UpperCAmelCase_ = "p059_cipher.txt"  )						->int:





  """simple docstring"""


  __UpperCAmelCase    :						list[int]
  __UpperCAmelCase    :						list[str]
  __UpperCAmelCase    :						str
  __UpperCAmelCase    :						str
  __UpperCAmelCase    :						str       =						Path(UpperCAmelCase_  ).parent.joinpath(UpperCAmelCase_  ).read_text(encoding='''utf-8'''  )

  __UpperCAmelCase    :						List[str]       =						[int(UpperCAmelCase_  ) for number in data.strip().split(''','''  )]

  __UpperCAmelCase    :						List[Any]       =						filter_valid_chars(UpperCAmelCase_  )
  for common_word in COMMON_WORDS:
    __UpperCAmelCase    :						Optional[int]       =						filter_common_word(UpperCAmelCase_  ,			UpperCAmelCase_  )
    if len(UpperCAmelCase_  ) == 1:
      break

  __UpperCAmelCase    :						Optional[Any]       =						possibles[0]
  return sum(ord(UpperCAmelCase_  ) for char in decoded_text  )


if __name__ == "__main__":
  print(f"""{solution() = }""")