import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch

				from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
				from transformers.models.esm.modeling_esm import (
				    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
				    EsmEmbeddings,
				    create_position_ids_from_input_ids,
				)



class   lowercase_       :



							def __init__(		self ,							lowercase_ ,							lowercase_=13 ,							lowercase_=7 ,							lowercase_=False ,							lowercase_=True ,							lowercase_=False ,							lowercase_=True ,							lowercase_=33 ,							lowercase_=32 ,							lowercase_=5 ,							lowercase_=4 ,							lowercase_=37 ,							lowercase_="gelu" ,							lowercase_=0.1 ,							lowercase_=0.1 ,							lowercase_=512 ,							lowercase_=16 ,							lowercase_=2 ,							lowercase_=0.02 ,							lowercase_=3 ,							lowercase_=4 ,							lowercase_=None ,							):
													_snake_case		:	Any							=       parent
													_snake_case		:	Optional[int]							=       batch_size
													_snake_case		:	Any							=       seq_length
													_snake_case		:	Dict							=       is_training
													_snake_case		:	Union[str, Any]							=       use_input_mask
													_snake_case		:	Union[str, Any]							=       use_token_type_ids
													_snake_case		:	Union[str, Any]							=       use_labels
													_snake_case		:	Dict							=       vocab_size
													_snake_case		:	List[str]							=       hidden_size
													_snake_case		:	Tuple							=       num_hidden_layers
													_snake_case		:	Optional[Any]							=       num_attention_heads
													_snake_case		:	List[str]							=       intermediate_size
													_snake_case		:	int							=       hidden_act
													_snake_case		:	Any							=       hidden_dropout_prob
													_snake_case		:	List[Any]							=       attention_probs_dropout_prob
													_snake_case		:	List[str]							=       max_position_embeddings
													_snake_case		:	List[str]							=       type_vocab_size
													_snake_case		:	Optional[Any]							=       type_sequence_label_size
													_snake_case		:	List[str]							=       initializer_range
													_snake_case		:	List[Any]							=       num_labels
													_snake_case		:	List[str]							=       num_choices
													_snake_case		:	Any							=       scope



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Optional[Any]							=       ids_tensor([self.batch_size, self.seq_length] ,							self.vocab_size							)

													_snake_case		:	Any							=       None
													if self.use_input_mask:
																			_snake_case		:	List[Any]							=       random_attention_mask([self.batch_size, self.seq_length]							)

													_snake_case		:	int							=       None
													_snake_case		:	int							=       None
													_snake_case		:	Optional[int]							=       None
													if self.use_labels:
																			_snake_case		:	Tuple							=       ids_tensor([self.batch_size] ,							self.type_sequence_label_size							)
																			_snake_case		:	Tuple							=       ids_tensor([self.batch_size, self.seq_length] ,							self.num_labels							)
																			_snake_case		:	Tuple							=       ids_tensor([self.batch_size] ,							self.num_choices							)

													_snake_case		:	Optional[Any]							=       self.get_config()

													return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



							def 			UpperCamelCase			(		self							):
													return EsmConfig(
													    vocab_size=self.vocab_size ,							hidden_size=self.hidden_size ,							pad_token_id=1 ,							num_hidden_layers=self.num_hidden_layers ,							num_attention_heads=self.num_attention_heads ,							intermediate_size=self.intermediate_size ,							hidden_act=self.hidden_act ,							hidden_dropout_prob=self.hidden_dropout_prob ,							attention_probs_dropout_prob=self.attention_probs_dropout_prob ,							max_position_embeddings=self.max_position_embeddings ,							type_vocab_size=self.type_vocab_size ,							initializer_range=self.initializer_range ,							)



							def 			UpperCamelCase			(		self ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_							):
													_snake_case		:	Optional[Any]							=       EsmModel(config=UpperCamelCase_							)
													model.to(UpperCamelCase_							)
													model.eval()
													_snake_case		:	str							=       model(UpperCamelCase_ ,							attention_mask=UpperCamelCase_							)
													_snake_case		:	List[Any]							=       model(UpperCamelCase_							)
													_snake_case		:	Union[str, Any]							=       model(UpperCamelCase_							)

													self.parent.assertEqual(result.last_hidden_state.shape ,							(self.batch_size, self.seq_length, self.hidden_size)							)
													self.parent.assertEqual(result.pooler_output.shape ,							(self.batch_size, self.hidden_size)							)



							def 			UpperCamelCase			(		self ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_							):
													_snake_case		:	List[str]							=       EsmForMaskedLM(config=UpperCamelCase_							)
													model.to(UpperCamelCase_							)
													model.eval()
													_snake_case		:	Dict							=       model(UpperCamelCase_ ,							attention_mask=UpperCamelCase_ ,							labels=UpperCamelCase_							)
													self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length, self.vocab_size)							)



							def 			UpperCamelCase			(		self ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_ ,							lowercase_							):
													_snake_case		:	Dict							=       self.num_labels
													_snake_case		:	int							=       EsmForTokenClassification(config=UpperCamelCase_							)
													model.to(UpperCamelCase_							)
													model.eval()
													_snake_case		:	str							=       model(UpperCamelCase_ ,							attention_mask=UpperCamelCase_ ,							labels=UpperCamelCase_							)
													self.parent.assertEqual(result.logits.shape ,							(self.batch_size, self.seq_length, self.num_labels)							)



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Dict							=       self.prepare_config_and_inputs()
													(
													    (
													    _snake_case  
													)							,(
													    _snake_case  
													)							,(
													    _snake_case  
													)							,(
													    _snake_case  
													)							,(
													    _snake_case  
													)							,(
													    _snake_case  
													)							,  
													)		:	int							=       config_and_inputs
													_snake_case		:	Any							=       {"input_ids": input_ids, "attention_mask": input_mask}
													return config, inputs_dict




@require_torch
class   lowercase_       (	a_		,					a_		,					unittest.TestCase						):
							_lowerCamelCase		=						False

							_lowerCamelCase		=						(
							    (
							        EsmForMaskedLM,
							        EsmModel,
							        EsmForSequenceClassification,
							        EsmForTokenClassification,
							    )
							    if is_torch_available()
							    else ()
							)
							_lowerCamelCase		=						()
							_lowerCamelCase		=						(
							    {
							        'feature-extraction': EsmModel,
							        'fill-mask': EsmForMaskedLM,
							        'text-classification': EsmForSequenceClassification,
							        'token-classification': EsmForTokenClassification,
							        'zero-shot': EsmForSequenceClassification,
							    }
							    if is_torch_available()
							    else {}
							)
							_lowerCamelCase		=						True



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Optional[Any]							=       EsmModelTester(self							)
													_snake_case		:	List[str]							=       ConfigTester(self ,							config_class=UpperCamelCase_ ,							hidden_size=37							)



							def 			UpperCamelCase			(		self							):
													self.config_tester.run_common_tests()



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Dict							=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*UpperCamelCase_							)



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Union[str, Any]							=       self.model_tester.prepare_config_and_inputs()
													for type in ["absolute", "relative_key", "relative_key_query"]:
																			_snake_case		:	Dict							=       type
																			self.model_tester.create_and_check_model(*UpperCamelCase_							)



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Union[str, Any]							=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase_							)



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Optional[Any]							=       self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_token_classification(*UpperCamelCase_							)



							@slow
							def 			UpperCamelCase			(		self							):

													for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			_snake_case		:	Tuple							=       EsmModel.from_pretrained(UpperCamelCase_							)
																			self.assertIsNotNone(UpperCamelCase_							)



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Optional[int]							=       self.model_tester.prepare_config_and_inputs()[0]
													_snake_case		:	List[str]							=       EsmEmbeddings(config=UpperCamelCase_							)

													_snake_case		:	Union[str, Any]							=       torch.as_tensor([[12, 31, 13, model.padding_idx]]							)
													_snake_case		:	str							=       torch.as_tensor(
													    [
													        [
													            0 + model.padding_idx + 1,
													            1 + model.padding_idx + 1,
													            2 + model.padding_idx + 1,
													            model.padding_idx,
													        ]
													    ]							)
													_snake_case		:	str							=       create_position_ids_from_input_ids(UpperCamelCase_ ,							model.padding_idx							)
													self.assertEqual(position_ids.shape ,							expected_positions.shape							)
													self.assertTrue(torch.all(torch.eq(UpperCamelCase_ ,							UpperCamelCase_							)							)							)



							def 			UpperCamelCase			(		self							):
													_snake_case		:	Optional[int]							=       self.model_tester.prepare_config_and_inputs()[0]
													_snake_case		:	Dict							=       EsmEmbeddings(config=UpperCamelCase_							)

													_snake_case		:	Union[str, Any]							=       torch.empty(2 ,							4 ,							30							)
													_snake_case		:	List[str]							=       [
													    0 + embeddings.padding_idx + 1,
													    1 + embeddings.padding_idx + 1,
													    2 + embeddings.padding_idx + 1,
													    3 + embeddings.padding_idx + 1,
													]
													_snake_case		:	List[Any]							=       torch.as_tensor([expected_single_positions, expected_single_positions]							)
													_snake_case		:	Optional[int]							=       embeddings.create_position_ids_from_inputs_embeds(UpperCamelCase_							)
													self.assertEqual(position_ids.shape ,							expected_positions.shape							)
													self.assertTrue(torch.all(torch.eq(UpperCamelCase_ ,							UpperCamelCase_							)							)							)



							@unittest.skip("Esm does not support embedding resizing"							)
							def 			UpperCamelCase			(		self							):
													pass



							@unittest.skip("Esm does not support embedding resizing"							)
							def 			UpperCamelCase			(		self							):
													pass



							@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."							)
							def 			UpperCamelCase			(		self							):
													pass




@require_torch
class   lowercase_       (	a_						):



							@slow
							def 			UpperCamelCase			(		self							):

													with torch.no_grad():
																			_snake_case		:	Optional[int]							=       EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"							)
																			model.eval()
																			_snake_case		:	str							=       torch.tensor([[0, 1, 2, 3, 4, 5]]							)
																			_snake_case		:	List[Any]							=       model(UpperCamelCase_							)[0]

																			_snake_case		:	Union[str, Any]							=       33

																			_snake_case		:	str							=       torch.Size((1, 6, vocab_size)							)
																			self.assertEqual(output.shape ,							UpperCamelCase_							)

																			_snake_case		:	int							=       torch.tensor(
																			    [[[8.9_215, -10.5_898, -6.4_671], [-6.3_967, -13.9_114, -1.1_212], [-7.7_812, -13.9_516, -3.7_406]]]							)
																			self.assertTrue(torch.allclose(output[:, :3, :3] ,							UpperCamelCase_ ,							atol=1e-4							)							)



							@slow
							def 			UpperCamelCase			(		self							):

													with torch.no_grad():
																			_snake_case		:	Optional[Any]							=       EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"							)
																			model.eval()

																			_snake_case		:	Optional[int]							=       torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]							)
																			_snake_case		:	int							=       model(UpperCamelCase_							)[0]
																			# compare the actual values for a slice.
																			_snake_case		:	Any							=       torch.tensor(
																			    [[[0.1_444, 0.5_413, 0.3_248], [0.3_034, 0.0_053, 0.3_108], [0.3_228, -0.2_499, 0.3_415]]]							)
																			self.assertTrue(torch.allclose(output[:, :3, :3] ,							UpperCamelCase_ ,							atol=1e-4							)							)







"""simple docstring"""
import os

import jsonlines
import numpy as np
from tqdm import tqdm


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


def 							A						(				__snake_case: Any  )    ->		Optional[int]:

       """simple docstring"""





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

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

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

       __magic_name__      				=    ['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 							A						(				__snake_case: Any ,     __snake_case: str=False  )    ->		Optional[Any]:

       """simple docstring"""





       __magic_name__      				=    _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
              __magic_name__      				=    example['document']['tokens']
              __magic_name__      				=    []
              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": -1_0_0,  # ignore index in cross-entropy
                      "end_token": -1_0_0,  # 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

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

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

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

       # checking above code
       if assertion:
              __magic_name__      				=    doc['is_html'][answer['start_token'] : answer['end_token']]
              __magic_name__      				=    doc['token'][answer['start_token'] : answer['end_token']]
              __magic_name__      				=    ' '.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 							A						(				__snake_case: List[Any] ,     __snake_case: int ,     __snake_case: Tuple=2_0_4_8 ,     __snake_case: List[str]=4_0_9_6 ,     __snake_case: Optional[int]=True  )    ->		Any:

       """simple docstring"""





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

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

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

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

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

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

       __magic_name__      				=    input_ids[:q_len]
       __magic_name__      				=    range(__snake_case ,     len(__snake_case  ) ,     max_length - doc_stride  )

       __magic_name__      				=    []
       __magic_name__      				=    []
       __magic_name__      				=    []
       __magic_name__      				=    []  # null, yes, no, long, short
       for i in doc_start_indices:
              __magic_name__      				=    i + max_length - q_len
              __magic_name__      				=    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:
                     __magic_name__      				=    start_token - i + q_len
                     __magic_name__      				=    end_token - i + q_len
                     answers_category.append(answer['category'][0]  )  # ["short"] -> "short"
              else:
                     __magic_name__      				=    -1_0_0
                     __magic_name__      				=    -1_0_0
                     answers_category.append('null'  )
              __magic_name__      				=    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 							A						(				__snake_case: Tuple ,     __snake_case: Dict ,     __snake_case: Any=2_0_4_8 ,     __snake_case: List[Any]=4_0_9_6 ,     __snake_case: Any=False  )    ->		Tuple:

       """simple docstring"""





       __magic_name__      				=    get_strided_contexts_and_ans(
           __snake_case ,     __snake_case ,     doc_stride=__snake_case ,     max_length=__snake_case ,     assertion=__snake_case ,     )

       return example


def 							A						(				__snake_case: Optional[int] ,     __snake_case: Union[str, Any]  )    ->		Dict:

       """simple docstring"""





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

 snake_case : Any     		=		load_dataset("""natural_questions""")
 snake_case : List[Any]     		=		BigBirdTokenizer.from_pretrained("""google/bigbird-roberta-base""")

 snake_case : Tuple     		=		data["""train""" if PROCESS_TRAIN == """true""" else """validation"""]

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

 np.random.seed(SEED)
 snake_case : Union[str, Any]     		=		"""nq-training.jsonl""" if PROCESS_TRAIN == """true""" else """nq-validation.jsonl"""
 save_to_disk(data, file_name=cache_file_name)