

import gc
import unittest

import numpy as np
import torch

from diffusers import DanceDiffusionPipeline, IPNDMScheduler, UNetaDModel
from diffusers.utils import slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu, skip_mps

from ..pipeline_params import UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS, UNCONDITIONAL_AUDIO_GENERATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()







class 					__lowercase  (						_SCREAMING_SNAKE_CASE		,       unittest.TestCase				):





							"""simple docstring"""

							UpperCAmelCase_						:		Dict								=	DanceDiffusionPipeline
							UpperCAmelCase_						:		Dict								=	UNCONDITIONAL_AUDIO_GENERATION_PARAMS
							UpperCAmelCase_						:		Tuple								=	PipelineTesterMixin.required_optional_params - {
							    '''callback''',
							    '''latents''',
							    '''callback_steps''',
							    '''output_type''',
							    '''num_images_per_prompt''',
							}
							UpperCAmelCase_						:		str								=	UNCONDITIONAL_AUDIO_GENERATION_BATCH_PARAMS
							UpperCAmelCase_						:		Optional[int]								=	False
							UpperCAmelCase_						:		Tuple								=	False




							def       snake_case			(					self						)		->       Any:
														torch.manual_seed(0						)
														A      :					Optional[Any]      		=					UNetaDModel(
														    block_out_channels=(32, 32, 64)    ,						extra_in_channels=16    ,						sample_size=5_12    ,						sample_rate=1_60_00    ,						in_channels=2    ,						out_channels=2    ,						flip_sin_to_cos=__UpperCAmelCase    ,						use_timestep_embedding=__UpperCAmelCase    ,						time_embedding_type='''fourier'''    ,						mid_block_type='''UNetMidBlock1D'''    ,						down_block_types=('''DownBlock1DNoSkip''', '''DownBlock1D''', '''AttnDownBlock1D''')    ,						up_block_types=('''AttnUpBlock1D''', '''UpBlock1D''', '''UpBlock1DNoSkip''')    ,						)
														A      :					Optional[int]      		=					IPNDMScheduler()

														A      :					Tuple      		=					{
														    '''unet''': unet,
														    '''scheduler''': scheduler,
														}
														return components




							def       snake_case			(					self    ,						__UpperCAmelCase    ,						__UpperCAmelCase=0						)		->       int:
														if str(__UpperCAmelCase						).startswith('''mps'''						):
																					A      :					Any      		=					torch.manual_seed(__UpperCAmelCase						)
														else:
																					A      :					Optional[Any]      		=					torch.Generator(device=__UpperCAmelCase						).manual_seed(__UpperCAmelCase						)
														A      :					Optional[Any]      		=					{
														    '''batch_size''': 1,
														    '''generator''': generator,
														    '''num_inference_steps''': 4,
														}
														return inputs




							def       snake_case			(					self						)		->       Tuple:
														A      :					List[str]      		=					'''cpu'''  # ensure determinism for the device-dependent torch.Generator
														A      :					Any      		=					self.get_dummy_components()
														A      :					str      		=					DanceDiffusionPipeline(**__UpperCAmelCase						)
														A      :					int      		=					pipe.to(__UpperCAmelCase						)
														pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

														A      :					List[str]      		=					self.get_dummy_inputs(__UpperCAmelCase						)
														A      :					Optional[Any]      		=					pipe(**__UpperCAmelCase						)
														A      :					Union[str, Any]      		=					output.audios

														A      :					Dict      		=					audio[0, -3:, -3:]

														assert audio.shape == (1, 2, components["unet"].sample_size)
														A      :					Optional[Any]      		=					np.array([-0.7_2_6_5, 1.0_0_0_0, -0.8_3_8_8, 0.1_1_7_5, 0.9_4_9_8, -1.0_0_0_0]						)
														assert np.abs(audio_slice.flatten() - expected_slice						).max() < 1E-2




							@skip_mps
							def       snake_case			(					self						)		->       Tuple:
														return super().test_save_load_local()




							@skip_mps
							def       snake_case			(					self						)		->       Optional[Any]:
														return super().test_dict_tuple_outputs_equivalent(expected_max_difference=3E-3						)




							@skip_mps
							def       snake_case			(					self						)		->       Optional[Any]:
														return super().test_save_load_optional_components()




							@skip_mps
							def       snake_case			(					self						)		->       str:
														return super().test_attention_slicing_forward_pass()






							def       snake_case			(					self						)		->       str:
														super().test_inference_batch_single_identical(expected_max_diff=3E-3						)







@slow
@require_torch_gpu
class 					__lowercase  (						unittest.TestCase				):





							"""simple docstring"""

							def       snake_case			(					self						)		->       List[Any]:
														# clean up the VRAM after each test
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()




							def       snake_case			(					self						)		->       Dict:
														A      :					int      		=					torch_device

														A      :					Any      		=					DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''						)
														A      :					Tuple      		=					pipe.to(__UpperCAmelCase						)
														pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

														A      :					Optional[Any]      		=					torch.manual_seed(0						)
														A      :					int      		=					pipe(generator=__UpperCAmelCase    ,						num_inference_steps=1_00    ,						audio_length_in_s=4.0_9_6						)
														A      :					Union[str, Any]      		=					output.audios

														A      :					str      		=					audio[0, -3:, -3:]

														assert audio.shape == (1, 2, pipe.unet.sample_size)
														A      :					Tuple      		=					np.array([-0.0_1_9_2, -0.0_2_3_1, -0.0_3_1_8, -0.0_0_5_9, 0.0_0_0_2, -0.0_0_2_0]						)

														assert np.abs(audio_slice.flatten() - expected_slice						).max() < 1E-2






							def       snake_case			(					self						)		->       Optional[int]:
														A      :					Any      		=					torch_device

														A      :					List[str]      		=					DanceDiffusionPipeline.from_pretrained('''harmonai/maestro-150k'''    ,						torch_dtype=torch.floataa						)
														A      :					Optional[int]      		=					pipe.to(__UpperCAmelCase						)
														pipe.set_progress_bar_config(disable=__UpperCAmelCase						)

														A      :					str      		=					torch.manual_seed(0						)
														A      :					Optional[int]      		=					pipe(generator=__UpperCAmelCase    ,						num_inference_steps=1_00    ,						audio_length_in_s=4.0_9_6						)
														A      :					Tuple      		=					output.audios

														A      :					Union[str, Any]      		=					audio[0, -3:, -3:]

														assert audio.shape == (1, 2, pipe.unet.sample_size)
														A      :					Union[str, Any]      		=					np.array([-0.0_3_6_7, -0.0_4_8_8, -0.0_7_7_1, -0.0_5_2_5, -0.0_4_4_4, -0.0_3_4_1]						)

														assert np.abs(audio_slice.flatten() - expected_slice						).max() < 1E-2









import argparse
import dataclasses
import json
import logging
import os
import shutil
from typing import List, Optional

import datasets
from accelerate import Accelerator
from datasets import load_dataset
from finetuning import finetune
from tqdm.auto import tqdm

import transformers
from transformers import AutoConfig, set_seed
from transformers.trainer_utils import IntervalStrategy


lowercase			:				Optional[int]   =     logging.getLogger(__name__)

lowercase			:				str   =     "pytorch_model.bin"







@dataclasses.dataclass
class 					__lowercase  :





							"""simple docstring"""

							UpperCAmelCase_						:		str								=	dataclasses.field(
							    metadata={'''help''': '''Path to pretrained model or model identifier from huggingface.co/models.'''}				)
							UpperCAmelCase_						:		Optional[str]								=	dataclasses.field(
							    default=_SCREAMING_SNAKE_CASE		,       metadata={'''help''': '''Where do you want to store the pretrained models downloaded from huggingface.co.'''}		,       )







@dataclasses.dataclass
class 					__lowercase  :





							"""simple docstring"""

							UpperCAmelCase_						:		str								=	dataclasses.field(metadata={'''help''': '''A csv or a json file containing the training data.'''}				)
							UpperCAmelCase_						:		str								=	dataclasses.field(metadata={'''help''': '''A csv or a json file containing the data to predict on.'''}				)
							UpperCAmelCase_						:		Optional[str]								=	dataclasses.field(
							    default=_SCREAMING_SNAKE_CASE		,       metadata={'''help''': '''A csv or a json file containing the validation data.'''}				)
							UpperCAmelCase_						:		Optional[str]								=	dataclasses.field(
							    default=_SCREAMING_SNAKE_CASE		,       metadata={'''help''': '''The name of the task to train on.'''}		,       )
							UpperCAmelCase_						:		Optional[List[str]]								=	dataclasses.field(
							    default=_SCREAMING_SNAKE_CASE		,       metadata={'''help''': '''The list of labels for the task.'''}				)







@dataclasses.dataclass
class 					__lowercase  :





							"""simple docstring"""

							UpperCAmelCase_						:		str								=	dataclasses.field(
							    metadata={'''help''': '''The output directory where the model predictions and checkpoints will be written.'''}				)
							UpperCAmelCase_						:		Optional[str]								=	dataclasses.field(
							    default='''accuracy'''		,       metadata={'''help''': '''The evaluation metric used for the task.'''}				)
							UpperCAmelCase_						:		Optional[str]								=	dataclasses.field(
							    default='''no'''		,       metadata={
							        '''help''': '''The evaluation strategy to adopt during training. Possible values are: ["no", "step", "epoch]'''
							    }		,       )
							UpperCAmelCase_						:		Optional[int]								=	dataclasses.field(
							    default=10		,       metadata={'''help''': '''Number of evaluation calls with no improvement after which training will be stopped.'''}		,       )
							UpperCAmelCase_						:		Optional[float]								=	dataclasses.field(
							    default=0.0		,       metadata={
							        '''help''': '''How much the specified evaluation metric must improve to satisfy early stopping conditions.'''
							    }		,       )
							UpperCAmelCase_						:		Optional[bool]								=	dataclasses.field(
							    default=_SCREAMING_SNAKE_CASE		,       metadata={'''help''': '''Whether to filter the pseudo-labeled data based on the confidence score.'''}		,       )
							UpperCAmelCase_						:		Optional[bool]								=	dataclasses.field(
							    default=_SCREAMING_SNAKE_CASE		,       metadata={'''help''': '''Whether to filter the pseudo-labeled data based on the validation performance.'''}		,       )
							UpperCAmelCase_						:		Optional[bool]								=	dataclasses.field(
							    default=_SCREAMING_SNAKE_CASE		,       metadata={'''help''': '''Whether to fine-tune on labeled data after pseudo training.'''}		,       )
							UpperCAmelCase_						:		Optional[float]								=	dataclasses.field(
							    default=0.0		,       metadata={'''help''': '''Confidence threshold for pseudo-labeled data filtering.'''}		,       )
							UpperCAmelCase_						:		Optional[int]								=	dataclasses.field(
							    default=1_00		,       metadata={'''help''': '''Number of evaluation calls with no improvement after which training will be stopped.'''}		,       )
							UpperCAmelCase_						:		Optional[int]								=	dataclasses.field(
							    default=_SCREAMING_SNAKE_CASE		,       metadata={'''help''': '''Random seed for initialization.'''}		,       )






def        snake_case__		(    lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_		):
							A      :					Dict      		=					datasets.concatenate_datasets([infer_input, infer_output]     ,				axis=1		)

							if args.do_filter_by_confidence:
														A      :					Union[str, Any]      		=					dataset.filter(lambda lowerCamelCase_		: example["probability"] > args.confidence_threshold		)

							if args.do_filter_by_val_performance:
														assert eval_result >= 0.0 and eval_result <= 1.0
														A      :					Optional[int]      		=					int(eval_result * len(lowerCamelCase_		)		)
														print(lowerCamelCase_		)
														A      :					Any      		=					dataset.sort('''probability'''     ,				reverse=lowerCamelCase_		)
														A      :					Any      		=					dataset.select(range(lowerCamelCase_		)		)

							A      :					Dict      		=					dataset.remove_columns(['''label''', '''probability''']		)
							A      :					List[str]      		=					dataset.rename_column('''prediction'''     ,				'''label'''		)
							A      :					Tuple      		=					dataset.map(lambda lowerCamelCase_		: {"label": idalabel[example["label"]]}		)
							A      :					Dict      		=					dataset.shuffle(seed=args.seed		)

							A      :					str      		=					os.path.join(lowerCamelCase_     ,				F'train_pseudo.{args.data_file_extension}'		)
							if args.data_file_extension == "csv":
														dataset.to_csv(lowerCamelCase_     ,				index=lowerCamelCase_		)
							else:
														dataset.to_json(lowerCamelCase_		)






def        snake_case__		(    lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				**lowerCamelCase_		):
							A      :					Optional[Any]      		=					Accelerator()
							# Make one log on every process with the configuration for debugging.
							logging.basicConfig(
							    format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s'''     ,				datefmt='''%m/%d/%Y %H:%M:%S'''     ,				level=logging.INFO     ,				)
							logger.info(accelerator.state		)

							# Setup logging, we only want one process per machine to log things on the
							# screen. accelerator.is_local_main_process is only True for one process per
							# machine.
							logger.setLevel(logging.INFO if accelerator.is_local_main_process else logging.ERROR		)

							if accelerator.is_local_main_process:
														datasets.utils.logging.set_verbosity_warning()
														transformers.utils.logging.set_verbosity_info()
							else:
														datasets.utils.logging.set_verbosity_error()
														transformers.utils.logging.set_verbosity_error()

							A      :					List[str]      		=					STModelArguments(model_name_or_path=lowerCamelCase_		)
							A      :					List[Any]      		=					STDataArguments(train_file=lowerCamelCase_     ,				infer_file=lowerCamelCase_		)
							A      :					int      		=					STTrainingArguments(output_dir=lowerCamelCase_		)
							A      :					Any      		=					argparse.Namespace()

							for arg_class in (model_args, data_args, training_args):
														for key, value in vars(lowerCamelCase_		).items():
																					setattr(lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_		)

							for key, value in kwargs.items():
														if hasattr(lowerCamelCase_     ,				lowerCamelCase_		):
																					setattr(lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_		)

    # Sanity checks
							A      :					Tuple      		=					{}
							A      :					str      		=					None

							# You need to provide the training data and the data to predict on
							assert args.train_file is not None
							assert args.infer_file is not None
							A      :					Dict      		=					args.train_file
							A      :					Any      		=					args.infer_file

							if args.evaluation_strategy != IntervalStrategy.NO.value:
														assert args.eval_file is not None
														A      :					Optional[int]      		=					args.eval_file

							for key in data_files:
														A      :					str      		=					data_files[key].split('''.'''		)[-1]
														assert extension in ["csv", "json"], F'`{key}_file` should be a csv or a json file.'
														if args.data_file_extension is None:
																					A      :					List[Any]      		=					extension
														else:
																					assert extension == args.data_file_extension, F'`{key}_file` should be a {args.data_file_extension} file`.'

							assert (
							    args.eval_metric in datasets.list_metrics()
							), F'{args.eval_metric} not in the list of supported metrics {datasets.list_metrics()}.'

							# If passed along, set the training seed now.
							if args.seed is not None:
														set_seed(args.seed		)

							logger.info('''Creating the initial data directory for self-training...'''		)
							A      :					Any      		=					F'{args.output_dir}/self-train_iter-{{}}'.format
							A      :					Tuple      		=					data_dir_format(0		)

							if accelerator.is_main_process:
														if args.output_dir is not None:
																					os.makedirs(args.output_dir     ,				exist_ok=lowerCamelCase_		)
																					os.makedirs(lowerCamelCase_     ,				exist_ok=lowerCamelCase_		)
							accelerator.wait_for_everyone()

							A      :					str      		=					None
							A      :					List[str]      		=					None
							A      :					List[str]      		=					0
							A      :					List[Any]      		=					False
							# Show the progress bar
							A      :					Any      		=					tqdm(range(args.max_selftrain_iterations		)     ,				disable=not accelerator.is_local_main_process		)

							# Self-train
							for iteration in range(0     ,				int(args.max_selftrain_iterations		)		):
														A      :					Union[str, Any]      		=					data_dir_format(lowerCamelCase_		)
														assert os.path.exists(lowerCamelCase_		)

														# Stage 1: initial fine-tuning for iteration = 0 or pseudo-training for
														# iteration > 0
														A      :					str      		=					os.path.join(lowerCamelCase_     ,				'''stage-1'''		)
														A      :					Optional[Any]      		=					{
														    '''accelerator''': accelerator,
														    '''model_name_or_path''': args.model_name_or_path,
														    '''cache_dir''': args.cache_dir,
														    '''do_train''': True,
														    '''train_file''': data_files['''train'''] if iteration == 0 else data_files['''train_pseudo'''],
														    '''do_eval''': True if args.eval_file is not None else False,
														    '''eval_file''': data_files['''eval'''],
														    '''do_predict''': True,
														    '''infer_file''': data_files['''infer'''],
														    '''task_name''': args.task_name,
														    '''label_list''': args.label_list,
														    '''output_dir''': current_output_dir,
														    '''eval_metric''': args.eval_metric,
														    '''evaluation_strategy''': args.evaluation_strategy,
														    '''early_stopping_patience''': args.early_stopping_patience,
														    '''early_stopping_threshold''': args.early_stopping_threshold,
														    '''seed''': args.seed,
														}
														# Add additional training arguments
														for key, value in kwargs.items():
																					if key not in arguments_dict and not hasattr(lowerCamelCase_     ,				lowerCamelCase_		):
																												arguments_dict.update({key: value}		)

														A      :					str      		=					os.path.join(lowerCamelCase_     ,				'''best-checkpoint'''     ,				lowerCamelCase_		)
														if os.path.exists(lowerCamelCase_		):
																					logger.info(
																					    '''Found existing model checkpoint at %s. Skipping self-training: iteration: %d, stage: 1.'''     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				)
														else:
																					logger.info('''***** Running self-training: iteration: %d, stage: 1 *****'''     ,				lowerCamelCase_		)
																					finetune(**lowerCamelCase_		)
																					accelerator.wait_for_everyone()
																					assert os.path.exists(lowerCamelCase_		)
																					logger.info('''Self-training job completed: iteration: %d, stage: 1.'''     ,				lowerCamelCase_		)

														if iteration > 0 and args.finetune_on_labeled_data:
																					# Stage 2 (optional): fine-tuning on the original labeled data
																					A      :					Any      		=					os.path.join(lowerCamelCase_     ,				'''best-checkpoint'''		)
																					A      :					int      		=					os.path.join(lowerCamelCase_     ,				'''stage-2'''		)
																					# Update arguments_dict
																					A      :					int      		=					model_path
																					A      :					Optional[Any]      		=					data_files['''train''']
																					A      :					Tuple      		=					current_output_dir

																					A      :					List[str]      		=					os.path.join(lowerCamelCase_     ,				'''best-checkpoint'''     ,				lowerCamelCase_		)
																					if os.path.exists(lowerCamelCase_		):
																												logger.info(
																												    '''Found existing model checkpoint at %s. Skipping self-training: iteration: %d, stage: 2.'''     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				)
																					else:
																												logger.info('''***** Running self-training: iteration: %d, stage: 2 *****'''     ,				lowerCamelCase_		)
																												finetune(**lowerCamelCase_		)
																												accelerator.wait_for_everyone()
																												assert os.path.exists(lowerCamelCase_		)
																												logger.info('''Self-training job completed: iteration: %d, stage: 2.'''     ,				lowerCamelCase_		)

														A      :					List[Any]      		=					iteration
														A      :					Optional[Any]      		=					data_dir_format(iteration + 1		)

														A      :					Tuple      		=					AutoConfig.from_pretrained(os.path.join(lowerCamelCase_     ,				'''best-checkpoint'''		)		)
														A      :					List[Any]      		=					config.idalabel
														A      :					Optional[int]      		=					os.path.join(lowerCamelCase_     ,				'''eval_results_best-checkpoint.json'''		)
														A      :					Union[str, Any]      		=					os.path.join(lowerCamelCase_     ,				'''test_results_best-checkpoint.json'''		)
														assert os.path.exists(lowerCamelCase_		)

														with open(lowerCamelCase_     ,				'''r'''		) as f:
																					A      :					Optional[int]      		=					float(json.load(lowerCamelCase_		)[args.eval_metric]		)
														A      :					Any      		=					os.path.join(lowerCamelCase_     ,				'''infer_output_best-checkpoint.csv'''		)
														assert os.path.exists(lowerCamelCase_		)
														# Loading the dataset from local csv or json files.
														A      :					str      		=					load_dataset(args.data_file_extension     ,				data_files={'''data''': data_files['''infer''']}		)['''data''']
														A      :					Tuple      		=					load_dataset('''csv'''     ,				data_files={'''data''': infer_output_file}		)['''data''']

														if accelerator.is_main_process:
																					os.makedirs(lowerCamelCase_     ,				exist_ok=lowerCamelCase_		)
																					shutil.copy(lowerCamelCase_     ,				os.path.join(lowerCamelCase_     ,				F'eval_results_iter-{iteration}.json'		)		)
																					if os.path.exists(lowerCamelCase_		):
																												shutil.copy(lowerCamelCase_     ,				os.path.join(lowerCamelCase_     ,				F'test_results_iter-{iteration}.json'		)		)
																					create_pseudo_labeled_data(lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_     ,				lowerCamelCase_		)
														accelerator.wait_for_everyone()

														A      :					Optional[int]      		=					os.path.join(lowerCamelCase_     ,				F'train_pseudo.{args.data_file_extension}'		)

														if args.evaluation_strategy != IntervalStrategy.NO.value:
																					A      :					Any      		=					eval_result

																					if best_iteration is None:
																												A      :					int      		=					new_iteration
																												A      :					Optional[Any]      		=					new_eval_result
																					else:
																												if new_eval_result - best_eval_result > args.early_stopping_threshold:
																																			A      :					Optional[int]      		=					new_iteration
																																			A      :					Any      		=					new_eval_result
																																			A      :					Any      		=					0
																												else:
																																			if new_eval_result == best_eval_result:
																																										A      :					Any      		=					new_iteration
																																										A      :					Any      		=					new_eval_result
																																			early_stopping_patience_counter += 1

																												if early_stopping_patience_counter >= args.early_stopping_patience:
																																			A      :					int      		=					True

														progress_bar.update(1		)

														if should_training_stop:
																					break

							if best_iteration is not None:
														# Save the best iteration
														logger.info('''Best iteration: %d'''     ,				lowerCamelCase_		)
														logger.info('''Best evaluation result: %s = %f'''     ,				args.eval_metric     ,				lowerCamelCase_		)
														accelerator.wait_for_everyone()
														if accelerator.is_main_process:
																					shutil.copy(
																					    os.path.join(lowerCamelCase_     ,				F'eval_results_iter-{iteration}.json'		)     ,				os.path.join(lowerCamelCase_     ,				'''eval_results_best-iteration.json'''		)     ,				)
							else:
														# Assume that the last iteration is the best
														logger.info('''Best iteration: %d'''     ,				args.max_selftrain_iterations - 1		)
														logger.info('''Best evaluation result: %s = %f'''     ,				args.eval_metric     ,				lowerCamelCase_		)
														accelerator.wait_for_everyone()
														if accelerator.is_main_process:
																					shutil.copy(
																					    os.path.join(lowerCamelCase_     ,				F'eval_results_iter-{args.max_selftrain_iterations - 1}.json'		)     ,				os.path.join(lowerCamelCase_     ,				'''eval_results_best-iteration.json'''		)     ,				)






