





import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


A_        =  namedtuple("ExampleNamedTuple", "a b c")

class SCREAMING_SNAKE_CASE_			(       unittest.TestCase ):





					"""simple docstring"""


					def        __magic_name__			(			self						):
										lowerCamelCase__      =							torch.randn(5						,	2						)
										lowerCamelCase__      =							torch.device("cuda"						) if torch.cuda.is_available() else torch.device("cpu"						)

										lowerCamelCase__      =							send_to_device(_lowerCAmelCase						,	_lowerCAmelCase						)
										self.assertTrue(torch.equal(resulta.cpu()						,	_lowerCAmelCase						)						)

										lowerCamelCase__      =							send_to_device((tensor, [tensor, tensor], 1)						,	_lowerCAmelCase						)
										self.assertIsInstance(_lowerCAmelCase						,	_lowerCAmelCase						)
										self.assertTrue(torch.equal(resulta[0].cpu()						,	_lowerCAmelCase						)						)
										self.assertIsInstance(resulta[1]						,	_lowerCAmelCase						)
										self.assertTrue(torch.equal(resulta[1][0].cpu()						,	_lowerCAmelCase						)						)
										self.assertTrue(torch.equal(resulta[1][1].cpu()						,	_lowerCAmelCase						)						)
										self.assertEqual(resulta[2]						,	1						)

										lowerCamelCase__      =							send_to_device({"a": tensor, "b": [tensor, tensor], "c": 1}						,	_lowerCAmelCase						)
										self.assertIsInstance(_lowerCAmelCase						,	_lowerCAmelCase						)
										self.assertTrue(torch.equal(resulta["a"].cpu()						,	_lowerCAmelCase						)						)
										self.assertIsInstance(resulta["b"]						,	_lowerCAmelCase						)
										self.assertTrue(torch.equal(resulta["b"][0].cpu()						,	_lowerCAmelCase						)						)
										self.assertTrue(torch.equal(resulta["b"][1].cpu()						,	_lowerCAmelCase						)						)
										self.assertEqual(resulta["c"]						,	1						)

										lowerCamelCase__      =							send_to_device(ExampleNamedTuple(a=_lowerCAmelCase						,	b=[tensor, tensor]						,	c=1						)						,	_lowerCAmelCase						)
										self.assertIsInstance(_lowerCAmelCase						,	_lowerCAmelCase						)
										self.assertTrue(torch.equal(resulta.a.cpu()						,	_lowerCAmelCase						)						)
										self.assertIsInstance(resulta.b						,	_lowerCAmelCase						)
										self.assertTrue(torch.equal(resulta.b[0].cpu()						,	_lowerCAmelCase						)						)
										self.assertTrue(torch.equal(resulta.b[1].cpu()						,	_lowerCAmelCase						)						)
										self.assertEqual(resulta.c						,	1						)

										lowerCamelCase__      =							send_to_device(UserDict({"a": tensor, "b": [tensor, tensor], "c": 1}						)						,	_lowerCAmelCase						)
										self.assertIsInstance(_lowerCAmelCase						,	_lowerCAmelCase						)
										self.assertTrue(torch.equal(resulta["a"].cpu()						,	_lowerCAmelCase						)						)
										self.assertIsInstance(resulta["b"]						,	_lowerCAmelCase						)
										self.assertTrue(torch.equal(resulta["b"][0].cpu()						,	_lowerCAmelCase						)						)
										self.assertTrue(torch.equal(resulta["b"][1].cpu()						,	_lowerCAmelCase						)						)
										self.assertEqual(resulta["c"]						,	1						)





					def        __magic_name__			(			self						):
										with self.assertRaises(_lowerCAmelCase						) as cm:
															lowerCamelCase__      =							recursively_apply(torch.tensor						,	(torch.tensor(1						), 1)						,	error_on_other_type=_lowerCAmelCase						)
										self.assertEqual(
										    str(cm.exception						)						,	"Unsupported types (<class 'int'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed."						,	)





					def        __magic_name__			(			self						):
										lowerCamelCase__      =							torch.tensor([1, 2, 3, 4, 5]						)
										self.assertEqual(listify(_lowerCAmelCase						)						,	[1, 2, 3, 4, 5]						)

										lowerCamelCase__      =							torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]						)
										self.assertEqual(listify(_lowerCAmelCase						)						,	[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]						)

										lowerCamelCase__      =							torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]						)
										self.assertEqual(
										    listify(_lowerCAmelCase						)						,	[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]], [[11, 12, 13, 14, 15], [16, 17, 18, 19, 20]]]						)





					def        __magic_name__			(			self						):
										with patch_environment(aa=1						,	BB=2						):
															self.assertEqual(os.environ.get("AA"						)						,	"1"						)
															self.assertEqual(os.environ.get("BB"						)						,	"2"						)

										self.assertNotIn("AA"						,	os.environ						)
										self.assertNotIn("BB"						,	os.environ						)





					def        __magic_name__			(			self						):
										lowerCamelCase__      =							RegressionModel()
										lowerCamelCase__      =							model.forward
										lowerCamelCase__      =							convert_outputs_to_fpaa(model.forward						)
										lowerCamelCase__      =							extract_model_from_parallel(_lowerCAmelCase						,	keep_fpaa_wrapper=_lowerCAmelCase						)
										lowerCamelCase__      =							pickle.dumps(_lowerCAmelCase						)





					@require_cuda
					def        __magic_name__			(			self						):
										lowerCamelCase__      =							RegressionModel()
										lowerCamelCase__      =							model.forward
										lowerCamelCase__      =							torch.cuda.amp.autocast(dtype=torch.floataa						)(model.forward						)
										lowerCamelCase__      =							convert_outputs_to_fpaa(model.forward						)
										lowerCamelCase__      =							extract_model_from_parallel(_lowerCAmelCase						,	keep_fpaa_wrapper=_lowerCAmelCase						)
										lowerCamelCase__      =							pickle.dumps(_lowerCAmelCase						)





					@require_cuda
					@require_torch_min_version(version="2.0"						)
					def        __magic_name__			(			self						):
										lowerCamelCase__      =							RegressionModel()
										lowerCamelCase__      =							model.forward
										lowerCamelCase__      =							torch.cuda.amp.autocast(dtype=torch.floataa						)(model.forward						)
										lowerCamelCase__      =							convert_outputs_to_fpaa(model.forward						)
										lowerCamelCase__      =							torch.compile(model.forward						,	backend="inductor"						)
										lowerCamelCase__      =							torch.randn(4						,	10						).cuda()
										lowerCamelCase__      =							model(_lowerCAmelCase						)





					def        __magic_name__			(			self						):
										lowerCamelCase__      =							RegressionModel()
										# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
										lowerCamelCase__      =							torch.nn.parallel.DataParallel(_lowerCAmelCase						)
										lowerCamelCase__      =							extract_model_from_parallel(_lowerCAmelCase						)

										self.assertEqual(_lowerCAmelCase						,	_lowerCAmelCase						)





					@require_torch_min_version(version="2.0"						)
					def        __magic_name__			(			self						):
										lowerCamelCase__      =							RegressionModel()
										lowerCamelCase__      =							torch.compile(_lowerCAmelCase						)

										# could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
										lowerCamelCase__      =							torch.nn.parallel.DataParallel(_lowerCAmelCase						)
										lowerCamelCase__      =							torch.compile(_lowerCAmelCase						)
										lowerCamelCase__      =							extract_model_from_parallel(_lowerCAmelCase						)

										self.assertEqual(compiled_model._orig_mod						,	compiled_model_unwrapped._orig_mod						)






					def        __magic_name__			(			self						):
										self.assertEqual(find_device([1, "a", torch.tensor([1, 2, 3]						)]						)						,	torch.device("cpu"						)						)
										self.assertEqual(find_device({"a": 1, "b": torch.tensor([1, 2, 3]						)}						)						,	torch.device("cpu"						)						)
										self.assertIsNone(find_device([1, "a"]						)						)








# this script builds a small sample spm file tests/fixtures/test_sentencepiece_no_bos.model, with features needed by pegasus

# 1. pip install sentencepiece
#
# 2. wget https://raw.githubusercontent.com/google/sentencepiece/master/data/botchan.txt

# 3. build
import sentencepiece as spm

# pegasus:
# 1. no bos
# 2. eos_id is 1
# 3. unk_id is 2
# build a sample spm file accordingly
spm.SentencePieceTrainer.train(
    "--input=botchan.txt --model_prefix=test_sentencepiece_no_bos --bos_id=-1 --unk_id=2  --eos_id=1  --vocab_size=1000"
)

# 4. now update the fixture
# mv test_sentencepiece_no_bos.model ../../tests/fixtures/

