

import unittest

from transformers import (
    MODEL_FOR_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
    TextGenerationPipeline,
    logging,
    pipeline,
)
from transformers.testing_utils import (
    CaptureLogger,
    is_pipeline_test,
    require_accelerate,
    require_tf,
    require_torch,
    require_torch_gpu,
    require_torch_or_tf,
)

from .test_pipelines_common import ANY




@is_pipeline_test
@require_torch_or_tf
class      UpperCamelCase (     unittest.TestCase   ):
    __UpperCamelCase         =MODEL_FOR_CAUSAL_LM_MAPPING
    __UpperCamelCase         =TF_MODEL_FOR_CAUSAL_LM_MAPPING

    @require_torch
    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			pipeline(task='text-generation'							,    model='sshleifer/tiny-ctrl'							,    framework='pt'    )
     # Using `do_sample=False` to force deterministic output
     SCREAMING_SNAKE_CASE								=			text_generator('This is a test'							,    do_sample=snake_case__    )
     self.assertEqual(
         snake_case__							,    [
             {
                 'generated_text': (
                     'This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope.'
                     ' oscope. FiliFili@@'
                 )
             }
         ]							,    )

     SCREAMING_SNAKE_CASE								=			text_generator(['This is a test', 'This is a second test']    )
     self.assertEqual(
         snake_case__							,    [
             [
                 {
                     'generated_text': (
                         'This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope.'
                         ' oscope. FiliFili@@'
                     )
                 }
             ],
             [
                 {
                     'generated_text': (
                         'This is a second test ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy'
                         ' oscope. oscope. FiliFili@@'
                     )
                 }
             ],
         ]							,    )

     SCREAMING_SNAKE_CASE								=			text_generator('This is a test'							,    do_sample=snake_case__							,    num_return_sequences=2							,    return_tensors=snake_case__    )
     self.assertEqual(
         snake_case__							,    [
             {'generated_token_ids': ANY(snake_case__    )},
             {'generated_token_ids': ANY(snake_case__    )},
         ]							,    )
     SCREAMING_SNAKE_CASE								=			text_generator.model.config.eos_token_id
     SCREAMING_SNAKE_CASE								=			'<pad>'
     SCREAMING_SNAKE_CASE								=			text_generator(
         ['This is a test', 'This is a second test']							,    do_sample=snake_case__							,    num_return_sequences=2							,    batch_size=2							,    return_tensors=snake_case__							,    )
     self.assertEqual(
         snake_case__							,    [
             [
                 {'generated_token_ids': ANY(snake_case__    )},
                 {'generated_token_ids': ANY(snake_case__    )},
             ],
             [
                 {'generated_token_ids': ANY(snake_case__    )},
                 {'generated_token_ids': ANY(snake_case__    )},
             ],
         ]							,    )

    @require_tf
    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			pipeline(task='text-generation'							,    model='sshleifer/tiny-ctrl'							,    framework='tf'    )

     # Using `do_sample=False` to force deterministic output
     SCREAMING_SNAKE_CASE								=			text_generator('This is a test'							,    do_sample=snake_case__    )
     self.assertEqual(
         snake_case__							,    [
             {
                 'generated_text': (
                     'This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵'
                     ' please,'
                 )
             }
         ]							,    )

     SCREAMING_SNAKE_CASE								=			text_generator(['This is a test', 'This is a second test']							,    do_sample=snake_case__    )
     self.assertEqual(
         snake_case__							,    [
             [
                 {
                     'generated_text': (
                         'This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵'
                         ' please,'
                     )
                 }
             ],
             [
                 {
                     'generated_text': (
                         'This is a second test Chieftain Chieftain prefecture prefecture prefecture Cannes Cannes'
                         ' Cannes 閲閲Cannes Cannes Cannes 攵 please,'
                     )
                 }
             ],
         ]							,    )

    def UpperCamelCase			(			self :      str							,    snake_case__ :      Optional[Any]							,    snake_case__ :      Tuple							,    snake_case__ :      str    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			TextGenerationPipeline(model=snake_case__							,    tokenizer=snake_case__    )
     return text_generator, ["This is a test", "Another test"]

    def UpperCamelCase			(			self :      Optional[int]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			'Hello I believe in'
     SCREAMING_SNAKE_CASE								=			pipeline('text-generation'							,    model='hf-internal-testing/tiny-random-gpt2'    )
     SCREAMING_SNAKE_CASE								=			text_generator(snake_case__    )
     self.assertEqual(
         snake_case__							,    [{'generated_text': 'Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe'}]							,    )

     SCREAMING_SNAKE_CASE								=			text_generator(snake_case__							,    stop_sequence=' fe'    )
     self.assertEqual(snake_case__							,    [{'generated_text': 'Hello I believe in fe'}]    )

    def UpperCamelCase			(			self :      Optional[int]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			text_generator.model
     SCREAMING_SNAKE_CASE								=			text_generator.tokenizer

     SCREAMING_SNAKE_CASE								=			text_generator('This is a test'    )
     self.assertEqual(snake_case__							,    [{'generated_text': ANY(snake_case__    )}]    )
     self.assertTrue(outputs[0]['generated_text'].startswith('This is a test'    )    )

     SCREAMING_SNAKE_CASE								=			text_generator('This is a test'							,    return_full_text=snake_case__    )
     self.assertEqual(snake_case__							,    [{'generated_text': ANY(snake_case__    )}]    )
     self.assertNotIn('This is a test'							,    outputs[0]['generated_text']    )

     SCREAMING_SNAKE_CASE								=			pipeline(task='text-generation'							,    model=snake_case__							,    tokenizer=snake_case__							,    return_full_text=snake_case__    )
     SCREAMING_SNAKE_CASE								=			text_generator('This is a test'    )
     self.assertEqual(snake_case__							,    [{'generated_text': ANY(snake_case__    )}]    )
     self.assertNotIn('This is a test'							,    outputs[0]['generated_text']    )

     SCREAMING_SNAKE_CASE								=			text_generator('This is a test'							,    return_full_text=snake_case__    )
     self.assertEqual(snake_case__							,    [{'generated_text': ANY(snake_case__    )}]    )
     self.assertTrue(outputs[0]['generated_text'].startswith('This is a test'    )    )

     SCREAMING_SNAKE_CASE								=			text_generator(['This is great !', 'Something else']							,    num_return_sequences=2							,    do_sample=snake_case__    )
     self.assertEqual(
         snake_case__							,    [
             [{'generated_text': ANY(snake_case__    )}, {'generated_text': ANY(snake_case__    )}],
             [{'generated_text': ANY(snake_case__    )}, {'generated_text': ANY(snake_case__    )}],
         ]							,    )

     if text_generator.tokenizer.pad_token is not None:
      SCREAMING_SNAKE_CASE								=			text_generator(
          ['This is great !', 'Something else']							,    num_return_sequences=2							,    batch_size=2							,    do_sample=snake_case__    )
      self.assertEqual(
          snake_case__							,    [
              [{'generated_text': ANY(snake_case__    )}, {'generated_text': ANY(snake_case__    )}],
              [{'generated_text': ANY(snake_case__    )}, {'generated_text': ANY(snake_case__    )}],
          ]							,    )

     with self.assertRaises(snake_case__    ):
      SCREAMING_SNAKE_CASE								=			text_generator('test'							,    return_full_text=snake_case__							,    return_text=snake_case__    )
     with self.assertRaises(snake_case__    ):
      SCREAMING_SNAKE_CASE								=			text_generator('test'							,    return_full_text=snake_case__							,    return_tensors=snake_case__    )
     with self.assertRaises(snake_case__    ):
      SCREAMING_SNAKE_CASE								=			text_generator('test'							,    return_text=snake_case__							,    return_tensors=snake_case__    )

     # Empty prompt is slighly special
     # it requires BOS token to exist.
     # Special case for Pegasus which will always append EOS so will
     # work even without BOS.
     if (
         text_generator.tokenizer.bos_token_id is not None
         or "Pegasus" in tokenizer.__class__.__name__
         or "Git" in model.__class__.__name__
     ):
      SCREAMING_SNAKE_CASE								=			text_generator(''    )
      self.assertEqual(snake_case__							,    [{'generated_text': ANY(snake_case__    )}]    )
     else:
      with self.assertRaises((ValueError, AssertionError)    ):
       SCREAMING_SNAKE_CASE								=			text_generator(''    )

     if text_generator.framework == "tf":
      # TF generation does not support max_new_tokens, and it's impossible
      # to control long generation with only max_length without
      # fancy calculation, dismissing tests for now.
      return
     # We don't care about infinite range models.
     # They already work.
     # Skip this test for XGLM, since it uses sinusoidal positional embeddings which are resized on-the-fly.
     SCREAMING_SNAKE_CASE								=			['RwkvForCausalLM', 'XGLMForCausalLM', 'GPTNeoXForCausalLM']
     if (
         tokenizer.model_max_length < 1_0_0_0_0
         and text_generator.model.__class__.__name__ not in EXTRA_MODELS_CAN_HANDLE_LONG_INPUTS
     ):
      # Handling of large generations
      with self.assertRaises((RuntimeError, IndexError, ValueError, AssertionError)    ):
       text_generator('This is a test' * 5_0_0							,    max_new_tokens=2_0    )

      SCREAMING_SNAKE_CASE								=			text_generator('This is a test' * 5_0_0							,    handle_long_generation='hole'							,    max_new_tokens=2_0    )
      # Hole strategy cannot work
      with self.assertRaises(snake_case__    ):
       text_generator(
           'This is a test' * 5_0_0							,    handle_long_generation='hole'							,    max_new_tokens=tokenizer.model_max_length + 1_0							,    )

    @require_torch
    @require_accelerate
    @require_torch_gpu
    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






     import torch

     # Classic `model_kwargs`
     SCREAMING_SNAKE_CASE								=			pipeline(
         model='hf-internal-testing/tiny-random-bloom'							,    model_kwargs={'device_map': 'auto', 'torch_dtype': torch.bfloataa}							,    )
     self.assertEqual(pipe.model.device							,    torch.device(0    )    )
     self.assertEqual(pipe.model.lm_head.weight.dtype							,    torch.bfloataa    )
     SCREAMING_SNAKE_CASE								=			pipe('This is a test'    )
     self.assertEqual(
         snake_case__							,    [
             {
                 'generated_text': (
                     'This is a test test test test test test test test test test test test test test test test'
                     ' test'
                 )
             }
         ]							,    )

     # Upgraded those two to real pipeline arguments (they just get sent for the model as they're unlikely to mean anything else.)
     SCREAMING_SNAKE_CASE								=			pipeline(model='hf-internal-testing/tiny-random-bloom'							,    device_map='auto'							,    torch_dtype=torch.bfloataa    )
     self.assertEqual(pipe.model.device							,    torch.device(0    )    )
     self.assertEqual(pipe.model.lm_head.weight.dtype							,    torch.bfloataa    )
     SCREAMING_SNAKE_CASE								=			pipe('This is a test'    )
     self.assertEqual(
         snake_case__							,    [
             {
                 'generated_text': (
                     'This is a test test test test test test test test test test test test test test test test'
                     ' test'
                 )
             }
         ]							,    )

     # torch_dtype will be automatically set to float32 if not provided - check: https://github.com/huggingface/transformers/pull/20602
     SCREAMING_SNAKE_CASE								=			pipeline(model='hf-internal-testing/tiny-random-bloom'							,    device_map='auto'    )
     self.assertEqual(pipe.model.device							,    torch.device(0    )    )
     self.assertEqual(pipe.model.lm_head.weight.dtype							,    torch.floataa    )
     SCREAMING_SNAKE_CASE								=			pipe('This is a test'    )
     self.assertEqual(
         snake_case__							,    [
             {
                 'generated_text': (
                     'This is a test test test test test test test test test test test test test test test test'
                     ' test'
                 )
             }
         ]							,    )

    @require_torch
    @require_torch_gpu
    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






     import torch

     SCREAMING_SNAKE_CASE								=			pipeline(model='hf-internal-testing/tiny-random-bloom'							,    device=0							,    torch_dtype=torch.floataa    )
     pipe('This is a test'    )

    @require_torch
    @require_accelerate
    @require_torch_gpu
    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






     import torch

     SCREAMING_SNAKE_CASE								=			pipeline(model='hf-internal-testing/tiny-random-bloom'							,    device_map='auto'							,    torch_dtype=torch.floataa    )
     pipe('This is a test'							,    do_sample=snake_case__							,    top_p=0.5    )




    def UpperCamelCase			(			self :      Dict    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			'Hello world'
     SCREAMING_SNAKE_CASE								=			pipeline('text-generation'							,    model='hf-internal-testing/tiny-random-gpt2'    )
     if text_generator.model.framework == "tf":
      SCREAMING_SNAKE_CASE								=			logging.get_logger('transformers.generation.tf_utils'    )
     else:
      SCREAMING_SNAKE_CASE								=			logging.get_logger('transformers.generation.utils'    )
     SCREAMING_SNAKE_CASE								=			'Both `max_new_tokens`'  # The beggining of the message to be checked in this test

     # Both are set by the user -> log warning
     with CaptureLogger(snake_case__    ) as cl:
      SCREAMING_SNAKE_CASE								=			text_generator(snake_case__							,    max_length=1_0							,    max_new_tokens=1    )
     self.assertIn(snake_case__							,    cl.out    )

     # The user only sets one -> no warning
     with CaptureLogger(snake_case__    ) as cl:
      SCREAMING_SNAKE_CASE								=			text_generator(snake_case__							,    max_new_tokens=1    )
     self.assertNotIn(snake_case__							,    cl.out    )

     with CaptureLogger(snake_case__    ) as cl:
      SCREAMING_SNAKE_CASE								=			text_generator(snake_case__							,    max_length=1_0    )
     self.assertNotIn(snake_case__							,    cl.out    )




import unittest

from transformers import AutoTokenizer, NystromformerConfig, 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, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          NystromformerForMaskedLM,
          NystromformerForMultipleChoice,
          NystromformerForQuestionAnswering,
          NystromformerForSequenceClassification,
          NystromformerForTokenClassification,
          NystromformerModel,
      )
      from transformers.models.nystromformer.modeling_nystromformer import NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST




class      UpperCamelCase :

    def __init__(			self :      List[str]							,    snake_case__ :      Dict							,    snake_case__ :      Optional[Any]=1_3							,    snake_case__ :      Union[str, Any]=7							,    snake_case__ :      List[str]=True							,    snake_case__ :      Any=True							,    snake_case__ :      List[str]=True							,    snake_case__ :      Optional[Any]=True							,    snake_case__ :      List[str]=9_9							,    snake_case__ :      str=3_2							,    snake_case__ :      Dict=5							,    snake_case__ :      str=4							,    snake_case__ :      int=3_7							,    snake_case__ :      Union[str, Any]="gelu"							,    snake_case__ :      Any=0.1							,    snake_case__ :      Dict=0.1							,    snake_case__ :      Optional[Any]=5_1_2							,    snake_case__ :      List[Any]=1_6							,    snake_case__ :      str=2							,    snake_case__ :      int=0.02							,    snake_case__ :      List[str]=3							,    snake_case__ :      Dict=4							,    snake_case__ :      str=None							,    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			parent
     SCREAMING_SNAKE_CASE								=			batch_size
     SCREAMING_SNAKE_CASE								=			seq_length
     SCREAMING_SNAKE_CASE								=			is_training
     SCREAMING_SNAKE_CASE								=			use_input_mask
     SCREAMING_SNAKE_CASE								=			use_token_type_ids
     SCREAMING_SNAKE_CASE								=			use_labels
     SCREAMING_SNAKE_CASE								=			vocab_size
     SCREAMING_SNAKE_CASE								=			hidden_size
     SCREAMING_SNAKE_CASE								=			num_hidden_layers
     SCREAMING_SNAKE_CASE								=			num_attention_heads
     SCREAMING_SNAKE_CASE								=			intermediate_size
     SCREAMING_SNAKE_CASE								=			hidden_act
     SCREAMING_SNAKE_CASE								=			hidden_dropout_prob
     SCREAMING_SNAKE_CASE								=			attention_probs_dropout_prob
     SCREAMING_SNAKE_CASE								=			max_position_embeddings
     SCREAMING_SNAKE_CASE								=			type_vocab_size
     SCREAMING_SNAKE_CASE								=			type_sequence_label_size
     SCREAMING_SNAKE_CASE								=			initializer_range
     SCREAMING_SNAKE_CASE								=			num_labels
     SCREAMING_SNAKE_CASE								=			num_choices
     SCREAMING_SNAKE_CASE								=			scope

    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






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

     SCREAMING_SNAKE_CASE								=			None
     if self.use_input_mask:
      SCREAMING_SNAKE_CASE								=			random_attention_mask([self.batch_size, self.seq_length]    )

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

     SCREAMING_SNAKE_CASE								=			None
     SCREAMING_SNAKE_CASE								=			None
     SCREAMING_SNAKE_CASE								=			None
     if self.use_labels:
      SCREAMING_SNAKE_CASE								=			ids_tensor([self.batch_size]							,    self.type_sequence_label_size    )
      SCREAMING_SNAKE_CASE								=			ids_tensor([self.batch_size, self.seq_length]							,    self.num_labels    )
      SCREAMING_SNAKE_CASE								=			ids_tensor([self.batch_size]							,    self.num_choices    )

     SCREAMING_SNAKE_CASE								=			self.get_config()

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

    def UpperCamelCase			(			self :      Optional[Any]    ):

     """simple docstring"""






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

    def UpperCamelCase			(			self :      List[str]							,    snake_case__ :      Any							,    snake_case__ :      Tuple							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      Dict							,    snake_case__ :      List[str]							,    snake_case__ :      int							,    snake_case__ :      Union[str, Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			NystromformerModel(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    attention_mask=snake_case__							,    token_type_ids=snake_case__    )
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    token_type_ids=snake_case__    )
     SCREAMING_SNAKE_CASE								=			model(snake_case__    )
     self.parent.assertEqual(result.last_hidden_state.shape							,    (self.batch_size, self.seq_length, self.hidden_size)    )

    def UpperCamelCase			(			self :      List[Any]							,    snake_case__ :      str							,    snake_case__ :      Dict							,    snake_case__ :      Dict							,    snake_case__ :      Dict							,    snake_case__ :      int							,    snake_case__ :      int							,    snake_case__ :      Tuple    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			NystromformerForMaskedLM(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    attention_mask=snake_case__							,    token_type_ids=snake_case__							,    labels=snake_case__    )
     self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.seq_length, self.vocab_size)    )

    def UpperCamelCase			(			self :      Union[str, Any]							,    snake_case__ :      Tuple							,    snake_case__ :      Optional[Any]							,    snake_case__ :      int							,    snake_case__ :      Optional[int]							,    snake_case__ :      str							,    snake_case__ :      List[str]							,    snake_case__ :      Optional[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			NystromformerForQuestionAnswering(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			model(
         snake_case__							,    attention_mask=snake_case__							,    token_type_ids=snake_case__							,    start_positions=snake_case__							,    end_positions=snake_case__							,    )
     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 UpperCamelCase			(			self :      List[str]							,    snake_case__ :      int							,    snake_case__ :      int							,    snake_case__ :      int							,    snake_case__ :      str							,    snake_case__ :      Optional[int]							,    snake_case__ :      Any							,    snake_case__ :      Union[str, Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.num_labels
     SCREAMING_SNAKE_CASE								=			NystromformerForSequenceClassification(snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    attention_mask=snake_case__							,    token_type_ids=snake_case__							,    labels=snake_case__    )
     self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.num_labels)    )

    def UpperCamelCase			(			self :      Tuple							,    snake_case__ :      int							,    snake_case__ :      Tuple							,    snake_case__ :      List[str]							,    snake_case__ :      List[Any]							,    snake_case__ :      Optional[Any]							,    snake_case__ :      List[Any]							,    snake_case__ :      Dict    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.num_labels
     SCREAMING_SNAKE_CASE								=			NystromformerForTokenClassification(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			model(snake_case__							,    attention_mask=snake_case__							,    token_type_ids=snake_case__							,    labels=snake_case__    )
     self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.seq_length, self.num_labels)    )

    def UpperCamelCase			(			self :      int							,    snake_case__ :      Optional[Any]							,    snake_case__ :      List[str]							,    snake_case__ :      Union[str, Any]							,    snake_case__ :      Tuple							,    snake_case__ :      Optional[int]							,    snake_case__ :      Any							,    snake_case__ :      List[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.num_choices
     SCREAMING_SNAKE_CASE								=			NystromformerForMultipleChoice(config=snake_case__    )
     model.to(snake_case__    )
     model.eval()
     SCREAMING_SNAKE_CASE								=			input_ids.unsqueeze(1    ).expand(-1							,    self.num_choices							,    -1    ).contiguous()
     SCREAMING_SNAKE_CASE								=			token_type_ids.unsqueeze(1    ).expand(-1							,    self.num_choices							,    -1    ).contiguous()
     SCREAMING_SNAKE_CASE								=			input_mask.unsqueeze(1    ).expand(-1							,    self.num_choices							,    -1    ).contiguous()
     SCREAMING_SNAKE_CASE								=			model(
         snake_case__							,    attention_mask=snake_case__							,    token_type_ids=snake_case__							,    labels=snake_case__							,    )
     self.parent.assertEqual(result.logits.shape							,    (self.batch_size, self.num_choices)    )




    def UpperCamelCase			(			self :      Union[str, Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.prepare_config_and_inputs()
     (
         (
         SCREAMING_SNAKE_CASE
     )						,	(
         SCREAMING_SNAKE_CASE
     )						,	(
         SCREAMING_SNAKE_CASE
     )						,	(
         SCREAMING_SNAKE_CASE
     )						,	(
         SCREAMING_SNAKE_CASE
     )						,	(
         SCREAMING_SNAKE_CASE
     )						,	(
         SCREAMING_SNAKE_CASE
     )						,	
     )								=			config_and_inputs
     SCREAMING_SNAKE_CASE								=			{'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
     return config, inputs_dict




@require_torch
class      UpperCamelCase (     SCREAMING_SNAKE_CASE	,			SCREAMING_SNAKE_CASE	,			unittest.TestCase   ):
    __UpperCamelCase         =(
        (
            NystromformerModel,
            NystromformerForMaskedLM,
            NystromformerForMultipleChoice,
            NystromformerForQuestionAnswering,
            NystromformerForSequenceClassification,
            NystromformerForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    __UpperCamelCase         =(
        {
            "feature-extraction": NystromformerModel,
            "fill-mask": NystromformerForMaskedLM,
            "question-answering": NystromformerForQuestionAnswering,
            "text-classification": NystromformerForSequenceClassification,
            "token-classification": NystromformerForTokenClassification,
            "zero-shot": NystromformerForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    __UpperCamelCase         =False
    __UpperCamelCase         =False

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			NystromformerModelTester(self    )
     SCREAMING_SNAKE_CASE								=			ConfigTester(self							,    config_class=snake_case__							,    hidden_size=3_7    )

    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     self.config_tester.run_common_tests()

    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*snake_case__    )

    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     for type in ["absolute", "relative_key", "relative_key_query"]:
      SCREAMING_SNAKE_CASE								=			type
      self.model_tester.create_and_check_model(*snake_case__    )

    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_masked_lm(*snake_case__    )

    def UpperCamelCase			(			self :      Union[str, Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_multiple_choice(*snake_case__    )

    def UpperCamelCase			(			self :      Tuple    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_question_answering(*snake_case__    )

    def UpperCamelCase			(			self :      Any    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_sequence_classification(*snake_case__    )

    def UpperCamelCase			(			self :      List[Any]    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_token_classification(*snake_case__    )




    @slow
    def UpperCamelCase			(			self :      List[str]    ):

     """simple docstring"""






     for model_name in NYSTROMFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
      SCREAMING_SNAKE_CASE								=			NystromformerModel.from_pretrained(snake_case__    )
      self.assertIsNotNone(snake_case__    )




@require_torch
class      UpperCamelCase (     unittest.TestCase   ):

    @slow
    def UpperCamelCase			(			self :      int    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			NystromformerModel.from_pretrained('uw-madison/nystromformer-512'    )
     SCREAMING_SNAKE_CASE								=			torch.tensor([[0, 1, 2, 3, 4, 5]]    )

     with torch.no_grad():
      SCREAMING_SNAKE_CASE								=			model(snake_case__    )[0]

     SCREAMING_SNAKE_CASE								=			torch.Size((1, 6, 7_6_8)    )
     self.assertEqual(output.shape							,    snake_case__    )

     SCREAMING_SNAKE_CASE								=			torch.tensor(
         [[[-0.4_532, -0.0_936, 0.5_137], [-0.2_676, 0.0_628, 0.6_186], [-0.3_629, -0.1_726, 0.4_716]]]    )

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




    @slow
    def UpperCamelCase			(			self :      Dict    ):

     """simple docstring"""






     SCREAMING_SNAKE_CASE								=			'the [MASK] of Belgium is Brussels'

     SCREAMING_SNAKE_CASE								=			AutoTokenizer.from_pretrained('uw-madison/nystromformer-512'    )
     SCREAMING_SNAKE_CASE								=			NystromformerForMaskedLM.from_pretrained('uw-madison/nystromformer-512'    )

     SCREAMING_SNAKE_CASE								=			tokenizer(snake_case__							,    return_tensors='pt'    )

     with torch.no_grad():
      SCREAMING_SNAKE_CASE								=			model(encoding.input_ids    ).logits

     SCREAMING_SNAKE_CASE								=			token_logits[:, 2, :].argmax(-1    )[0]

     self.assertEqual(tokenizer.decode(snake_case__    )							,    'capital'    )

