"""simple docstring"""
import os

import jsonlines
import numpy as np
from tqdm import tqdm


lowercase__ 	=     2048
lowercase__ 	=     4096
lowercase__ 	=     42
lowercase__ 	=     os.environ.pop('PROCESS_TRAIN', 'false')
lowercase__ 	=     {'null': 0, 'short': 1, 'long': 2, 'yes': 3, 'no': 4}







def 						__a       (     _SCREAMING_SNAKE_CASE       )     ->Optional[Any]:


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

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

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

  a__:					Any      							=     ["""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       (     _SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE=False       )     ->List[str]:
  a__:					List[Any]      							=     _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
    a__:					Optional[Any]      							=     example["""document"""]["""tokens"""]
    a__:					Union[str, Any]      							=     []
    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

  a__:					Union[str, 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

  a__:					str      							=     example["""document"""]["""tokens"""]
  a__:					Any      							=     answer["""start_token"""]
  a__:					Optional[int]      							=     answer["""end_token"""]

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

  # checking above code
  if assertion:
    a__:					str      							=     doc["""is_html"""][answer["""start_token"""] : answer["""end_token"""]]
    a__:					Union[str, Any]      							=     doc["""token"""][answer["""start_token"""] : answer["""end_token"""]]
    a__:					Dict      							=     """ """.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       (     _SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE=2048		,					_SCREAMING_SNAKE_CASE=4096		,					_SCREAMING_SNAKE_CASE=True       )     ->Dict:
  a__:					List[Any]      							=     get_context_and_ans(_lowerCamelCase		,					assertion=_lowerCamelCase       )
  a__:					Dict      							=     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__:					List[Any]      							=     tokenizer(example['question']['text']		,					out['context']       ).input_ids
  a__:					Optional[int]      							=     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__:					int      							=     []
    a__:					List[str]      							=     []
    a__:					Optional[int]      							=     input_ids[:q_len]
    a__:					Dict      							=     range(_lowerCamelCase		,					len(_lowerCamelCase       )		,					max_length - doc_stride       )
    for i in doc_start_indices:
      a__:					Dict      							=     i + max_length - q_len
      a__:					Tuple      							=     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,
        },
    }

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

  a__:					Union[str, Any]      							=     input_ids[answer["""start_token"""] : answer["""end_token"""] + 1]  # right & left are inclusive
  a__:					Optional[int]      							=     answer["""start_token"""]
  a__:					List[Any]      							=     answer["""end_token"""]

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

  a__:					List[Any]      							=     input_ids[:q_len]
  a__:					Tuple      							=     range(_lowerCamelCase		,					len(_lowerCamelCase       )		,					max_length - doc_stride       )

  a__:					Optional[Any]      							=     []
  a__:					List[Any]      							=     []
  a__:					Optional[int]      							=     []
  a__:					Tuple      							=     []  # null, yes, no, long, short
  for i in doc_start_indices:
    a__:					Dict      							=     i + max_length - q_len
    a__:					Union[str, 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:
      a__:					str      							=     start_token - i + q_len
      a__:					int      							=     end_token - i + q_len
      answers_category.append(answer['category'][0]       )  # ["short"] -> "short"
    else:
      a__:					Optional[Any]      							=     -100
      a__:					Optional[Any]      							=     -100
      answers_category.append('null'       )
    a__:					List[str]      							=     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       (     _SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE=2048		,					_SCREAMING_SNAKE_CASE=4096		,					_SCREAMING_SNAKE_CASE=False       )     ->int:
  a__:					Union[str, Any]      							=     get_strided_contexts_and_ans(
      _lowerCamelCase		,					_lowerCamelCase		,					doc_stride=_lowerCamelCase		,					max_length=_lowerCamelCase		,					assertion=_lowerCamelCase		,					)

  return example







def 						__a       (     _SCREAMING_SNAKE_CASE		,					_SCREAMING_SNAKE_CASE       )     ->Optional[Any]:


  with jsonlines.open(_lowerCamelCase		,					'a'       ) as writer:
    for example in tqdm(_lowerCamelCase		,					total=len(_lowerCamelCase       )		,					desc='Saving samples ... '       ):
      a__:					int      							=     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__ 	=     load_dataset('natural_questions')
      lowercase__ 	=     BigBirdTokenizer.from_pretrained('google/bigbird-roberta-base')

      lowercase__ 	=     data['train' if PROCESS_TRAIN == 'true' else 'validation']

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

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



"""simple docstring"""
import inspect
import unittest

from transformers import RegNetConfig
from transformers.file_utils import cached_property, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

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


if is_torch_available():
 import torch
 from torch import nn

 from transformers import RegNetForImageClassification, RegNetModel
 from transformers.models.regnet.modeling_regnet import REGNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
 from PIL import Image

 from transformers import AutoImageProcessor




class        __snake_case	:



      def __init__(						self   ,      lowercase   ,      lowercase=3   ,      lowercase=32   ,      lowercase=3   ,      lowercase=10   ,      lowercase=[10, 20, 30, 40]   ,      lowercase=[1, 1, 2, 1]   ,      lowercase=True   ,      lowercase=True   ,      lowercase="relu"   ,      lowercase=3   ,      lowercase=None   ,      )     ->      Any:





        '''simple docstring'''





        a__:					Union[str, Any]      							=     parent
        a__:					int      							=     batch_size
        a__:					List[Any]      							=     image_size
        a__:					Any      							=     num_channels
        a__:					Dict      							=     embeddings_size
        a__:					str      							=     hidden_sizes
        a__:					List[Any]      							=     depths
        a__:					Optional[int]      							=     is_training
        a__:					Optional[int]      							=     use_labels
        a__:					Tuple      							=     hidden_act
        a__:					Any      							=     num_labels
        a__:					Union[str, Any]      							=     scope
        a__:					Tuple      							=     len(lowercase)



      def   lowerCamelCase_       (						self)     ->      int:





        '''simple docstring'''





        a__:					Optional[int]      							=     floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

        a__:					str      							=     None
        if self.use_labels:
          a__:					Dict      							=     ids_tensor([self.batch_size]   ,      self.num_labels)

        a__:					int      							=     self.get_config()

        return config, pixel_values, labels



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





        '''simple docstring'''





        return RegNetConfig(
            num_channels=self.num_channels   ,      embeddings_size=self.embeddings_size   ,      hidden_sizes=self.hidden_sizes   ,      depths=self.depths   ,      hidden_act=self.hidden_act   ,      num_labels=self.num_labels   ,      )



      def   lowerCamelCase_       (						self   ,      lowercase   ,      lowercase   ,      lowercase)     ->      Any:





        '''simple docstring'''





        a__:					Dict      							=     RegNetModel(config=lowercase)
        model.to(lowercase)
        model.eval()
        a__:					List[str]      							=     model(lowercase)
        # expected last hidden states: B, C, H // 32, W // 32
        self.parent.assertEqual(
            result.last_hidden_state.shape   ,      (self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)   ,      )



      def   lowerCamelCase_       (						self   ,      lowercase   ,      lowercase   ,      lowercase)     ->      List[Any]:





        '''simple docstring'''





        a__:					Tuple      							=     self.num_labels
        a__:					Tuple      							=     RegNetForImageClassification(lowercase)
        model.to(lowercase)
        model.eval()
        a__:					Any      							=     model(lowercase   ,      labels=lowercase)
        self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.num_labels))



      def   lowerCamelCase_       (						self)     ->      str:





        '''simple docstring'''





        a__:					int      							=     self.prepare_config_and_inputs()
        a__							,   a__							,   a__:					List[Any]      							=     config_and_inputs
        a__:					Union[str, Any]      							=     {'pixel_values': pixel_values}
        return config, inputs_dict







@require_torch
class        __snake_case	(					__lowerCAmelCase	,				__lowerCAmelCase	,				unittest.TestCase    ):
      a__				=		(RegNetModel, RegNetForImageClassification) if is_torch_available() else ()
      a__				=		(
          {"""feature-extraction""": RegNetModel, """image-classification""": RegNetForImageClassification}
          if is_torch_available()
          else {}
      )

      a__				=		False
      a__				=		False
      a__				=		False
      a__				=		False



      def   lowerCamelCase_       (						self)     ->      Dict:





        '''simple docstring'''





        a__:					Dict      							=     RegNetModelTester(self)
        a__:					Dict      							=     ConfigTester(self   ,      config_class=lowercase   ,      has_text_modality=lowercase)



      def   lowerCamelCase_       (						self)     ->      Dict:





        '''simple docstring'''





        self.create_and_test_config_common_properties()
        self.config_tester.create_and_test_config_to_json_string()
        self.config_tester.create_and_test_config_to_json_file()
        self.config_tester.create_and_test_config_from_and_save_pretrained()
        self.config_tester.create_and_test_config_with_num_labels()
        self.config_tester.check_config_can_be_init_without_params()
        self.config_tester.check_config_arguments_init()



      def   lowerCamelCase_       (						self)     ->      Optional[Any]:





        '''simple docstring'''





        return



      @unittest.skip(reason='RegNet does not use inputs_embeds')
      def   lowerCamelCase_       (						self)     ->      List[str]:





        '''simple docstring'''





        pass



      @unittest.skip(reason='RegNet does not support input and output embeddings')
      def   lowerCamelCase_       (						self)     ->      Optional[Any]:





        '''simple docstring'''





        pass



      def   lowerCamelCase_       (						self)     ->      List[Any]:





        '''simple docstring'''





        a__							,   a__:					str      							=     self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          a__:					int      							=     model_class(lowercase)
          a__:					str      							=     inspect.signature(model.forward)
          # signature.parameters is an OrderedDict => so arg_names order is deterministic
          a__:					int      							=     [*signature.parameters.keys()]

          a__:					List[str]      							=     ['pixel_values']
          self.assertListEqual(arg_names[:1]   ,      lowercase)



      def   lowerCamelCase_       (						self)     ->      Optional[Any]:





        '''simple docstring'''





        a__:					Any      							=     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase)



      def   lowerCamelCase_       (						self)     ->      Dict:





        '''simple docstring'''





        a__							,   a__:					int      							=     self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
          a__:					int      							=     model_class(config=lowercase)
          for name, module in model.named_modules():
            if isinstance(lowercase   ,      (nn.BatchNormad, nn.GroupNorm)):
              self.assertTrue(
                  torch.all(module.weight == 1)   ,      msg=f'Parameter {name} of model {model_class} seems not properly initialized'   ,      )
              self.assertTrue(
                  torch.all(module.bias == 0)   ,      msg=f'Parameter {name} of model {model_class} seems not properly initialized'   ,      )



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





        '''simple docstring'''





        def check_hidden_states_output(lowercase   ,      lowercase   ,      lowercase):
          a__:					int      							=     model_class(lowercase)
          model.to(lowercase)
          model.eval()

          with torch.no_grad():
            a__:					int      							=     model(**self._prepare_for_class(lowercase   ,      lowercase))

          a__:					Any      							=     outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

          a__:					List[str]      							=     self.model_tester.num_stages
          self.assertEqual(len(lowercase)   ,      expected_num_stages + 1)

          # RegNet's feature maps are of shape (batch_size, num_channels, height, width)
          self.assertListEqual(
              list(hidden_states[0].shape[-2:])   ,      [self.model_tester.image_size // 2, self.model_tester.image_size // 2]   ,      )

        a__							,   a__:					Optional[Any]      							=     self.model_tester.prepare_config_and_inputs_for_common()
        a__:					Dict      							=     ['basic', 'bottleneck']
        for model_class in self.all_model_classes:
          for layer_type in layers_type:
            a__:					Dict      							=     layer_type
            a__:					Dict      							=     True
            check_hidden_states_output(lowercase   ,      lowercase   ,      lowercase)

            # check that output_hidden_states also work using config
            del inputs_dict["output_hidden_states"]
            a__:					str      							=     True

            check_hidden_states_output(lowercase   ,      lowercase   ,      lowercase)



      def   lowerCamelCase_       (						self)     ->      List[Any]:





        '''simple docstring'''





        a__:					Optional[Any]      							=     self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*lowercase)



      @slow
      def   lowerCamelCase_       (						self)     ->      str:





        '''simple docstring'''





        for model_name in REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          a__:					Tuple      							=     RegNetModel.from_pretrained(lowercase)
          self.assertIsNotNone(lowercase)







def 						__a       (     )     ->Dict:
  a__:					str      							=     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'       )
  return image




@require_torch
@require_vision
class        __snake_case	(					unittest.TestCase    ):



      @cached_property
      def   lowerCamelCase_       (						self)     ->      Any:





        '''simple docstring'''





        return (
            AutoImageProcessor.from_pretrained(REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0])
            if is_vision_available()
            else None
        )



      @slow
      def   lowerCamelCase_       (						self)     ->      Optional[int]:





        '''simple docstring'''





        a__:					Tuple      							=     RegNetForImageClassification.from_pretrained(REGNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]).to(lowercase)

        a__:					Tuple      							=     self.default_image_processor
        a__:					str      							=     prepare_img()
        a__:					Optional[int]      							=     image_processor(images=lowercase   ,      return_tensors='pt').to(lowercase)

        # forward pass
        with torch.no_grad():
          a__:					List[Any]      							=     model(**lowercase)

        # verify the logits
        a__:					Tuple      							=     torch.Size((1, 10_00))
        self.assertEqual(outputs.logits.shape   ,      lowercase)

        a__:					Optional[Any]      							=     torch.tensor([-0.4180, -1.5051, -3.4836]).to(lowercase)

        self.assertTrue(torch.allclose(outputs.logits[0, :3]   ,      lowercase   ,      atol=1e-4))


