




import os
import tempfile
import unittest

from transformers import DistilBertConfig, is_torch_available
from transformers.testing_utils import require_torch, require_torch_gpu, 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 (
      DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      DistilBertForMaskedLM,
      DistilBertForMultipleChoice,
      DistilBertForQuestionAnswering,
      DistilBertForSequenceClassification,
      DistilBertForTokenClassification,
      DistilBertModel,
  )


class 				UpperCamelCase_    ( UpperCAmelCase__     ):

    '''simple docstring'''






    def __init__(			self      :	int					, UpperCAmelCase__      :	Any					, UpperCAmelCase__      :	Union[str, Any]=13					, UpperCAmelCase__      :	List[str]=7					, UpperCAmelCase__      :	Any=True					, UpperCAmelCase__      :	Tuple=True					, UpperCAmelCase__      :	str=False					, UpperCAmelCase__      :	List[str]=True					, UpperCAmelCase__      :	Optional[Any]=99					, UpperCAmelCase__      :	Dict=32					, UpperCAmelCase__      :	Union[str, Any]=5					, UpperCAmelCase__      :	Optional[Any]=4					, UpperCAmelCase__      :	int=37					, UpperCAmelCase__      :	str="gelu"					, UpperCAmelCase__      :	Optional[Any]=0.1					, UpperCAmelCase__      :	Optional[int]=0.1					, UpperCAmelCase__      :	List[str]=512					, UpperCAmelCase__      :	List[Any]=16					, UpperCAmelCase__      :	Union[str, Any]=2					, UpperCAmelCase__      :	Optional[int]=0.02					, UpperCAmelCase__      :	List[Any]=3					, UpperCAmelCase__      :	Union[str, Any]=4					, UpperCAmelCase__      :	List[Any]=None					, )						->Dict:





           '''simple docstring'''

           A__       =    parent
           A__       =    batch_size
           A__       =    seq_length
           A__       =    is_training
           A__       =    use_input_mask
           A__       =    use_token_type_ids
           A__       =    use_labels
           A__       =    vocab_size
           A__       =    hidden_size
           A__       =    num_hidden_layers
           A__       =    num_attention_heads
           A__       =    intermediate_size
           A__       =    hidden_act
           A__       =    hidden_dropout_prob
           A__       =    attention_probs_dropout_prob
           A__       =    max_position_embeddings
           A__       =    type_vocab_size
           A__       =    type_sequence_label_size
           A__       =    initializer_range
           A__       =    num_labels
           A__       =    num_choices
           A__       =    scope


    def        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any])						->Optional[int]:





           '''simple docstring'''

           A__       =    ids_tensor([self.batch_size, self.seq_length]					, self.vocab_size)

           A__       =    None
           if self.use_input_mask:
                  A__       =    random_attention_mask([self.batch_size, self.seq_length])

           A__       =    None
           A__       =    None
           A__       =    None
           if self.use_labels:
                  A__       =    ids_tensor([self.batch_size]					, self.type_sequence_label_size)
                  A__       =    ids_tensor([self.batch_size, self.seq_length]					, self.num_labels)
                  A__       =    ids_tensor([self.batch_size]					, self.num_choices)

           A__       =    self.get_config()

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


    def        SCREAMING_SNAKE_CASE							(			self      :	str)						->str:





           '''simple docstring'''

           return DistilBertConfig(
               vocab_size=self.vocab_size					, dim=self.hidden_size					, n_layers=self.num_hidden_layers					, n_heads=self.num_attention_heads					, hidden_dim=self.intermediate_size					, hidden_act=self.hidden_act					, dropout=self.hidden_dropout_prob					, attention_dropout=self.attention_probs_dropout_prob					, max_position_embeddings=self.max_position_embeddings					, initializer_range=self.initializer_range					, )


    def        SCREAMING_SNAKE_CASE							(			self      :	int					, UpperCAmelCase__      :	int					, UpperCAmelCase__      :	Optional[Any]					, UpperCAmelCase__      :	Optional[int]					, UpperCAmelCase__      :	Union[str, Any]					, UpperCAmelCase__      :	Tuple					, UpperCAmelCase__      :	List[Any])						->Dict:





           '''simple docstring'''

           A__       =    DistilBertModel(config=UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()
           A__       =    model(UpperCAmelCase__					, UpperCAmelCase__)
           A__       =    model(UpperCAmelCase__)
           self.parent.assertEqual(result.last_hidden_state.shape					, (self.batch_size, self.seq_length, self.hidden_size))


    def        SCREAMING_SNAKE_CASE							(			self      :	str					, UpperCAmelCase__      :	str					, UpperCAmelCase__      :	Union[str, Any]					, UpperCAmelCase__      :	Optional[int]					, UpperCAmelCase__      :	str					, UpperCAmelCase__      :	List[str]					, UpperCAmelCase__      :	Any)						->int:





           '''simple docstring'''

           A__       =    DistilBertForMaskedLM(config=UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()
           A__       =    model(UpperCAmelCase__					, attention_mask=UpperCAmelCase__					, labels=UpperCAmelCase__)
           self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.seq_length, self.vocab_size))


    def        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any]					, UpperCAmelCase__      :	Tuple					, UpperCAmelCase__      :	List[Any]					, UpperCAmelCase__      :	str					, UpperCAmelCase__      :	Optional[int]					, UpperCAmelCase__      :	List[Any]					, UpperCAmelCase__      :	Any)						->List[str]:





           '''simple docstring'''

           A__       =    DistilBertForQuestionAnswering(config=UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()
           A__       =    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        SCREAMING_SNAKE_CASE							(			self      :	List[Any]					, UpperCAmelCase__      :	Any					, UpperCAmelCase__      :	List[Any]					, UpperCAmelCase__      :	str					, UpperCAmelCase__      :	Optional[Any]					, UpperCAmelCase__      :	List[Any]					, UpperCAmelCase__      :	int)						->Union[str, Any]:





           '''simple docstring'''

           A__       =    self.num_labels
           A__       =    DistilBertForSequenceClassification(UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()
           A__       =    model(UpperCAmelCase__					, attention_mask=UpperCAmelCase__					, labels=UpperCAmelCase__)
           self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.num_labels))


    def        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any]					, UpperCAmelCase__      :	Union[str, Any]					, UpperCAmelCase__      :	Dict					, UpperCAmelCase__      :	Union[str, Any]					, UpperCAmelCase__      :	Optional[int]					, UpperCAmelCase__      :	Optional[Any]					, UpperCAmelCase__      :	Union[str, Any])						->Optional[Any]:





           '''simple docstring'''

           A__       =    self.num_labels
           A__       =    DistilBertForTokenClassification(config=UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()

           A__       =    model(UpperCAmelCase__					, attention_mask=UpperCAmelCase__					, labels=UpperCAmelCase__)
           self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.seq_length, self.num_labels))


    def        SCREAMING_SNAKE_CASE							(			self      :	str					, UpperCAmelCase__      :	Union[str, Any]					, UpperCAmelCase__      :	Optional[int]					, UpperCAmelCase__      :	Tuple					, UpperCAmelCase__      :	Any					, UpperCAmelCase__      :	Optional[int]					, UpperCAmelCase__      :	Dict)						->Dict:





           '''simple docstring'''

           A__       =    self.num_choices
           A__       =    DistilBertForMultipleChoice(config=UpperCAmelCase__)
           model.to(UpperCAmelCase__)
           model.eval()
           A__       =    input_ids.unsqueeze(1).expand(-1					, self.num_choices					, -1).contiguous()
           A__       =    input_mask.unsqueeze(1).expand(-1					, self.num_choices					, -1).contiguous()
           A__       =    model(
               UpperCAmelCase__					, attention_mask=UpperCAmelCase__					, labels=UpperCAmelCase__					, )
           self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.num_choices))





    def        SCREAMING_SNAKE_CASE							(			self      :	Dict)						->List[Any]:





           '''simple docstring'''

           A__       =    self.prepare_config_and_inputs()
           ((A__) ,							(A__) ,							(A__) ,							(A__) ,							(A__) ,							(A__))       =    config_and_inputs
           A__       =    {'''input_ids''': input_ids, '''attention_mask''': input_mask}
           return config, inputs_dict


@require_torch
class 				UpperCamelCase_    ( UpperCAmelCase__		,					UpperCAmelCase__		,					unittest.TestCase     ):

    '''simple docstring'''






    UpperCAmelCase__							  =      (
        (
            DistilBertModel,
            DistilBertForMaskedLM,
            DistilBertForMultipleChoice,
            DistilBertForQuestionAnswering,
            DistilBertForSequenceClassification,
            DistilBertForTokenClassification,
        )
        if is_torch_available()
        else None
    )
    UpperCAmelCase__							  =      (
        {
            '''feature-extraction''': DistilBertModel,
            '''fill-mask''': DistilBertForMaskedLM,
            '''question-answering''': DistilBertForQuestionAnswering,
            '''text-classification''': DistilBertForSequenceClassification,
            '''token-classification''': DistilBertForTokenClassification,
            '''zero-shot''': DistilBertForSequenceClassification,
        }
        if is_torch_available()
        else {}
    )
    UpperCAmelCase__							  =      True
    UpperCAmelCase__							  =      True
    UpperCAmelCase__							  =      True
    UpperCAmelCase__							  =      True


    def        SCREAMING_SNAKE_CASE							(			self      :	str)						->List[str]:





           '''simple docstring'''

           A__       =    DistilBertModelTester(self)
           A__       =    ConfigTester(self					, config_class=UpperCAmelCase__					, dim=37)


    def        SCREAMING_SNAKE_CASE							(			self      :	Union[str, Any])						->List[Any]:





           '''simple docstring'''

           self.config_tester.run_common_tests()


    def        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any])						->Optional[int]:





           '''simple docstring'''

           A__       =    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_model(*UpperCAmelCase__)


    def        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any])						->Union[str, Any]:





           '''simple docstring'''

           A__       =    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_masked_lm(*UpperCAmelCase__)


    def        SCREAMING_SNAKE_CASE							(			self      :	List[str])						->Any:





           '''simple docstring'''

           A__       =    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_question_answering(*UpperCAmelCase__)


    def        SCREAMING_SNAKE_CASE							(			self      :	Any)						->Dict:





           '''simple docstring'''

           A__       =    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_sequence_classification(*UpperCAmelCase__)


    def        SCREAMING_SNAKE_CASE							(			self      :	int)						->List[Any]:





           '''simple docstring'''

           A__       =    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_token_classification(*UpperCAmelCase__)


    def        SCREAMING_SNAKE_CASE							(			self      :	Tuple)						->str:





           '''simple docstring'''

           A__       =    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_multiple_choice(*UpperCAmelCase__)


    @slow
    def        SCREAMING_SNAKE_CASE							(			self      :	Optional[int])						->Tuple:





           '''simple docstring'''

           for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A__       =    DistilBertModel.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)





    @slow
    @require_torch_gpu
    def        SCREAMING_SNAKE_CASE							(			self      :	Union[str, Any])						->Dict:





           '''simple docstring'''

           A__ ,							A__       =    self.model_tester.prepare_config_and_inputs_for_common()
           for model_class in self.all_model_classes:
                  # BertForMultipleChoice behaves incorrectly in JIT environments.
                  if model_class == DistilBertForMultipleChoice:
                         return

                  A__       =    True
                  A__       =    model_class(config=UpperCAmelCase__)

                  A__       =    self._prepare_for_class(UpperCAmelCase__					, UpperCAmelCase__)
                  A__       =    torch.jit.trace(
                      UpperCAmelCase__					, (inputs_dict['''input_ids'''].to('''cpu'''), inputs_dict['''attention_mask'''].to('''cpu''')))

                  with tempfile.TemporaryDirectory() as tmp:
                         torch.jit.save(UpperCAmelCase__					, os.path.join(UpperCAmelCase__					, '''traced_model.pt'''))
                         A__       =    torch.jit.load(os.path.join(UpperCAmelCase__					, '''traced_model.pt''')					, map_location=UpperCAmelCase__)
                         loaded(inputs_dict['''input_ids'''].to(UpperCAmelCase__)					, inputs_dict['''attention_mask'''].to(UpperCAmelCase__))


@require_torch
class 				UpperCamelCase_    ( unittest.TestCase     ):

    '''simple docstring'''






    @slow
    def        SCREAMING_SNAKE_CASE							(			self      :	Any)						->int:





           '''simple docstring'''

           A__       =    DistilBertModel.from_pretrained('''distilbert-base-uncased''')
           A__       =    torch.tensor([[0, 345, 232, 328, 740, 140, 1_695, 69, 6_078, 1_588, 2]])
           A__       =    torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
           with torch.no_grad():
                  A__       =    model(UpperCAmelCase__					, attention_mask=UpperCAmelCase__)[0]
           A__       =    torch.Size((1, 11, 768))
           self.assertEqual(output.shape					, UpperCAmelCase__)
           A__       =    torch.tensor(
               [[[-0.1639, 0.3299, 0.1648], [-0.1746, 0.3289, 0.1710], [-0.1884, 0.3357, 0.1810]]])

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








import os

import pytest

from transformers.dynamic_module_utils import get_imports


_lowerCamelCase	:		Any							       =   """
import os
"""

_lowerCamelCase	:		Optional[int]							       =   """
def foo():
    import os
    return False
"""

_lowerCamelCase	:		List[Any]							       =   """
def foo():
    def bar():
        if True:
            import os
        return False
    return bar()
"""

_lowerCamelCase	:		List[Any]							       =   """
import os

try:
    import bar
except ImportError:
    raise ValueError()
"""

_lowerCamelCase	:		Union[str, Any]							       =   """
import os

def foo():
    try:
        import bar
    except ImportError:
        raise ValueError()
"""

_lowerCamelCase	:		List[Any]							       =   """
import os

try:
    import bar
except (ImportError, AttributeError):
    raise ValueError()
"""

_lowerCamelCase	:		List[Any]							       =   """
import os

try:
    import bar
except ImportError as e:
    raise ValueError()
"""

_lowerCamelCase	:		str							       =   """
import os

try:
    import bar
except:
    raise ValueError()
"""

_lowerCamelCase	:		Optional[Any]							       =   """
import os

try:
    import bar
    import baz
except ImportError:
    raise ValueError()
"""

_lowerCamelCase	:		Any							       =   """
import os

try:
    import bar
    import baz
except ImportError:
    x = 1
    raise ValueError()
"""

_lowerCamelCase	:		Dict							       =   [
    TOP_LEVEL_IMPORT,
    IMPORT_IN_FUNCTION,
    DEEPLY_NESTED_IMPORT,
    TOP_LEVEL_TRY_IMPORT,
    GENERIC_EXCEPT_IMPORT,
    MULTILINE_TRY_IMPORT,
    MULTILINE_BOTH_IMPORT,
    MULTIPLE_EXCEPTS_IMPORT,
    EXCEPT_AS_IMPORT,
    TRY_IMPORT_IN_FUNCTION,
]







@pytest.mark.parametrize('''case''' ,				lowercase_						)
def SCREAMING_SNAKE_CASE				(       lowercase_ ,				lowercase_						)   ->       int:



       """simple docstring"""

       A__       =    os.path.join(lowercase_ ,				'''test_file.py'''						)
       with open(lowercase_ ,				'''w'''						) as _tmp_file:
              _tmp_file.write(lowercase_						)

       A__       =    get_imports(lowercase_						)
       assert parsed_imports == ["os"]


