import os
import tempfile
import unittest

import numpy as np
import torch

from diffusers.models import ModelMixin, UNetaDConditionModel
from diffusers.models.attention_processor import AttnProcessor, LoRAAttnProcessor
from diffusers.utils import (
    floats_tensor,
    logging,
    skip_mps,
    torch_device,
)
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism

from .test_modeling_common import ModelTesterMixin, UNetTesterMixin


enable_full_determinism()

UpperCAmelCase     						=			logging.get_logger(__name__)


def SCREAMING_SNAKE_CASE			(       lowercase_:		List[str]						,  lowercase_:		str = True			):
       _UpperCamelCase       :    List[str]					  =					{}
       for name in model.attn_processors.keys():
              _UpperCamelCase       :    Any					  =					name.endswith('attn2.processor'			) and not (
                  name.startswith('transformer_in'			) or 'temp_attentions' in name.split('.'			)
              )
              _UpperCamelCase       :    str					  =					model.config.cross_attention_dim if has_cross_attention else None
              if name.startswith('mid_block'			):
                     _UpperCamelCase       :    Optional[int]					  =					model.config.block_out_channels[-1]
              elif name.startswith('up_blocks'			):
                     _UpperCamelCase       :    Optional[int]					  =					int(name[len('up_blocks.'			)]			)
                     _UpperCamelCase       :    Tuple					  =					list(reversed(model.config.block_out_channels			)			)[block_id]
              elif name.startswith('down_blocks'			):
                     _UpperCamelCase       :    Optional[int]					  =					int(name[len('down_blocks.'			)]			)
                     _UpperCamelCase       :    int					  =					model.config.block_out_channels[block_id]
              elif name.startswith('transformer_in'			):
                     # Note that the `8 * ...` comes from: https://github.com/huggingface/diffusers/blob/7139f0e874f10b2463caa8cbd585762a309d12d6/src/diffusers/models/unet_3d_condition.py#L148
                     _UpperCamelCase       :    Any					  =					8 * model.config.attention_head_dim

              _UpperCamelCase       :    Dict					  =					LoRAAttnProcessor(hidden_size=lowercase_						,  cross_attention_dim=lowercase_			)
              _UpperCamelCase       :    Union[str, Any]					  =					lora_attn_procs[name].to(model.device			)

              if mock_weights:
                     # add 1 to weights to mock trained weights
                     with torch.no_grad():
                            lora_attn_procs[name].to_q_lora.up.weight += 1
                            lora_attn_procs[name].to_k_lora.up.weight += 1
                            lora_attn_procs[name].to_v_lora.up.weight += 1
                            lora_attn_procs[name].to_out_lora.up.weight += 1

       return lora_attn_procs





@skip_mps
class 				__lowerCAmelCase					(    snake_case__,		snake_case__,		unittest.TestCase  ):



 '''simple docstring'''
 _lowerCamelCase  :		Dict			      =	UNetaDConditionModel
 _lowerCamelCase  :		Tuple			      =	'''sample'''

 @property
 def    __snake_case   (     self    :						List[Any]					)	->							Tuple:
        _UpperCamelCase       :    str					  =					4
        _UpperCamelCase       :    Tuple					  =					4
        _UpperCamelCase       :    Any					  =					4
        _UpperCamelCase       :    Dict					  =					(3_2, 3_2)

        _UpperCamelCase       :    Tuple					  =					floats_tensor((batch_size, num_channels, num_frames) + sizes					).to(_A					)
        _UpperCamelCase       :    List[Any]					  =					torch.tensor([1_0]					).to(_A					)
        _UpperCamelCase       :    str					  =					floats_tensor((batch_size, 4, 3_2)					).to(_A					)

        return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states}

 @property
 def    __snake_case   (     self    :						List[str]					)	->							Any:
        return (4, 4, 3_2, 3_2)

 @property
 def    __snake_case   (     self    :						Union[str, Any]					)	->							int:
        return (4, 4, 3_2, 3_2)

 def    __snake_case   (     self    :						Tuple					)	->							List[Any]:
        _UpperCamelCase       :    Union[str, Any]					  =					{
            'block_out_channels': (3_2, 6_4),
            'down_block_types': (
                'CrossAttnDownBlock3D',
                'DownBlock3D',
            ),
            'up_block_types': ('UpBlock3D', 'CrossAttnUpBlock3D'),
            'cross_attention_dim': 3_2,
            'attention_head_dim': 8,
            'out_channels': 4,
            'in_channels': 4,
            'layers_per_block': 1,
            'sample_size': 3_2,
        }
        _UpperCamelCase       :    Tuple					  =					self.dummy_input
        return init_dict, inputs_dict

 @unittest.skipIf(
     torch_device != 'cuda' or not is_xformers_available()    ,     reason='XFormers attention is only available with CUDA and `xformers` installed'    ,     )
 def    __snake_case   (     self    :						Tuple					)	->							Optional[Any]:
        _UpperCamelCase       :    List[str]					  =					self.prepare_init_args_and_inputs_for_common()
        _UpperCamelCase       :    str					  =					self.model_class(**_A					)

        model.enable_xformers_memory_efficient_attention()

        assert (
            model.mid_block.attentions[0].transformer_blocks[0].attna.processor.__class__.__name__
            == "XFormersAttnProcessor"
        ), "xformers is not enabled"

 def    __snake_case   (     self    :						Optional[int]					)	->							Dict:
        _UpperCamelCase       :    Union[str, Any]					  =					self.prepare_init_args_and_inputs_for_common()

        _UpperCamelCase       :    Tuple					  =					3_2

        _UpperCamelCase       :    Any					  =					self.model_class(**_A					)
        model.to(_A					)
        model.eval()

        with torch.no_grad():
               _UpperCamelCase       :    List[str]					  =					model(**_A					)

               if isinstance(_A    ,     _A					):
                      _UpperCamelCase       :    Dict					  =					output.sample

        self.assertIsNotNone(_A					)
        _UpperCamelCase       :    Dict					  =					inputs_dict['sample'].shape
        self.assertEqual(output.shape    ,     _A    ,     'Input and output shapes do not match'					)

 def    __snake_case   (     self    :						Optional[int]					)	->							Dict:
        _UpperCamelCase       :    Optional[Any]					  =					self.prepare_init_args_and_inputs_for_common()
        _UpperCamelCase       :    Optional[Any]					  =					self.model_class(**_A					)
        model.to(_A					)
        model.eval()

        with torch.no_grad():
               # Warmup pass when using mps (see #372)
               if torch_device == "mps" and isinstance(_A    ,     _A					):
                      model(**self.dummy_input					)

               _UpperCamelCase       :    Optional[Any]					  =					model(**_A					)
               if isinstance(_A    ,     _A					):
                      _UpperCamelCase       :    Any					  =					first.sample

               _UpperCamelCase       :    int					  =					model(**_A					)
               if isinstance(_A    ,     _A					):
                      _UpperCamelCase       :    str					  =					second.sample

        _UpperCamelCase       :    Optional[int]					  =					first.cpu().numpy()
        _UpperCamelCase       :    Tuple					  =					second.cpu().numpy()
        _UpperCamelCase       :    str					  =					out_a[~np.isnan(_A					)]
        _UpperCamelCase       :    str					  =					out_a[~np.isnan(_A					)]
        _UpperCamelCase       :    Dict					  =					np.amax(np.abs(out_a - out_a					)					)
        self.assertLessEqual(_A    ,     1e-5					)

 def    __snake_case   (     self    :						Optional[int]					)	->							List[str]:
        _UpperCamelCase       :    List[str]					  =					self.prepare_init_args_and_inputs_for_common()

        _UpperCamelCase       :    Tuple					  =					8

        _UpperCamelCase       :    Optional[int]					  =					self.model_class(**_A					)
        model.to(_A					)
        model.eval()

        model.set_attention_slice('auto'					)
        with torch.no_grad():
               _UpperCamelCase       :    Any					  =					model(**_A					)
        assert output is not None

        model.set_attention_slice('max'					)
        with torch.no_grad():
               _UpperCamelCase       :    Optional[Any]					  =					model(**_A					)
        assert output is not None

        model.set_attention_slice(2					)
        with torch.no_grad():
               _UpperCamelCase       :    Optional[Any]					  =					model(**_A					)
        assert output is not None

 def    __snake_case   (     self    :						List[str]					)	->							List[Any]:
        _UpperCamelCase       :    Any					  =					self.prepare_init_args_and_inputs_for_common()

        _UpperCamelCase       :    str					  =					8

        _UpperCamelCase       :    Union[str, Any]					  =					self.model_class(**_A					)
        model.to(_A					)

        with torch.no_grad():
               _UpperCamelCase       :    Optional[int]					  =					model(**_A					).sample

        _UpperCamelCase       :    Optional[Any]					  =					create_lora_layers(_A					)

        # make sure we can set a list of attention processors
        model.set_attn_processor(_A					)
        model.to(_A					)

        # test that attn processors can be set to itself
        model.set_attn_processor(model.attn_processors					)

        with torch.no_grad():
               _UpperCamelCase       :    str					  =					model(**_A    ,     cross_attention_kwargs={'scale': 0.0}					).sample
               _UpperCamelCase       :    Any					  =					model(**_A    ,     cross_attention_kwargs={'scale': 0.5}					).sample
               _UpperCamelCase       :    Any					  =					model(**_A    ,     cross_attention_kwargs={'scale': 0.5}					).sample

        assert (samplea - samplea).abs().max() < 3e-3
        assert (samplea - samplea).abs().max() < 3e-3

        # sample 2 and sample 3 should be different
        assert (samplea - samplea).abs().max() > 3e-3

 def    __snake_case   (     self    :						Any					)	->							int:
        _UpperCamelCase       :    Union[str, Any]					  =					self.prepare_init_args_and_inputs_for_common()

        _UpperCamelCase       :    Any					  =					8

        torch.manual_seed(0					)
        _UpperCamelCase       :    Optional[Any]					  =					self.model_class(**_A					)
        model.to(_A					)

        with torch.no_grad():
               _UpperCamelCase       :    str					  =					model(**_A					).sample

        _UpperCamelCase       :    Optional[int]					  =					create_lora_layers(_A					)
        model.set_attn_processor(_A					)

        with torch.no_grad():
               _UpperCamelCase       :    List[str]					  =					model(**_A    ,     cross_attention_kwargs={'scale': 0.5}					).sample

        with tempfile.TemporaryDirectory() as tmpdirname:
               model.save_attn_procs(_A					)
               self.assertTrue(os.path.isfile(os.path.join(_A    ,     'pytorch_lora_weights.bin'					)					)					)
               torch.manual_seed(0					)
               _UpperCamelCase       :    Tuple					  =					self.model_class(**_A					)
               new_model.to(_A					)
               new_model.load_attn_procs(_A					)

        with torch.no_grad():
               _UpperCamelCase       :    List[str]					  =					new_model(**_A    ,     cross_attention_kwargs={'scale': 0.5}					).sample

        assert (sample - new_sample).abs().max() < 1e-3

        # LoRA and no LoRA should NOT be the same
        assert (sample - old_sample).abs().max() > 1e-4

 def    __snake_case   (     self    :						List[str]					)	->							List[Any]:
        _UpperCamelCase       :    str					  =					self.prepare_init_args_and_inputs_for_common()

        _UpperCamelCase       :    str					  =					8

        torch.manual_seed(0					)
        _UpperCamelCase       :    List[Any]					  =					self.model_class(**_A					)
        model.to(_A					)

        with torch.no_grad():
               _UpperCamelCase       :    Optional[Any]					  =					model(**_A					).sample

        _UpperCamelCase       :    Any					  =					create_lora_layers(_A					)
        model.set_attn_processor(_A					)

        with torch.no_grad():
               _UpperCamelCase       :    List[Any]					  =					model(**_A    ,     cross_attention_kwargs={'scale': 0.5}					).sample

        with tempfile.TemporaryDirectory() as tmpdirname:
               model.save_attn_procs(_A    ,     safe_serialization=_A					)
               self.assertTrue(os.path.isfile(os.path.join(_A    ,     'pytorch_lora_weights.safetensors'					)					)					)
               torch.manual_seed(0					)
               _UpperCamelCase       :    Dict					  =					self.model_class(**_A					)
               new_model.to(_A					)
               new_model.load_attn_procs(_A					)

        with torch.no_grad():
               _UpperCamelCase       :    Tuple					  =					new_model(**_A    ,     cross_attention_kwargs={'scale': 0.5}					).sample

        assert (sample - new_sample).abs().max() < 3e-3

        # LoRA and no LoRA should NOT be the same
        assert (sample - old_sample).abs().max() > 1e-4

 def    __snake_case   (     self    :						Dict					)	->							int:
        # enable deterministic behavior for gradient checkpointing
        _UpperCamelCase       :    Union[str, Any]					  =					self.prepare_init_args_and_inputs_for_common()

        _UpperCamelCase       :    int					  =					8

        torch.manual_seed(0					)
        _UpperCamelCase       :    int					  =					self.model_class(**_A					)
        model.to(_A					)

        _UpperCamelCase       :    Tuple					  =					create_lora_layers(_A    ,     mock_weights=_A					)
        model.set_attn_processor(_A					)
        # Saving as torch, properly reloads with directly filename
        with tempfile.TemporaryDirectory() as tmpdirname:
               model.save_attn_procs(_A					)
               self.assertTrue(os.path.isfile(os.path.join(_A    ,     'pytorch_lora_weights.bin'					)					)					)
               torch.manual_seed(0					)
               _UpperCamelCase       :    Optional[int]					  =					self.model_class(**_A					)
               new_model.to(_A					)
               new_model.load_attn_procs(_A    ,     weight_name='pytorch_lora_weights.bin'					)

 def    __snake_case   (     self    :						str					)	->							Optional[Any]:
        _UpperCamelCase       :    Dict					  =					self.prepare_init_args_and_inputs_for_common()

        _UpperCamelCase       :    List[str]					  =					8

        torch.manual_seed(0					)
        _UpperCamelCase       :    Union[str, Any]					  =					self.model_class(**_A					)
        model.to(_A					)

        _UpperCamelCase       :    Union[str, Any]					  =					create_lora_layers(_A    ,     mock_weights=_A					)
        model.set_attn_processor(_A					)
        # Saving as torch, properly reloads with directly filename
        with tempfile.TemporaryDirectory() as tmpdirname:
               model.save_attn_procs(_A					)
               self.assertTrue(os.path.isfile(os.path.join(_A    ,     'pytorch_lora_weights.bin'					)					)					)
               torch.manual_seed(0					)
               _UpperCamelCase       :    List[str]					  =					self.model_class(**_A					)
               new_model.to(_A					)
               with self.assertRaises(_A					) as e:
                      new_model.load_attn_procs(_A    ,     use_safetensors=_A					)
               self.assertIn('Error no file named pytorch_lora_weights.safetensors'    ,     str(e.exception					)					)

 def    __snake_case   (     self    :						Union[str, Any]					)	->							Any:
        _UpperCamelCase       :    Optional[int]					  =					self.prepare_init_args_and_inputs_for_common()

        _UpperCamelCase       :    List[str]					  =					8

        torch.manual_seed(0					)
        _UpperCamelCase       :    int					  =					self.model_class(**_A					)
        model.to(_A					)

        with torch.no_grad():
               _UpperCamelCase       :    Optional[int]					  =					model(**_A					).sample

        _UpperCamelCase       :    List[Any]					  =					create_lora_layers(_A					)
        model.set_attn_processor(_A					)

        with torch.no_grad():
               _UpperCamelCase       :    Any					  =					model(**_A    ,     cross_attention_kwargs={'scale': 0.0}					).sample

        model.set_attn_processor(AttnProcessor()					)

        with torch.no_grad():
               _UpperCamelCase       :    Dict					  =					model(**_A					).sample

        assert (sample - new_sample).abs().max() < 1e-4
        assert (sample - old_sample).abs().max() < 3e-3

 @unittest.skipIf(
     torch_device != 'cuda' or not is_xformers_available()    ,     reason='XFormers attention is only available with CUDA and `xformers` installed'    ,     )
 def    __snake_case   (     self    :						Optional[Any]					)	->							List[Any]:
        # enable deterministic behavior for gradient checkpointing
        _UpperCamelCase       :    Union[str, Any]					  =					self.prepare_init_args_and_inputs_for_common()

        _UpperCamelCase       :    Dict					  =					4

        torch.manual_seed(0					)
        _UpperCamelCase       :    Union[str, Any]					  =					self.model_class(**_A					)
        model.to(_A					)
        _UpperCamelCase       :    Union[str, Any]					  =					create_lora_layers(_A					)
        model.set_attn_processor(_A					)

        # default
        with torch.no_grad():
               _UpperCamelCase       :    Optional[int]					  =					model(**_A					).sample

               model.enable_xformers_memory_efficient_attention()
               _UpperCamelCase       :    Tuple					  =					model(**_A					).sample

               model.disable_xformers_memory_efficient_attention()
               _UpperCamelCase       :    Tuple					  =					model(**_A					).sample

        assert (sample - on_sample).abs().max() < 1e-4
        assert (sample - off_sample).abs().max() < 1e-4

 def    __snake_case   (     self    :						Optional[Any]					)	->							List[str]:
        _UpperCamelCase       :    int					  =					self.prepare_init_args_and_inputs_for_common()
        _UpperCamelCase       :    Optional[int]					  =					3_2

        _UpperCamelCase       :    Tuple					  =					self.model_class(**_A					)
        model.to(_A					)
        model.eval()

        with torch.no_grad():
               _UpperCamelCase       :    str					  =					model(**_A					)[0]

        model.enable_forward_chunking()
        with torch.no_grad():
               _UpperCamelCase       :    Tuple					  =					model(**_A					)[0]

        self.assertEqual(output.shape    ,     output_a.shape    ,     'Shape doesn\'t match'					)
        assert np.abs(output.cpu() - output_a.cpu()					).max() < 1e-2


# (todo: sayakpaul) implement SLOW tests.




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


UpperCAmelCase     						=			{
    '''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''': [5_15]},
}


@pytest.mark.parametrize(
    'task_cls'						,  [
        AudioClassification,
        AutomaticSpeechRecognition,
        ImageClassification,
        LanguageModeling,
        QuestionAnsweringExtractive,
        Summarization,
        TextClassification,
    ]						,  )
def SCREAMING_SNAKE_CASE			(       lowercase_:		Optional[int]			):
       _UpperCamelCase       :    int					  =					task_cls()
       _UpperCamelCase       :    List[Any]					  =					asdict(lowercase_			)
       _UpperCamelCase       :    Union[str, Any]					  =					task_template_from_dict(lowercase_			)
       assert task == reloaded




class 				__lowerCAmelCase					:



 '''simple docstring'''

 def    __snake_case   (     self    :						Dict					)	->							Any:
        _UpperCamelCase       :    List[Any]					  =					LanguageModeling(text_column='input_text'					)
        assert {"input_text": "text"} == task.column_mapping

 def    __snake_case   (     self    :						List[Any]					)	->							List[Any]:
        _UpperCamelCase       :    List[str]					  =					Features({'text': Value('string'					)}					)
        _UpperCamelCase       :    Optional[int]					  =					{'text_column': 'input_text'}
        _UpperCamelCase       :    Optional[Any]					  =					LanguageModeling.from_dict(lowerCamelCase					)
        assert "language-modeling" == task.task
        assert input_schema == task.input_schema







class 				__lowerCAmelCase					(    a_  ):



 '''simple docstring'''

 def    __snake_case   (     self    :						List[Any]					)	->							Union[str, Any]:
        _UpperCamelCase       :    Union[str, Any]					  =					sorted(['pos', 'neg']					)

 def    __snake_case   (     self    :						List[str]					)	->							Dict:
        _UpperCamelCase       :    List[str]					  =					TextClassification(text_column='input_text'    ,     label_column='input_label'					)
        self.assertDictEqual({'input_text': 'text', 'input_label': 'labels'}    ,     task.column_mapping					)

 def    __snake_case   (     self    :						Optional[Any]					)	->							Dict:
        _UpperCamelCase       :    Optional[Any]					  =					Features({'text': Value('string'					)}					)
        # Labels are cast to tuple during `TextClassification.__post_init__`, so we do the same here
        _UpperCamelCase       :    List[Any]					  =					Features({'labels': ClassLabel}					)
        _UpperCamelCase       :    Optional[int]					  =					{'text_column': 'input_text', 'label_column': 'input_labels'}
        _UpperCamelCase       :    str					  =					TextClassification.from_dict(lowerCamelCase					)
        self.assertEqual('text-classification'    ,     task.task					)
        self.assertEqual(lowerCamelCase    ,     task.input_schema					)
        self.assertEqual(lowerCamelCase    ,     task.label_schema					)

 def    __snake_case   (     self    :						str					)	->							Optional[int]:
        _UpperCamelCase       :    Any					  =					TextClassification(text_column='input_text'    ,     label_column='input_label'					)
        self.assertEqual(task.label_schema['labels']    ,     lowerCamelCase					)
        _UpperCamelCase       :    str					  =					task.align_with_features(Features({'input_label': ClassLabel(names=self.labels					)}					)					)
        self.assertEqual(task.label_schema['labels']    ,     ClassLabel(names=self.labels					)					)







class 				__lowerCAmelCase					(    a_  ):



 '''simple docstring'''

 def    __snake_case   (     self    :						Dict					)	->							str:
        _UpperCamelCase       :    Dict					  =					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    __snake_case   (     self    :						Union[str, Any]					)	->							Tuple:
        _UpperCamelCase       :    Optional[int]					  =					Features({'question': Value('string'					), 'context': Value('string'					)}					)
        _UpperCamelCase       :    Any					  =					Features(
            {
                'answers': Sequence(
                    {
                        'text': Value('string'					),
                        'answer_start': Value('int32'					),
                    }					)
            }					)
        _UpperCamelCase       :    str					  =					{
            'context_column': 'input_input_context',
            'question_column': 'input_question',
            'answers_column': 'input_answers',
        }
        _UpperCamelCase       :    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					(    a_  ):



 '''simple docstring'''

 def    __snake_case   (     self    :						Optional[Any]					)	->							List[str]:
        _UpperCamelCase       :    List[str]					  =					Summarization(text_column='input_text'    ,     summary_column='input_summary'					)
        self.assertDictEqual({'input_text': 'text', 'input_summary': 'summary'}    ,     task.column_mapping					)

 def    __snake_case   (     self    :						int					)	->							Dict:
        _UpperCamelCase       :    Tuple					  =					Features({'text': Value('string'					)}					)
        _UpperCamelCase       :    List[Any]					  =					Features({'summary': Value('string'					)}					)
        _UpperCamelCase       :    Any					  =					{'text_column': 'input_text', 'summary_column': 'input_summary'}
        _UpperCamelCase       :    Any					  =					Summarization.from_dict(lowerCamelCase					)
        self.assertEqual('summarization'    ,     task.task					)
        self.assertEqual(lowerCamelCase    ,     task.input_schema					)
        self.assertEqual(lowerCamelCase    ,     task.label_schema					)







class 				__lowerCAmelCase					(    a_  ):



 '''simple docstring'''

 def    __snake_case   (     self    :						Any					)	->							Optional[int]:
        _UpperCamelCase       :    Dict					  =					AutomaticSpeechRecognition(audio_column='input_audio'    ,     transcription_column='input_transcription'					)
        self.assertDictEqual({'input_audio': 'audio', 'input_transcription': 'transcription'}    ,     task.column_mapping					)

 def    __snake_case   (     self    :						List[Any]					)	->							int:
        _UpperCamelCase       :    Union[str, Any]					  =					Features({'audio': Audio()}					)
        _UpperCamelCase       :    Any					  =					Features({'transcription': Value('string'					)}					)
        _UpperCamelCase       :    Optional[int]					  =					{
            'audio_column': 'input_audio',
            'transcription_column': 'input_transcription',
        }
        _UpperCamelCase       :    int					  =					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					(    a_  ):



 '''simple docstring'''

 def    __snake_case   (     self    :						Union[str, Any]					)	->							Dict:
        _UpperCamelCase       :    Dict					  =					sorted(['pos', 'neg']					)

 def    __snake_case   (     self    :						Optional[int]					)	->							Optional[int]:
        _UpperCamelCase       :    Union[str, Any]					  =					AudioClassification(audio_column='input_audio'    ,     label_column='input_label'					)
        self.assertDictEqual({'input_audio': 'audio', 'input_label': 'labels'}    ,     task.column_mapping					)

 def    __snake_case   (     self    :						Tuple					)	->							Dict:
        _UpperCamelCase       :    List[str]					  =					Features({'audio': Audio()}					)
        _UpperCamelCase       :    Union[str, Any]					  =					Features({'labels': ClassLabel}					)
        _UpperCamelCase       :    Union[str, Any]					  =					{
            'audio_column': 'input_image',
            'label_column': 'input_label',
        }
        _UpperCamelCase       :    Tuple					  =					AudioClassification.from_dict(lowerCamelCase					)
        self.assertEqual('audio-classification'    ,     task.task					)
        self.assertEqual(lowerCamelCase    ,     task.input_schema					)
        self.assertEqual(lowerCamelCase    ,     task.label_schema					)

 def    __snake_case   (     self    :						Any					)	->							List[Any]:
        _UpperCamelCase       :    Any					  =					AudioClassification(audio_column='input_audio'    ,     label_column='input_label'					)
        self.assertEqual(task.label_schema['labels']    ,     lowerCamelCase					)
        _UpperCamelCase       :    Union[str, Any]					  =					task.align_with_features(Features({'input_label': ClassLabel(names=self.labels					)}					)					)
        self.assertEqual(task.label_schema['labels']    ,     ClassLabel(names=self.labels					)					)







class 				__lowerCAmelCase					(    a_  ):



 '''simple docstring'''

 def    __snake_case   (     self    :						int					)	->							str:
        _UpperCamelCase       :    Tuple					  =					sorted(['pos', 'neg']					)

 def    __snake_case   (     self    :						Optional[Any]					)	->							Any:
        _UpperCamelCase       :    str					  =					ImageClassification(image_column='input_image'    ,     label_column='input_label'					)
        self.assertDictEqual({'input_image': 'image', 'input_label': 'labels'}    ,     task.column_mapping					)

 def    __snake_case   (     self    :						int					)	->							Dict:
        _UpperCamelCase       :    Dict					  =					Features({'image': Image()}					)
        _UpperCamelCase       :    Tuple					  =					Features({'labels': ClassLabel}					)
        _UpperCamelCase       :    Optional[int]					  =					{
            'image_column': 'input_image',
            'label_column': 'input_label',
        }
        _UpperCamelCase       :    List[str]					  =					ImageClassification.from_dict(lowerCamelCase					)
        self.assertEqual('image-classification'    ,     task.task					)
        self.assertEqual(lowerCamelCase    ,     task.input_schema					)
        self.assertEqual(lowerCamelCase    ,     task.label_schema					)

 def    __snake_case   (     self    :						List[str]					)	->							Union[str, Any]:
        _UpperCamelCase       :    Any					  =					ImageClassification(image_column='input_image'    ,     label_column='input_label'					)
        self.assertEqual(task.label_schema['labels']    ,     lowerCamelCase					)
        _UpperCamelCase       :    Optional[Any]					  =					task.align_with_features(Features({'input_label': ClassLabel(names=self.labels					)}					)					)
        self.assertEqual(task.label_schema['labels']    ,     ClassLabel(names=self.labels					)					)







class 				__lowerCAmelCase					(    a_  ):



 '''simple docstring'''

 def    __snake_case   (     self    :						List[str]					)	->							Dict:
        _UpperCamelCase       :    List[str]					  =					DatasetInfo(task_templates=QuestionAnsweringExtractive()					)
        _UpperCamelCase       :    Tuple					  =					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    :						Dict					):
               return x

        def dont_keep_task(lowerCamelCase    :						List[Any]					):
               _UpperCamelCase       :    Union[str, Any]					  =					deepcopy(lowerCamelCase					)
               _UpperCamelCase       :    List[str]					  =					0
               return out

        _UpperCamelCase       :    Any					  =					dataset.map(lowerCamelCase					)
        assert mapped_dataset.info.task_templates == dataset.info.task_templates
        # reload from cache
        _UpperCamelCase       :    Any					  =					dataset.map(lowerCamelCase					)
        assert mapped_dataset.info.task_templates == dataset.info.task_templates

        _UpperCamelCase       :    int					  =					dataset.map(lowerCamelCase					)
        assert mapped_dataset.info.task_templates == []
        # reload from cache
        _UpperCamelCase       :    int					  =					dataset.map(lowerCamelCase					)
        assert mapped_dataset.info.task_templates == []

 def    __snake_case   (     self    :						Union[str, Any]					)	->							Any:
        _UpperCamelCase       :    List[str]					  =					Features({'text': Value('string'					), 'label': ClassLabel(names=('pos', 'neg')					)}					)
        _UpperCamelCase       :    List[str]					  =					TextClassification(text_column='text'    ,     label_column='label'					)
        _UpperCamelCase       :    List[str]					  =					DatasetInfo(features=lowerCamelCase    ,     task_templates=lowerCamelCase					)
        _UpperCamelCase       :    Any					  =					Dataset.from_dict({'text': ['A sentence.'], 'label': ['pos']}    ,     info=lowerCamelCase					)

        def process(lowerCamelCase    :						Any					):
               return example

        _UpperCamelCase       :    Dict					  =					dataset.remove_columns('label'					)
        _UpperCamelCase       :    Union[str, Any]					  =					modified_dataset.map(lowerCamelCase					)
        assert mapped_dataset.info.task_templates == []



