





import os

import jsonlines
import numpy as np
from tqdm import tqdm


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







def    __lowerCamelCase					(							__a   :Optional[Any]    )		->				Dict:






		"""simple docstring"""

		def choose_first(__a   :int    ,     __a   :List[str]=False    ):
				assert isinstance(__a    ,     __a    )
				if len(__a    ) == 1:
						A__        =   answer[0]
						return {k: [answer[k]] for k in answer} if is_long_answer else answer
				for a in answer:
						if is_long_answer:
								A__        =   {k: [a[k]] for k in a}
						if len(a["""start_token"""]    ) > 0:
								break
				return a

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

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

		A__        =   ["""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    __lowerCamelCase					(							__a   :Union[str, Any]    ,     __a   :Tuple=False    )		->				Union[str, Any]:






		"""simple docstring"""

		A__        =   _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
				A__        =   example["""document"""]["""tokens"""]
				A__        =   []
				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_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

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

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

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

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






		"""simple docstring"""

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

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

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

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

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

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

		A__        =   input_ids[:q_len]
		A__        =   range(__a    ,     len(__a    )    ,     max_length - doc_stride    )

		A__        =   []
		A__        =   []
		A__        =   []
		A__        =   []  # null, yes, no, long, short
		for i in doc_start_indices:
				A__        =   i + max_length - q_len
				A__        =   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:
						A__        =   start_token - i + q_len
						A__        =   end_token - i + q_len
						answers_category.append(answer["""category"""][0]    )  # ["short"] -> "short"
				else:
						A__        =   -1_0_0
						A__        =   -1_0_0
						answers_category.append("""null"""    )
				A__        =   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    __lowerCamelCase					(							__a   :int    ,     __a   :Any    ,     __a   :Optional[int]=2_0_4_8    ,     __a   :Dict=4_0_9_6    ,     __a   :int=False    )		->				Union[str, Any]:






		"""simple docstring"""

		A__        =   get_strided_contexts_and_ans(
		    __a    ,     __a    ,     doc_stride=__a    ,     max_length=__a    ,     assertion=__a    ,     )

		return example







def    __lowerCamelCase					(							__a   :List[str]    ,     __a   :Any    )		->				Dict:






		"""simple docstring"""

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

					A    :					Optional[int] =      load_dataset('''natural_questions''')
					A    :					Tuple =      BigBirdTokenizer.from_pretrained('''google/bigbird-roberta-base''')

					A    :					Optional[int] =      data['''train''' if PROCESS_TRAIN == '''true''' else '''validation''']

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

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










import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
					import jax.numpy as jnp

					from transformers.models.albert.modeling_flax_albert import (
					    FlaxAlbertForMaskedLM,
					    FlaxAlbertForMultipleChoice,
					    FlaxAlbertForPreTraining,
					    FlaxAlbertForQuestionAnswering,
					    FlaxAlbertForSequenceClassification,
					    FlaxAlbertForTokenClassification,
					    FlaxAlbertModel,
					)


class   A			(unittest.TestCase    ):
			'''simple docstring'''




			def __init__(			self   :    List[str]							,	__lowerCAmelCase   :    int							,	__lowerCAmelCase   :    List[str]=13							,	__lowerCAmelCase   :    int=7							,	__lowerCAmelCase   :    Tuple=True							,	__lowerCAmelCase   :    int=True							,	__lowerCAmelCase   :    Dict=True							,	__lowerCAmelCase   :    Union[str, Any]=True							,	__lowerCAmelCase   :    List[Any]=99							,	__lowerCAmelCase   :    Optional[Any]=32							,	__lowerCAmelCase   :    Optional[Any]=5							,	__lowerCAmelCase   :    Tuple=4							,	__lowerCAmelCase   :    Any=37							,	__lowerCAmelCase   :    str="gelu"							,	__lowerCAmelCase   :    Optional[Any]=0.1							,	__lowerCAmelCase   :    Union[str, Any]=0.1							,	__lowerCAmelCase   :    List[str]=5_12							,	__lowerCAmelCase   :    int=16							,	__lowerCAmelCase   :    Optional[int]=2							,	__lowerCAmelCase   :    List[Any]=0.0_2							,	__lowerCAmelCase   :    Tuple=4							,	)	->	Dict:







					"""simple docstring"""






					A__        =   parent
					A__        =   batch_size
					A__        =   seq_length
					A__        =   is_training
					A__        =   use_attention_mask
					A__        =   use_token_type_ids
					A__        =   use_labels
					A__        =   vocab_size
					A__        =   hidden_size
					A__        =   num_hidden_layers
					A__        =   num_attention_heads
					A__        =   intermediate_size
					A__        =   hidden_act
					A__        =   hidden_dropout_prob
					A__        =   attention_probs_dropout_prob
					A__        =   max_position_embeddings
					A__        =   type_vocab_size
					A__        =   type_sequence_label_size
					A__        =   initializer_range
					A__        =   num_choices






			def    a_       (			self   :    Any   )	->	str:







					"""simple docstring"""






					A__        =   ids_tensor([self.batch_size, self.seq_length]							,	self.vocab_size   )

					A__        =   None
					if self.use_attention_mask:
							A__        =   random_attention_mask([self.batch_size, self.seq_length]   )

					A__        =   None
					if self.use_token_type_ids:
							A__        =   ids_tensor([self.batch_size, self.seq_length]							,	self.type_vocab_size   )

					A__        =   AlbertConfig(
					    vocab_size=self.vocab_size							,	hidden_size=self.hidden_size							,	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							,	is_decoder=__lowerCAmelCase							,	initializer_range=self.initializer_range							,	)

					return config, input_ids, token_type_ids, attention_mask







			def    a_       (			self   :    Optional[int]   )	->	Optional[int]:







					"""simple docstring"""






					A__        =   self.prepare_config_and_inputs()
					A__      ,			A__      ,			A__      ,			A__        =   config_and_inputs
					A__        =   {"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": attention_mask}
					return config, inputs_dict


@require_flax
class   A			(SCREAMING_SNAKE_CASE   ,     unittest.TestCase    ):
			'''simple docstring'''




			__lowerCamelCase	:       str							=    (
			    (
			        FlaxAlbertModel,
			        FlaxAlbertForPreTraining,
			        FlaxAlbertForMaskedLM,
			        FlaxAlbertForMultipleChoice,
			        FlaxAlbertForQuestionAnswering,
			        FlaxAlbertForSequenceClassification,
			        FlaxAlbertForTokenClassification,
			        FlaxAlbertForQuestionAnswering,
			    )
			    if is_flax_available()
			    else ()
			)






			def    a_       (			self   :    str   )	->	Optional[int]:







					"""simple docstring"""






					A__        =   FlaxAlbertModelTester(self   )







			@slow
			def    a_       (			self   :    int   )	->	Tuple:







					"""simple docstring"""






					for model_class_name in self.all_model_classes:
							A__        =   model_class_name.from_pretrained("""albert-base-v2"""   )
							A__        =   model(np.ones((1, 1)   )   )
							self.assertIsNotNone(__lowerCAmelCase   )


@require_flax
class   A			(unittest.TestCase    ):
			'''simple docstring'''




			@slow
			def    a_       (			self   :    Dict   )	->	List[Any]:







					"""simple docstring"""






					A__        =   FlaxAlbertModel.from_pretrained("""albert-base-v2"""   )
					A__        =   np.array([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]   )
					A__        =   np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]   )
					A__        =   model(__lowerCAmelCase							,	attention_mask=__lowerCAmelCase   )[0]
					A__        =   (1, 11, 7_68)
					self.assertEqual(output.shape							,	__lowerCAmelCase   )
					A__        =   np.array(
					    [[[-0.6_5_1_3, 1.5_0_3_5, -0.2_7_6_6], [-0.6_5_1_5, 1.5_0_4_6, -0.2_7_8_0], [-0.6_5_1_2, 1.5_0_4_9, -0.2_7_8_4]]]   )

					self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]							,	__lowerCAmelCase							,	atol=1e-4   )   )



