"""simple docstring"""
import unittest

from transformers import YosoConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

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


if is_torch_available():
 import torch

 from transformers import (
     YosoForMaskedLM,
     YosoForMultipleChoice,
     YosoForQuestionAnswering,
     YosoForSequenceClassification,
     YosoForTokenClassification,
     YosoModel,
 )
 from transformers.models.yoso.modeling_yoso import YOSO_PRETRAINED_MODEL_ARCHIVE_LIST




class   __snake_case			:



      def __init__(		self       ,     lowercase       ,     lowercase=13       ,     lowercase=7       ,     lowercase=True       ,     lowercase=True       ,     lowercase=True       ,     lowercase=True       ,     lowercase=99       ,     lowercase=32       ,     lowercase=5       ,     lowercase=4       ,     lowercase=37       ,     lowercase="gelu"       ,     lowercase=0.1       ,     lowercase=0.1       ,     lowercase=5_12       ,     lowercase=16       ,     lowercase=2       ,     lowercase=0.02       ,     lowercase=3       ,     lowercase=4       ,     lowercase=None       ,     )      ->     List[str]:





        '''simple docstring'''





        a__:	int     				=      parent
        a__:	List[Any]     				=      batch_size
        a__:	List[str]     				=      seq_length
        a__:	Any     				=      is_training
        a__:	str     				=      use_input_mask
        a__:	List[str]     				=      use_token_type_ids
        a__:	Optional[Any]     				=      use_labels
        a__:	Optional[Any]     				=      vocab_size
        a__:	int     				=      hidden_size
        a__:	Any     				=      num_hidden_layers
        a__:	Optional[Any]     				=      num_attention_heads
        a__:	Optional[int]     				=      intermediate_size
        a__:	Dict     				=      hidden_act
        a__:	List[str]     				=      hidden_dropout_prob
        a__:	Optional[int]     				=      attention_probs_dropout_prob
        a__:	Optional[int]     				=      max_position_embeddings
        a__:	Union[str, Any]     				=      type_vocab_size
        a__:	List[Any]     				=      type_sequence_label_size
        a__:	Tuple     				=      initializer_range
        a__:	List[Any]     				=      num_labels
        a__:	Dict     				=      num_choices
        a__:	Optional[Any]     				=      scope



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





        '''simple docstring'''





        a__:	int     				=      ids_tensor([self.batch_size, self.seq_length]       ,     self.vocab_size)

        a__:	Dict     				=      None
        if self.use_input_mask:
          a__:	str     				=      random_attention_mask([self.batch_size, self.seq_length])

        a__:	int     				=      None
        if self.use_token_type_ids:
          a__:	Optional[Any]     				=      ids_tensor([self.batch_size, self.seq_length]       ,     self.type_vocab_size)

        a__:	List[Any]     				=      None
        a__:	List[Any]     				=      None
        a__:	Optional[int]     				=      None
        if self.use_labels:
          a__:	Union[str, Any]     				=      ids_tensor([self.batch_size]       ,     self.type_sequence_label_size)
          a__:	List[str]     				=      ids_tensor([self.batch_size, self.seq_length]       ,     self.num_labels)
          a__:	Dict     				=      ids_tensor([self.batch_size]       ,     self.num_choices)

        a__:	Tuple     				=      self.get_config()

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



      def     lowerCamelCase_  (		self)      ->     int:





        '''simple docstring'''





        return YosoConfig(
            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=lowercase       ,     initializer_range=self.initializer_range       ,     )



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





        '''simple docstring'''





        a__:	Any     				=      self.get_config()
        a__:	Dict     				=      3_00
        return config



      def     lowerCamelCase_  (		self)      ->     Tuple:





        '''simple docstring'''





        (
            (
            a__
        )				,       (
            a__
        )				,       (
            a__
        )				,       (
            a__
        )				,       (
            a__
        )				,       (
            a__
        )				,       (
            a__
        )				,       
        ):	str     				=      self.prepare_config_and_inputs()

        a__:	Any     				=      True
        a__:	Dict     				=      floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
        a__:	Dict     				=      ids_tensor([self.batch_size, self.seq_length]       ,     vocab_size=2)

        return (
            config,
            input_ids,
            token_type_ids,
            input_mask,
            sequence_labels,
            token_labels,
            choice_labels,
            encoder_hidden_states,
            encoder_attention_mask,
        )



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





        '''simple docstring'''





        a__:	List[str]     				=      YosoModel(config=lowercase)
        model.to(lowercase)
        model.eval()
        a__:	Union[str, Any]     				=      model(lowercase       ,     attention_mask=lowercase       ,     token_type_ids=lowercase)
        a__:	Any     				=      model(lowercase       ,     token_type_ids=lowercase)
        a__:	str     				=      model(lowercase)
        self.parent.assertEqual(result.last_hidden_state.shape       ,     (self.batch_size, self.seq_length, self.hidden_size))



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





        '''simple docstring'''





        a__:	int     				=      True
        a__:	Any     				=      YosoModel(lowercase)
        model.to(lowercase)
        model.eval()
        a__:	Tuple     				=      model(
            lowercase       ,     attention_mask=lowercase       ,     token_type_ids=lowercase       ,     encoder_hidden_states=lowercase       ,     encoder_attention_mask=lowercase       ,     )
        a__:	List[Any]     				=      model(
            lowercase       ,     attention_mask=lowercase       ,     token_type_ids=lowercase       ,     encoder_hidden_states=lowercase       ,     )
        a__:	Dict     				=      model(lowercase       ,     attention_mask=lowercase       ,     token_type_ids=lowercase)
        self.parent.assertEqual(result.last_hidden_state.shape       ,     (self.batch_size, self.seq_length, self.hidden_size))



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





        '''simple docstring'''





        a__:	int     				=      YosoForMaskedLM(config=lowercase)
        model.to(lowercase)
        model.eval()
        a__:	List[Any]     				=      model(lowercase       ,     attention_mask=lowercase       ,     token_type_ids=lowercase       ,     labels=lowercase)
        self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.seq_length, self.vocab_size))



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





        '''simple docstring'''





        a__:	Any     				=      YosoForQuestionAnswering(config=lowercase)
        model.to(lowercase)
        model.eval()
        a__:	str     				=      model(
            lowercase       ,     attention_mask=lowercase       ,     token_type_ids=lowercase       ,     start_positions=lowercase       ,     end_positions=lowercase       ,     )
        self.parent.assertEqual(result.start_logits.shape       ,     (self.batch_size, self.seq_length))
        self.parent.assertEqual(result.end_logits.shape       ,     (self.batch_size, self.seq_length))



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





        '''simple docstring'''





        a__:	Dict     				=      self.num_labels
        a__:	Dict     				=      YosoForSequenceClassification(lowercase)
        model.to(lowercase)
        model.eval()
        a__:	Optional[Any]     				=      model(lowercase       ,     attention_mask=lowercase       ,     token_type_ids=lowercase       ,     labels=lowercase)
        self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.num_labels))



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





        '''simple docstring'''





        a__:	Any     				=      self.num_labels
        a__:	int     				=      YosoForTokenClassification(config=lowercase)
        model.to(lowercase)
        model.eval()
        a__:	List[Any]     				=      model(lowercase       ,     attention_mask=lowercase       ,     token_type_ids=lowercase       ,     labels=lowercase)
        self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.seq_length, self.num_labels))



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





        '''simple docstring'''





        a__:	Optional[int]     				=      self.num_choices
        a__:	List[str]     				=      YosoForMultipleChoice(config=lowercase)
        model.to(lowercase)
        model.eval()
        a__:	Any     				=      input_ids.unsqueeze(1).expand(-1       ,     self.num_choices       ,     -1).contiguous()
        a__:	Tuple     				=      token_type_ids.unsqueeze(1).expand(-1       ,     self.num_choices       ,     -1).contiguous()
        a__:	List[Any]     				=      input_mask.unsqueeze(1).expand(-1       ,     self.num_choices       ,     -1).contiguous()
        a__:	int     				=      model(
            lowercase       ,     attention_mask=lowercase       ,     token_type_ids=lowercase       ,     labels=lowercase       ,     )
        self.parent.assertEqual(result.logits.shape       ,     (self.batch_size, self.num_choices))



      def     lowerCamelCase_  (		self)      ->     Tuple:





        '''simple docstring'''





        a__:	int     				=      self.prepare_config_and_inputs()
        (
            (
            a__
        )				,       (
            a__
        )				,       (
            a__
        )				,       (
            a__
        )				,       (
            a__
        )				,       (
            a__
        )				,       (
            a__
        )				,       
        ):	Tuple     				=      config_and_inputs
        a__:	Any     				=      {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
        return config, inputs_dict







@require_torch
class   __snake_case			(	__lowerCAmelCase			, __lowerCAmelCase			, unittest.TestCase							):
      a__								=						(
          (
              YosoModel,
              YosoForMaskedLM,
              YosoForMultipleChoice,
              YosoForQuestionAnswering,
              YosoForSequenceClassification,
              YosoForTokenClassification,
          )
          if is_torch_available()
          else ()
      )
      a__								=						False
      a__								=						False
      a__								=						False

      a__								=						()
      a__								=						(
          {
              """feature-extraction""": YosoModel,
              """fill-mask""": YosoForMaskedLM,
              """question-answering""": YosoForQuestionAnswering,
              """text-classification""": YosoForSequenceClassification,
              """token-classification""": YosoForTokenClassification,
              """zero-shot""": YosoForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )



      def     lowerCamelCase_  (		self)      ->     int:





        '''simple docstring'''





        a__:	List[str]     				=      YosoModelTester(self)
        a__:	Optional[int]     				=      ConfigTester(self       ,     config_class=lowercase       ,     hidden_size=37)



      def     lowerCamelCase_  (		self)      ->     Tuple:





        '''simple docstring'''





        self.config_tester.run_common_tests()



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





        '''simple docstring'''





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



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





        '''simple docstring'''





        a__:	Union[str, Any]     				=      self.model_tester.prepare_config_and_inputs()
        for type in ["absolute", "relative_key", "relative_key_query"]:
          a__:	Optional[Any]     				=      type
          self.model_tester.create_and_check_model(*lowercase)



      def     lowerCamelCase_  (		self)      ->     Tuple:





        '''simple docstring'''





        a__:	Tuple     				=      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*lowercase)



      def     lowerCamelCase_  (		self)      ->     Tuple:





        '''simple docstring'''





        a__:	Union[str, Any]     				=      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*lowercase)



      def     lowerCamelCase_  (		self)      ->     str:





        '''simple docstring'''





        a__:	List[str]     				=      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*lowercase)



      def     lowerCamelCase_  (		self)      ->     Union[str, Any]:





        '''simple docstring'''





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



      def     lowerCamelCase_  (		self)      ->     str:





        '''simple docstring'''





        a__:	Dict     				=      self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*lowercase)



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





        '''simple docstring'''





        for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          a__:	Optional[Any]     				=      YosoModel.from_pretrained(lowercase)
          self.assertIsNotNone(lowercase)



      def     lowerCamelCase_  (		self)      ->     Dict:





        '''simple docstring'''





        return







@require_torch
class   __snake_case			(	unittest.TestCase							):



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





        '''simple docstring'''





        a__:	str     				=      YosoModel.from_pretrained('uw-madison/yoso-4096')
        a__:	Any     				=      torch.tensor([[0, 1, 2, 3, 4, 5]])

        with torch.no_grad():
          a__:	Tuple     				=      model(lowercase)[0]

        a__:	Dict     				=      torch.Size((1, 6, 7_68))
        self.assertEqual(output.shape       ,     lowercase)

        a__:	Optional[Any]     				=      torch.tensor(
            [[[-0.0611, 0.1242, 0.0840], [0.0280, -0.0048, 0.1125], [0.0106, 0.0226, 0.0751]]])

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



      @slow
      def     lowerCamelCase_  (		self)      ->     Any:





        '''simple docstring'''





        a__:	str     				=      YosoForMaskedLM.from_pretrained('uw-madison/yoso-4096')
        a__:	List[Any]     				=      torch.tensor([[0, 1, 2, 3, 4, 5]])

        with torch.no_grad():
          a__:	Optional[int]     				=      model(lowercase)[0]

        a__:	Dict     				=      5_02_65

        a__:	int     				=      torch.Size((1, 6, vocab_size))
        self.assertEqual(output.shape       ,     lowercase)

        a__:	List[str]     				=      torch.tensor(
            [[[-2.1313, -3.7285, -2.2407], [-2.7047, -3.3314, -2.6408], [0.0629, -2.5166, -0.3356]]])

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



      @slow
      def     lowerCamelCase_  (		self)      ->     int:





        '''simple docstring'''





        a__:	int     				=      YosoForMaskedLM.from_pretrained('uw-madison/yoso-4096')
        a__:	Union[str, Any]     				=      torch.arange(40_96).unsqueeze(0)

        with torch.no_grad():
          a__:	List[Any]     				=      model(lowercase)[0]

        a__:	Tuple     				=      5_02_65

        a__:	List[str]     				=      torch.Size((1, 40_96, vocab_size))
        self.assertEqual(output.shape       ,     lowercase)

        a__:	List[Any]     				=      torch.tensor(
            [[[-2.3914, -4.3742, -5.0956], [-4.0988, -4.2384, -7.0406], [-3.1427, -3.7192, -6.6800]]])

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




# Copyright 2022 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from typing import List



def 			UpperCamelCase       (				lowerCAmelCase__					,     lowerCAmelCase__						):
					'''simple docstring'''
					if name != "training_function" and name != "main":
										raise ValueError(f'Incorrect function name passed: {name}, choose either \'main\' or \'training_function\''						)
					lowercase      ,		lowercase		      =		[], False
					for line in lines:
										if not found_start and f'def {name}' in line:
															lowercase		      =		True
															good_lines.append(lowerCAmelCase__						)
															continue
										if found_start:
															if name == "training_function" and "def main" in line:
																				return good_lines
															if name == "main" and "if __name__" in line:
																				return good_lines
															good_lines.append(lowerCAmelCase__						)



def 			UpperCamelCase       (				lowerCAmelCase__						):
					'''simple docstring'''
					return [line for line in lines if not line.lstrip().startswith('''#'''						) and line != "\n"]



def 			UpperCamelCase       (				lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__					,     lowerCAmelCase__ = None						):
					'''simple docstring'''
					with open(lowerCAmelCase__					,     '''r'''						) as f:
										lowercase		      =		f.readlines()
					with open(os.path.abspath(os.path.join('''examples'''					,     '''nlp_example.py'''						)						)					,     '''r'''						) as f:
										lowercase		      =		f.readlines()
					with open(lowerCAmelCase__					,     '''r'''						) as f:
										lowercase		      =		f.readlines()
					if secondary_filename is not None:
										with open(lowerCAmelCase__					,     '''r'''						) as f:
															lowercase		      =		f.readlines()

    # This is our base, we remove all the code from here in our `full_filename` and `feature_filename` to find the new content
					if parser_only:
										lowercase		      =		clean_lines(get_function_contents_by_name(lowerCAmelCase__					,     '''main'''						)						)
										lowercase		      =		clean_lines(get_function_contents_by_name(lowerCAmelCase__					,     '''main'''						)						)
										lowercase		      =		clean_lines(get_function_contents_by_name(lowerCAmelCase__					,     '''main'''						)						)
										if secondary_filename is not None:
															lowercase		      =		clean_lines(get_function_contents_by_name(lowerCAmelCase__					,     '''main'''						)						)
					else:
										lowercase		      =		clean_lines(get_function_contents_by_name(lowerCAmelCase__					,     '''training_function'''						)						)
										lowercase		      =		clean_lines(get_function_contents_by_name(lowerCAmelCase__					,     '''training_function'''						)						)
										lowercase		      =		clean_lines(get_function_contents_by_name(lowerCAmelCase__					,     '''training_function'''						)						)
										if secondary_filename is not None:
															lowercase		      =		clean_lines(
															    get_function_contents_by_name(lowerCAmelCase__					,     '''training_function'''						)						)

					lowercase		      =		'''train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size)\n'''

					# Specific code in our script that differs from the full version, aka what is new
					lowercase		      =		[]
					lowercase		      =		[]  # We keep track of the idxs just in case it's a repeated statement
					lowercase		      =		iter(lowerCAmelCase__						)
					for i in range(len(lowerCAmelCase__						) - 1						):
										if i not in passed_idxs:
															lowercase		      =		next(lowerCAmelCase__						)
															if (line not in full_file_func) and (line.lstrip() != _dl_line):
																				if "TESTING_MOCKED_DATALOADERS" not in line:
																									new_feature_code.append(lowerCAmelCase__						)
																									passed_idxs.append(lowerCAmelCase__						)
																				else:
																									# Skip over the `config['num_epochs'] = 2` statement
																									lowercase		      =		next(lowerCAmelCase__						)

    # Extract out just the new parts from the full_file_training_func
					lowercase		      =		[]
					lowercase		      =		[]  # We keep track of the idxs just in case it's a repeated statement
					for i, line in enumerate(lowerCAmelCase__						):
										if i not in passed_idxs:
															if (line not in full_file_func) and (line.lstrip() != _dl_line):
																				if "TESTING_MOCKED_DATALOADERS" not in line:
																									new_full_example_parts.append(lowerCAmelCase__						)
																									passed_idxs.append(lowerCAmelCase__						)

    # Finally, get the overall diff
					lowercase		      =		[line for line in new_feature_code if line not in new_full_example_parts]
					if secondary_filename is not None:
										lowercase		      =		[line for line in full_file_contents if line not in secondary_file_func]
										lowercase		      =		[line for line in diff_from_example if line not in diff_from_two]

					return diff_from_example




