




import unittest

from transformers import ErnieMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           ErnieMForInformationExtraction,
           ErnieMForMultipleChoice,
           ErnieMForQuestionAnswering,
           ErnieMForSequenceClassification,
           ErnieMForTokenClassification,
           ErnieMModel,
       )
       from transformers.models.ernie_m.modeling_ernie_m import ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST




class 		UpperCamelCase__     :

    '''simple docstring'''



    def __init__(   self					,  UpperCamelCase__					,  UpperCamelCase__=13					,  UpperCamelCase__=7					,  UpperCamelCase__=True					,  UpperCamelCase__=True					,  UpperCamelCase__=True					,  UpperCamelCase__=99					,  UpperCamelCase__=32					,  UpperCamelCase__=5					,  UpperCamelCase__=4					,  UpperCamelCase__=37					,  UpperCamelCase__="gelu"					,  UpperCamelCase__=0.1					,  UpperCamelCase__=0.1					,  UpperCamelCase__=512					,  UpperCamelCase__=16					,  UpperCamelCase__=2					,  UpperCamelCase__=0.02					,  UpperCamelCase__=3					,  UpperCamelCase__=4					,  UpperCamelCase__=None					,  )					->				List[str]:
           lowerCamelCase       :				Any      			= parent
           lowerCamelCase       :				Optional[int]      			= batch_size
           lowerCamelCase       :				Dict      			= seq_length
           lowerCamelCase       :				int      			= is_training
           lowerCamelCase       :				Dict      			= use_input_mask
           lowerCamelCase       :				str      			= use_labels
           lowerCamelCase       :				List[str]      			= vocab_size
           lowerCamelCase       :				Tuple      			= hidden_size
           lowerCamelCase       :				int      			= num_hidden_layers
           lowerCamelCase       :				Optional[int]      			= num_attention_heads
           lowerCamelCase       :				List[str]      			= intermediate_size
           lowerCamelCase       :				Optional[int]      			= hidden_act
           lowerCamelCase       :				Union[str, Any]      			= hidden_dropout_prob
           lowerCamelCase       :				str      			= attention_probs_dropout_prob
           lowerCamelCase       :				List[Any]      			= max_position_embeddings
           lowerCamelCase       :				int      			= type_vocab_size
           lowerCamelCase       :				Optional[int]      			= type_sequence_label_size
           lowerCamelCase       :				str      			= initializer_range
           lowerCamelCase       :				Any      			= num_labels
           lowerCamelCase       :				Any      			= num_choices
           lowerCamelCase       :				Optional[Any]      			= scope







    def    _lowercase					(   self					)					->				List[Any]:
           lowerCamelCase       :				Optional[Any]      			= ids_tensor([self.batch_size, self.seq_length]					,  self.vocab_size					)

           lowerCamelCase       :				Optional[int]      			= None
           if self.use_input_mask:
                  lowerCamelCase       :				List[str]      			= random_attention_mask([self.batch_size, self.seq_length]					)

           lowerCamelCase       :				str      			= None
           lowerCamelCase       :				List[str]      			= None
           lowerCamelCase       :				Union[str, Any]      			= None
           if self.use_labels:
                  lowerCamelCase       :				Optional[int]      			= ids_tensor([self.batch_size]					,  self.type_sequence_label_size					)
                  lowerCamelCase       :				Any      			= ids_tensor([self.batch_size, self.seq_length]					,  self.num_labels					)
                  lowerCamelCase       :				Union[str, Any]      			= ids_tensor([self.batch_size]					,  self.num_choices					)

           lowerCamelCase       :				Any      			= self.get_config()

           return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels







    def    _lowercase					(   self					)					->				str:
           lowerCamelCase       :				Tuple      			= ids_tensor([self.batch_size, self.seq_length]					,  self.vocab_size					)

           lowerCamelCase       :				str      			= None
           if self.use_input_mask:
                  lowerCamelCase       :				Dict      			= random_attention_mask([self.batch_size, self.seq_length]					)
           lowerCamelCase       :				int      			= self.get_config()

           return config, input_ids, input_mask







    def    _lowercase					(   self					)					->				Optional[int]:
           return ErnieMConfig(
               vocab_size=self.vocab_size					,  hidden_size=self.hidden_size					,  num_hidden_layers=self.num_hidden_layers					,  num_attention_heads=self.num_attention_heads					,  intermediate_size=self.intermediate_size					,  hidden_act=self.hidden_act					,  hidden_dropout_prob=self.hidden_dropout_prob					,  attention_probs_dropout_prob=self.attention_probs_dropout_prob					,  max_position_embeddings=self.max_position_embeddings					,  type_vocab_size=self.type_vocab_size					,  initializer_range=self.initializer_range					,  )







    def    _lowercase					(   self					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					)					->				Optional[Any]:
           lowerCamelCase       :				List[Any]      			= ErnieMModel(config=UpperCamelCase__					)
           model.to(UpperCamelCase__					)
           model.eval()
           lowerCamelCase       :				Optional[int]      			= model(UpperCamelCase__					,  return_dict=UpperCamelCase__					)
           self.parent.assertEqual(result.last_hidden_state.shape					,  (self.batch_size, self.seq_length, self.hidden_size)					)







    def    _lowercase					(   self					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					)					->				Dict:
           lowerCamelCase       :				Optional[int]      			= ErnieMForQuestionAnswering(config=UpperCamelCase__					)
           model.to(UpperCamelCase__					)
           model.eval()
           lowerCamelCase       :				Optional[int]      			= model(
               UpperCamelCase__					,  attention_mask=UpperCamelCase__					,  start_positions=UpperCamelCase__					,  end_positions=UpperCamelCase__					,  )
           self.parent.assertEqual(result.start_logits.shape					,  (self.batch_size, self.seq_length)					)
           self.parent.assertEqual(result.end_logits.shape					,  (self.batch_size, self.seq_length)					)







    def    _lowercase					(   self					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					)					->				Optional[Any]:
           lowerCamelCase       :				Optional[int]      			= ErnieMForInformationExtraction(config=UpperCamelCase__					)
           model.to(UpperCamelCase__					)
           model.eval()
           lowerCamelCase       :				List[str]      			= torch.ones_like(UpperCamelCase__					,  dtype=torch.floataa					)
           lowerCamelCase       :				int      			= model(
               UpperCamelCase__					,  attention_mask=UpperCamelCase__					,  start_positions=UpperCamelCase__					,  end_positions=UpperCamelCase__					,  )
           self.parent.assertEqual(result.start_logits.shape					,  (self.batch_size, self.seq_length)					)
           self.parent.assertEqual(result.end_logits.shape					,  (self.batch_size, self.seq_length)					)







    def    _lowercase					(   self					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					)					->				List[Any]:
           lowerCamelCase       :				List[str]      			= self.num_labels
           lowerCamelCase       :				str      			= ErnieMForSequenceClassification(UpperCamelCase__					)
           model.to(UpperCamelCase__					)
           model.eval()
           lowerCamelCase       :				Union[str, Any]      			= model(UpperCamelCase__					,  attention_mask=UpperCamelCase__					,  labels=UpperCamelCase__					)
           self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.num_labels)					)







    def    _lowercase					(   self					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					)					->				Optional[int]:
           lowerCamelCase       :				List[Any]      			= self.num_labels
           lowerCamelCase       :				str      			= ErnieMForTokenClassification(config=UpperCamelCase__					)
           model.to(UpperCamelCase__					)
           model.eval()
           input_ids.to(UpperCamelCase__					)
           input_mask.to(UpperCamelCase__					)
           token_labels.to(UpperCamelCase__					)

           lowerCamelCase       :				Dict      			= model(UpperCamelCase__					,  attention_mask=UpperCamelCase__					,  labels=UpperCamelCase__					)

           self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.seq_length, self.num_labels)					)







    def    _lowercase					(   self					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					)					->				Optional[int]:
           lowerCamelCase       :				Tuple      			= self.num_choices
           lowerCamelCase       :				Optional[Any]      			= ErnieMForMultipleChoice(config=UpperCamelCase__					)
           model.to(UpperCamelCase__					)
           model.eval()
           lowerCamelCase       :				Optional[Any]      			= input_ids.unsqueeze(1					).expand(-1					,  self.num_choices					,  -1					).contiguous()
           lowerCamelCase       :				Optional[int]      			= input_mask.unsqueeze(1					).expand(-1					,  self.num_choices					,  -1					).contiguous()
           lowerCamelCase       :				int      			= model(
               UpperCamelCase__					,  attention_mask=UpperCamelCase__					,  labels=UpperCamelCase__					,  )
           self.parent.assertEqual(result.logits.shape					,  (self.batch_size, self.num_choices)					)







    def    _lowercase					(   self					)					->				Any:
           lowerCamelCase       :				Optional[int]      			= self.prepare_config_and_inputs()
           (
               (
               lowerCamelCase
           )			,				(
               lowerCamelCase
           )			,				(
               lowerCamelCase
           )			,				(
               lowerCamelCase
           )			,				(
               lowerCamelCase
           )			,				(
               lowerCamelCase
           )			,				
           )       :				int      			= config_and_inputs
           lowerCamelCase       :				str      			= {"input_ids": input_ids, "attention_mask": input_mask}
           return config, inputs_dict




@require_torch
class 		UpperCamelCase__     (lowerCAmelCase__   ,    lowerCAmelCase__   ,    unittest.TestCase					):

    '''simple docstring'''



    lowerCamelCase_				:     Dict     					=				(
        (
            ErnieMModel,
            ErnieMForMultipleChoice,
            ErnieMForQuestionAnswering,
            ErnieMForSequenceClassification,
            ErnieMForTokenClassification,
        )
        if is_torch_available()
        else ()
    )
    lowerCamelCase_				:     Union[str, Any]     					=				()
    lowerCamelCase_				:     str     					=				(
        {
            """feature-extraction""": ErnieMModel,
            """question-answering""": ErnieMForQuestionAnswering,
            """text-classification""": ErnieMForSequenceClassification,
            """token-classification""": ErnieMForTokenClassification,
            """zero-shot""": ErnieMForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    lowerCamelCase_				:     str     					=				False







    def    _lowercase					(   self					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					,  UpperCamelCase__					)					->				List[Any]:
           if pipeline_test_casse_name == "QAPipelineTests":
                  return True

           return False







    def    _lowercase					(   self					)					->				Tuple:
           lowerCamelCase       :				Tuple      			= ErnieMModelTester(self					)
           lowerCamelCase       :				List[Any]      			= ConfigTester(self					,  config_class=UpperCamelCase__					,  hidden_size=37					)







    def    _lowercase					(   self					)					->				List[Any]:
           self.config_tester.run_common_tests()







    def    _lowercase					(   self					)					->				List[str]:
           lowerCamelCase       :				Optional[Any]      			= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase__					)







    def    _lowercase					(   self					)					->				int:
           lowerCamelCase       :				List[Any]      			= self.model_tester.prepare_config_and_inputs()
           for type in ["absolute", "relative_key", "relative_key_query"]:
                  lowerCamelCase       :				int      			= type
                  self.model_tester.create_and_check_model(*UpperCamelCase__					)







    def    _lowercase					(   self					)					->				Dict:
           lowerCamelCase       :				int      			= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_multiple_choice(*UpperCamelCase__					)







    def    _lowercase					(   self					)					->				List[Any]:
           lowerCamelCase       :				Optional[Any]      			= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_question_answering(*UpperCamelCase__					)







    def    _lowercase					(   self					)					->				Any:
           lowerCamelCase       :				Tuple      			= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_sequence_classification(*UpperCamelCase__					)







    def    _lowercase					(   self					)					->				Dict:
           lowerCamelCase       :				Union[str, Any]      			= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_information_extraction(*UpperCamelCase__					)







    def    _lowercase					(   self					)					->				Dict:
           lowerCamelCase       :				Any      			= self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__					)







    @slow
    def    _lowercase					(   self					)					->				Optional[Any]:
           for model_name in ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowerCamelCase       :				Union[str, Any]      			= ErnieMModel.from_pretrained(UpperCamelCase__					)
                  self.assertIsNotNone(UpperCamelCase__					)




@require_torch
class 		UpperCamelCase__     (unittest.TestCase					):

    '''simple docstring'''



    @slow
    def    _lowercase					(   self					)					->				Dict:
           lowerCamelCase       :				Tuple      			= ErnieMModel.from_pretrained("susnato/ernie-m-base_pytorch"					)
           model.eval()
           lowerCamelCase       :				Optional[Any]      			= torch.tensor([[0, 1, 2, 3, 4, 5]]					)
           lowerCamelCase       :				Dict      			= model(UpperCamelCase__					)[0]

           # TODO Replace vocab size
           lowerCamelCase       :				Optional[int]      			= 768

           lowerCamelCase       :				Optional[Any]      			= torch.Size((1, 6, hidden_size)					)
           self.assertEqual(output.shape					,  UpperCamelCase__					)

           lowerCamelCase       :				Dict      			= torch.tensor(
               [[[-0.0012, 0.1245, -0.0214], [-0.0742, 0.0244, -0.0771], [-0.0333, 0.1164, -0.1554]]]					)

           self.assertTrue(torch.allclose(output[:, :3, :3]					,  UpperCamelCase__					,  atol=1e-3					)					)












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






