





"""simple docstring"""




import unittest
from queue import Empty
from threading import Thread

from transformers import AutoTokenizer, TextIteratorStreamer, TextStreamer, is_torch_available
from transformers.testing_utils import CaptureStdout, require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
      import torch

      from transformers import AutoModelForCausalLM

@require_torch
class 		UpperCamelCase__			(				unittest.TestCase       ):



     '''simple docstring'''

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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	)
           SCREAMING_SNAKE_CASE		=						AutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	).to(__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						-1

           SCREAMING_SNAKE_CASE		=						ids_tensor((1, 5)						,vocab_size=model.config.vocab_size	).to(__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						model.generate(__lowerCAmelCase						,max_new_tokens=10						,do_sample=__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						tokenizer.decode(greedy_ids[0]	)

           with CaptureStdout() as cs:
                 SCREAMING_SNAKE_CASE		=						TextStreamer(__lowerCAmelCase	)
                 model.generate(__lowerCAmelCase						,max_new_tokens=10						,do_sample=__lowerCAmelCase						,streamer=__lowerCAmelCase	)
           # The greedy text should be printed to stdout, except for the final "\n" in the streamer
           SCREAMING_SNAKE_CASE		=						cs.out[:-1]

           self.assertEqual(__lowerCAmelCase						,__lowerCAmelCase	)






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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	)
           SCREAMING_SNAKE_CASE		=						AutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	).to(__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						-1

           SCREAMING_SNAKE_CASE		=						ids_tensor((1, 5)						,vocab_size=model.config.vocab_size	).to(__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						model.generate(__lowerCAmelCase						,max_new_tokens=10						,do_sample=__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						tokenizer.decode(greedy_ids[0]	)

           SCREAMING_SNAKE_CASE		=						TextIteratorStreamer(__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						{"""input_ids""": input_ids, """max_new_tokens""": 10, """do_sample""": False, """streamer""": streamer}
           SCREAMING_SNAKE_CASE		=						Thread(target=model.generate						,kwargs=__lowerCAmelCase	)
           thread.start()
           SCREAMING_SNAKE_CASE		=						""""""
           for new_text in streamer:
                 streamer_text += new_text

           self.assertEqual(__lowerCAmelCase						,__lowerCAmelCase	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]	)					->		str:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	)
           SCREAMING_SNAKE_CASE		=						AutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	).to(__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						-1

           SCREAMING_SNAKE_CASE		=						ids_tensor((1, 5)						,vocab_size=model.config.vocab_size	).to(__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						model.generate(__lowerCAmelCase						,max_new_tokens=10						,do_sample=__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						greedy_ids[:, input_ids.shape[1] :]
           SCREAMING_SNAKE_CASE		=						tokenizer.decode(new_greedy_ids[0]	)

           with CaptureStdout() as cs:
                 SCREAMING_SNAKE_CASE		=						TextStreamer(__lowerCAmelCase						,skip_prompt=__lowerCAmelCase	)
                 model.generate(__lowerCAmelCase						,max_new_tokens=10						,do_sample=__lowerCAmelCase						,streamer=__lowerCAmelCase	)
           # The greedy text should be printed to stdout, except for the final "\n" in the streamer
           SCREAMING_SNAKE_CASE		=						cs.out[:-1]

           self.assertEqual(__lowerCAmelCase						,__lowerCAmelCase	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Union[str, Any]	)					->		Optional[int]:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						AutoTokenizer.from_pretrained("""distilgpt2"""	)
           SCREAMING_SNAKE_CASE		=						AutoModelForCausalLM.from_pretrained("""distilgpt2"""	).to(__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						-1

           SCREAMING_SNAKE_CASE		=						torch.ones((1, 5)						,device=__lowerCAmelCase	).long() * model.config.bos_token_id
           with CaptureStdout() as cs:
                 SCREAMING_SNAKE_CASE		=						TextStreamer(__lowerCAmelCase						,skip_special_tokens=__lowerCAmelCase	)
                 model.generate(__lowerCAmelCase						,max_new_tokens=1						,do_sample=__lowerCAmelCase						,streamer=__lowerCAmelCase	)

           # The prompt contains a special token, so the streamer should not print it. As such, the output text, when
           # re-tokenized, must only contain one token
           SCREAMING_SNAKE_CASE		=						cs.out[:-1]  # Remove the final "\n"
           SCREAMING_SNAKE_CASE		=						tokenizer(__lowerCAmelCase						,return_tensors="""pt"""	)
           self.assertEqual(streamer_text_tokenized.input_ids.shape						,(1, 1)	)






     def     SCREAMING_SNAKE_CASE__  (							self						:   Optional[Any]	)					->		int:




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						AutoTokenizer.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	)
           SCREAMING_SNAKE_CASE		=						AutoModelForCausalLM.from_pretrained("""hf-internal-testing/tiny-random-gpt2"""	).to(__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						-1

           SCREAMING_SNAKE_CASE		=						ids_tensor((1, 5)						,vocab_size=model.config.vocab_size	).to(__lowerCAmelCase	)
           SCREAMING_SNAKE_CASE		=						TextIteratorStreamer(__lowerCAmelCase						,timeout=0.001	)
           SCREAMING_SNAKE_CASE		=						{"""input_ids""": input_ids, """max_new_tokens""": 10, """do_sample""": False, """streamer""": streamer}
           SCREAMING_SNAKE_CASE		=						Thread(target=model.generate						,kwargs=__lowerCAmelCase	)
           thread.start()

           # The streamer will timeout after 0.001 seconds, so an exception will be raised
           with self.assertRaises(__lowerCAmelCase	):
                 SCREAMING_SNAKE_CASE		=						""""""
                 for new_text in streamer:
                       streamer_text += new_text









import unittest
from dataclasses import dataclass

import pytest

from accelerate.commands.config.config_args import SageMakerConfig
from accelerate.utils import ComputeEnvironment
from accelerate.utils.launch import _convert_nargs_to_dict

@dataclass
class 		UpperCamelCase__			(				lowerCAmelCase_       ):



     '''simple docstring'''

     __snake_case   :     int		     =      ComputeEnvironment.AMAZON_SAGEMAKER
     __snake_case   :     List[Any]		     =      True
     __snake_case   :     Optional[int]		     =      "ml.p3.2xlarge"
     __snake_case   :     List[str]		     =      "accelerate_sagemaker_execution_role"
     __snake_case   :     Tuple		     =      "hf-sm"
     __snake_case   :     Any		     =      "us-east-1"
     __snake_case   :     Union[str, Any]		     =      1
     __snake_case   :     Dict		     =      "accelerate-sagemaker-1"
     __snake_case   :     Tuple		     =      "1.6"
     __snake_case   :     List[str]		     =      "4.4"
     __snake_case   :     str		     =      "train.py"
     __snake_case   :     List[str]		     =      [
         "--model_name_or_path",
         "bert",
         "--do_train",
         "False",
         "--epochs",
         "3",
         "--learning_rate",
         "5e-5",
         "--max_steps",
         "50.5",
     ]
     __snake_case   :     Optional[int]		     =      [
         "--model_name_or_path",
         "bert",
         "--do_train",
         "--do_test",
         "False",
         "--do_predict",
         "--epochs",
         "3",
         "--learning_rate",
         "5e-5",
         "--max_steps",
         "50.5",
     ]

class 		UpperCamelCase__			(				unittest.TestCase       ):



     '''simple docstring'''

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




           '''simple docstring'''


           SCREAMING_SNAKE_CASE		=						_convert_nargs_to_dict(MockLaunchConfig.success_training_script_args	)
           assert isinstance(converted_args["""model_name_or_path"""]						,lowerCamelCase__	)
           assert isinstance(converted_args["""do_train"""]						,lowerCamelCase__	)
           assert isinstance(converted_args["""epochs"""]						,lowerCamelCase__	)
           assert isinstance(converted_args["""learning_rate"""]						,lowerCamelCase__	)
           assert isinstance(converted_args["""max_steps"""]						,lowerCamelCase__	)

           with pytest.raises(lowerCamelCase__	):
                 _convert_nargs_to_dict(MockLaunchConfig.fail_training_script_args	)


