import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTNeoConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
    import jax
    import jax.numpy as jnp

    from transformers.modeling_flax_pytorch_utils import (
        convert_pytorch_state_dict_to_flax,
        load_flax_weights_in_pytorch_model,
    )
    from transformers.models.gpt_neo.modeling_flax_gpt_neo import FlaxGPTNeoForCausalLM, FlaxGPTNeoModel

if is_torch_available():
    import torch
class 				lowerCamelCase			:






      def __init__(      self  ,		__lowerCamelCase  ,		__lowerCamelCase=14  ,		__lowerCamelCase=7  ,		__lowerCamelCase=True  ,		__lowerCamelCase=True  ,		__lowerCamelCase=False  ,		__lowerCamelCase=True  ,		__lowerCamelCase=99  ,		__lowerCamelCase=32  ,		__lowerCamelCase=4  ,		__lowerCamelCase=4  ,		__lowerCamelCase=[[["global", "local"], 2]]  ,		__lowerCamelCase=37  ,		__lowerCamelCase="gelu"  ,		__lowerCamelCase=0.1  ,		__lowerCamelCase=0.1  ,		__lowerCamelCase=5_12  ,		__lowerCamelCase=7  ,		__lowerCamelCase=0.02  ,		)     ->				Union[str, Any]:


          '''simple docstring'''







          snake_case:						Dict         =   parent
          snake_case:						int         =   batch_size
          snake_case:						List[str]         =   seq_length
          snake_case:						List[str]         =   is_training
          snake_case:						Optional[int]         =   use_input_mask
          snake_case:						List[str]         =   use_token_type_ids
          snake_case:						Tuple         =   use_labels
          snake_case:						List[Any]         =   vocab_size
          snake_case:						int         =   hidden_size
          snake_case:						Union[str, Any]         =   num_hidden_layers
          snake_case:						Dict         =   num_attention_heads
          snake_case:						str         =   attention_types
          snake_case:						Any         =   intermediate_size
          snake_case:						Tuple         =   hidden_act
          snake_case:						Union[str, Any]         =   hidden_dropout_prob
          snake_case:						str         =   attention_probs_dropout_prob
          snake_case:						str         =   max_position_embeddings
          snake_case:						str         =   window_size
          snake_case:						Any         =   initializer_range
          snake_case:						Optional[int]         =   None
          snake_case:						Union[str, Any]         =   vocab_size - 1
          snake_case:						Union[str, Any]         =   vocab_size - 1
          snake_case:						str         =   vocab_size - 1






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


          '''simple docstring'''







          snake_case:						List[str]         =   ids_tensor([self.batch_size, self.seq_length]  ,		self.vocab_size						)

          snake_case:						Optional[int]         =   None
          if self.use_input_mask:
              snake_case:						str         =   random_attention_mask([self.batch_size, self.seq_length]						)

          snake_case:						Dict         =   GPTNeoConfig(
              vocab_size=self.vocab_size  ,		hidden_size=self.hidden_size  ,		num_layers=self.num_hidden_layers  ,		num_heads=self.num_attention_heads  ,		max_position_embeddings=self.max_position_embeddings  ,		use_cache=__lowerCamelCase  ,		bos_token_id=self.bos_token_id  ,		eos_token_id=self.eos_token_id  ,		pad_token_id=self.pad_token_id  ,		window_size=self.window_size  ,		attention_types=self.attention_types  ,		)

          return (config, input_ids, input_mask)






      def 			lowerCAmelCase_   (      self						)     ->				Dict:


          '''simple docstring'''







          snake_case:						Optional[Any]         =   self.prepare_config_and_inputs()
          snake_case    , snake_case    , snake_case:						Any         =   config_and_inputs
          snake_case:						List[str]         =   {"""input_ids""": input_ids, """attention_mask""": attention_mask}
          return config, inputs_dict






      def 			lowerCAmelCase_   (      self  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase						)     ->				str:


          '''simple docstring'''







          snake_case:						Optional[int]         =   20
          snake_case:						Optional[Any]         =   model_class_name(__lowerCamelCase						)

          snake_case:						List[str]         =   model.init_cache(input_ids.shape[0]  ,		__lowerCamelCase						)
          snake_case:						Dict         =   jnp.ones((input_ids.shape[0], max_decoder_length)  ,		dtype="""i4"""						)

          snake_case:						List[str]         =   jnp.broadcast_to(
              jnp.arange(input_ids.shape[-1] - 1						)[None, :]  ,		(input_ids.shape[0], input_ids.shape[-1] - 1)						)
          snake_case:						int         =   model(
              input_ids[:, :-1]  ,		attention_mask=__lowerCamelCase  ,		past_key_values=__lowerCamelCase  ,		position_ids=__lowerCamelCase  ,		)

          snake_case:						List[str]         =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]  ,		dtype="""i4"""						)
          snake_case:						List[str]         =   model(
              input_ids[:, -1:]  ,		attention_mask=__lowerCamelCase  ,		past_key_values=outputs_cache.past_key_values  ,		position_ids=__lowerCamelCase  ,		)

          snake_case:						Dict         =   model(__lowerCamelCase						)

          snake_case:						List[Any]         =   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])						)						)
          self.parent.assertTrue(diff < 1e-3  ,		msg=F"Max diff is {diff}"						)






      def 			lowerCAmelCase_   (      self  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase						)     ->				Optional[Any]:


          '''simple docstring'''







          snake_case:						Dict         =   20
          snake_case:						str         =   model_class_name(__lowerCamelCase						)

          snake_case:						int         =   jnp.concatenate(
              [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])						)]  ,		axis=-1  ,		)

          snake_case:						Optional[Any]         =   model.init_cache(input_ids.shape[0]  ,		__lowerCamelCase						)
          snake_case:						List[Any]         =   jnp.broadcast_to(
              jnp.arange(input_ids.shape[-1] - 1						)[None, :]  ,		(input_ids.shape[0], input_ids.shape[-1] - 1)						)

          snake_case:						Dict         =   model(
              input_ids[:, :-1]  ,		attention_mask=__lowerCamelCase  ,		past_key_values=__lowerCamelCase  ,		position_ids=__lowerCamelCase  ,		)
          snake_case:						Optional[int]         =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]  ,		dtype="""i4"""						)
          snake_case:						int         =   model(
              input_ids[:, -1:]  ,		past_key_values=outputs_cache.past_key_values  ,		attention_mask=__lowerCamelCase  ,		position_ids=__lowerCamelCase  ,		)

          snake_case:						Optional[Any]         =   model(__lowerCamelCase  ,		attention_mask=__lowerCamelCase						)

          snake_case:						Any         =   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])						)						)
          self.parent.assertTrue(diff < 1e-3  ,		msg=F"Max diff is {diff}"						)





@require_flax
class 				lowerCamelCase			(							__snake_case				,     __snake_case				,     unittest.TestCase					):
      __lowerCamelCase							=					(FlaxGPTNeoModel, FlaxGPTNeoForCausalLM) if is_flax_available() else ()
      __lowerCamelCase							=					(FlaxGPTNeoForCausalLM,) if is_flax_available() else ()






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


          '''simple docstring'''







          snake_case:						Dict         =   FlaxGPTNeoModelTester(self						)






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


          '''simple docstring'''







          for model_class_name in self.all_model_classes:
              snake_case    , snake_case    , snake_case:						str         =   self.model_tester.prepare_config_and_inputs()
              self.model_tester.check_use_cache_forward(__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase						)






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


          '''simple docstring'''







          for model_class_name in self.all_model_classes:
              snake_case    , snake_case    , snake_case:						List[Any]         =   self.model_tester.prepare_config_and_inputs()
              self.model_tester.check_use_cache_forward_with_attn_mask(
                  __lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase  ,		__lowerCamelCase						)






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


          '''simple docstring'''







          snake_case:						List[str]         =   GPTaTokenizer.from_pretrained("""gpt2"""  ,		pad_token="""<|endoftext|>"""  ,		padding_side="""left"""						)
          snake_case:						int         =   tokenizer(["""Hello this is a long string""", """Hey"""]  ,		return_tensors="""np"""  ,		padding=__lowerCamelCase  ,		truncation=__lowerCamelCase						)

          snake_case:						str         =   FlaxGPTNeoForCausalLM.from_pretrained("""EleutherAI/gpt-neo-125M"""						)
          snake_case:						int         =   False
          snake_case:						List[Any]         =   model.config.eos_token_id

          snake_case:						Any         =   jax.jit(model.generate						)

          snake_case:						Union[str, Any]         =   jit_generate(
              inputs["""input_ids"""]  ,		attention_mask=inputs["""attention_mask"""]  ,		pad_token_id=tokenizer.pad_token_id						).sequences

          snake_case:						int         =   tokenizer.batch_decode(__lowerCamelCase  ,		skip_special_tokens=__lowerCamelCase						)

          snake_case:						int         =   [
              """Hello this is a long string of text.\n\nI'm trying to get the text of the""",
              """Hey, I'm a little late to the party. I'm going to""",
          ]

          self.assertListEqual(__lowerCamelCase  ,		__lowerCamelCase						)






      @is_pt_flax_cross_test
      def 			lowerCAmelCase_   (      self						)     ->				Optional[Any]:


          '''simple docstring'''







          snake_case    , snake_case:						Tuple         =   self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__						):
                  # prepare inputs
                  snake_case:						Union[str, Any]         =   self._prepare_for_class(__lowerCamelCase  ,		__lowerCamelCase						)
                  snake_case:						Optional[Any]         =   {k: torch.tensor(v.tolist()						) for k, v in prepared_inputs_dict.items()}

                  # load corresponding PyTorch class
                  snake_case:						Dict         =   model_class.__name__[4:]  # Skip the "Flax" at the beginning
                  snake_case:						Optional[Any]         =   getattr(__lowerCamelCase  ,		__lowerCamelCase						)

                  snake_case    , snake_case:						Union[str, Any]         =   pt_inputs["""input_ids"""].shape
                  snake_case:						Union[str, Any]         =   np.random.randint(0  ,		seq_length - 1  ,		size=(batch_size,)						)
                  for batch_idx, start_index in enumerate(__lowerCamelCase						):
                      snake_case:						List[str]         =   0
                      snake_case:						str         =   1
                      snake_case:						Tuple         =   0
                      snake_case:						List[Any]         =   1
                  snake_case:						List[str]         =   pt_model_class(__lowerCamelCase						).eval()
                  snake_case:						Optional[Any]         =   model_class(__lowerCamelCase  ,		dtype=jnp.floataa						)

                  snake_case:						List[Any]         =   convert_pytorch_state_dict_to_flax(pt_model.state_dict()  ,		__lowerCamelCase						)
                  snake_case:						Union[str, Any]         =   fx_state

                  with torch.no_grad():
                      snake_case:						Any         =   pt_model(**__lowerCamelCase						).to_tuple()

                  snake_case:						Dict         =   fx_model(**__lowerCamelCase						).to_tuple()
                  self.assertEqual(len(__lowerCamelCase						)  ,		len(__lowerCamelCase						)  ,		"""Output lengths differ between Flax and PyTorch"""						)
                  for fx_output, pt_output in zip(__lowerCamelCase  ,		__lowerCamelCase						):
                      self.assert_almost_equals(fx_output[:, -1]  ,		pt_output[:, -1].numpy()  ,		4e-2						)

                  with tempfile.TemporaryDirectory() as tmpdirname:
                      pt_model.save_pretrained(__lowerCamelCase						)
                      snake_case:						Tuple         =   model_class.from_pretrained(__lowerCamelCase  ,		from_pt=__lowerCamelCase						)

                  snake_case:						Union[str, Any]         =   fx_model_loaded(**__lowerCamelCase						).to_tuple()
                  self.assertEqual(
                      len(__lowerCamelCase						)  ,		len(__lowerCamelCase						)  ,		"""Output lengths differ between Flax and PyTorch"""						)
                  for fx_output_loaded, pt_output in zip(__lowerCamelCase  ,		__lowerCamelCase						):
                      self.assert_almost_equals(fx_output_loaded[:, -1]  ,		pt_output[:, -1].numpy()  ,		4e-2						)






      @is_pt_flax_cross_test
      def 			lowerCAmelCase_   (      self						)     ->				str:


          '''simple docstring'''







          snake_case    , snake_case:						Tuple         =   self.model_tester.prepare_config_and_inputs_for_common()
          for model_class in self.all_model_classes:
              with self.subTest(model_class.__name__						):
                  # prepare inputs
                  snake_case:						int         =   self._prepare_for_class(__lowerCamelCase  ,		__lowerCamelCase						)
                  snake_case:						Any         =   {k: torch.tensor(v.tolist()						) for k, v in prepared_inputs_dict.items()}

                  # load corresponding PyTorch class
                  snake_case:						List[Any]         =   model_class.__name__[4:]  # Skip the "Flax" at the beginning
                  snake_case:						Union[str, Any]         =   getattr(__lowerCamelCase  ,		__lowerCamelCase						)

                  snake_case:						Optional[Any]         =   pt_model_class(__lowerCamelCase						).eval()
                  snake_case:						Dict         =   model_class(__lowerCamelCase  ,		dtype=jnp.floataa						)

                  snake_case:						str         =   load_flax_weights_in_pytorch_model(__lowerCamelCase  ,		fx_model.params						)
                  snake_case    , snake_case:						Optional[int]         =   pt_inputs["""input_ids"""].shape
                  snake_case:						Optional[int]         =   np.random.randint(0  ,		seq_length - 1  ,		size=(batch_size,)						)
                  for batch_idx, start_index in enumerate(__lowerCamelCase						):
                      snake_case:						Dict         =   0
                      snake_case:						Any         =   1
                      snake_case:						int         =   0
                      snake_case:						str         =   1

                  # make sure weights are tied in PyTorch
                  pt_model.tie_weights()

                  with torch.no_grad():
                      snake_case:						Optional[Any]         =   pt_model(**__lowerCamelCase						).to_tuple()

                  snake_case:						str         =   fx_model(**__lowerCamelCase						).to_tuple()
                  self.assertEqual(len(__lowerCamelCase						)  ,		len(__lowerCamelCase						)  ,		"""Output lengths differ between Flax and PyTorch"""						)
                  for fx_output, pt_output in zip(__lowerCamelCase  ,		__lowerCamelCase						):
                      self.assert_almost_equals(fx_output[:, -1]  ,		pt_output[:, -1].numpy()  ,		4e-2						)

                  with tempfile.TemporaryDirectory() as tmpdirname:
                      fx_model.save_pretrained(__lowerCamelCase						)
                      snake_case:						Tuple         =   pt_model_class.from_pretrained(__lowerCamelCase  ,		from_flax=__lowerCamelCase						)

                  with torch.no_grad():
                      snake_case:						List[str]         =   pt_model_loaded(**__lowerCamelCase						).to_tuple()

                  self.assertEqual(
                      len(__lowerCamelCase						)  ,		len(__lowerCamelCase						)  ,		"""Output lengths differ between Flax and PyTorch"""						)
                  for fx_output, pt_output in zip(__lowerCamelCase  ,		__lowerCamelCase						):
                      self.assert_almost_equals(fx_output[:, -1]  ,		pt_output[:, -1].numpy()  ,		4e-2						)






      @slow
      def 			lowerCAmelCase_   (      self						)     ->				Tuple:


          '''simple docstring'''







          for model_class_name in self.all_model_classes:
              snake_case:						int         =   model_class_name.from_pretrained("""EleutherAI/gpt-neo-125M"""						)
              snake_case:						Union[str, Any]         =   model(np.ones((1, 1)						)						)
              self.assertIsNotNone(__lowerCamelCase						)



from copy import deepcopy
from unittest.case import TestCase

import pytest

from datasets.arrow_dataset import Dataset
from datasets.features import Audio, ClassLabel, Features, Image, Sequence, Value
from datasets.info import DatasetInfo
from datasets.tasks import (
    AudioClassification,
    AutomaticSpeechRecognition,
    ImageClassification,
    LanguageModeling,
    QuestionAnsweringExtractive,
    Summarization,
    TextClassification,
    task_template_from_dict,
)
from datasets.utils.py_utils import asdict


__lowerCAmelCase    :  int										= {
    'id': '5733be284776f41900661182',
    'title': 'University_of_Notre_Dame',
    'context': 'Architecturally, the school has a Catholic character. Atop the Main Building\'s gold dome is a golden statue of the Virgin Mary. Immediately in front of the Main Building and facing it, is a copper statue of Christ with arms upraised with the legend "Venite Ad Me Omnes". Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. It is a replica of the grotto at Lourdes, France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous in 1858. At the end of the main drive (and in a direct line that connects through 3 statues and the Gold Dome), is a simple, modern stone statue of Mary.',
    'question': 'To whom did the Virgin Mary allegedly appear in 1858 in Lourdes France?',
    'answers': {'text': ['Saint Bernadette Soubirous'], 'answer_start': [515]},
}

@pytest.mark.parametrize(
    """task_cls"""			,	[
        AudioClassification,
        AutomaticSpeechRecognition,
        ImageClassification,
        LanguageModeling,
        QuestionAnsweringExtractive,
        Summarization,
        TextClassification,
    ]			,	)
def   a_					(_lowerCAmelCase  :   List[Any]       )->				Tuple:
    snake_case:						Dict         =   task_cls()
    snake_case:						Any         =   asdict(_lowerCAmelCase       )
    snake_case:						Any         =   task_template_from_dict(_lowerCAmelCase       )
    assert task == reloaded
class 				lowerCamelCase			:






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


          '''simple docstring'''







          snake_case:						str         =   LanguageModeling(text_column="""input_text"""						)
          assert {"input_text": "text"} == task.column_mapping






      def 			lowerCAmelCase_   (      self						)     ->				Tuple:


          '''simple docstring'''







          snake_case:						List[Any]         =   Features({"""text""": Value("""string"""						)}						)
          snake_case:						Optional[int]         =   {"""text_column""": """input_text"""}
          snake_case:						List[Any]         =   LanguageModeling.from_dict(__lowerCamelCase						)
          assert "language-modeling" == task.task
          assert input_schema == task.input_schema





class 				lowerCamelCase			(							__snake_case					):






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


          '''simple docstring'''







          snake_case:						Optional[int]         =   sorted(["""pos""", """neg"""]						)






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


          '''simple docstring'''







          snake_case:						str         =   TextClassification(text_column="""input_text"""  ,		label_column="""input_label"""						)
          self.assertDictEqual({"""input_text""": """text""", """input_label""": """labels"""}  ,		task.column_mapping						)






      def 			lowerCAmelCase_   (      self						)     ->				Any:


          '''simple docstring'''







          snake_case:						Any         =   Features({"""text""": Value("""string"""						)}						)
          # Labels are cast to tuple during `TextClassification.__post_init__`, so we do the same here
          snake_case:						List[str]         =   Features({"""labels""": ClassLabel}						)
          snake_case:						Any         =   {"""text_column""": """input_text""", """label_column""": """input_labels"""}
          snake_case:						List[str]         =   TextClassification.from_dict(__lowerCamelCase						)
          self.assertEqual("""text-classification"""  ,		task.task						)
          self.assertEqual(__lowerCamelCase  ,		task.input_schema						)
          self.assertEqual(__lowerCamelCase  ,		task.label_schema						)






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


          '''simple docstring'''







          snake_case:						str         =   TextClassification(text_column="""input_text"""  ,		label_column="""input_label"""						)
          self.assertEqual(task.label_schema["""labels"""]  ,		__lowerCamelCase						)
          snake_case:						Any         =   task.align_with_features(Features({"""input_label""": ClassLabel(names=self.labels						)}						)						)
          self.assertEqual(task.label_schema["""labels"""]  ,		ClassLabel(names=self.labels						)						)





class 				lowerCamelCase			(							__snake_case					):






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


          '''simple docstring'''







          snake_case:						List[Any]         =   QuestionAnsweringExtractive(
              context_column="""input_context"""  ,		question_column="""input_question"""  ,		answers_column="""input_answers"""						)
          self.assertDictEqual(
              {"""input_context""": """context""", """input_question""": """question""", """input_answers""": """answers"""}  ,		task.column_mapping						)






      def 			lowerCAmelCase_   (      self						)     ->				Any:


          '''simple docstring'''







          snake_case:						Any         =   Features({"""question""": Value("""string"""						), """context""": Value("""string"""						)}						)
          snake_case:						Tuple         =   Features(
              {
                  """answers""": Sequence(
                      {
                          """text""": Value("""string"""						),
                          """answer_start""": Value("""int32"""						),
                      }						)
              }						)
          snake_case:						Tuple         =   {
              """context_column""": """input_input_context""",
              """question_column""": """input_question""",
              """answers_column""": """input_answers""",
          }
          snake_case:						Optional[Any]         =   QuestionAnsweringExtractive.from_dict(__lowerCamelCase						)
          self.assertEqual("""question-answering-extractive"""  ,		task.task						)
          self.assertEqual(__lowerCamelCase  ,		task.input_schema						)
          self.assertEqual(__lowerCamelCase  ,		task.label_schema						)





class 				lowerCamelCase			(							__snake_case					):






      def 			lowerCAmelCase_   (      self						)     ->				Tuple:


          '''simple docstring'''







          snake_case:						Any         =   Summarization(text_column="""input_text"""  ,		summary_column="""input_summary"""						)
          self.assertDictEqual({"""input_text""": """text""", """input_summary""": """summary"""}  ,		task.column_mapping						)






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


          '''simple docstring'''







          snake_case:						List[Any]         =   Features({"""text""": Value("""string"""						)}						)
          snake_case:						Dict         =   Features({"""summary""": Value("""string"""						)}						)
          snake_case:						Tuple         =   {"""text_column""": """input_text""", """summary_column""": """input_summary"""}
          snake_case:						List[str]         =   Summarization.from_dict(__lowerCamelCase						)
          self.assertEqual("""summarization"""  ,		task.task						)
          self.assertEqual(__lowerCamelCase  ,		task.input_schema						)
          self.assertEqual(__lowerCamelCase  ,		task.label_schema						)





class 				lowerCamelCase			(							__snake_case					):






      def 			lowerCAmelCase_   (      self						)     ->				Dict:


          '''simple docstring'''







          snake_case:						str         =   AutomaticSpeechRecognition(audio_column="""input_audio"""  ,		transcription_column="""input_transcription"""						)
          self.assertDictEqual({"""input_audio""": """audio""", """input_transcription""": """transcription"""}  ,		task.column_mapping						)






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


          '''simple docstring'''







          snake_case:						Union[str, Any]         =   Features({"""audio""": Audio()}						)
          snake_case:						Union[str, Any]         =   Features({"""transcription""": Value("""string"""						)}						)
          snake_case:						Any         =   {
              """audio_column""": """input_audio""",
              """transcription_column""": """input_transcription""",
          }
          snake_case:						str         =   AutomaticSpeechRecognition.from_dict(__lowerCamelCase						)
          self.assertEqual("""automatic-speech-recognition"""  ,		task.task						)
          self.assertEqual(__lowerCamelCase  ,		task.input_schema						)
          self.assertEqual(__lowerCamelCase  ,		task.label_schema						)





class 				lowerCamelCase			(							__snake_case					):






      def 			lowerCAmelCase_   (      self						)     ->				int:


          '''simple docstring'''







          snake_case:						Union[str, Any]         =   sorted(["""pos""", """neg"""]						)






      def 			lowerCAmelCase_   (      self						)     ->				Dict:


          '''simple docstring'''







          snake_case:						Dict         =   AudioClassification(audio_column="""input_audio"""  ,		label_column="""input_label"""						)
          self.assertDictEqual({"""input_audio""": """audio""", """input_label""": """labels"""}  ,		task.column_mapping						)






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


          '''simple docstring'''







          snake_case:						Optional[Any]         =   Features({"""audio""": Audio()}						)
          snake_case:						int         =   Features({"""labels""": ClassLabel}						)
          snake_case:						int         =   {
              """audio_column""": """input_image""",
              """label_column""": """input_label""",
          }
          snake_case:						Optional[Any]         =   AudioClassification.from_dict(__lowerCamelCase						)
          self.assertEqual("""audio-classification"""  ,		task.task						)
          self.assertEqual(__lowerCamelCase  ,		task.input_schema						)
          self.assertEqual(__lowerCamelCase  ,		task.label_schema						)






      def 			lowerCAmelCase_   (      self						)     ->				str:


          '''simple docstring'''







          snake_case:						List[str]         =   AudioClassification(audio_column="""input_audio"""  ,		label_column="""input_label"""						)
          self.assertEqual(task.label_schema["""labels"""]  ,		__lowerCamelCase						)
          snake_case:						int         =   task.align_with_features(Features({"""input_label""": ClassLabel(names=self.labels						)}						)						)
          self.assertEqual(task.label_schema["""labels"""]  ,		ClassLabel(names=self.labels						)						)





class 				lowerCamelCase			(							__snake_case					):






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


          '''simple docstring'''







          snake_case:						Dict         =   sorted(["""pos""", """neg"""]						)






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


          '''simple docstring'''







          snake_case:						Optional[Any]         =   ImageClassification(image_column="""input_image"""  ,		label_column="""input_label"""						)
          self.assertDictEqual({"""input_image""": """image""", """input_label""": """labels"""}  ,		task.column_mapping						)






      def 			lowerCAmelCase_   (      self						)     ->				int:


          '''simple docstring'''







          snake_case:						Dict         =   Features({"""image""": Image()}						)
          snake_case:						str         =   Features({"""labels""": ClassLabel}						)
          snake_case:						List[str]         =   {
              """image_column""": """input_image""",
              """label_column""": """input_label""",
          }
          snake_case:						Any         =   ImageClassification.from_dict(__lowerCamelCase						)
          self.assertEqual("""image-classification"""  ,		task.task						)
          self.assertEqual(__lowerCamelCase  ,		task.input_schema						)
          self.assertEqual(__lowerCamelCase  ,		task.label_schema						)






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


          '''simple docstring'''







          snake_case:						Union[str, Any]         =   ImageClassification(image_column="""input_image"""  ,		label_column="""input_label"""						)
          self.assertEqual(task.label_schema["""labels"""]  ,		__lowerCamelCase						)
          snake_case:						str         =   task.align_with_features(Features({"""input_label""": ClassLabel(names=self.labels						)}						)						)
          self.assertEqual(task.label_schema["""labels"""]  ,		ClassLabel(names=self.labels						)						)





class 				lowerCamelCase			(							__snake_case					):






      def 			lowerCAmelCase_   (      self						)     ->				str:


          '''simple docstring'''







          snake_case:						List[str]         =   DatasetInfo(task_templates=QuestionAnsweringExtractive()						)
          snake_case:						int         =   Dataset.from_dict({k: [v] for k, v in SAMPLE_QUESTION_ANSWERING_EXTRACTIVE.items()}  ,		info=__lowerCamelCase						)
          assert isinstance(dataset.info.task_templates  ,		__lowerCamelCase						)
          assert len(dataset.info.task_templates						) == 1

          def keep_task(__lowerCamelCase						):
              return x

          def dont_keep_task(__lowerCamelCase						):
              snake_case:						Union[str, Any]         =   deepcopy(__lowerCamelCase						)
              snake_case:						int         =   0
              return out

          snake_case:						int         =   dataset.map(__lowerCamelCase						)
          assert mapped_dataset.info.task_templates == dataset.info.task_templates
          # reload from cache
          snake_case:						Optional[Any]         =   dataset.map(__lowerCamelCase						)
          assert mapped_dataset.info.task_templates == dataset.info.task_templates

          snake_case:						Optional[int]         =   dataset.map(__lowerCamelCase						)
          assert mapped_dataset.info.task_templates == []
          # reload from cache
          snake_case:						Union[str, Any]         =   dataset.map(__lowerCamelCase						)
          assert mapped_dataset.info.task_templates == []






      def 			lowerCAmelCase_   (      self						)     ->				int:


          '''simple docstring'''







          snake_case:						Any         =   Features({"""text""": Value("""string"""						), """label""": ClassLabel(names=("""pos""", """neg""")						)}						)
          snake_case:						Optional[int]         =   TextClassification(text_column="""text"""  ,		label_column="""label"""						)
          snake_case:						str         =   DatasetInfo(features=__lowerCamelCase  ,		task_templates=__lowerCamelCase						)
          snake_case:						Optional[Any]         =   Dataset.from_dict({"""text""": ["""A sentence."""], """label""": ["""pos"""]}  ,		info=__lowerCamelCase						)

          def process(__lowerCamelCase						):
              return example

          snake_case:						Optional[Any]         =   dataset.remove_columns("""label"""						)
          snake_case:						Tuple         =   modified_dataset.map(__lowerCamelCase						)
          assert mapped_dataset.info.task_templates == []


