



"""simple docstring"""







import argparse
import logging

import numpy as np
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer


logging.basicConfig(
    format="""%(asctime)s - %(levelname)s - %(name)s - %(message)s""",
    datefmt="""%m/%d/%Y %H:%M:%S""",
    level=logging.INFO,
)
UpperCAmelCase:  Any			     =				logging.getLogger(__name__)


def  __SCREAMING_SNAKE_CASE (  __UpperCAmelCase      ):
		np.random.seed(args.seed      )
		torch.manual_seed(args.seed      )
		if args.n_gpu > 0:
				torch.cuda.manual_seed_all(args.seed      )


def  __SCREAMING_SNAKE_CASE (  ):
		_lowercase				: Any     =   argparse.ArgumentParser()
		parser.add_argument(
		    """--model_name_or_path"""   ,				default=__UpperCAmelCase   ,				type=__UpperCAmelCase   ,				required=__UpperCAmelCase   ,				)
		parser.add_argument("""--prompt"""   ,				type=__UpperCAmelCase   ,				default=""""""      )
		parser.add_argument("""--length"""   ,				type=__UpperCAmelCase   ,				default=20      )
		parser.add_argument("""--stop_token"""   ,				type=__UpperCAmelCase   ,				default=__UpperCAmelCase   ,				help="""Token at which text generation is stopped"""      )
		parser.add_argument(
		    """--temperature"""   ,				type=__UpperCAmelCase   ,				default=1.0   ,				help="""temperature of 1.0 has no effect, lower tend toward greedy sampling"""   ,				)
		parser.add_argument(
		    """--repetition_penalty"""   ,				type=__UpperCAmelCase   ,				default=1.0   ,				help="""primarily useful for CTRL model; in that case, use 1.2"""      )
		parser.add_argument("""--k"""   ,				type=__UpperCAmelCase   ,				default=0      )
		parser.add_argument("""--penalty_alpha"""   ,				type=__UpperCAmelCase   ,				default=0.0      )
		parser.add_argument("""--p"""   ,				type=__UpperCAmelCase   ,				default=0.9      )

		parser.add_argument("""--prefix"""   ,				type=__UpperCAmelCase   ,				default=""""""   ,				help="""Text added prior to input."""      )
		parser.add_argument("""--padding_text"""   ,				type=__UpperCAmelCase   ,				default=""""""   ,				help="""Deprecated, the use of `--prefix` is preferred."""      )
		parser.add_argument("""--xlm_language"""   ,				type=__UpperCAmelCase   ,				default=""""""   ,				help="""Optional language when used with the XLM model."""      )

		parser.add_argument("""--seed"""   ,				type=__UpperCAmelCase   ,				default=42   ,				help="""random seed for initialization"""      )
		parser.add_argument("""--no_cuda"""   ,				action="""store_true"""   ,				help="""Avoid using CUDA when available"""      )
		parser.add_argument(
		    """--fp16"""   ,				action="""store_true"""   ,				help="""Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit"""   ,				)
		_lowercase				: str     =   parser.parse_args()

		_lowercase				: str     =   torch.device("""cuda""" if torch.cuda.is_available() and not args.no_cuda else """cpu"""      )
		_lowercase				: List[str]     =   0 if args.no_cuda else torch.cuda.device_count()

		logger.warning(F"""device: {args.device}, n_gpu: {args.n_gpu}, 16-bits training: {args.fpaa}"""      )

		set_seed(__UpperCAmelCase      )

		# Initialize the model and tokenizer
		_lowercase				: List[str]     =   AutoTokenizer.from_pretrained(args.model_name_or_path      )
		_lowercase				: Union[str, Any]     =   AutoModelForCausalLM.from_pretrained(args.model_name_or_path      )

		# tokenizer = GPT2Tokenizer.from_pretrained(args.model_name_or_path)
		# model = OPTForCausalLM.from_pretrained(args.model_name_or_path)
		model.to(args.device      )

		if args.fpaa:
				model.half()

		logger.info(__UpperCAmelCase      )
		_lowercase				: Any     =   args.prompt if args.prompt else input("""Model prompt >>> """      )

		_lowercase				: Optional[Any]     =   tokenizer(__UpperCAmelCase   ,				return_tensors="""pt"""   ,				add_special_tokens=__UpperCAmelCase      )
		_lowercase				: Dict     =   {key: value.to(args.device      ) for key, value in inputs.items()}

		_lowercase				: Optional[Any]     =   model.generate(
		    **__UpperCAmelCase   ,				max_length=args.length + len(inputs["""input_ids"""][0]      )   ,				penalty_alpha=args.penalty_alpha   ,				top_k=args.k   ,				)

		_lowercase				: Tuple     =   []
		for generated_sequence_idx, generated_sequence in enumerate(__UpperCAmelCase      ):
				print(F"""=== GENERATED SEQUENCE {generated_sequence_idx + 1} ==="""      )
				_lowercase				: List[Any]     =   generated_sequence.tolist()

				# Decode text
				_lowercase				: Optional[int]     =   tokenizer.decode(__UpperCAmelCase   ,				clean_up_tokenization_spaces=__UpperCAmelCase   ,				add_special_tokens=__UpperCAmelCase      )

				# Remove all text after the stop token
				_lowercase				: Dict     =   text[: text.find(args.stop_token      ) if args.stop_token else None]

				# Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
				_lowercase				: Tuple     =   (
				    prompt_text + text[len(tokenizer.decode(inputs["""input_ids"""][0]   ,				clean_up_tokenization_spaces=__UpperCAmelCase      )      ) :]
				)

				generated_sequences.append(__UpperCAmelCase      )
				print(__UpperCAmelCase      )

		return generated_sequences


if __name__ == "__main__":
		main()








import shutil
import tempfile
import unittest

import pytest

from transformers import WhisperTokenizer, is_speech_available
from transformers.testing_utils import require_sentencepiece, require_torch, require_torchaudio

from .test_feature_extraction_whisper import floats_list


if is_speech_available():
		from transformers import WhisperFeatureExtractor, WhisperProcessor


__UpperCAmelCase    	=  5_03_58
__UpperCAmelCase    	=  5_03_62


@require_torch
@require_torchaudio
@require_sentencepiece
class 						lowerCamelCase__			(		unittest.TestCase ):


						def 	_lowerCamelCase      (			self    :	List[Any]      ):
													a__:							str									="openai/whisper-small.en"
													a__:							Union[str, Any]									=tempfile.mkdtemp()


						def 	_lowerCamelCase      (			self    :	Optional[Any]   ,							**_a    :	Optional[int]      ):
													return WhisperTokenizer.from_pretrained(self.checkpoint   ,							**_a      )


						def 	_lowerCamelCase      (			self    :	str   ,							**_a    :	Tuple      ):
													return WhisperFeatureExtractor.from_pretrained(self.checkpoint   ,							**_a      )


						def 	_lowerCamelCase      (			self    :	Union[str, Any]      ):
													shutil.rmtree(self.tmpdirname      )


						def 	_lowerCamelCase      (			self    :	int      ):
													a__:							str									=self.get_tokenizer()
													a__:							Union[str, Any]									=self.get_feature_extractor()

													a__:							int									=WhisperProcessor(tokenizer=_a   ,							feature_extractor=_a      )

													processor.save_pretrained(self.tmpdirname      )
													a__:							List[Any]									=WhisperProcessor.from_pretrained(self.tmpdirname      )

													self.assertEqual(processor.tokenizer.get_vocab()   ,							tokenizer.get_vocab()      )
													self.assertIsInstance(processor.tokenizer   ,							_a      )

													self.assertEqual(processor.feature_extractor.to_json_string()   ,							feature_extractor.to_json_string()      )
													self.assertIsInstance(processor.feature_extractor   ,							_a      )


						def 	_lowerCamelCase      (			self    :	Dict      ):
													a__:							str									=WhisperProcessor(tokenizer=self.get_tokenizer()   ,							feature_extractor=self.get_feature_extractor()      )
													processor.save_pretrained(self.tmpdirname      )

													a__:							Tuple									=self.get_tokenizer(bos_token="(BOS)"   ,							eos_token="(EOS)"      )
													a__:							Optional[int]									=self.get_feature_extractor(do_normalize=_a   ,							padding_value=1.0      )

													a__:							Any									=WhisperProcessor.from_pretrained(
													    self.tmpdirname   ,							bos_token="(BOS)"   ,							eos_token="(EOS)"   ,							do_normalize=_a   ,							padding_value=1.0      )

													self.assertEqual(processor.tokenizer.get_vocab()   ,							tokenizer_add_kwargs.get_vocab()      )
													self.assertIsInstance(processor.tokenizer   ,							_a      )

													self.assertEqual(processor.feature_extractor.to_json_string()   ,							feature_extractor_add_kwargs.to_json_string()      )
													self.assertIsInstance(processor.feature_extractor   ,							_a      )


						def 	_lowerCamelCase      (			self    :	Tuple      ):
													a__:							Optional[int]									=self.get_feature_extractor()
													a__:							str									=self.get_tokenizer()

													a__:							int									=WhisperProcessor(tokenizer=_a   ,							feature_extractor=_a      )

													a__:							List[str]									=floats_list((3, 1_0_0_0)      )

													a__:							Optional[int]									=feature_extractor(_a   ,							return_tensors="np"      )
													a__:							Optional[int]									=processor(_a   ,							return_tensors="np"      )

													for key in input_feat_extract.keys():
																				self.assertAlmostEqual(input_feat_extract[key].sum()   ,							input_processor[key].sum()   ,							delta=1e-2      )


						def 	_lowerCamelCase      (			self    :	Optional[Any]      ):
													a__:							Dict									=self.get_feature_extractor()
													a__:							int									=self.get_tokenizer()

													a__:							List[Any]									=WhisperProcessor(tokenizer=_a   ,							feature_extractor=_a      )

													a__:							List[Any]									="This is a test string"

													a__:							Tuple									=processor(text=_a      )

													a__:							int									=tokenizer(_a      )

													for key in encoded_tok.keys():
																				self.assertListEqual(encoded_tok[key]   ,							encoded_processor[key]      )


						def 	_lowerCamelCase      (			self    :	Any      ):
													a__:							int									=self.get_feature_extractor()
													a__:							List[str]									=self.get_tokenizer()

													a__:							List[str]									=WhisperProcessor(tokenizer=_a   ,							feature_extractor=_a      )

													a__:							List[Any]									=[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

													a__:							List[Any]									=processor.batch_decode(_a      )
													a__:							int									=tokenizer.batch_decode(_a      )

													self.assertListEqual(_a   ,							_a      )


						def 	_lowerCamelCase      (			self    :	List[str]      ):
													a__:							List[Any]									=self.get_feature_extractor()
													a__:							Any									=self.get_tokenizer()

													a__:							Optional[int]									=WhisperProcessor(tokenizer=_a   ,							feature_extractor=_a      )

													self.assertListEqual(
													    processor.model_input_names   ,							feature_extractor.model_input_names   ,							msg="`processor` and `feature_extractor` model input names do not match"   ,							)


						def 	_lowerCamelCase      (			self    :	List[Any]      ):
													a__:							Any									=self.get_feature_extractor()
													a__:							str									=self.get_tokenizer()

													a__:							Dict									=WhisperProcessor(tokenizer=_a   ,							feature_extractor=_a      )
													a__:							List[Any]									=processor.get_decoder_prompt_ids(task="transcribe"   ,							no_timestamps=_a      )

													self.assertIsInstance(_a   ,							_a      )
													for ids in forced_decoder_ids:
																				self.assertIsInstance(_a   ,							(list, tuple)      )

													a__:							Any									=[TRANSCRIBE, NOTIMESTAMPS]
													self.assertListEqual([ids[-1] for ids in forced_decoder_ids]   ,							_a      )


						def 	_lowerCamelCase      (			self    :	str      ):
													a__:							Tuple									=WhisperProcessor(tokenizer=self.get_tokenizer()   ,							feature_extractor=self.get_feature_extractor()      )
													a__:							List[Any]									=processor.get_prompt_ids("Mr. Quilter"      )
													a__:							Optional[Any]									=processor.tokenizer.decode(_a      )

													self.assertListEqual(prompt_ids.tolist()   ,							[5_0_3_6_0, 1_7_7_0, 1_3, 2_2_6_4, 3_4_6, 3_5_3]      )
													self.assertEqual(_a   ,							"<|startofprev|> Mr. Quilter"      )


						def 	_lowerCamelCase      (			self    :	List[str]      ):
													a__:							int									=WhisperProcessor(tokenizer=self.get_tokenizer()   ,							feature_extractor=self.get_feature_extractor()      )
													a__:							List[str]									=processor.get_prompt_ids(""      )
													a__:							Any									=processor.tokenizer.decode(_a      )

													self.assertListEqual(prompt_ids.tolist()   ,							[5_0_3_6_0, 2_2_0]      )
													self.assertEqual(_a   ,							"<|startofprev|> "      )







						def 	_lowerCamelCase      (			self    :	Any      ):
													a__:							Tuple									=WhisperProcessor(tokenizer=self.get_tokenizer()   ,							feature_extractor=self.get_feature_extractor()      )

													def _test_prompt_error_raised_helper(_a    :	Union[str, Any]   ,							_a    :	List[str]      ):
																				with pytest.raises(_a      ) as excinfo:
																											processor.get_prompt_ids(_a      )
																				a__:							List[Any]									=F"Encountered text in the prompt corresponding to disallowed special token: {special_token}."
																				self.assertEqual(_a   ,							str(excinfo.value      )      )

													_test_prompt_error_raised_helper("<|startofprev|> test"   ,							"<|startofprev|>"      )
													_test_prompt_error_raised_helper("test <|notimestamps|>"   ,							"<|notimestamps|>"      )
													_test_prompt_error_raised_helper("test <|zh|> test <|transcribe|>"   ,							"<|zh|>"      )





