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





   '''simple docstring'''


   @slow
   @unittest.skip('''Model is not available.'''  )
   def 				__magic_name__       (      self  :							Union[str, Any]  ):

      """simple docstring"""
      _A:						List[Any]	    = '''nvidia/megatron-gpt2-345m/'''
      if "MYDIR" in os.environ:
         _A:						int	    = os.path.join(os.environ['''MYDIR''']    ,    lowerCAmelCase_  )
      _A:						int	    = GPTaLMHeadModel.from_pretrained(lowerCAmelCase_  )
      model.to(lowerCAmelCase_  )
      model.half()

      _A:						Tuple	    = torch.tensor(
          [[1_0_1, 7_1_1_0, 1_0_0_5, 1_0_5_6, 2_0_2_3, 1_1_3_3_3, 1_7_4_1_3, 1_0_2_9, 1_0_2]]    ,    device=lowerCAmelCase_    ,    dtype=torch.long    ,    )

      with torch.no_grad():
         _A:						Dict	    = model(lowerCAmelCase_  ).logits

      _A:						Dict	    = torch.Size((1, 9, 5_0_2_5_7)  )
      self.assertEqual(output.shape    ,    lowerCAmelCase_  )

      _A:						str	    = torch.tensor(
          [
              4.9414,
              -0.2920,
              -1.2148,
              -4.0273,
              -0.5161,
              -5.2109,
              -1.2412,
              -1.8301,
              -1.7734,
              -4.7148,
              -0.2317,
              -1.0811,
              -2.1777,
              0.4141,
              -3.7969,
              -4.0586,
              -2.5332,
              -3.3809,
              4.3867,
          ]    ,    device=lowerCAmelCase_    ,    dtype=torch.half    ,    )

      for i in range(1_9  ):
         _A:						Optional[Any]	    = 8 * i // 1_7, 2_7_9_2 * i  # along the diagonal
         _A:						List[Any]	    = output[0, r, c], expected_diag[i]
         _A:						Union[str, Any]	    = F"""row={r} col={c} computed={computed} expected={expected}"""
         self.assertAlmostEqual(lowerCAmelCase_    ,    lowerCAmelCase_    ,    delta=1e-4    ,    msg=lowerCAmelCase_  )





# 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/



