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 				a_ :






 '''simple docstring'''






 def __init__(			self					:  Tuple			,    lowercase__					:  Optional[int]			,    lowercase__					:  Optional[int]=13			,    lowercase__					:  int=7			,    lowercase__					:  str=True			,    lowercase__					:  str=True			,    lowercase__					:  List[str]=True			,    lowercase__					:  List[Any]=True			,    lowercase__					:  Tuple=99			,    lowercase__					:  Any=32			,    lowercase__					:  Tuple=5			,    lowercase__					:  Dict=4			,    lowercase__					:  Optional[Any]=37			,    lowercase__					:  Tuple="gelu"			,    lowercase__					:  Tuple=0.1			,    lowercase__					:  Any=0.1			,    lowercase__					:  str=512			,    lowercase__					:  int=16			,    lowercase__					:  List[Any]=2			,    lowercase__					:  str=0.02			,    lowercase__					:  Union[str, Any]=3			,    lowercase__					:  List[str]=4			,    lowercase__					:  Dict=None			,    ):







        '''simple docstring'''



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




 def 		__snake_case						(			self					:  Optional[int]):







        '''simple docstring'''



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

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

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

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

        lowerCAmelCase__     	=					self.get_config()

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




 def 		__snake_case						(			self					:  Dict):







        '''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 		__snake_case						(			self					:  Union[str, Any]):







        '''simple docstring'''



        lowerCAmelCase__     	=					self.get_config()
        lowerCAmelCase__     	=					300
        return config




 def 		__snake_case						(			self					:  str):







        '''simple docstring'''



        (
            (
            lowerCAmelCase__
        )	,					(
            lowerCAmelCase__
        )	,					(
            lowerCAmelCase__
        )	,					(
            lowerCAmelCase__
        )	,					(
            lowerCAmelCase__
        )	,					(
            lowerCAmelCase__
        )	,					(
            lowerCAmelCase__
        )	,					
        )     	=					self.prepare_config_and_inputs()

        lowerCAmelCase__     	=					True
        lowerCAmelCase__     	=					floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
        lowerCAmelCase__     	=					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 		__snake_case						(			self					:  int			,    lowercase__					:  Tuple			,    lowercase__					:  int			,    lowercase__					:  Any			,    lowercase__					:  Optional[Any]			,    lowercase__					:  str			,    lowercase__					:  Tuple			,    lowercase__					:  List[Any]):







        '''simple docstring'''



        lowerCAmelCase__     	=					YosoModel(config=lowercase__)
        model.to(lowercase__)
        model.eval()
        lowerCAmelCase__     	=					model(lowercase__			,    attention_mask=lowercase__			,    token_type_ids=lowercase__)
        lowerCAmelCase__     	=					model(lowercase__			,    token_type_ids=lowercase__)
        lowerCAmelCase__     	=					model(lowercase__)
        self.parent.assertEqual(result.last_hidden_state.shape			,    (self.batch_size, self.seq_length, self.hidden_size))




 def 		__snake_case						(			self					:  Tuple			,    lowercase__					:  Optional[int]			,    lowercase__					:  List[Any]			,    lowercase__					:  Any			,    lowercase__					:  Optional[Any]			,    lowercase__					:  Dict			,    lowercase__					:  List[str]			,    lowercase__					:  Optional[Any]			,    lowercase__					:  Tuple			,    lowercase__					:  int			,    ):







        '''simple docstring'''



        lowerCAmelCase__     	=					True
        lowerCAmelCase__     	=					YosoModel(lowercase__)
        model.to(lowercase__)
        model.eval()
        lowerCAmelCase__     	=					model(
            lowercase__			,    attention_mask=lowercase__			,    token_type_ids=lowercase__			,    encoder_hidden_states=lowercase__			,    encoder_attention_mask=lowercase__			,    )
        lowerCAmelCase__     	=					model(
            lowercase__			,    attention_mask=lowercase__			,    token_type_ids=lowercase__			,    encoder_hidden_states=lowercase__			,    )
        lowerCAmelCase__     	=					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 		__snake_case						(			self					:  Dict			,    lowercase__					:  Any			,    lowercase__					:  List[Any]			,    lowercase__					:  Any			,    lowercase__					:  Optional[Any]			,    lowercase__					:  Dict			,    lowercase__					:  List[Any]			,    lowercase__					:  Tuple):







        '''simple docstring'''



        lowerCAmelCase__     	=					YosoForMaskedLM(config=lowercase__)
        model.to(lowercase__)
        model.eval()
        lowerCAmelCase__     	=					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 		__snake_case						(			self					:  Any			,    lowercase__					:  Tuple			,    lowercase__					:  Optional[Any]			,    lowercase__					:  int			,    lowercase__					:  List[Any]			,    lowercase__					:  List[str]			,    lowercase__					:  List[str]			,    lowercase__					:  Optional[int]):







        '''simple docstring'''



        lowerCAmelCase__     	=					YosoForQuestionAnswering(config=lowercase__)
        model.to(lowercase__)
        model.eval()
        lowerCAmelCase__     	=					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 		__snake_case						(			self					:  Optional[int]			,    lowercase__					:  List[Any]			,    lowercase__					:  Optional[Any]			,    lowercase__					:  Union[str, Any]			,    lowercase__					:  Any			,    lowercase__					:  Optional[int]			,    lowercase__					:  Tuple			,    lowercase__					:  Union[str, Any]):







        '''simple docstring'''



        lowerCAmelCase__     	=					self.num_labels
        lowerCAmelCase__     	=					YosoForSequenceClassification(lowercase__)
        model.to(lowercase__)
        model.eval()
        lowerCAmelCase__     	=					model(lowercase__			,    attention_mask=lowercase__			,    token_type_ids=lowercase__			,    labels=lowercase__)
        self.parent.assertEqual(result.logits.shape			,    (self.batch_size, self.num_labels))




 def 		__snake_case						(			self					:  Union[str, Any]			,    lowercase__					:  Optional[Any]			,    lowercase__					:  Optional[Any]			,    lowercase__					:  Union[str, Any]			,    lowercase__					:  List[Any]			,    lowercase__					:  int			,    lowercase__					:  List[str]			,    lowercase__					:  Any):







        '''simple docstring'''



        lowerCAmelCase__     	=					self.num_labels
        lowerCAmelCase__     	=					YosoForTokenClassification(config=lowercase__)
        model.to(lowercase__)
        model.eval()
        lowerCAmelCase__     	=					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 		__snake_case						(			self					:  Any			,    lowercase__					:  Tuple			,    lowercase__					:  int			,    lowercase__					:  List[str]			,    lowercase__					:  Union[str, Any]			,    lowercase__					:  List[str]			,    lowercase__					:  Union[str, Any]			,    lowercase__					:  int):







        '''simple docstring'''



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




 def 		__snake_case						(			self					:  str):







        '''simple docstring'''



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






@require_torch
class 				a_ (						SCREAMING_SNAKE_CASE       ,     SCREAMING_SNAKE_CASE       ,     unittest.TestCase						):






 '''simple docstring'''






 UpperCAmelCase_	 =					(
     (
         YosoModel,
         YosoForMaskedLM,
         YosoForMultipleChoice,
         YosoForQuestionAnswering,
         YosoForSequenceClassification,
         YosoForTokenClassification,
     )
     if is_torch_available()
     else ()
 )
 UpperCAmelCase_	 =					False
 UpperCAmelCase_	 =					False
 UpperCAmelCase_	 =					False

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




 def 		__snake_case						(			self					:  str):







        '''simple docstring'''



        lowerCAmelCase__     	=					YosoModelTester(self)
        lowerCAmelCase__     	=					ConfigTester(self			,    config_class=lowercase__			,    hidden_size=37)




 def 		__snake_case						(			self					:  Tuple):







        '''simple docstring'''



        self.config_tester.run_common_tests()




 def 		__snake_case						(			self					:  List[str]):







        '''simple docstring'''



        lowerCAmelCase__     	=					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*lowercase__)




 def 		__snake_case						(			self					:  List[str]):







        '''simple docstring'''



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




 def 		__snake_case						(			self					:  str):







        '''simple docstring'''



        lowerCAmelCase__     	=					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*lowercase__)




 def 		__snake_case						(			self					:  int):







        '''simple docstring'''



        lowerCAmelCase__     	=					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_multiple_choice(*lowercase__)




 def 		__snake_case						(			self					:  int):







        '''simple docstring'''



        lowerCAmelCase__     	=					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_question_answering(*lowercase__)




 def 		__snake_case						(			self					:  int):







        '''simple docstring'''



        lowerCAmelCase__     	=					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_sequence_classification(*lowercase__)




 def 		__snake_case						(			self					:  List[Any]):







        '''simple docstring'''



        lowerCAmelCase__     	=					self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*lowercase__)




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







        '''simple docstring'''



        for model_name in YOSO_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               lowerCAmelCase__     	=					YosoModel.from_pretrained(lowercase__)
               self.assertIsNotNone(lowercase__)




 def 		__snake_case						(			self					:  List[str]):







        '''simple docstring'''



        return






@require_torch
class 				a_ (						unittest.TestCase						):






 '''simple docstring'''






 @slow
 def 		__snake_case						(			self					:  str):







        '''simple docstring'''



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

        with torch.no_grad():
               lowerCAmelCase__     	=					model(lowercase__)[0]

        lowerCAmelCase__     	=					torch.Size((1, 6, 768))
        self.assertEqual(output.shape			,    lowercase__)

        lowerCAmelCase__     	=					torch.tensor(
            [[[-0.0_611, 0.1_242, 0.0_840], [0.0_280, -0.0_048, 0.1_125], [0.0_106, 0.0_226, 0.0_751]]])

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




 @slow
 def 		__snake_case						(			self					:  List[Any]):







        '''simple docstring'''



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

        with torch.no_grad():
               lowerCAmelCase__     	=					model(lowercase__)[0]

        lowerCAmelCase__     	=					50_265

        lowerCAmelCase__     	=					torch.Size((1, 6, vocab_size))
        self.assertEqual(output.shape			,    lowercase__)

        lowerCAmelCase__     	=					torch.tensor(
            [[[-2.1_313, -3.7_285, -2.2_407], [-2.7_047, -3.3_314, -2.6_408], [0.0_629, -2.5_166, -0.3_356]]])

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




 @slow
 def 		__snake_case						(			self					:  Union[str, Any]):







        '''simple docstring'''



        lowerCAmelCase__     	=					YosoForMaskedLM.from_pretrained('uw-madison/yoso-4096')
        lowerCAmelCase__     	=					torch.arange(4_096).unsqueeze(0)

        with torch.no_grad():
               lowerCAmelCase__     	=					model(lowercase__)[0]

        lowerCAmelCase__     	=					50_265

        lowerCAmelCase__     	=					torch.Size((1, 4_096, vocab_size))
        self.assertEqual(output.shape			,    lowercase__)

        lowerCAmelCase__     	=					torch.tensor(
            [[[-2.3_914, -4.3_742, -5.0_956], [-4.0_988, -4.2_384, -7.0_406], [-3.1_427, -3.7_192, -6.6_800]]])

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

import itertools
from dataclasses import dataclass
from typing import Optional

import pyarrow as pa

import datasets
from datasets.table import table_cast


lowerCAmelCase__							= datasets.utils.logging.get_logger(__name__)
@dataclass
class 				a_ (						datasets.BuilderConfig						):






 '''simple docstring'''






 UpperCAmelCase_	 =					None






class 				a_ (						datasets.ArrowBasedBuilder						):






 '''simple docstring'''






 UpperCAmelCase_	 =					ArrowConfig




 def 		__snake_case						(			self					:  Any):







        '''simple docstring'''



        return datasets.DatasetInfo(features=self.config.features)




 def 		__snake_case						(			self					:  Dict			,    lowercase__					:  str):







        '''simple docstring'''



        if not self.config.data_files:
               raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}""")
        lowerCAmelCase__     	=					dl_manager.download_and_extract(self.config.data_files)
        if isinstance(lowercase__			,    (str, list, tuple)):
               lowerCAmelCase__     	=					data_files
               if isinstance(lowercase__			,    lowercase__):
                      lowerCAmelCase__     	=					[files]
               # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
               lowerCAmelCase__     	=					[dl_manager.iter_files(lowercase__) for file in files]
               return [datasets.SplitGenerator(name=datasets.Split.TRAIN			,    gen_kwargs={'files': files})]
        lowerCAmelCase__     	=					[]
        for split_name, files in data_files.items():
               if isinstance(lowercase__			,    lowercase__):
                      lowerCAmelCase__     	=					[files]
               # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
               lowerCAmelCase__     	=					[dl_manager.iter_files(lowercase__) for file in files]
               # Infer features is they are stoed in the arrow schema
               if self.info.features is None:
                      for file in itertools.chain.from_iterable(lowercase__):
                             with open(lowercase__			,    'rb') as f:
                                    lowerCAmelCase__     	=					datasets.Features.from_arrow_schema(pa.ipc.open_stream(lowercase__).schema)
                             break
               splits.append(datasets.SplitGenerator(name=lowercase__			,    gen_kwargs={'files': files}))
        return splits




 def 		__snake_case						(			self					:  Dict			,    lowercase__					:  pa.Table):







        '''simple docstring'''



        if self.info.features is not None:
               # more expensive cast to support nested features with keys in a different order
               # allows str <-> int/float or str to Audio for example
               lowerCAmelCase__     	=					table_cast(lowercase__			,    self.info.features.arrow_schema)
        return pa_table




 def 		__snake_case						(			self					:  List[Any]			,    lowercase__					:  str):







        '''simple docstring'''



        for file_idx, file in enumerate(itertools.chain.from_iterable(lowercase__)):
               with open(lowercase__			,    'rb') as f:
                      try:
                             for batch_idx, record_batch in enumerate(pa.ipc.open_stream(lowercase__)):
                                    lowerCAmelCase__     	=					pa.Table.from_batches([record_batch])
                                    # Uncomment for debugging (will print the Arrow table size and elements)
                                    # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                                    # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                                    yield F"""{file_idx}_{batch_idx}""", self._cast_table(lowercase__)
                      except ValueError as e:
                             logger.error(F"""Failed to read file '{file}' with error {type(lowercase__)}: {e}""")
                             raise
