


'''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		):





     def 							__UpperCAmelCase      ( self       :      Tuple    )				->  Optional[int]:
        lowerCAmelCase     =			AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'    )
        lowerCAmelCase     =			AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'    ).to(UpperCAmelCase__    )
        lowerCAmelCase     =			-1

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

        with CaptureStdout() as cs:
           lowerCAmelCase     =			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
        lowerCAmelCase     =			cs.out[:-1]

        self.assertEqual(UpperCAmelCase__    ,      UpperCAmelCase__    )





     def 							__UpperCAmelCase      ( self       :      Any    )				->  Any:
        lowerCAmelCase     =			AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'    )
        lowerCAmelCase     =			AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'    ).to(UpperCAmelCase__    )
        lowerCAmelCase     =			-1

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

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

        self.assertEqual(UpperCAmelCase__    ,      UpperCAmelCase__    )





     def 							__UpperCAmelCase      ( self       :      Dict    )				->  str:
        lowerCAmelCase     =			AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'    )
        lowerCAmelCase     =			AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'    ).to(UpperCAmelCase__    )
        lowerCAmelCase     =			-1

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

        with CaptureStdout() as cs:
           lowerCAmelCase     =			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
        lowerCAmelCase     =			cs.out[:-1]

        self.assertEqual(UpperCAmelCase__    ,      UpperCAmelCase__    )





     def 							__UpperCAmelCase      ( self       :      List[Any]    )				->  Optional[int]:
        # 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
        lowerCAmelCase     =			AutoTokenizer.from_pretrained('distilgpt2'    )
        lowerCAmelCase     =			AutoModelForCausalLM.from_pretrained('distilgpt2'    ).to(UpperCAmelCase__    )
        lowerCAmelCase     =			-1

        lowerCAmelCase     =			torch.ones((1, 5)    ,      device=UpperCAmelCase__    ).long() * model.config.bos_token_id
        with CaptureStdout() as cs:
           lowerCAmelCase     =			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
        lowerCAmelCase     =			cs.out[:-1]  # Remove the final "\n"
        lowerCAmelCase     =			tokenizer(UpperCAmelCase__    ,      return_tensors='pt'    )
        self.assertEqual(streamer_text_tokenized.input_ids.shape    ,      (1, 1)    )





     def 							__UpperCAmelCase      ( self       :      Optional[Any]    )				->  List[Any]:
        lowerCAmelCase     =			AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'    )
        lowerCAmelCase     =			AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'    ).to(UpperCAmelCase__    )
        lowerCAmelCase     =			-1

        lowerCAmelCase     =			ids_tensor((1, 5)    ,      vocab_size=model.config.vocab_size    ).to(UpperCAmelCase__    )
        lowerCAmelCase     =			TextIteratorStreamer(UpperCAmelCase__    ,      timeout=0.001    )
        lowerCAmelCase     =			{'input_ids': input_ids, 'max_new_tokens': 1_0, 'do_sample': False, 'streamer': streamer}
        lowerCAmelCase     =			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__    ):
           lowerCAmelCase     =			''
           for new_text in streamer:
              streamer_text += new_text







# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

import torch

from ...utils import is_npu_available, is_xpu_available
from .config_args import ClusterConfig, default_json_config_file
from .config_utils import SubcommandHelpFormatter


UpperCAmelCase	      = '''Create a default config file for Accelerate with only a few flags set.'''






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE="no" ,							__SCREAMING_SNAKE_CASE = default_json_config_file ,							__SCREAMING_SNAKE_CASE = False							):
	lowercase       =    Path(__SCREAMING_SNAKE_CASE							)
	path.parent.mkdir(parents=__SCREAMING_SNAKE_CASE ,							exist_ok=__SCREAMING_SNAKE_CASE							)
	if path.exists():
		print(
		    F'''Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`.'''							)
		return False
	lowercase       =    mixed_precision.lower()
	if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
		raise ValueError(
		    F'''`mixed_precision` should be one of \'no\', \'fp16\', \'bf16\', or \'fp8\'. Received {mixed_precision}'''							)
	lowercase       =    {
	    'compute_environment': 'LOCAL_MACHINE',
	    'mixed_precision': mixed_precision,
	}
	if torch.cuda.is_available():
		lowercase       =    torch.cuda.device_count()
		lowercase       =    num_gpus
		lowercase       =    False
		if num_gpus > 1:
			lowercase       =    'MULTI_GPU'
		else:
			lowercase       =    'NO'
	elif is_xpu_available() and use_xpu:
		lowercase       =    torch.xpu.device_count()
		lowercase       =    num_xpus
		lowercase       =    False
		if num_xpus > 1:
			lowercase       =    'MULTI_XPU'
		else:
			lowercase       =    'NO'
	elif is_npu_available():
		lowercase       =    torch.npu.device_count()
		lowercase       =    num_npus
		lowercase       =    False
		if num_npus > 1:
			lowercase       =    'MULTI_NPU'
		else:
			lowercase       =    'NO'
	else:
		lowercase       =    0
		lowercase       =    True
		lowercase       =    1
		lowercase       =    'NO'
	lowercase       =    ClusterConfig(**__SCREAMING_SNAKE_CASE							)
	config.to_json_file(__SCREAMING_SNAKE_CASE							)
	return path






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
	lowercase       =    parser.add_parser('default' ,							parents=__SCREAMING_SNAKE_CASE ,							help=__SCREAMING_SNAKE_CASE ,							formatter_class=__SCREAMING_SNAKE_CASE							)
	parser.add_argument(
	    '--config_file' ,							default=__SCREAMING_SNAKE_CASE ,							help=(
	        'The path to use to store the config file. Will default to a file named default_config.yaml in the cache '
	        'location, which is the content of the environment `HF_HOME` suffixed with \'accelerate\', or if you don\'t have '
	        'such an environment variable, your cache directory (\'~/.cache\' or the content of `XDG_CACHE_HOME`) suffixed '
	        'with \'huggingface\'.'
	    ) ,							dest='save_location' ,							)

	parser.add_argument(
	    '--mixed_precision' ,							choices=['no', 'fp16', 'bf16'] ,							type=__SCREAMING_SNAKE_CASE ,							help='Whether or not to use mixed precision training. '
	    'Choose between FP16 and BF16 (bfloat16) training. '
	    'BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.' ,							default='no' ,							)
	parser.set_defaults(func=__SCREAMING_SNAKE_CASE							)
	return parser






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE							):
	lowercase       =    write_basic_config(args.mixed_precision ,							args.save_location							)
	if config_file:
		print(F'''accelerate configuration saved at {config_file}'''							)

