




'''simple docstring'''
import argparse
import logging
import os
import sys
import tempfile
from pathlib import Path

import lightning_base
import pytest
import pytorch_lightning as pl
import torch
from convert_pl_checkpoint_to_hf import convert_pl_to_hf
from distillation import distill_main
from finetune import SummarizationModule, main
from huggingface_hub import list_models
from parameterized import parameterized
from run_eval import generate_summaries_or_translations
from torch import nn

from transformers import AutoConfig, AutoModelForSeqaSeqLM
from transformers.testing_utils import CaptureStderr, CaptureStdout, TestCasePlus, require_torch_gpu, slow
from utils import label_smoothed_nll_loss, lmap, load_json


logging.basicConfig(level=logging.DEBUG)

A							:				Union[str, Any]			     =	logging.getLogger()
A							:				str			     =	torch.cuda.is_available()
A							:				Optional[Any]			     =	{
    """max_tokens_per_batch""": None,
    """supervise_forward""": True,
    """normalize_hidden""": True,
    """label_smoothing""": 0.2,
    """eval_max_gen_length""": None,
    """eval_beams""": 1,
    """val_metric""": """loss""",
    """save_top_k""": 1,
    """adafactor""": True,
    """early_stopping_patience""": 2,
    """logger_name""": """default""",
    """length_penalty""": 0.5,
    """cache_dir""": """""",
    """task""": """summarization""",
    """num_workers""": 2,
    """alpha_hid""": 0,
    """freeze_embeds""": True,
    """enc_only""": False,
    """tgt_suffix""": """""",
    """resume_from_checkpoint""": None,
    """sortish_sampler""": True,
    """student_decoder_layers""": 1,
    """val_check_interval""": 1.0,
    """output_dir""": """""",
    """fp16""": False,  # TODO(SS): set this to CUDA_AVAILABLE if ci installs apex or start using native amp
    """no_teacher""": False,
    """fp16_opt_level""": """O1""",
    """gpus""": 1 if CUDA_AVAILABLE else 0,
    """n_tpu_cores""": 0,
    """max_grad_norm""": 1.0,
    """do_train""": True,
    """do_predict""": True,
    """accumulate_grad_batches""": 1,
    """server_ip""": """""",
    """server_port""": """""",
    """seed""": 42,
    """model_name_or_path""": """sshleifer/bart-tiny-random""",
    """config_name""": """""",
    """tokenizer_name""": """facebook/bart-large""",
    """do_lower_case""": False,
    """learning_rate""": 0.3,
    """lr_scheduler""": """linear""",
    """weight_decay""": 0.0,
    """adam_epsilon""": 1E-08,
    """warmup_steps""": 0,
    """max_epochs""": 1,
    """train_batch_size""": 2,
    """eval_batch_size""": 2,
    """max_source_length""": 12,
    """max_target_length""": 12,
    """val_max_target_length""": 12,
    """test_max_target_length""": 12,
    """fast_dev_run""": False,
    """no_cache""": False,
    """n_train""": -1,
    """n_val""": -1,
    """n_test""": -1,
    """student_encoder_layers""": 1,
    """freeze_encoder""": False,
    """auto_scale_batch_size""": False,
    """overwrite_output_dir""": False,
    """student""": None,
}
def     snake_case_  (       a__		:     Path   ,a__		:     list      ):



    """simple docstring"""




    __lowercase          =						"""\n""".join(a__      )
    Path(a__      ).open("""w"""      ).writelines(a__      )


A							:				List[Any]			     =	[""" Sam ate lunch today.""", """Sams lunch ingredients."""]
A							:				int			     =	["""A very interesting story about what I ate for lunch.""", """Avocado, celery, turkey, coffee"""]
A							:				Dict			     =	"""patrickvonplaten/t5-tiny-random"""
A							:				Tuple			     =	"""sshleifer/t5-tinier-random"""
A							:				int			     =	"""sshleifer/bart-tiny-random"""
A							:				List[Any]			     =	"""sshleifer/tiny-mbart"""
A							:				Optional[int]			     =	"""sshleifer/tiny-marian-en-de"""
A							:				Dict			     =	"""stas/tiny-wmt19-en-de"""


A							:				Tuple			     =	logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)
logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks
def     snake_case_  (       a__		:     str      ):



    """simple docstring"""




    for split in ["train", "val", "test"]:
        _dump_articles(os.path.join(a__   ,f'{split}.source'      )   ,a__      )
        _dump_articles(os.path.join(a__   ,f'{split}.target'      )   ,a__      )
    return tmp_dir






class 	SCREAMING_SNAKE_CASE(    __A					):

       @classmethod
       def    snake_case__      (   cls					)   -> Optional[int]:

           """simple docstring"""



           logging.disable(logging.CRITICAL					)  # remove noisy download output from tracebacks
           return cls

       @slow
       @require_torch_gpu
       def    snake_case__      (   self					)   -> Tuple:

           """simple docstring"""



           __lowercase          =						list_models()
           __lowercase          =						"""sshleifer"""
           __lowercase          =						[x.modelId for x in model_list if x.modelId.startswith(lowerCamelCase__					)]
           __lowercase          =						["""sshleifer/blenderbot-3B""", """sshleifer/blenderbot-90M"""]
           __lowercase          =						[]
           for m in model_ids:
               if m in allowed_to_be_broken:
                   continue
               try:
                   AutoConfig.from_pretrained(lowerCamelCase__					)
               except Exception:
                   failures.append(lowerCamelCase__					)
           assert not failures, F'The following models could not be loaded through AutoConfig: {failures}'

       def    snake_case__      (   self					)   -> Any:

           """simple docstring"""



           __lowercase          =						{"""student_encoder_layers""": 2, """student_decoder_layers""": 1, """no_teacher""": True}
           self._test_distiller_cli(lowerCamelCase__					)

       def    snake_case__      (   self					)   -> Any:

           """simple docstring"""



           __lowercase          =						{
               """student_encoder_layers""": 2,
               """student_decoder_layers""": 1,
               """max_epochs""": 4,
               """val_check_interval""": 0.25,
               """alpha_hid""": 2.0,
               """model_name_or_path""": """IGNORE_THIS_IT_DOESNT_GET_USED""",
           }
           __lowercase          =						self._test_distiller_cli(lowerCamelCase__ ,		check_contents=lowerCamelCase__					)

           __lowercase          =						list(Path(model.output_dir					).glob("""*.ckpt"""					)					)
           self.assertEqual(1 ,		len(lowerCamelCase__					)					)
           __lowercase          =						list(Path(model.output_dir					).glob("""**/*.bin"""					)					)
           self.assertEqual(len(lowerCamelCase__					) ,		2					)
           __lowercase          =						lmap(str.strip ,		Path(model.hparams.data_dir					).joinpath("""test.source"""					).open().readlines()					)
           __lowercase          =						tempfile.mktemp()  # XXX: not being cleaned up
           generate_summaries_or_translations(lowerCamelCase__ ,		lowerCamelCase__ ,		str(model.output_dir / """best_tfmr"""					)					)
           self.assertTrue(Path(lowerCamelCase__					).exists()					)

           __lowercase          =						self.get_auto_remove_tmp_dir()
           convert_pl_to_hf(ckpts[0] ,		transformer_ckpts[0].parent ,		lowerCamelCase__					)
           assert os.path.exists(os.path.join(lowerCamelCase__ ,		"""pytorch_model.bin"""					)					)

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

           """simple docstring"""



           __lowercase          =						AutoModelForSeqaSeqLM.from_pretrained(lowerCamelCase__					)
           __lowercase    ,__lowercase          =						model.dummy_inputs["""input_ids"""], model.dummy_inputs["""attention_mask"""]
           __lowercase          =						torch.tensor([[0, 4, 8, 2], [0, 8, 2, 1]] ,		dtype=torch.long ,		device=model.device					)
           __lowercase          =						target_ids[:, :-1].contiguous()  # Why this line?
           __lowercase          =						target_ids[:, 1:].clone()  # why clone?
           __lowercase          =						model(
               lowerCamelCase__ ,		attention_mask=lowerCamelCase__ ,		decoder_input_ids=lowerCamelCase__ ,		labels=lowerCamelCase__ ,		use_cache=lowerCamelCase__					).loss

           __lowercase          =						model(lowerCamelCase__ ,		attention_mask=lowerCamelCase__ ,		decoder_input_ids=lowerCamelCase__ ,		use_cache=lowerCamelCase__					).logits

           __lowercase          =						nn.functional.log_softmax(lowerCamelCase__ ,		dim=-1					)
           __lowercase    ,__lowercase          =						label_smoothed_nll_loss(
               lowerCamelCase__ ,		lowerCamelCase__ ,		0.1 ,		ignore_index=model.config.pad_token_id					)
           with self.assertRaises(lowerCamelCase__					):
               # TODO: understand why this breaks
               self.assertEqual(lowerCamelCase__ ,		lowerCamelCase__					)

       def    snake_case__      (   self					)   -> Tuple:

           """simple docstring"""



           __lowercase          =						{
               """student_encoder_layers""": 2,
               """student_decoder_layers""": 1,
               """num_train_epochs""": 4,
               """val_check_interval""": 0.25,
               """alpha_hid""": 2.0,
               """task""": """translation""",
               """model_name_or_path""": """IGNORE_THIS_IT_DOESNT_GET_USED""",
               """tokenizer_name""": MBART_TINY,
               """teacher""": MBART_TINY,
               """src_lang""": """en_XX""",
               """tgt_lang""": """ro_RO""",
           }
           __lowercase          =						self._test_distiller_cli(lowerCamelCase__ ,		check_contents=lowerCamelCase__					)
           assert model.model.config.model_type == "mbart"

           __lowercase          =						list(Path(model.output_dir					).glob("""*.ckpt"""					)					)
           self.assertEqual(1 ,		len(lowerCamelCase__					)					)
           __lowercase          =						list(Path(model.output_dir					).glob("""**/*.bin"""					)					)
           __lowercase          =						list(Path(model.output_dir					).glob("""best_tfmr/*"""					)					)
           assert len(lowerCamelCase__					) > 2
           self.assertEqual(len(lowerCamelCase__					) ,		2					)

       def    snake_case__      (   self					)   -> Dict:

           """simple docstring"""



           __lowercase          =						{
               """student_encoder_layers""": 1,
               """student_decoder_layers""": 1,
               """alpha_hid""": 2.0,
               """teacher""": T5_TINY,
               """model_name_or_path""": T5_TINY,
               """tokenizer_name""": T5_TINY,
           }
           self._test_distiller_cli(lowerCamelCase__					)

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

           """simple docstring"""



           __lowercase          =						{
               """teacher""": T5_TINY,
               """student""": T5_TINIER,
               """model_name_or_path""": T5_TINIER,
               """tokenizer_name""": T5_TINIER,
           }
           self._test_distiller_cli(lowerCamelCase__					)




       def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__=True					)   -> int:

           """simple docstring"""



           __lowercase          =						{
               """label_smoothing""": 0.0,
               """early_stopping_patience""": -1,
               """train_batch_size""": 1,
               """eval_batch_size""": 2,
               """max_epochs""": 2,
               """alpha_mlm""": 0.2,
               """alpha_ce""": 0.8,
               """do_predict""": True,
               """model_name_or_path""": """sshleifer/tinier_bart""",
               """teacher""": CHEAP_ARGS["""model_name_or_path"""],
               """val_check_interval""": 0.5,
           }
           default_updates.update(lowerCamelCase__					)
           __lowercase          =						CHEAP_ARGS.copy()
           __lowercase          =						make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()					)
           __lowercase          =						self.get_auto_remove_tmp_dir()

           args_d.update(data_dir=lowerCamelCase__ ,		output_dir=lowerCamelCase__ ,		**lowerCamelCase__					)
           __lowercase          =						distill_main(argparse.Namespace(**lowerCamelCase__					)					)
           if not check_contents:
               return model
           __lowercase          =						os.listdir(lowerCamelCase__					)
           __lowercase          =						{os.path.basename(lowerCamelCase__					) for p in contents}
           __lowercase          =						[p for p in contents if p.endswith("""ckpt"""					)]
           assert len(lowerCamelCase__					) > 0

           self.assertIn("""test_generations.txt""" ,		lowerCamelCase__					)
           self.assertIn("""test_results.txt""" ,		lowerCamelCase__					)

           __lowercase          =						load_json(model.metrics_save_path					)
           __lowercase          =						metrics["""val"""][-1]
           self.assertGreaterEqual(last_step_stats["""val_avg_gen_time"""] ,		0.01					)
           self.assertGreaterEqual(1.0 ,		last_step_stats["""val_avg_gen_time"""]					)
           self.assertIsInstance(last_step_stats[F'val_avg_{model.val_metric}'] ,		lowerCamelCase__					)
           __lowercase          =						int(args_d["""max_epochs"""] * (1 / args_d["""val_check_interval"""]) + 1					)
           self.assertEqual(len(metrics["""val"""]					) ,		lowerCamelCase__					)
           self.assertEqual(len(metrics["""test"""]					) ,		1					)
           return model






class 	SCREAMING_SNAKE_CASE(    __A					):

       @parameterized.expand(
           [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY, FSMT_TINY] ,		)
       def    snake_case__      (   self ,		lowerCamelCase__					)   -> str:

           """simple docstring"""



           __lowercase          =						CHEAP_ARGS.copy()
           __lowercase          =						"""translation""" if model in [MBART_TINY, MARIAN_TINY, FSMT_TINY] else """summarization"""
           __lowercase          =						0.1 if task == """translation""" else 0

           __lowercase          =						make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()					)
           __lowercase          =						self.get_auto_remove_tmp_dir()
           args_d.update(
               data_dir=lowerCamelCase__ ,		model_name_or_path=lowerCamelCase__ ,		tokenizer_name=lowerCamelCase__ ,		train_batch_size=2 ,		eval_batch_size=2 ,		output_dir=lowerCamelCase__ ,		do_predict=lowerCamelCase__ ,		task=lowerCamelCase__ ,		src_lang="""en_XX""" ,		tgt_lang="""ro_RO""" ,		freeze_encoder=lowerCamelCase__ ,		freeze_embeds=lowerCamelCase__ ,		)
           assert "n_train" in args_d
           __lowercase          =						argparse.Namespace(**lowerCamelCase__					)
           __lowercase          =						main(lowerCamelCase__					)

           __lowercase          =						module.model.get_input_embeddings()
           assert not input_embeds.weight.requires_grad
           if model == T5_TINY:
               __lowercase          =						module.model.lm_head
               assert not lm_head.weight.requires_grad
               assert (lm_head.weight == input_embeds.weight).all().item()
           elif model == FSMT_TINY:
               __lowercase          =						module.model.model
               __lowercase          =						fsmt.decoder.embed_positions
               assert not embed_pos.weight.requires_grad
               assert not fsmt.decoder.embed_tokens.weight.requires_grad
               # check that embeds are not the same
               assert fsmt.decoder.embed_tokens != fsmt.encoder.embed_tokens
           else:
               __lowercase          =						module.model.model
               __lowercase          =						bart.decoder.embed_positions
               assert not embed_pos.weight.requires_grad
               assert not bart.shared.weight.requires_grad
               # check that embeds are the same
               assert bart.decoder.embed_tokens == bart.encoder.embed_tokens
               assert bart.decoder.embed_tokens == bart.shared

           __lowercase          =						load_json(module.output_dir / """text_batch.json"""					)
           assert isinstance(lowerCamelCase__ ,		lowerCamelCase__					)
           assert len(lowerCamelCase__					) >= 4

       def    snake_case__      (   self					)   -> str:

           """simple docstring"""



           __lowercase          =						CHEAP_ARGS.copy()

           __lowercase          =						"""summarization"""
           __lowercase          =						make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()					)

           args_d.update(
               data_dir=lowerCamelCase__ ,		tokenizer_name=lowerCamelCase__ ,		train_batch_size=2 ,		eval_batch_size=2 ,		do_predict=lowerCamelCase__ ,		task=lowerCamelCase__ ,		src_lang="""en_XX""" ,		tgt_lang="""ro_RO""" ,		freeze_encoder=lowerCamelCase__ ,		freeze_embeds=lowerCamelCase__ ,		)

           # test models whose config includes the extra_model_args
           __lowercase          =						BART_TINY
           __lowercase          =						self.get_auto_remove_tmp_dir()
           __lowercase          =						args_d.copy()
           args_da.update(
               model_name_or_path=lowerCamelCase__ ,		output_dir=lowerCamelCase__ ,		)
           __lowercase          =						("""encoder_layerdrop""", """decoder_layerdrop""", """dropout""", """attention_dropout""")
           for p in extra_model_params:
               __lowercase          =						0.5
           __lowercase          =						argparse.Namespace(**lowerCamelCase__					)
           __lowercase          =						main(lowerCamelCase__					)
           for p in extra_model_params:
               assert getattr(model.config ,		lowerCamelCase__					) == 0.5, F'failed to override the model config for param {p}'

           # test models whose config doesn't include the extra_model_args
           __lowercase          =						T5_TINY
           __lowercase          =						self.get_auto_remove_tmp_dir()
           __lowercase          =						args_d.copy()
           args_da.update(
               model_name_or_path=lowerCamelCase__ ,		output_dir=lowerCamelCase__ ,		)
           __lowercase          =						"""encoder_layerdrop"""
           __lowercase          =						0.5
           __lowercase          =						argparse.Namespace(**lowerCamelCase__					)
           with pytest.raises(lowerCamelCase__					) as excinfo:
               __lowercase          =						main(lowerCamelCase__					)
           assert str(excinfo.value					) == F'model config doesn\'t have a `{unsupported_param}` attribute'




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

           """simple docstring"""



           __lowercase          =						CHEAP_ARGS.copy()

           __lowercase          =						"""summarization"""
           __lowercase          =						make_test_data_dir(tmp_dir=self.get_auto_remove_tmp_dir()					)

           __lowercase          =						BART_TINY
           __lowercase          =						self.get_auto_remove_tmp_dir()

           args_d.update(
               data_dir=lowerCamelCase__ ,		model_name_or_path=lowerCamelCase__ ,		output_dir=lowerCamelCase__ ,		tokenizer_name=lowerCamelCase__ ,		train_batch_size=2 ,		eval_batch_size=2 ,		do_predict=lowerCamelCase__ ,		task=lowerCamelCase__ ,		src_lang="""en_XX""" ,		tgt_lang="""ro_RO""" ,		freeze_encoder=lowerCamelCase__ ,		freeze_embeds=lowerCamelCase__ ,		)

           # emulate finetune.py
           __lowercase          =						argparse.ArgumentParser()
           __lowercase          =						pl.Trainer.add_argparse_args(lowerCamelCase__					)
           __lowercase          =						SummarizationModule.add_model_specific_args(lowerCamelCase__ ,		os.getcwd()					)
           __lowercase          =						{"""--help""": True}

           # --help test
           with pytest.raises(lowerCamelCase__					) as excinfo:
               with CaptureStdout() as cs:
                   __lowercase          =						parser.parse_args(lowerCamelCase__					)
               assert False, "--help is expected to sys.exit"
           assert excinfo.type == SystemExit
           __lowercase          =						lightning_base.arg_to_scheduler_metavar
           assert expected in cs.out, "--help is expected to list the supported schedulers"

           # --lr_scheduler=non_existing_scheduler test
           __lowercase          =						"""non_existing_scheduler"""
           __lowercase          =						{F'--lr_scheduler={unsupported_param}'}
           with pytest.raises(lowerCamelCase__					) as excinfo:
               with CaptureStderr() as cs:
                   __lowercase          =						parser.parse_args(lowerCamelCase__					)
               assert False, "invalid argument is expected to sys.exit"
           assert excinfo.type == SystemExit
           __lowercase          =						F'invalid choice: \'{unsupported_param}\''
           assert expected in cs.err, F'should have bailed on invalid choice of scheduler {unsupported_param}'

           # --lr_scheduler=existing_scheduler test
           __lowercase          =						"""cosine"""
           __lowercase          =						args_d.copy()
           __lowercase          =						supported_param
           __lowercase          =						argparse.Namespace(**lowerCamelCase__					)
           __lowercase          =						main(lowerCamelCase__					)
           assert (
               getattr(model.hparams ,		"""lr_scheduler"""					) == supported_param
           ), F'lr_scheduler={supported_param} shouldn\'t fail'










'''simple docstring'''
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


A							:				Dict			     =	{
    """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     snake_case_  (       a__		:     str      ):



    """simple docstring"""




    __lowercase          =						task_cls()
    __lowercase          =						asdict(a__      )
    __lowercase          =						task_template_from_dict(a__      )
    assert task == reloaded






class 	SCREAMING_SNAKE_CASE:

       def    snake_case__      (   self					)   -> Any:

           """simple docstring"""



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




       def    snake_case__      (   self					)   -> Any:

           """simple docstring"""



           __lowercase          =						Features({"""text""": Value("""string"""					)}					)
           __lowercase          =						{"""text_column""": """input_text"""}
           __lowercase          =						LanguageModeling.from_dict(lowerCamelCase__					)
           assert "language-modeling" == task.task
           assert input_schema == task.input_schema






class 	SCREAMING_SNAKE_CASE(    __A					):

       def    snake_case__      (   self					)   -> Dict:

           """simple docstring"""



           __lowercase          =						sorted(["""pos""", """neg"""]					)

       def    snake_case__      (   self					)   -> Tuple:

           """simple docstring"""



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

       def    snake_case__      (   self					)   -> Dict:

           """simple docstring"""



           __lowercase          =						Features({"""text""": Value("""string"""					)}					)
           # Labels are cast to tuple during `TextClassification.__post_init__`, so we do the same here
           __lowercase          =						Features({"""labels""": ClassLabel}					)
           __lowercase          =						{"""text_column""": """input_text""", """label_column""": """input_labels"""}
           __lowercase          =						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:

           """simple docstring"""



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






class 	SCREAMING_SNAKE_CASE(    __A					):

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

           """simple docstring"""



           __lowercase          =						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					)   -> Dict:

           """simple docstring"""



           __lowercase          =						Features({"""question""": Value("""string"""					), """context""": Value("""string"""					)}					)
           __lowercase          =						Features(
               {
                   """answers""": Sequence(
                       {
                           """text""": Value("""string"""					),
                           """answer_start""": Value("""int32"""					),
                       }					)
               }					)
           __lowercase          =						{
               """context_column""": """input_input_context""",
               """question_column""": """input_question""",
               """answers_column""": """input_answers""",
           }
           __lowercase          =						QuestionAnsweringExtractive.from_dict(lowerCamelCase__					)
           self.assertEqual("""question-answering-extractive""" ,		task.task					)
           self.assertEqual(lowerCamelCase__ ,		task.input_schema					)
           self.assertEqual(lowerCamelCase__ ,		task.label_schema					)






class 	SCREAMING_SNAKE_CASE(    __A					):

       def    snake_case__      (   self					)   -> int:

           """simple docstring"""



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




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

           """simple docstring"""



           __lowercase          =						Features({"""text""": Value("""string"""					)}					)
           __lowercase          =						Features({"""summary""": Value("""string"""					)}					)
           __lowercase          =						{"""text_column""": """input_text""", """summary_column""": """input_summary"""}
           __lowercase          =						Summarization.from_dict(lowerCamelCase__					)
           self.assertEqual("""summarization""" ,		task.task					)
           self.assertEqual(lowerCamelCase__ ,		task.input_schema					)
           self.assertEqual(lowerCamelCase__ ,		task.label_schema					)






class 	SCREAMING_SNAKE_CASE(    __A					):

       def    snake_case__      (   self					)   -> str:

           """simple docstring"""



           __lowercase          =						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[str]:

           """simple docstring"""



           __lowercase          =						Features({"""audio""": Audio()}					)
           __lowercase          =						Features({"""transcription""": Value("""string"""					)}					)
           __lowercase          =						{
               """audio_column""": """input_audio""",
               """transcription_column""": """input_transcription""",
           }
           __lowercase          =						AutomaticSpeechRecognition.from_dict(lowerCamelCase__					)
           self.assertEqual("""automatic-speech-recognition""" ,		task.task					)
           self.assertEqual(lowerCamelCase__ ,		task.input_schema					)
           self.assertEqual(lowerCamelCase__ ,		task.label_schema					)






class 	SCREAMING_SNAKE_CASE(    __A					):

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

           """simple docstring"""



           __lowercase          =						sorted(["""pos""", """neg"""]					)

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

           """simple docstring"""



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

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

           """simple docstring"""



           __lowercase          =						Features({"""audio""": Audio()}					)
           __lowercase          =						Features({"""labels""": ClassLabel}					)
           __lowercase          =						{
               """audio_column""": """input_image""",
               """label_column""": """input_label""",
           }
           __lowercase          =						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					)   -> str:

           """simple docstring"""



           __lowercase          =						AudioClassification(audio_column="""input_audio""" ,		label_column="""input_label"""					)
           self.assertEqual(task.label_schema["""labels"""] ,		lowerCamelCase__					)
           __lowercase          =						task.align_with_features(Features({"""input_label""": ClassLabel(names=self.labels					)}					)					)
           self.assertEqual(task.label_schema["""labels"""] ,		ClassLabel(names=self.labels					)					)






class 	SCREAMING_SNAKE_CASE(    __A					):

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

           """simple docstring"""



           __lowercase          =						sorted(["""pos""", """neg"""]					)

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

           """simple docstring"""



           __lowercase          =						ImageClassification(image_column="""input_image""" ,		label_column="""input_label"""					)
           self.assertDictEqual({"""input_image""": """image""", """input_label""": """labels"""} ,		task.column_mapping					)

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

           """simple docstring"""



           __lowercase          =						Features({"""image""": Image()}					)
           __lowercase          =						Features({"""labels""": ClassLabel}					)
           __lowercase          =						{
               """image_column""": """input_image""",
               """label_column""": """input_label""",
           }
           __lowercase          =						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					)   -> Tuple:

           """simple docstring"""



           __lowercase          =						ImageClassification(image_column="""input_image""" ,		label_column="""input_label"""					)
           self.assertEqual(task.label_schema["""labels"""] ,		lowerCamelCase__					)
           __lowercase          =						task.align_with_features(Features({"""input_label""": ClassLabel(names=self.labels					)}					)					)
           self.assertEqual(task.label_schema["""labels"""] ,		ClassLabel(names=self.labels					)					)






class 	SCREAMING_SNAKE_CASE(    __A					):

       def    snake_case__      (   self					)   -> Tuple:

           """simple docstring"""



           __lowercase          =						DatasetInfo(task_templates=QuestionAnsweringExtractive()					)
           __lowercase          =						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__					):
               __lowercase          =						deepcopy(lowerCamelCase__					)
               __lowercase          =						0
               return out

           __lowercase          =						dataset.map(lowerCamelCase__					)
           assert mapped_dataset.info.task_templates == dataset.info.task_templates
           # reload from cache
           __lowercase          =						dataset.map(lowerCamelCase__					)
           assert mapped_dataset.info.task_templates == dataset.info.task_templates

           __lowercase          =						dataset.map(lowerCamelCase__					)
           assert mapped_dataset.info.task_templates == []
           # reload from cache
           __lowercase          =						dataset.map(lowerCamelCase__					)
           assert mapped_dataset.info.task_templates == []




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

           """simple docstring"""



           __lowercase          =						Features({"""text""": Value("""string"""					), """label""": ClassLabel(names=("""pos""", """neg""")					)}					)
           __lowercase          =						TextClassification(text_column="""text""" ,		label_column="""label"""					)
           __lowercase          =						DatasetInfo(features=lowerCamelCase__ ,		task_templates=lowerCamelCase__					)
           __lowercase          =						Dataset.from_dict({"""text""": ["""A sentence."""], """label""": ["""pos"""]} ,		info=lowerCamelCase__					)

           def process(lowerCamelCase__					):
               return example

           __lowercase          =						dataset.remove_columns("""label"""					)
           __lowercase          =						modified_dataset.map(lowerCamelCase__					)
           assert mapped_dataset.info.task_templates == []




