





import os
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_sentencepiece, require_tokenizers, require_torch, slow, torch_device


if is_torch_available():
				import torch

				from transformers import GPTaLMHeadModel





@require_torch
@require_sentencepiece
@require_tokenizers
class 						_A( unittest.TestCase       ):


				"""simple docstring"""


				@slow
				@unittest.skip('Model is not available.'					)
				def  UpperCAmelCase_		(      self					):
									__A     :						Union[str, Any]          =    'nvidia/megatron-gpt2-345m/'
									if "MYDIR" in os.environ:
														__A     :						Dict          =    os.path.join(os.environ['MYDIR']					, _A					)
									__A     :						str          =    GPTaLMHeadModel.from_pretrained(_A					)
									model.to(_A					)
									model.half()

									__A     :						Dict          =    torch.tensor(
									    [[101, 7110, 1005, 1056, 2023, 11333, 17413, 1029, 102]]					, device=_A					, dtype=torch.long					, )

									with torch.no_grad():
														__A     :						Union[str, Any]          =    model(_A					).logits

									__A     :						Union[str, Any]          =    torch.Size((1, 9, 50257)					)
									self.assertEqual(output.shape					, _A					)

									__A     :						List[Any]          =    torch.tensor(
									    [
									        4.9_4_1_4,
									        -0.2_9_2_0,
									        -1.2_1_4_8,
									        -4.0_2_7_3,
									        -0.5_1_6_1,
									        -5.2_1_0_9,
									        -1.2_4_1_2,
									        -1.8_3_0_1,
									        -1.7_7_3_4,
									        -4.7_1_4_8,
									        -0.2_3_1_7,
									        -1.0_8_1_1,
									        -2.1_7_7_7,
									        0.4_1_4_1,
									        -3.7_9_6_9,
									        -4.0_5_8_6,
									        -2.5_3_3_2,
									        -3.3_8_0_9,
									        4.3_8_6_7,
									    ]					, device=_A					, dtype=torch.half					, )

									for i in range(19					):
														__A   ,    __A     :						str          =    8 * i // 17, 2792 * i  # along the diagonal
														__A   ,    __A     :						Any          =    output[0, r, c], expected_diag[i]
														__A     :						Optional[int]          =    F"""row={r} col={c} computed={computed} expected={expected}"""
														self.assertAlmostEqual(_A					, _A					, delta=1e-4					, msg=_A					)











from .hash_table import HashTable





class 						_A( snake_case__       ):


				"""simple docstring"""


				def __init__(      self					, *_A					, **_A					):
									super().__init__(*_A					, **_A					)





				def  UpperCAmelCase_		(      self					, _A					, _A=None					):
									__A     :						Optional[int]          =    1
									__A     :						Union[str, Any]          =    self.hash_function(key + i * i					)

									while self.values[new_key] is not None and self.values[new_key] != key:
														i += 1
														__A     :						Tuple          =    (
														    self.hash_function(key + i * i					)
														    if not self.balanced_factor() >= self.lim_charge
														    else None
														)

														if new_key is None:
																			break

									return new_key




