

"""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 					_lowerCAmelCase  (	unittest.TestCase   ):







 """simple docstring"""







 def 					_lowercase (			self							:			List[Any]	):
    __lowercase								=						AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2"	)
    __lowercase								=						AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"	).to(UpperCAmelCase__	)
    __lowercase								=						-1

    __lowercase								=						ids_tensor((1, 5),					vocab_size=model.config.vocab_size	).to(UpperCAmelCase__	)
    __lowercase								=						model.generate(UpperCAmelCase__,					max_new_tokens=1_0,					do_sample=UpperCAmelCase__	)
    __lowercase								=						tokenizer.decode(greedy_ids[0]	)

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

    self.assertEqual(UpperCAmelCase__,					UpperCAmelCase__	)


 def 					_lowercase (			self							:			List[str]	):
    __lowercase								=						AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2"	)
    __lowercase								=						AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"	).to(UpperCAmelCase__	)
    __lowercase								=						-1

    __lowercase								=						ids_tensor((1, 5),					vocab_size=model.config.vocab_size	).to(UpperCAmelCase__	)
    __lowercase								=						model.generate(UpperCAmelCase__,					max_new_tokens=1_0,					do_sample=UpperCAmelCase__	)
    __lowercase								=						tokenizer.decode(greedy_ids[0]	)

    __lowercase								=						TextIteratorStreamer(UpperCAmelCase__	)
    __lowercase								=						{"input_ids": input_ids, "max_new_tokens": 1_0, "do_sample": False, "streamer": streamer}
    __lowercase								=						Thread(target=model.generate,					kwargs=UpperCAmelCase__	)
    thread.start()
    __lowercase								=						""
    for new_text in streamer:
       streamer_text += new_text

    self.assertEqual(UpperCAmelCase__,					UpperCAmelCase__	)


 def 					_lowercase (			self							:			Union[str, Any]	):
    __lowercase								=						AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2"	)
    __lowercase								=						AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"	).to(UpperCAmelCase__	)
    __lowercase								=						-1

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

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

    self.assertEqual(UpperCAmelCase__,					UpperCAmelCase__	)


 def 					_lowercase (			self							:			Optional[Any]	):
    # Tests that we can pass `decode_kwargs` to the streamer to control how the tokens are decoded. Must be tested
    # with actual models -- the dummy models' tokenizers are not aligned with their models, and
    # `skip_special_tokens=True` has no effect on them
    __lowercase								=						AutoTokenizer.from_pretrained("distilgpt2"	)
    __lowercase								=						AutoModelForCausalLM.from_pretrained("distilgpt2"	).to(UpperCAmelCase__	)
    __lowercase								=						-1

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

    # 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
    __lowercase								=						cs.out[:-1]  # Remove the final "\n"
    __lowercase								=						tokenizer(UpperCAmelCase__,					return_tensors="pt"	)
    self.assertEqual(streamer_text_tokenized.input_ids.shape,					(1, 1)	)



 def 					_lowercase (			self							:			Any	):
    __lowercase								=						AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2"	)
    __lowercase								=						AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2"	).to(UpperCAmelCase__	)
    __lowercase								=						-1

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

    # The streamer will timeout after 0.001 seconds, so an exception will be raised
    with self.assertRaises(UpperCAmelCase__	):
       __lowercase								=						""
       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 __snake_case				(						lowerCamelCase_		):
   lowerCAmelCase_														=							ComputeEnvironment.AMAZON_SAGEMAKER
   lowerCAmelCase_														=							True
   lowerCAmelCase_														=							"ml.p3.2xlarge"
   lowerCAmelCase_														=							"accelerate_sagemaker_execution_role"
   lowerCAmelCase_														=							"hf-sm"
   lowerCAmelCase_														=							"us-east-1"
   lowerCAmelCase_														=							1
   lowerCAmelCase_														=							"accelerate-sagemaker-1"
   lowerCAmelCase_														=							"1.6"
   lowerCAmelCase_														=							"4.4"
   lowerCAmelCase_														=							"train.py"
   lowerCAmelCase_														=							[
       "--model_name_or_path",
       "bert",
       "--do_train",
       "False",
       "--epochs",
       "3",
       "--learning_rate",
       "5e-5",
       "--max_steps",
       "50.5",
   ]
   lowerCAmelCase_														=							[
       "--model_name_or_path",
       "bert",
       "--do_train",
       "--do_test",
       "False",
       "--do_predict",
       "--epochs",
       "3",
       "--learning_rate",
       "5e-5",
       "--max_steps",
       "50.5",
   ]







class __snake_case				(						unittest.TestCase		):




   def 	__a						(    self			:				List[Any]      ):



     """simple docstring"""






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

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