
import math
import os
import re
import sys
import unittest
from pathlib import Path
from typing import Tuple
from unittest.mock import patch

from parameterized import parameterized

from transformers.testing_utils import (
    CaptureStderr,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_torch_dist_unique_port,
    require_apex,
    require_bitsandbytes,
    require_fairscale,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    require_torch_non_multi_gpu,
    slow,
)
from transformers.trainer_callback import TrainerState
from transformers.trainer_utils import set_seed


lowerCamelCase__          =		os.path.abspath(os.path.dirname(__file__))
with ExtendSysPath(F"""{bindir}/../../examples/pytorch/translation"""):
					from run_translation import main  # noqa


set_seed(42)
lowerCamelCase__          =		"""sshleifer/student_marian_en_ro_6_1"""
lowerCamelCase__          =		"""sshleifer/tiny-mbart"""







@require_torch
class      A__							(     __magic_name__	):






				def 							_lowerCamelCase		(						self					:      Any    ,     a					:      Dict=False    ,     a					:      List[Any]=None    ,     a					:      Dict=True    ,     a					:      Tuple=True    ,     a					:      Union[str, Any]=True    ,     a					:      int=True    ,     ):






							'''simple docstring'''







							lowerCAmelCase__      :      Union[str, Any]         =		self.run_trainer(
							    eval_steps=1    ,     max_len=12    ,     model_name=a    ,     num_train_epochs=1    ,     distributed=a    ,     extra_args_str=a    ,     predict_with_generate=a    ,     do_train=a    ,     do_eval=a    ,     do_predict=a    ,     )
							lowerCAmelCase__      :      Tuple         =		TrainerState.load_from_json(os.path.join(a    ,     'trainer_state.json'   )   ).log_history

							if not do_eval:
										return

							lowerCAmelCase__      :      Dict         =		[log for log in logs if 'eval_loss' in log.keys()]

							lowerCAmelCase__      :      List[Any]         =		eval_metrics[0]
							if predict_with_generate:
										assert "eval_bleu" in first_step_stats

										lowerCAmelCase__      :      Tuple         =		eval_metrics[-1]
										assert isinstance(last_step_stats['eval_bleu']    ,     a   )
										assert not math.isnan(float(last_step_stats['eval_loss']   )   ), "eval_loss must not be `nan`"






				@require_torch_non_multi_gpu
				def 							_lowerCamelCase		(						self					:      Union[str, Any]   ):






							'''simple docstring'''







							self.run_seqaseq_quick()






				@require_torch_multi_gpu
				def 							_lowerCamelCase		(						self					:      Dict   ):






							'''simple docstring'''







							self.run_seqaseq_quick(distributed=a   )






				@require_torch_multi_gpu
				def 							_lowerCamelCase		(						self					:      Any   ):






							'''simple docstring'''







							self.run_seqaseq_quick(distributed=a   )






				@unittest.skip('Requires an update of the env running those tests'   )
				@require_torch_multi_gpu
				@require_fairscale
				def 							_lowerCamelCase		(						self					:      List[str]   ):






							'''simple docstring'''







							self.run_seqaseq_quick(distributed=a    ,     extra_args_str='--sharded_ddp simple'   )






				@unittest.skip('Requires an update of the env running those tests'   )
				@require_torch_multi_gpu
				@require_fairscale
				def 							_lowerCamelCase		(						self					:      Dict   ):






							'''simple docstring'''







							self.run_seqaseq_quick(distributed=a    ,     extra_args_str='--sharded_ddp simple --fp16'   )






				@unittest.skip('Requires an update of the env running those tests'   )
				@require_torch_multi_gpu
				@require_fairscale
				def 							_lowerCamelCase		(						self					:      Optional[int]   ):






							'''simple docstring'''







							self.run_seqaseq_quick(distributed=a    ,     extra_args_str='--sharded_ddp zero_dp_2'    ,     predict_with_generate=a   )






				@unittest.skip('Requires an update of the env running those tests'   )
				@require_torch_multi_gpu
				@require_fairscale
				def 							_lowerCamelCase		(						self					:      int   ):






							'''simple docstring'''







							self.run_seqaseq_quick(
							    distributed=a    ,     extra_args_str='--sharded_ddp zero_dp_2 --fp16'    ,     predict_with_generate=a   )






				@require_apex
				@require_torch_gpu
				def 							_lowerCamelCase		(						self					:      Dict   ):






							'''simple docstring'''







							self.run_seqaseq_quick(distributed=a    ,     extra_args_str='--fp16 --fp16_backend=apex'   )
							# test 2nd time - was getting eval_loss': nan'
							# to reproduce the problem set distributed=False
							self.run_seqaseq_quick(distributed=a    ,     extra_args_str='--fp16 --fp16_backend=apex'   )






				@parameterized.expand(['base', 'low', 'high', 'mixed']   )
				@require_torch_multi_gpu
				def 							_lowerCamelCase		(						self					:      Optional[int]    ,     a					:      Tuple   ):






							'''simple docstring'''







							lowerCAmelCase__      :      str         =		{
							    # test with the default log_level - should be info and thus log info once
							    'base': {'extra_args_str': '', 'n_matches': 1},
							    # test with low log_level and log_level_replica - should be noisy on all processes
							    # now the info string should appear twice on 2 processes
							    'low': {'extra_args_str': '--log_level debug --log_level_replica debug', 'n_matches': 2},
							    # test with high log_level and low log_level_replica
							    # now the info string should appear once only on the replica
							    'high': {'extra_args_str': '--log_level error --log_level_replica debug', 'n_matches': 1},
							    # test with high log_level and log_level_replica - should be quiet on all processes
							    'mixed': {'extra_args_str': '--log_level error --log_level_replica error', 'n_matches': 0},
							}

							lowerCAmelCase__      :      str         =		experiments[experiment_id]
							lowerCAmelCase__      :      Any         =		{'distributed': True, 'predict_with_generate': False, 'do_eval': False, 'do_predict': False}
							lowerCAmelCase__      :      int         =		'Running training'
							with CaptureStderr() as cl:
										self.run_seqaseq_quick(**a    ,     extra_args_str=data['extra_args_str']   )
							lowerCAmelCase__      :      List[Any]         =		len(re.findall(a    ,     cl.err   )   )
							self.assertEqual(a    ,     data['n_matches']   )






				@slow
				def 							_lowerCamelCase		(						self					:      int   ):






							'''simple docstring'''







							lowerCAmelCase__      :      List[str]         =		self.run_trainer(
							    eval_steps=2    ,     max_len=128    ,     model_name=a    ,     learning_rate=3E-4    ,     num_train_epochs=10    ,     distributed=a    ,     )

							# Check metrics
							lowerCAmelCase__      :      Dict         =		TrainerState.load_from_json(os.path.join(a    ,     'trainer_state.json'   )   ).log_history
							lowerCAmelCase__      :      str         =		[log for log in logs if 'eval_loss' in log.keys()]
							lowerCAmelCase__      :      str         =		eval_metrics[0]
							lowerCAmelCase__      :      Tuple         =		eval_metrics[-1]

							assert first_step_stats["eval_loss"] > last_step_stats["eval_loss"], "model learned nothing"
							assert isinstance(last_step_stats['eval_bleu']    ,     a   )

							# test if do_predict saves generations and metrics
							lowerCAmelCase__      :      List[Any]         =		os.listdir(a   )
							lowerCAmelCase__      :      List[str]         =		{os.path.basename(a   ) for p in contents}
							assert "generated_predictions.txt" in contents
							assert "predict_results.json" in contents






				@slow
				@require_bitsandbytes
				def 							_lowerCamelCase		(						self					:      Tuple   ):






							'''simple docstring'''







							from transformers.training_args import OptimizerNames

							def train_and_return_metrics(a					:      str   ) -> Tuple[int, float]:
										lowerCAmelCase__      :      Optional[int]         =		'--skip_memory_metrics 0'

										lowerCAmelCase__      :      List[Any]         =		self.run_trainer(
										    max_len=128    ,     model_name=a    ,     learning_rate=3E-4    ,     num_train_epochs=1    ,     optim=a    ,     distributed=a    ,     extra_args_str=a    ,     do_eval=a    ,     do_predict=a    ,     n_gpus_to_use=1    ,     )

										# Check metrics
										lowerCAmelCase__      :      Union[str, Any]         =		TrainerState.load_from_json(Path(a    ,     'trainer_state.json'   )   ).log_history
										lowerCAmelCase__      :      List[str]         =		int(logs[0]['train_mem_gpu_peaked_delta'] / 2**20   )
										lowerCAmelCase__      :      Union[str, Any]         =		int(logs[0]['train_mem_gpu_alloc_delta'] / 2**20   )

										lowerCAmelCase__      :      Optional[int]         =		logs[0]['train_loss']
										return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

							lowerCAmelCase__      ,				lowerCAmelCase__      ,				lowerCAmelCase__      :      str         =		train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value   )
							lowerCAmelCase__      ,				lowerCAmelCase__      ,				lowerCAmelCase__      :      int         =		train_and_return_metrics(OptimizerNames.ADAMW_BNB.value   )

							lowerCAmelCase__      :      Dict         =		gpu_alloc_mem_orig - gpu_alloc_mem_bnb

							lowerCAmelCase__      :      Any         =		gpu_peak_mem_orig + gpu_alloc_mem_orig
							lowerCAmelCase__      :      Dict         =		gpu_peak_mem_bnb + gpu_alloc_mem_bnb
							lowerCAmelCase__      :      Optional[Any]         =		gpu_total_mem_orig - gpu_total_mem_bnb

							# sshleifer/student_marian_en_ro_6_1 has 54M parameter, 29M of which is `nn.Embedding` which
							# doesn't get quantized and remains in fp32. Therefore we only have 25M parameters quantized
							# in 2 bytes and the diff in optim memory usage is derived as so:
							#
							# - normal 25*8=~200MB (8 bytes per param)
							# - bnb    25*2= ~50MB (2 bytes per param)
							#
							# Thus we should expect ~150MB total memory saved.
							#
							# Peak memory should be the same - the total should be different by about that same margin
							#
							# After leaving a small margin to accommodate for differences between gpus let's check
							# that we have at least 120MB in savings
							lowerCAmelCase__      :      Dict         =		120

							# uncomment the following if this test starts failing - requires py38 for a new print feature
							# gpu_peak_mem_diff = gpu_peak_mem_orig - gpu_peak_mem_bnb
							# print(f"{gpu_alloc_mem_orig=}MB {gpu_peak_mem_orig=}MB {gpu_alloc_mem_orig+gpu_peak_mem_orig=}MB")
							# print(f" {gpu_alloc_mem_bnb=}MB  {gpu_peak_mem_bnb=}MB  {gpu_alloc_mem_bnb+gpu_peak_mem_bnb=}MB")
							# print(f"{gpu_alloc_mem_diff=}MB")
							# print(f"{gpu_peak_mem_diff=}MB")
							# print(f"{gpu_total_mem_orig=}MB, {gpu_total_mem_bnb=}MB")
							# print(f"{gpu_total_mem_diff=}MB, {gpu_total_mem_diff=}MB")

							self.assertGreater(
							    a    ,     a    ,     'should use ~150MB less alloc gpu memory with BNB, compared to without it for this model but got'
							    f''' a difference of {gpu_alloc_mem_diff}MB, with gpu_alloc_mem_orig={gpu_alloc_mem_orig}MB and'''
							    f''' gpu_alloc_mem_bnb={gpu_alloc_mem_bnb}MB'''    ,     )

							self.assertGreater(
							    a    ,     a    ,     'should use ~150MB less total gpu memory with BNB, compared to without it for this model but got'
							    f''' a difference of {gpu_total_mem_diff}MB, with gpu_total_mem_orig={gpu_total_mem_orig}MB and'''
							    f''' gpu_total_mem_bnb={gpu_total_mem_bnb}MB'''    ,     )

							self.assertEqual(
							    a    ,     a    ,     f'''loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}'''   )






				def 							_lowerCamelCase		(						self					:      Any    ,     a					:      int    ,     a					:      str    ,     a					:      int    ,     a					:      float = 3E-3    ,     a					:      str = "adafactor"    ,     a					:      bool = False    ,     a					:      str = None    ,     a					:      int = 0    ,     a					:      bool = True    ,     a					:      bool = True    ,     a					:      bool = True    ,     a					:      bool = True    ,     a					:      int = None    ,     ):






							'''simple docstring'''







							lowerCAmelCase__      :      Optional[int]         =		self.test_file_dir / '../fixtures/tests_samples/wmt_en_ro'
							lowerCAmelCase__      :      List[str]         =		self.get_auto_remove_tmp_dir()
							lowerCAmelCase__      :      List[str]         =		f'''
            --model_name_or_path {model_name}
            --train_file {data_dir}/train.json
            --validation_file {data_dir}/val.json
            --test_file {data_dir}/test.json
            --output_dir {output_dir}
            --overwrite_output_dir
            --max_train_samples 8
            --max_source_length {max_len}
            --max_target_length {max_len}
            --do_train
            --num_train_epochs {str(a   )}
            --per_device_train_batch_size 4
            --learning_rate {learning_rate}
            --warmup_steps 8
            --logging_steps 0
            --logging_strategy no
            --save_steps {str(a   )}
            --group_by_length
            --label_smoothing_factor 0.1
            --target_lang ro_RO
            --source_lang en_XX
        '''.split()

							lowerCAmelCase__      :      int         =		f'''
            --do_eval
            --per_device_eval_batch_size 4
            --max_eval_samples 8
            --val_max_target_length {max_len}
            --evaluation_strategy steps
            --eval_steps {str(a   )}
        '''.split()

							lowerCAmelCase__      :      Any         =		'\n            --do_predict\n        '.split()

							lowerCAmelCase__      :      Tuple         =		[]
							if do_train:
										args += args_train

							if do_eval:
										args += args_eval

							if do_predict:
										args += args_predict

							if predict_with_generate:
										args += "--predict_with_generate".split()

							if do_train:
										if optim == "adafactor":
													args += "--adafactor".split()
										else:
													args += f'''--optim {optim}'''.split()

							if extra_args_str is not None:
										args += extra_args_str.split()

							if distributed:
										if n_gpus_to_use is None:
													lowerCAmelCase__      :      Tuple         =		get_gpu_count()
										lowerCAmelCase__      :      List[str]         =		get_torch_dist_unique_port()
										lowerCAmelCase__      :      Any         =		f'''
                -m torch.distributed.run
                --nproc_per_node={n_gpus_to_use}
                --master_port={master_port}
                {self.examples_dir_str}/pytorch/translation/run_translation.py
            '''.split()
										lowerCAmelCase__      :      int         =		[sys.executable] + distributed_args + args
										# keep for quick debug
										# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
										execute_subprocess_async(a    ,     env=self.get_env()   )
							else:
										lowerCAmelCase__      :      Optional[Any]         =		['run_translation.py'] + args
										with patch.object(a    ,     'argv'    ,     a   ):
													main()

							return output_dir

import argparse
import re
from typing import Dict

import torch
from datasets import Audio, Dataset, load_dataset, load_metric

from transformers import AutoFeatureExtractor, pipeline
def  lowerCAmelCase__			(  SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_     )				->		List[str]:
			lowerCAmelCase__      :      Optional[Any]         =		args.log_outputs
			lowerCAmelCase__      :      Union[str, Any]         =		'_'.join(args.dataset.split('/'     ) + [args.config, args.split]     )

			# load metric
			lowerCAmelCase__      :      Dict         =		load_metric('wer'     )
			lowerCAmelCase__      :      Tuple         =		load_metric('cer'     )

			# compute metrics
			lowerCAmelCase__      :      Dict         =		wer.compute(references=result['target']       ,							predictions=result['prediction']     )
			lowerCAmelCase__      :      int         =		cer.compute(references=result['target']       ,							predictions=result['prediction']     )

			# print & log results
			lowerCAmelCase__      :      Optional[int]         =		F'''WER: {wer_result}\nCER: {cer_result}'''
			print(SCREAMING_SNAKE_CASE_     )

			with open(F'''{dataset_id}_eval_results.txt'''       ,							'w'     ) as f:
						f.write(SCREAMING_SNAKE_CASE_     )

			# log all results in text file. Possibly interesting for analysis
			if log_outputs is not None:
						lowerCAmelCase__      :      List[str]         =		F'''log_{dataset_id}_predictions.txt'''
						lowerCAmelCase__      :      Union[str, Any]         =		F'''log_{dataset_id}_targets.txt'''

						with open(SCREAMING_SNAKE_CASE_       ,							'w'     ) as p, open(SCREAMING_SNAKE_CASE_       ,							'w'     ) as t:
									# mapping function to write output
									def write_to_file(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_     ):
												p.write(F'''{i}''' + '\n'     )
												p.write(batch['prediction'] + '\n'     )
												t.write(F'''{i}''' + '\n'     )
												t.write(batch['target'] + '\n'     )

									result.map(SCREAMING_SNAKE_CASE_       ,							with_indices=SCREAMING_SNAKE_CASE_     )
def  lowerCAmelCase__			(  SCREAMING_SNAKE_CASE_     )				->		str:
			lowerCAmelCase__      :      List[str]         =		'[,?.!\-\;\:"“%‘”�—’…–]'  # noqa: W605 IMPORTANT: this should correspond to the chars that were ignored during training

			lowerCAmelCase__      :      Union[str, Any]         =		re.sub(SCREAMING_SNAKE_CASE_       ,							''       ,							text.lower()     )

			# In addition, we can normalize the target text, e.g. removing new lines characters etc...
			# note that order is important here!
			lowerCAmelCase__      :      List[Any]         =		['\n\n', '\n', '   ', '  ']

			for t in token_sequences_to_ignore:
						lowerCAmelCase__      :      Optional[int]         =		' '.join(text.split(SCREAMING_SNAKE_CASE_     )     )

			return text
def  lowerCAmelCase__			(  SCREAMING_SNAKE_CASE_     )				->		List[Any]:
			# load dataset
			lowerCAmelCase__      :      Tuple         =		load_dataset(args.dataset       ,							args.config       ,							split=args.split       ,							use_auth_token=SCREAMING_SNAKE_CASE_     )

			# for testing: only process the first two examples as a test
			# dataset = dataset.select(range(10))

			# load processor
			lowerCAmelCase__      :      List[Any]         =		AutoFeatureExtractor.from_pretrained(args.model_id     )
			lowerCAmelCase__      :      Union[str, Any]         =		feature_extractor.sampling_rate

			# resample audio
			lowerCAmelCase__      :      Union[str, Any]         =		dataset.cast_column('audio'       ,							Audio(sampling_rate=SCREAMING_SNAKE_CASE_     )     )

			# load eval pipeline
			if args.device is None:
						lowerCAmelCase__      :      List[Any]         =		0 if torch.cuda.is_available() else -1
			lowerCAmelCase__      :      List[Any]         =		pipeline('automatic-speech-recognition'       ,							model=args.model_id       ,							device=args.device     )

			# map function to decode audio
			def map_to_pred(SCREAMING_SNAKE_CASE_     ):
						lowerCAmelCase__      :      Optional[int]         =		asr(
						    batch['audio']['array']       ,							chunk_length_s=args.chunk_length_s       ,							stride_length_s=args.stride_length_s     )

						lowerCAmelCase__      :      int         =		prediction['text']
						lowerCAmelCase__      :      Optional[int]         =		normalize_text(batch['sentence']     )
						return batch

			# run inference on all examples
			lowerCAmelCase__      :      Dict         =		dataset.map(SCREAMING_SNAKE_CASE_       ,							remove_columns=dataset.column_names     )

			# compute and log_results
			# do not change function below
			log_results(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_     )


if __name__ == "__main__":
					lowerCamelCase__          =		argparse.ArgumentParser()

					parser.add_argument(
					    """--model_id""", type=str, required=True, help="""Model identifier. Should be loadable with 🤗 Transformers"""
					)
					parser.add_argument(
					    """--dataset""",
					    type=str,
					    required=True,
					    help="""Dataset name to evaluate the `model_id`. Should be loadable with 🤗 Datasets""",
					)
					parser.add_argument(
					    """--config""", type=str, required=True, help="""Config of the dataset. *E.g.* `'en'`  for Common Voice"""
					)
					parser.add_argument("""--split""", type=str, required=True, help="""Split of the dataset. *E.g.* `'test'`""")
					parser.add_argument(
					    """--chunk_length_s""", type=float, default=None, help="""Chunk length in seconds. Defaults to 5 seconds."""
					)
					parser.add_argument(
					    """--stride_length_s""", type=float, default=None, help="""Stride of the audio chunks. Defaults to 1 second."""
					)
					parser.add_argument(
					    """--log_outputs""", action="""store_true""", help="""If defined, write outputs to log file for analysis."""
					)
					parser.add_argument(
					    """--device""",
					    type=int,
					    default=None,
					    help="""The device to run the pipeline on. -1 for CPU (default), 0 for the first GPU and so on.""",
					)
					lowerCamelCase__          =		parser.parse_args()

					main(args)