



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 					_A (			__UpperCAmelCase					):



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


           '''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 	_lowerCamelCase    (	self							:  Optional[Any]):


           '''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 	_lowerCamelCase    (	self							:  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 	_lowerCamelCase    (	self							:  int  ,	__SCREAMING_SNAKE_CASE							:  Optional[int]  ,	__SCREAMING_SNAKE_CASE							:  List[Any]  ,	__SCREAMING_SNAKE_CASE							:  Optional[int]  ,	__SCREAMING_SNAKE_CASE							:  Optional[int]  ,	__SCREAMING_SNAKE_CASE							:  int  ,	__SCREAMING_SNAKE_CASE							:  int):


           '''simple docstring'''

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



       def 	_lowerCamelCase    (	self							:  Optional[Any]  ,	__SCREAMING_SNAKE_CASE							:  Union[str, Any]  ,	__SCREAMING_SNAKE_CASE							:  str  ,	__SCREAMING_SNAKE_CASE							:  str  ,	__SCREAMING_SNAKE_CASE							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  Dict  ,	__SCREAMING_SNAKE_CASE							:  Tuple):


           '''simple docstring'''

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



       def 	_lowerCamelCase    (	self							:  Optional[Any]  ,	__SCREAMING_SNAKE_CASE							:  Dict  ,	__SCREAMING_SNAKE_CASE							:  List[Any]  ,	__SCREAMING_SNAKE_CASE							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  Any  ,	__SCREAMING_SNAKE_CASE							:  Any  ,	__SCREAMING_SNAKE_CASE							:  Optional[Any]):


           '''simple docstring'''

           __a		 =       DistilBertForQuestionAnswering(config=__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           model.eval()
           __a		 =       model(
               __SCREAMING_SNAKE_CASE  ,	attention_mask=__SCREAMING_SNAKE_CASE  ,	start_positions=__SCREAMING_SNAKE_CASE  ,	end_positions=__SCREAMING_SNAKE_CASE)
           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 	_lowerCamelCase    (	self							:  Optional[Any]  ,	__SCREAMING_SNAKE_CASE							:  Union[str, Any]  ,	__SCREAMING_SNAKE_CASE							:  List[str]  ,	__SCREAMING_SNAKE_CASE							:  Any  ,	__SCREAMING_SNAKE_CASE							:  str  ,	__SCREAMING_SNAKE_CASE							:  List[str]  ,	__SCREAMING_SNAKE_CASE							:  List[str]):


           '''simple docstring'''

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



       def 	_lowerCamelCase    (	self							:  Optional[Any]  ,	__SCREAMING_SNAKE_CASE							:  str  ,	__SCREAMING_SNAKE_CASE							:  Any  ,	__SCREAMING_SNAKE_CASE							:  str  ,	__SCREAMING_SNAKE_CASE							:  List[str]  ,	__SCREAMING_SNAKE_CASE							:  Optional[Any]  ,	__SCREAMING_SNAKE_CASE							:  List[Any]):


           '''simple docstring'''

           __a		 =       self.num_labels
           __a		 =       DistilBertForTokenClassification(config=__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           model.eval()

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



       def 	_lowerCamelCase    (	self							:  Optional[int]  ,	__SCREAMING_SNAKE_CASE							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  List[Any]  ,	__SCREAMING_SNAKE_CASE							:  Tuple  ,	__SCREAMING_SNAKE_CASE							:  Any  ,	__SCREAMING_SNAKE_CASE							:  int  ,	__SCREAMING_SNAKE_CASE							:  Union[str, Any]):


           '''simple docstring'''

           __a		 =       self.num_choices
           __a		 =       DistilBertForMultipleChoice(config=__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           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(
               __SCREAMING_SNAKE_CASE  ,	attention_mask=__SCREAMING_SNAKE_CASE  ,	labels=__SCREAMING_SNAKE_CASE  ,	)
           self.parent.assertEqual(result.logits.shape  ,	(self.batch_size, self.num_choices))





       def 	_lowerCamelCase    (	self							:  Optional[int]):


           '''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 					_A (			__UpperCAmelCase   ,__UpperCAmelCase   ,unittest.TestCase					):
       UpperCamelCase__						:					Optional[int]		      =		(
           (
               DistilBertModel,
               DistilBertForMaskedLM,
               DistilBertForMultipleChoice,
               DistilBertForQuestionAnswering,
               DistilBertForSequenceClassification,
               DistilBertForTokenClassification,
           )
           if is_torch_available()
           else None
       )
       UpperCamelCase__						:					List[str]		      =		(
           {
               '''feature-extraction''': DistilBertModel,
               '''fill-mask''': DistilBertForMaskedLM,
               '''question-answering''': DistilBertForQuestionAnswering,
               '''text-classification''': DistilBertForSequenceClassification,
               '''token-classification''': DistilBertForTokenClassification,
               '''zero-shot''': DistilBertForSequenceClassification,
           }
           if is_torch_available()
           else {}
       )
       UpperCamelCase__						:					Any		      =		True
       UpperCamelCase__						:					List[Any]		      =		True
       UpperCamelCase__						:					str		      =		True
       UpperCamelCase__						:					Optional[int]		      =		True



       def 	_lowerCamelCase    (	self							:  List[str]):


           '''simple docstring'''

           __a		 =       DistilBertModelTester(self)
           __a		 =       ConfigTester(self  ,	config_class=__SCREAMING_SNAKE_CASE  ,	dim=37)



       def 	_lowerCamelCase    (	self							:  Optional[int]):


           '''simple docstring'''

           self.config_tester.run_common_tests()



       def 	_lowerCamelCase    (	self							:  Optional[Any]):


           '''simple docstring'''

           __a		 =       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_model(*__SCREAMING_SNAKE_CASE)



       def 	_lowerCamelCase    (	self							:  Union[str, Any]):


           '''simple docstring'''

           __a		 =       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_masked_lm(*__SCREAMING_SNAKE_CASE)



       def 	_lowerCamelCase    (	self							:  Union[str, Any]):


           '''simple docstring'''

           __a		 =       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_question_answering(*__SCREAMING_SNAKE_CASE)



       def 	_lowerCamelCase    (	self							:  Dict):


           '''simple docstring'''

           __a		 =       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_sequence_classification(*__SCREAMING_SNAKE_CASE)



       def 	_lowerCamelCase    (	self							:  Tuple):


           '''simple docstring'''

           __a		 =       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_token_classification(*__SCREAMING_SNAKE_CASE)



       def 	_lowerCamelCase    (	self							:  int):


           '''simple docstring'''

           __a		 =       self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_distilbert_for_multiple_choice(*__SCREAMING_SNAKE_CASE)



       @slow
       def 	_lowerCamelCase    (	self							:  Dict):


           '''simple docstring'''

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





       @slow
       @require_torch_gpu
       def 	_lowerCamelCase    (	self							:  Optional[Any]):


           '''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=__SCREAMING_SNAKE_CASE)

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

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


@require_torch
class 					_A (			unittest.TestCase					):



       @slow
       def 	_lowerCamelCase    (	self							:  Tuple):


           '''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(__SCREAMING_SNAKE_CASE  ,	attention_mask=__SCREAMING_SNAKE_CASE)[0]
           __a		 =       torch.Size((1, 11, 768))
           self.assertEqual(output.shape  ,	__SCREAMING_SNAKE_CASE)
           __a		 =       torch.tensor(
               [[[-0.16_39, 0.32_99, 0.16_48], [-0.17_46, 0.32_89, 0.17_10], [-0.18_84, 0.33_57, 0.18_10]]])

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










__snake_case		:str   	=  [
    [0, 16, 13, 0, 0, 0],
    [0, 0, 10, 12, 0, 0],
    [0, 4, 0, 0, 14, 0],
    [0, 0, 9, 0, 0, 20],
    [0, 0, 0, 7, 0, 4],
    [0, 0, 0, 0, 0, 0],
]





def   __snake_case			(    _UpperCAmelCase     ,				_UpperCAmelCase     ,				_UpperCAmelCase     ,				_UpperCAmelCase       ):
    # Return True if there is node that has not iterated.
    __a		 =       [False] * len(_UpperCAmelCase       )
    __a		 =       [s]
    __a		 =       True

    while queue:
        __a		 =       queue.pop(0       )
        for ind in range(len(graph[u]       )       ):
            if visited[ind] is False and graph[u][ind] > 0:
                queue.append(_UpperCAmelCase       )
                __a		 =       True
                __a		 =       u

    return visited[t]





def   __snake_case			(    _UpperCAmelCase     ,				_UpperCAmelCase     ,				_UpperCAmelCase       ):
    __a		 =       [-1] * (len(_UpperCAmelCase       ))
    __a		 =       0
    __a		 =       []
    __a		 =       [i[:] for i in graph]  # Record original cut, copy.
    while bfs(_UpperCAmelCase     ,				_UpperCAmelCase     ,				_UpperCAmelCase     ,				_UpperCAmelCase       ):
        __a		 =       float('''Inf'''       )
        __a		 =       sink

        while s != source:
            # Find the minimum value in select path
            __a		 =       min(_UpperCAmelCase     ,				graph[parent[s]][s]       )
            __a		 =       parent[s]

        max_flow += path_flow
        __a		 =       sink

        while v != source:
            __a		 =       parent[v]
            graph[u][v] -= path_flow
            graph[v][u] += path_flow
            __a		 =       parent[v]

    for i in range(len(_UpperCAmelCase       )       ):
        for j in range(len(graph[0]       )       ):
            if graph[i][j] == 0 and temp[i][j] > 0:
                res.append((i, j)       )

    return res


if __name__ == "__main__":
 print(mincut(test_graph, source=0, sink=5))





