



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  lowerCamelCase__    (							lowercase_):




		'''simple docstring'''



		def __init__(self	,__lowerCamelCase	,__lowerCamelCase=13	,__lowerCamelCase=7	,__lowerCamelCase=True	,__lowerCamelCase=True	,__lowerCamelCase=False	,__lowerCamelCase=True	,__lowerCamelCase=99	,__lowerCamelCase=32	,__lowerCamelCase=5	,__lowerCamelCase=4	,__lowerCamelCase=37	,__lowerCamelCase="gelu"	,__lowerCamelCase=0.1	,__lowerCamelCase=0.1	,__lowerCamelCase=5_12	,__lowerCamelCase=16	,__lowerCamelCase=2	,__lowerCamelCase=0.02	,__lowerCamelCase=3	,__lowerCamelCase=4	,__lowerCamelCase=None	,)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       str				=  parent
						lowerCAmelCase__  :       List[Any]				=  batch_size
						lowerCAmelCase__  :       int				=  seq_length
						lowerCAmelCase__  :       str				=  is_training
						lowerCAmelCase__  :       Optional[int]				=  use_input_mask
						lowerCAmelCase__  :       Optional[int]				=  use_token_type_ids
						lowerCAmelCase__  :       List[Any]				=  use_labels
						lowerCAmelCase__  :       List[Any]				=  vocab_size
						lowerCAmelCase__  :       Dict				=  hidden_size
						lowerCAmelCase__  :       str				=  num_hidden_layers
						lowerCAmelCase__  :       int				=  num_attention_heads
						lowerCAmelCase__  :       Union[str, Any]				=  intermediate_size
						lowerCAmelCase__  :       Union[str, Any]				=  hidden_act
						lowerCAmelCase__  :       List[str]				=  hidden_dropout_prob
						lowerCAmelCase__  :       Optional[int]				=  attention_probs_dropout_prob
						lowerCAmelCase__  :       str				=  max_position_embeddings
						lowerCAmelCase__  :       List[Any]				=  type_vocab_size
						lowerCAmelCase__  :       List[str]				=  type_sequence_label_size
						lowerCAmelCase__  :       Optional[Any]				=  initializer_range
						lowerCAmelCase__  :       Dict				=  num_labels
						lowerCAmelCase__  :       Optional[Any]				=  num_choices
						lowerCAmelCase__  :       List[str]				=  scope



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       List[str]				=  ids_tensor([self.batch_size, self.seq_length]	,self.vocab_size				)

						lowerCAmelCase__  :       List[Any]				=  None
						if self.use_input_mask:
										lowerCAmelCase__  :       List[Any]				=  random_attention_mask([self.batch_size, self.seq_length]				)

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

						lowerCAmelCase__  :       Tuple				=  self.get_config()

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



		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""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	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    List[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Any				=  DistilBertModel(config=__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()
						lowerCAmelCase__  :       Dict				=  model(__lowerCamelCase	,__lowerCamelCase				)
						lowerCAmelCase__  :       int				=  model(__lowerCamelCase				)
						self.parent.assertEqual(result.last_hidden_state.shape	,(self.batch_size, self.seq_length, self.hidden_size)				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    List[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  DistilBertForMaskedLM(config=__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()
						lowerCAmelCase__  :       Optional[Any]				=  model(__lowerCamelCase	,attention_mask=__lowerCamelCase	,labels=__lowerCamelCase				)
						self.parent.assertEqual(result.logits.shape	,(self.batch_size, self.seq_length, self.vocab_size)				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    int:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  DistilBertForQuestionAnswering(config=__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()
						lowerCAmelCase__  :       Optional[Any]				=  model(
						    __lowerCamelCase	,attention_mask=__lowerCamelCase	,start_positions=__lowerCamelCase	,end_positions=__lowerCamelCase				)
						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	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[Any]				=  self.num_labels
						lowerCAmelCase__  :       List[Any]				=  DistilBertForSequenceClassification(__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()
						lowerCAmelCase__  :       Union[str, Any]				=  model(__lowerCamelCase	,attention_mask=__lowerCamelCase	,labels=__lowerCamelCase				)
						self.parent.assertEqual(result.logits.shape	,(self.batch_size, self.num_labels)				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[Any]				=  self.num_labels
						lowerCAmelCase__  :       Union[str, Any]				=  DistilBertForTokenClassification(config=__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()

						lowerCAmelCase__  :       Any				=  model(__lowerCamelCase	,attention_mask=__lowerCamelCase	,labels=__lowerCamelCase				)
						self.parent.assertEqual(result.logits.shape	,(self.batch_size, self.seq_length, self.num_labels)				)



		def 	lowerCAmelCase__   (self	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase	,__lowerCamelCase				)  ->    List[str]:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  self.num_choices
						lowerCAmelCase__  :       Optional[Any]				=  DistilBertForMultipleChoice(config=__lowerCamelCase				)
						model.to(__lowerCamelCase				)
						model.eval()
						lowerCAmelCase__  :       int				=  input_ids.unsqueeze(1				).expand(-1	,self.num_choices	,-1				).contiguous()
						lowerCAmelCase__  :       str				=  input_mask.unsqueeze(1				).expand(-1	,self.num_choices	,-1				).contiguous()
						lowerCAmelCase__  :       Union[str, Any]				=  model(
						    __lowerCamelCase	,attention_mask=__lowerCamelCase	,labels=__lowerCamelCase	,)
						self.parent.assertEqual(result.logits.shape	,(self.batch_size, self.num_choices)				)



		def 	lowerCAmelCase__   (self				)  ->    Optional[Any]:

						"""simple docstring"""






						lowerCAmelCase__  :       List[Any]				=  self.prepare_config_and_inputs()
						((lowerCAmelCase__)		,	(lowerCAmelCase__)		,	(lowerCAmelCase__)		,	(lowerCAmelCase__)		,	(lowerCAmelCase__)		,	(lowerCAmelCase__))  :       Optional[int]				=  config_and_inputs
						lowerCAmelCase__  :       List[Any]				=  {'''input_ids''': input_ids, '''attention_mask''': input_mask}
						return config, inputs_dict







@require_torch
class  lowerCamelCase__    (							lowercase_   ,       lowercase_   ,       unittest.TestCase):




		'''simple docstring'''



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



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       int				=  DistilBertModelTester(self				)
						lowerCAmelCase__  :       str				=  ConfigTester(self	,config_class=__lowerCamelCase	,dim=37				)



		def 	lowerCAmelCase__   (self				)  ->    Optional[Any]:

						"""simple docstring"""






						self.config_tester.run_common_tests()



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       Optional[int]				=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_distilbert_model(*__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    Dict:

						"""simple docstring"""






						lowerCAmelCase__  :       Dict				=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_distilbert_for_masked_lm(*__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       List[str]				=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_distilbert_for_question_answering(*__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__  :       Union[str, Any]				=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_distilbert_for_sequence_classification(*__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						lowerCAmelCase__  :       List[Any]				=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_distilbert_for_token_classification(*__lowerCamelCase				)



		def 	lowerCAmelCase__   (self				)  ->    Tuple:

						"""simple docstring"""






						lowerCAmelCase__  :       Tuple				=  self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_distilbert_for_multiple_choice(*__lowerCamelCase				)



		@slow
		def 	lowerCAmelCase__   (self				)  ->    str:

						"""simple docstring"""






						for model_name in DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										lowerCAmelCase__  :       Union[str, Any]				=  DistilBertModel.from_pretrained(__lowerCamelCase				)
										self.assertIsNotNone(__lowerCamelCase				)



		@slow
		@require_torch_gpu
		def 	lowerCAmelCase__   (self				)  ->    Any:

						"""simple docstring"""






						lowerCAmelCase__		,	lowerCAmelCase__  :       int				=  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

										lowerCAmelCase__  :       Optional[Any]				=  True
										lowerCAmelCase__  :       List[Any]				=  model_class(config=__lowerCamelCase				)

										lowerCAmelCase__  :       Any				=  self._prepare_for_class(__lowerCamelCase	,__lowerCamelCase				)
										lowerCAmelCase__  :       Optional[int]				=  torch.jit.trace(
										    __lowerCamelCase	,(inputs_dict['''input_ids'''].to('''cpu'''				), inputs_dict['''attention_mask'''].to('''cpu'''				))				)

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







@require_torch
class  lowerCamelCase__    (							unittest.TestCase):




		'''simple docstring'''



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

						"""simple docstring"""






						lowerCAmelCase__  :       Tuple				=  DistilBertModel.from_pretrained('''distilbert-base-uncased'''				)
						lowerCAmelCase__  :       Any				=  torch.tensor([[0, 3_45, 2_32, 3_28, 7_40, 1_40, 16_95, 69, 60_78, 15_88, 2]]				)
						lowerCAmelCase__  :       Any				=  torch.tensor([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]				)
						with torch.no_grad():
										lowerCAmelCase__  :       Dict				=  model(__lowerCamelCase	,attention_mask=__lowerCamelCase				)[0]
						lowerCAmelCase__  :       Tuple				=  torch.Size((1, 11, 7_68)				)
						self.assertEqual(output.shape	,__lowerCamelCase				)
						lowerCAmelCase__  :       List[Any]				=  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]	,__lowerCamelCase	,atol=1e-4				)				)






# This script creates a super tiny model that is useful inside tests, when we just want to test that
# the machinery works, without needing to the check the quality of the outcomes.
#
# This version creates a tiny vocab first, and then a tiny model - so the outcome is truly tiny -
# all files ~60KB. As compared to taking a full-size model, reducing to the minimum its layers and
# emb dimensions, but keeping the full vocab + merges files, leading to ~3MB in total for all files.
# The latter is done by `fsmt-make-super-tiny-model.py`.
#
# It will be used then as "stas/tiny-wmt19-en-ru"

from pathlib import Path
import json
import tempfile

from transformers import FSMTTokenizer, FSMTConfig, FSMTForConditionalGeneration
from transformers.models.fsmt.tokenization_fsmt import VOCAB_FILES_NAMES

__snake_case       :	List[str]            ='tiny-wmt19-en-ru'

# Build

# borrowed from a test
__snake_case       :	Any            =[
    'l',
    'o',
    'w',
    'e',
    'r',
    's',
    't',
    'i',
    'd',
    'n',
    'w</w>',
    'r</w>',
    't</w>',
    'lo',
    'low',
    'er</w>',
    'low</w>',
    'lowest</w>',
    'newer</w>',
    'wider</w>',
    '<unk>',
]
__snake_case       :	Union[str, Any]            =dict(zip(vocab, range(len(vocab))))
__snake_case       :	List[Any]            =['l o 123', 'lo w 1456', 'e r</w> 1789', '']

with tempfile.TemporaryDirectory() as tmpdirname:
							__snake_case       :	Optional[Any]            =Path(tmpdirname)
							__snake_case       :	int            =build_dir / VOCAB_FILES_NAMES['src_vocab_file']
							__snake_case       :	str            =build_dir / VOCAB_FILES_NAMES['tgt_vocab_file']
							__snake_case       :	Optional[int]            =build_dir / VOCAB_FILES_NAMES['merges_file']
							with open(src_vocab_file, 'w') as fp:
														fp.write(json.dumps(vocab_tokens))
							with open(tgt_vocab_file, 'w') as fp:
														fp.write(json.dumps(vocab_tokens))
							with open(merges_file, 'w') as fp:
														fp.write('\n'.join(merges))

							__snake_case       :	Union[str, Any]            =FSMTTokenizer(
							    langs=['en', 'ru'],
							    src_vocab_size=len(vocab),
							    tgt_vocab_size=len(vocab),
							    src_vocab_file=src_vocab_file,
							    tgt_vocab_file=tgt_vocab_file,
							    merges_file=merges_file,
							)

__snake_case       :	Dict            =FSMTConfig(
    langs=['ru', 'en'],
    src_vocab_size=1_0_0_0,
    tgt_vocab_size=1_0_0_0,
    d_model=4,
    encoder_layers=1,
    decoder_layers=1,
    encoder_ffn_dim=4,
    decoder_ffn_dim=4,
    encoder_attention_heads=1,
    decoder_attention_heads=1,
)

__snake_case       :	Optional[Any]            =FSMTForConditionalGeneration(config)
print(f"""num of params {tiny_model.num_parameters()}""")

# Test
__snake_case       :	Dict            =tokenizer(['Making tiny model'], return_tensors='pt')
__snake_case       :	Dict            =tiny_model(**batch)

print('test output:', len(outputs.logits[0]))

# Save
tiny_model.half()  # makes it smaller
tiny_model.save_pretrained(mname_tiny)
tokenizer.save_pretrained(mname_tiny)

print(f"""Generated {mname_tiny}""")

# Upload
# transformers-cli upload tiny-wmt19-en-ru

