


'''simple docstring'''

import collections
import inspect
import unittest

from transformers import NatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
     import torch
     from torch import nn

     from transformers import NatBackbone, NatForImageClassification, NatModel
     from transformers.models.nat.modeling_nat import NAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
     from PIL import Image

     from transformers import AutoImageProcessor





class 					snake_case   :





  """simple docstring"""




  def __init__(   self      ,  UpperCamelCase      ,  UpperCamelCase=13      ,  UpperCamelCase=64      ,  UpperCamelCase=4      ,  UpperCamelCase=3      ,  UpperCamelCase=16      ,  UpperCamelCase=[1, 2, 1]      ,  UpperCamelCase=[2, 4, 8]      ,  UpperCamelCase=3      ,  UpperCamelCase=2.0      ,  UpperCamelCase=True      ,  UpperCamelCase=0.0      ,  UpperCamelCase=0.0      ,  UpperCamelCase=0.1      ,  UpperCamelCase="gelu"      ,  UpperCamelCase=True      ,  UpperCamelCase=0.02      ,  UpperCamelCase=1e-5      ,  UpperCamelCase=True      ,  UpperCamelCase=None      ,  UpperCamelCase=True      ,  UpperCamelCase=10      ,  UpperCamelCase=["stage1", "stage2"]      ,  UpperCamelCase=[1, 2]      ,  ):







   """simple docstring"""
   lowerCamelCase_								=  parent
   lowerCamelCase_								=  batch_size
   lowerCamelCase_								=  image_size
   lowerCamelCase_								=  patch_size
   lowerCamelCase_								=  num_channels
   lowerCamelCase_								=  embed_dim
   lowerCamelCase_								=  depths
   lowerCamelCase_								=  num_heads
   lowerCamelCase_								=  kernel_size
   lowerCamelCase_								=  mlp_ratio
   lowerCamelCase_								=  qkv_bias
   lowerCamelCase_								=  hidden_dropout_prob
   lowerCamelCase_								=  attention_probs_dropout_prob
   lowerCamelCase_								=  drop_path_rate
   lowerCamelCase_								=  hidden_act
   lowerCamelCase_								=  patch_norm
   lowerCamelCase_								=  layer_norm_eps
   lowerCamelCase_								=  initializer_range
   lowerCamelCase_								=  is_training
   lowerCamelCase_								=  scope
   lowerCamelCase_								=  use_labels
   lowerCamelCase_								=  num_labels
   lowerCamelCase_								=  out_features
   lowerCamelCase_								=  out_indices






  def   snake_case		(   self					):







   """simple docstring"""
   lowerCamelCase_								=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]					)

   lowerCamelCase_								=  None
   if self.use_labels:
    lowerCamelCase_								=  ids_tensor([self.batch_size]      ,  self.num_labels					)

   lowerCamelCase_								=  self.get_config()

   return config, pixel_values, labels






  def   snake_case		(   self					):







   """simple docstring"""
   return NatConfig(
       num_labels=self.num_labels      ,  image_size=self.image_size      ,  patch_size=self.patch_size      ,  num_channels=self.num_channels      ,  embed_dim=self.embed_dim      ,  depths=self.depths      ,  num_heads=self.num_heads      ,  kernel_size=self.kernel_size      ,  mlp_ratio=self.mlp_ratio      ,  qkv_bias=self.qkv_bias      ,  hidden_dropout_prob=self.hidden_dropout_prob      ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,  drop_path_rate=self.drop_path_rate      ,  hidden_act=self.hidden_act      ,  patch_norm=self.patch_norm      ,  layer_norm_eps=self.layer_norm_eps      ,  initializer_range=self.initializer_range      ,  out_features=self.out_features      ,  out_indices=self.out_indices      ,  )






  def   snake_case		(   self      ,  UpperCamelCase      ,  UpperCamelCase      ,  UpperCamelCase					):







   """simple docstring"""
   lowerCamelCase_								=  NatModel(config=UpperCamelCase					)
   model.to(UpperCamelCase					)
   model.eval()
   lowerCamelCase_								=  model(UpperCamelCase					)

   lowerCamelCase_								=  lowerCamelCase_								=  (config.image_size // config.patch_size) // (2 ** (len(config.depths					) - 1))
   lowerCamelCase_								=  int(config.embed_dim * 2 ** (len(config.depths					) - 1)					)

   self.parent.assertEqual(
       result.last_hidden_state.shape      ,  (self.batch_size, expected_height, expected_width, expected_dim)					)






  def   snake_case		(   self      ,  UpperCamelCase      ,  UpperCamelCase      ,  UpperCamelCase					):







   """simple docstring"""
   lowerCamelCase_								=  NatForImageClassification(UpperCamelCase					)
   model.to(UpperCamelCase					)
   model.eval()
   lowerCamelCase_								=  model(UpperCamelCase      ,  labels=UpperCamelCase					)
   self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.num_labels)					)

   # test greyscale images
   lowerCamelCase_								=  1
   lowerCamelCase_								=  NatForImageClassification(UpperCamelCase					)
   model.to(UpperCamelCase					)
   model.eval()

   lowerCamelCase_								=  floats_tensor([self.batch_size, 1, self.image_size, self.image_size]					)
   lowerCamelCase_								=  model(UpperCamelCase					)
   self.parent.assertEqual(result.logits.shape      ,  (self.batch_size, self.num_labels)					)






  def   snake_case		(   self      ,  UpperCamelCase      ,  UpperCamelCase      ,  UpperCamelCase					):







   """simple docstring"""
   lowerCamelCase_								=  NatBackbone(config=UpperCamelCase					)
   model.to(UpperCamelCase					)
   model.eval()
   lowerCamelCase_								=  model(UpperCamelCase					)

   # verify hidden states
   self.parent.assertEqual(len(result.feature_maps					)      ,  len(config.out_features					)					)
   self.parent.assertListEqual(list(result.feature_maps[0].shape					)      ,  [self.batch_size, model.channels[0], 16, 16]					)

   # verify channels
   self.parent.assertEqual(len(model.channels					)      ,  len(config.out_features					)					)

   # verify backbone works with out_features=None
   lowerCamelCase_								=  None
   lowerCamelCase_								=  NatBackbone(config=UpperCamelCase					)
   model.to(UpperCamelCase					)
   model.eval()
   lowerCamelCase_								=  model(UpperCamelCase					)

   # verify feature maps
   self.parent.assertEqual(len(result.feature_maps					)      ,  1					)
   self.parent.assertListEqual(list(result.feature_maps[0].shape					)      ,  [self.batch_size, model.channels[-1], 4, 4]					)

   # verify channels
   self.parent.assertEqual(len(model.channels					)      ,  1					)







  def   snake_case		(   self					):







   """simple docstring"""
   lowerCamelCase_								=  self.prepare_config_and_inputs()
   lowerCamelCase_			,lowerCamelCase_			,lowerCamelCase_								=  config_and_inputs
   lowerCamelCase_								=  {"pixel_values": pixel_values}
   return config, inputs_dict





@require_natten
@require_torch
class 					snake_case   (   lowercase      ,						lowercase      ,						unittest.TestCase      ):





  """simple docstring"""




  _lowerCamelCase		   =  (
      (
          NatModel,
          NatForImageClassification,
          NatBackbone,
      )
      if is_torch_available()
      else ()
  )
  _lowerCamelCase		   =  (
      {"feature-extraction": NatModel, "image-classification": NatForImageClassification}
      if is_torch_available()
      else {}
  )
  _lowerCamelCase		   =  False

  _lowerCamelCase		   =  False
  _lowerCamelCase		   =  False
  _lowerCamelCase		   =  False
  _lowerCamelCase		   =  False






  def   snake_case		(   self					):







   """simple docstring"""
   lowerCamelCase_								=  NatModelTester(self					)
   lowerCamelCase_								=  ConfigTester(self      ,  config_class=UpperCamelCase      ,  embed_dim=37					)






  def   snake_case		(   self					):







   """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   snake_case		(   self					):







   """simple docstring"""
   return






  def   snake_case		(   self					):







   """simple docstring"""
   lowerCamelCase_								=  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_model(*UpperCamelCase					)






  def   snake_case		(   self					):







   """simple docstring"""
   lowerCamelCase_								=  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_for_image_classification(*UpperCamelCase					)






  def   snake_case		(   self					):







   """simple docstring"""
   lowerCamelCase_								=  self.model_tester.prepare_config_and_inputs()
   self.model_tester.create_and_check_backbone(*UpperCamelCase					)






  @unittest.skip(reason="Nat does not use inputs_embeds"					)
  def   snake_case		(   self					):







   """simple docstring"""
   pass






  @unittest.skip(reason="Nat does not use feedforward chunking"					)
  def   snake_case		(   self					):







   """simple docstring"""
   pass






  def   snake_case		(   self					):







   """simple docstring"""
   lowerCamelCase_			,lowerCamelCase_								=  self.model_tester.prepare_config_and_inputs_for_common()

   for model_class in self.all_model_classes:
    lowerCamelCase_								=  model_class(UpperCamelCase					)
    self.assertIsInstance(model.get_input_embeddings()      ,  (nn.Module)					)
    lowerCamelCase_								=  model.get_output_embeddings()
    self.assertTrue(x is None or isinstance(UpperCamelCase      ,  nn.Linear					)					)






  def   snake_case		(   self					):







   """simple docstring"""
   lowerCamelCase_			,lowerCamelCase_								=  self.model_tester.prepare_config_and_inputs_for_common()

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

    lowerCamelCase_								=  ["pixel_values"]
    self.assertListEqual(arg_names[:1]      ,  UpperCamelCase					)






  def   snake_case		(   self					):







   """simple docstring"""
   self.skipTest("Nat's attention operation is handled entirely by NATTEN."					)






  def   snake_case		(   self      ,  UpperCamelCase      ,  UpperCamelCase      ,  UpperCamelCase      ,  UpperCamelCase					):







   """simple docstring"""
   lowerCamelCase_								=  model_class(UpperCamelCase					)
   model.to(UpperCamelCase					)
   model.eval()

   with torch.no_grad():
    lowerCamelCase_								=  model(**self._prepare_for_class(UpperCamelCase      ,  UpperCamelCase					)					)

   lowerCamelCase_								=  outputs.hidden_states

   lowerCamelCase_								=  getattr(
       self.model_tester      ,  "expected_num_hidden_layers"      ,  len(self.model_tester.depths					) + 1					)
   self.assertEqual(len(UpperCamelCase					)      ,  UpperCamelCase					)

   # Nat has a different seq_length
   lowerCamelCase_								=  (
       config.patch_size
       if isinstance(config.patch_size      ,  collections.abc.Iterable					)
       else (config.patch_size, config.patch_size)
   )

   lowerCamelCase_								=  image_size[0] // patch_size[0]
   lowerCamelCase_								=  image_size[1] // patch_size[1]

   self.assertListEqual(
       list(hidden_states[0].shape[-3:]					)      ,  [height, width, self.model_tester.embed_dim]      ,  )

   if model_class.__name__ != "NatBackbone":
    lowerCamelCase_								=  outputs.reshaped_hidden_states
    self.assertEqual(len(UpperCamelCase					)      ,  UpperCamelCase					)

    lowerCamelCase_			,lowerCamelCase_			,lowerCamelCase_			,lowerCamelCase_								=  reshaped_hidden_states[0].shape
    lowerCamelCase_								=  (
        reshaped_hidden_states[0].view(UpperCamelCase      ,  UpperCamelCase      ,  UpperCamelCase      ,  UpperCamelCase					).permute(0      ,  2      ,  3      ,  1					)
    )
    self.assertListEqual(
        list(reshaped_hidden_states.shape[-3:]					)      ,  [height, width, self.model_tester.embed_dim]      ,  )






  def   snake_case		(   self					):







   """simple docstring"""
   lowerCamelCase_			,lowerCamelCase_								=  self.model_tester.prepare_config_and_inputs_for_common()

   lowerCamelCase_								=  (
       self.model_tester.image_size
       if isinstance(self.model_tester.image_size      ,  collections.abc.Iterable					)
       else (self.model_tester.image_size, self.model_tester.image_size)
   )

   for model_class in self.all_model_classes:
    lowerCamelCase_								=  True
    self.check_hidden_states_output(UpperCamelCase      ,  UpperCamelCase      ,  UpperCamelCase      ,  UpperCamelCase					)

    # check that output_hidden_states also work using config
    del inputs_dict["output_hidden_states"]
    lowerCamelCase_								=  True

    self.check_hidden_states_output(UpperCamelCase      ,  UpperCamelCase      ,  UpperCamelCase      ,  UpperCamelCase					)






  @slow
  def   snake_case		(   self					):







   """simple docstring"""
   for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    lowerCamelCase_								=  NatModel.from_pretrained(UpperCamelCase					)
    self.assertIsNotNone(UpperCamelCase					)







  def   snake_case		(   self					):







   """simple docstring"""
   lowerCamelCase_			,lowerCamelCase_								=  self.model_tester.prepare_config_and_inputs_for_common()

   lowerCamelCase_								=  _config_zero_init(UpperCamelCase					)
   for model_class in self.all_model_classes:
    lowerCamelCase_								=  model_class(config=UpperCamelCase					)
    for name, param in model.named_parameters():
     if "embeddings" not in name and param.requires_grad:
      self.assertIn(
          ((param.data.mean() * 1e9).round() / 1e9).item()      ,  [0.0, 1.0]      ,  msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''      ,  )





@require_natten
@require_vision
@require_torch
class 					snake_case   (   unittest.TestCase      ):





  """simple docstring"""




  @cached_property
  def   snake_case		(   self					):







   """simple docstring"""
   return AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224"					) if is_vision_available() else None







  @slow
  def   snake_case		(   self					):







   """simple docstring"""
   lowerCamelCase_								=  NatForImageClassification.from_pretrained("shi-labs/nat-mini-in1k-224"					).to(UpperCamelCase					)
   lowerCamelCase_								=  self.default_image_processor

   lowerCamelCase_								=  Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"					)
   lowerCamelCase_								=  image_processor(images=UpperCamelCase      ,  return_tensors="pt"					).to(UpperCamelCase					)

   # forward pass
   with torch.no_grad():
    lowerCamelCase_								=  model(**UpperCamelCase					)

   # verify the logits
   lowerCamelCase_								=  torch.Size((1, 1000)					)
   self.assertEqual(outputs.logits.shape      ,  UpperCamelCase					)
   lowerCamelCase_								=  torch.tensor([0.3_805, -0.8_676, -0.3_912]					).to(UpperCamelCase					)
   self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,  UpperCamelCase      ,  atol=1e-4					)					)





@require_torch
@require_natten
class 					snake_case   (   unittest.TestCase      ,						lowercase      ):





  """simple docstring"""




  _lowerCamelCase		   =  (NatBackbone,) if is_torch_available() else ()
  _lowerCamelCase		   =  NatConfig






  def   snake_case		(   self					):







   """simple docstring"""
   lowerCamelCase_								=  NatModelTester(self					)







'''simple docstring'''

import argparse
import json
import re
import string
import sys
from collections import Counter





def __snake_case (				UpperCAmelCase_						:						str		):

 def remove_articles(UpperCAmelCase_						:						Optional[Any]		):
  return re.sub(r"\b(a|an|the)\b" ,    " " ,    UpperCAmelCase_		)

 def white_space_fix(UpperCAmelCase_						:						List[str]		):
  return " ".join(text.split()		)

 def remove_punc(UpperCAmelCase_						:						Any		):
  lowerCamelCase_								=  set(string.punctuation		)
  return "".join(ch for ch in text if ch not in exclude		)

 def lower(UpperCAmelCase_						:						List[str]		):
  return text.lower()

 return white_space_fix(remove_articles(remove_punc(lower(UpperCAmelCase_		)		)		)		)





def __snake_case (				UpperCAmelCase_						:						Tuple ,    UpperCAmelCase_						:						int		):
 lowerCamelCase_								=  normalize_answer(UpperCAmelCase_		).split()
 lowerCamelCase_								=  normalize_answer(UpperCAmelCase_		).split()
 lowerCamelCase_								=  Counter(UpperCAmelCase_		) & Counter(UpperCAmelCase_		)
 lowerCamelCase_								=  sum(common.values()		)
 if num_same == 0:
  return 0
 lowerCamelCase_								=  1.0 * num_same / len(UpperCAmelCase_		)
 lowerCamelCase_								=  1.0 * num_same / len(UpperCAmelCase_		)
 lowerCamelCase_								=  (2 * precision * recall) / (precision + recall)
 return fa





def __snake_case (				UpperCAmelCase_						:						Tuple ,    UpperCAmelCase_						:						List[Any]		):
 return normalize_answer(UpperCAmelCase_		) == normalize_answer(UpperCAmelCase_		)





def __snake_case (				UpperCAmelCase_						:						Any ,    UpperCAmelCase_						:						Dict ,    UpperCAmelCase_						:						Union[str, Any]		):
 lowerCamelCase_								=  []
 for ground_truth in ground_truths:
  lowerCamelCase_								=  metric_fn(UpperCAmelCase_ ,    UpperCAmelCase_		)
  scores_for_ground_truths.append(UpperCAmelCase_		)
 return max(UpperCAmelCase_		)





def __snake_case (				UpperCAmelCase_						:						Optional[int] ,    UpperCAmelCase_						:						Dict		):
 lowerCamelCase_								=  lowerCamelCase_								=  lowerCamelCase_								=  0
 lowerCamelCase_								=  []
 for passage in dataset:
  for qa in passage["qas"]:
   total += 1
   if qa["id"] not in predictions:
    lowerCamelCase_								=  F'''Unanswered question {qa["id"]} will receive score 0.'''
    print(UpperCAmelCase_ ,    file=sys.stderr		)
    continue

   lowerCamelCase_								=  [x["text"] for x in qa["answers"]]
   lowerCamelCase_								=  predictions[qa["id"]]

   lowerCamelCase_								=  metric_max_over_ground_truths(UpperCAmelCase_ ,    UpperCAmelCase_ ,    UpperCAmelCase_		)
   if int(_exact_match		) == 1:
    correct_ids.append(qa["id"]		)
   exact_match += _exact_match

   fa += metric_max_over_ground_truths(UpperCAmelCase_ ,    UpperCAmelCase_ ,    UpperCAmelCase_		)

 lowerCamelCase_								=  exact_match / total
 lowerCamelCase_								=  fa / total

 return {"exact_match": exact_match, "f1": fa}, correct_ids


if __name__ == "__main__":
     a_  :    Optional[int]      =			"""1.0"""
     a_  :    Dict      =			argparse.ArgumentParser("""Official evaluation script for ReCoRD v1.0.""")
     parser.add_argument("""data_file""", help="""The dataset file in JSON format.""")
     parser.add_argument("""pred_file""", help="""The model prediction file in JSON format.""")
     parser.add_argument("""--output_correct_ids""", action="""store_true""", help="""Output the correctly answered query IDs.""")
     a_  :    Optional[int]      =			parser.parse_args()

     with open(args.data_file) as data_file:
          a_  :    Any      =			json.load(data_file)
          if dataset_json["version"] != expected_version:
               print(
                   f'''Evaluation expects v-{expected_version}, but got dataset with v-{dataset_json['version']}''',
                   file=sys.stderr,
               )
          a_  :    Any      =			dataset_json["""data"""]

     with open(args.pred_file) as pred_file:
          a_  :    List[str]      =			json.load(pred_file)

     a_     ,       a_  :    int      =			evaluate(dataset, predictions)

     if args.output_correct_ids:
          print(f'''Output {len(correct_ids)} correctly answered question IDs.''')
          with open("""correct_ids.json""", """w""") as f:
               json.dump(correct_ids, f)



