import itertools
import json
import linecache
import os
import pickle
import re
import socket
import string
from collections import Counter
from logging import getLogger
from pathlib import Path
from typing import Callable, Dict, Iterable, List

import git
import torch
from torch.utils.data import Dataset

from transformers import BartTokenizer, RagTokenizer, TaTokenizer




def        lowerCAmelCase_			(					__A,  __A,  __A,  __A,  __A=True,  __A="pt"							)					->    Optional[Any]:







						'''simple docstring'''


						UpperCAmelCase__								=  {"add_prefix_space": True} if isinstance(__A,  __A							) and not line.startswith(" "							) else {}
						UpperCAmelCase__								=  padding_side
						return tokenizer(
						    [line],  max_length=__A,  padding="max_length" if pad_to_max_length else None,  truncation=__A,  return_tensors=__A,  add_special_tokens=__A,  **__A,  )




def        lowerCAmelCase_			(					__A,  __A,  __A=None,  )					->    Any:







						'''simple docstring'''


						UpperCAmelCase__								=  input_ids.ne(__A							).any(dim=0							)
						if attention_mask is None:
												return input_ids[:, keep_column_mask]
						else:
												return (input_ids[:, keep_column_mask], attention_mask[:, keep_column_mask])





class 					A							(    UpperCAmelCase_					):


			def __init__(self					:						str	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Any="train"	,  __UpperCAmelCase					:						str=None	,  __UpperCAmelCase					:						Optional[Any]=None	,  __UpperCAmelCase					:						Any=None	,  __UpperCAmelCase					:						Any=""	,  )						->    Optional[Any]:
									"""simple docstring"""



									super().__init__()
									UpperCAmelCase__								=  Path(__UpperCAmelCase		).joinpath(type_path + ".source"		)
									UpperCAmelCase__								=  Path(__UpperCAmelCase		).joinpath(type_path + ".target"		)
									UpperCAmelCase__								=  self.get_char_lens(self.src_file		)
									UpperCAmelCase__								=  max_source_length
									UpperCAmelCase__								=  max_target_length
									assert min(self.src_lens		) > 0, f"""found empty line in {self.src_file}"""
									UpperCAmelCase__								=  tokenizer
									UpperCAmelCase__								=  prefix
									if n_obs is not None:
															UpperCAmelCase__								=  self.src_lens[:n_obs]
									UpperCAmelCase__								=  src_lang
									UpperCAmelCase__								=  tgt_lang


			def __len__(self					:						List[str]		)						->    Any:
									"""simple docstring"""



									return len(self.src_lens		)


			def __getitem__(self					:						str	,  __UpperCAmelCase					:						Union[str, Any]		)						->    Dict[str, torch.Tensor]:
									"""simple docstring"""



									UpperCAmelCase__								=  index + 1  # linecache starts at 1
									UpperCAmelCase__								=  self.prefix + linecache.getline(str(self.src_file		)	,  __UpperCAmelCase		).rstrip("\n"		)
									UpperCAmelCase__								=  linecache.getline(str(self.tgt_file		)	,  __UpperCAmelCase		).rstrip("\n"		)
									assert source_line, f"""empty source line for index {index}"""
									assert tgt_line, f"""empty tgt line for index {index}"""

									# Need to add eos token manually for T5
									if isinstance(self.tokenizer	,  __UpperCAmelCase		):
															source_line += self.tokenizer.eos_token
															tgt_line += self.tokenizer.eos_token

									# Pad source and target to the right
									UpperCAmelCase__								=  (
									    self.tokenizer.question_encoder if isinstance(self.tokenizer	,  __UpperCAmelCase		) else self.tokenizer
									)
									UpperCAmelCase__								=  self.tokenizer.generator if isinstance(self.tokenizer	,  __UpperCAmelCase		) else self.tokenizer

									UpperCAmelCase__								=  encode_line(__UpperCAmelCase	,  __UpperCAmelCase	,  self.max_source_length	,  "right"		)
									UpperCAmelCase__								=  encode_line(__UpperCAmelCase	,  __UpperCAmelCase	,  self.max_target_length	,  "right"		)

									UpperCAmelCase__								=  source_inputs["input_ids"].squeeze()
									UpperCAmelCase__								=  target_inputs["input_ids"].squeeze()
									UpperCAmelCase__								=  source_inputs["attention_mask"].squeeze()
									return {
									    "input_ids": source_ids,
									    "attention_mask": src_mask,
									    "decoder_input_ids": target_ids,
									}


			@staticmethod
			def       lowercase_				(__UpperCAmelCase					:						str		)						->    Tuple:
									"""simple docstring"""



									return [len(__UpperCAmelCase		) for x in Path(__UpperCAmelCase		).open().readlines()]





			def       lowercase_				(self					:						Optional[Any]	,  __UpperCAmelCase					:						List[str]		)						->    Dict[str, torch.Tensor]:
									"""simple docstring"""



									UpperCAmelCase__								=  torch.stack([x["input_ids"] for x in batch]		)
									UpperCAmelCase__								=  torch.stack([x["attention_mask"] for x in batch]		)
									UpperCAmelCase__								=  torch.stack([x["decoder_input_ids"] for x in batch]		)
									UpperCAmelCase__								=  (
									    self.tokenizer.generator.pad_token_id
									    if isinstance(self.tokenizer	,  __UpperCAmelCase		)
									    else self.tokenizer.pad_token_id
									)
									UpperCAmelCase__								=  (
									    self.tokenizer.question_encoder.pad_token_id
									    if isinstance(self.tokenizer	,  __UpperCAmelCase		)
									    else self.tokenizer.pad_token_id
									)
									UpperCAmelCase__								=  trim_batch(__UpperCAmelCase	,  __UpperCAmelCase		)
									UpperCAmelCase__		,				UpperCAmelCase__								=  trim_batch(__UpperCAmelCase	,  __UpperCAmelCase	,  attention_mask=__UpperCAmelCase		)
									UpperCAmelCase__								=  {
									    "input_ids": source_ids,
									    "attention_mask": source_mask,
									    "decoder_input_ids": y,
									}
									return batch




UpperCamelCase__    							= getLogger(__name__)




def        lowerCAmelCase_			(					__A							)					->    Dict:







						'''simple docstring'''


						return list(itertools.chain.from_iterable(__A							)							)




def        lowerCAmelCase_			(					__A							)					->    None:







						'''simple docstring'''


						UpperCAmelCase__								=  get_git_info()
						save_json(__A,  os.path.join(__A,  "git_log.json"							)							)




def        lowerCAmelCase_			(					__A,  __A,  __A=4,  **__A							)					->    List[str]:







						'''simple docstring'''


						with open(__A,  "w"							) as f:
												json.dump(__A,  __A,  indent=__A,  **__A							)




def        lowerCAmelCase_			(					__A							)					->    Optional[int]:







						'''simple docstring'''


						with open(__A							) as f:
												return json.load(__A							)




def        lowerCAmelCase_			(					)					->    List[Any]:







						'''simple docstring'''


						UpperCAmelCase__								=  git.Repo(search_parent_directories=__A							)
						UpperCAmelCase__								=  {
						    "repo_id": str(__A							),
						    "repo_sha": str(repo.head.object.hexsha							),
						    "repo_branch": str(repo.active_branch							),
						    "hostname": str(socket.gethostname()							),
						}
						return repo_infos




def        lowerCAmelCase_			(					__A,  __A							)					->    List:







						'''simple docstring'''


						return list(map(__A,  __A							)							)




def        lowerCAmelCase_			(					__A,  __A							)					->    List[Any]:







						'''simple docstring'''


						with open(__A,  "wb"							) as f:
												return pickle.dump(__A,  __A							)




def        lowerCAmelCase_			(					__A							)					->    Any:







						'''simple docstring'''


						def remove_articles(__A							):
												return re.sub(r"\b(a|an|the)\b",  " ",  __A							)

						def white_space_fix(__A							):
												return " ".join(text.split()							)

						def remove_punc(__A							):
												UpperCAmelCase__								=  set(string.punctuation							)
												return "".join(ch for ch in text if ch not in exclude							)

						def lower(__A							):
												return text.lower()

						return white_space_fix(remove_articles(remove_punc(lower(__A							)							)							)							)




def        lowerCAmelCase_			(					__A,  __A							)					->    Any:







						'''simple docstring'''


						UpperCAmelCase__								=  normalize_answer(__A							).split()
						UpperCAmelCase__								=  normalize_answer(__A							).split()
						UpperCAmelCase__								=  Counter(__A							) & Counter(__A							)
						UpperCAmelCase__								=  sum(common.values()							)
						if num_same == 0:
												return 0
						UpperCAmelCase__								=  1.0 * num_same / len(__A							)
						UpperCAmelCase__								=  1.0 * num_same / len(__A							)
						UpperCAmelCase__								=  (2 * precision * recall) / (precision + recall)
						return fa




def        lowerCAmelCase_			(					__A,  __A							)					->    Dict:







						'''simple docstring'''


						return normalize_answer(__A							) == normalize_answer(__A							)




def        lowerCAmelCase_			(					__A,  __A							)					->    Dict:







						'''simple docstring'''


						assert len(__A							) == len(__A							)
						UpperCAmelCase__								=  0
						for hypo, pred in zip(__A,  __A							):
												em += exact_match_score(__A,  __A							)
						if len(__A							) > 0:
												em /= len(__A							)
						return {"em": em}




def        lowerCAmelCase_			(					__A							)					->    int:







						'''simple docstring'''


						return model_prefix.startswith("rag"							)







def        lowerCAmelCase_			(					__A,  __A,  __A							)					->    Any:







						'''simple docstring'''


						UpperCAmelCase__								=  {p: p for p in extra_params}
						# T5 models don't have `dropout` param, they have `dropout_rate` instead
						UpperCAmelCase__								=  "dropout_rate"
						for p in extra_params:
												if getattr(__A,  __A,  __A							):
																		if not hasattr(__A,  __A							) and not hasattr(__A,  equivalent_param[p]							):
																								logger.info("config doesn't have a `{}` attribute".format(__A							)							)
																								delattr(__A,  __A							)
																								continue
																		UpperCAmelCase__								=  p if hasattr(__A,  __A							) else equivalent_param[p]
																		setattr(__A,  __A,  getattr(__A,  __A							)							)
																		delattr(__A,  __A							)
						return hparams, config




from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
      import tensorflow as tf

      from transformers.models.distilbert.modeling_tf_distilbert import (
          TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFDistilBertForMaskedLM,
          TFDistilBertForMultipleChoice,
          TFDistilBertForQuestionAnswering,
          TFDistilBertForSequenceClassification,
          TFDistilBertForTokenClassification,
          TFDistilBertModel,
      )





class 					A							:


   def __init__(self					:						Tuple	,  __UpperCAmelCase					:						str	,  )						->    Dict:
         """simple docstring"""



         UpperCAmelCase__								=  parent
         UpperCAmelCase__								=  1_3
         UpperCAmelCase__								=  7
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  False
         UpperCAmelCase__								=  True
         UpperCAmelCase__								=  9_9
         UpperCAmelCase__								=  3_2
         UpperCAmelCase__								=  2
         UpperCAmelCase__								=  4
         UpperCAmelCase__								=  3_7
         UpperCAmelCase__								=  "gelu"
         UpperCAmelCase__								=  0.1
         UpperCAmelCase__								=  0.1
         UpperCAmelCase__								=  5_1_2
         UpperCAmelCase__								=  1_6
         UpperCAmelCase__								=  2
         UpperCAmelCase__								=  0.02
         UpperCAmelCase__								=  3
         UpperCAmelCase__								=  4
         UpperCAmelCase__								=  None


   def       lowercase_				(self					:						int		)						->    Union[str, Any]:
         """simple docstring"""



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

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

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

         UpperCAmelCase__								=  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	,  )

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


   def       lowercase_				(self					:						int	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Tuple		)						->    Tuple:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertModel(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}

         UpperCAmelCase__								=  model(__UpperCAmelCase		)

         UpperCAmelCase__								=  [input_ids, input_mask]

         UpperCAmelCase__								=  model(__UpperCAmelCase		)

         self.parent.assertEqual(result.last_hidden_state.shape	,  (self.batch_size, self.seq_length, self.hidden_size)		)


   def       lowercase_				(self					:						str	,  __UpperCAmelCase					:						int	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Union[str, Any]	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Dict		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertForMaskedLM(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.seq_length, self.vocab_size)		)


   def       lowercase_				(self					:						int	,  __UpperCAmelCase					:						Dict	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						List[str]		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertForQuestionAnswering(config=__UpperCAmelCase		)
         UpperCAmelCase__								=  {
             "input_ids": input_ids,
             "attention_mask": input_mask,
         }
         UpperCAmelCase__								=  model(__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					:						Optional[int]	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						Dict	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						List[str]		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_labels
         UpperCAmelCase__								=  TFDistilBertForSequenceClassification(__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.num_labels)		)


   def       lowercase_				(self					:						Optional[int]	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						Optional[Any]	,  __UpperCAmelCase					:						Dict	,  __UpperCAmelCase					:						List[Any]	,  __UpperCAmelCase					:						str	,  __UpperCAmelCase					:						Any		)						->    int:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_choices
         UpperCAmelCase__								=  TFDistilBertForMultipleChoice(__UpperCAmelCase		)
         UpperCAmelCase__								=  tf.tile(tf.expand_dims(__UpperCAmelCase	,  1		)	,  (1, self.num_choices, 1)		)
         UpperCAmelCase__								=  tf.tile(tf.expand_dims(__UpperCAmelCase	,  1		)	,  (1, self.num_choices, 1)		)
         UpperCAmelCase__								=  {
             "input_ids": multiple_choice_inputs_ids,
             "attention_mask": multiple_choice_input_mask,
         }
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.num_choices)		)


   def       lowercase_				(self					:						Dict	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						Optional[int]	,  __UpperCAmelCase					:						Tuple	,  __UpperCAmelCase					:						List[str]	,  __UpperCAmelCase					:						Any	,  __UpperCAmelCase					:						Optional[int]		)						->    List[Any]:
         """simple docstring"""



         UpperCAmelCase__								=  self.num_labels
         UpperCAmelCase__								=  TFDistilBertForTokenClassification(__UpperCAmelCase		)
         UpperCAmelCase__								=  {"input_ids": input_ids, "attention_mask": input_mask}
         UpperCAmelCase__								=  model(__UpperCAmelCase		)
         self.parent.assertEqual(result.logits.shape	,  (self.batch_size, self.seq_length, self.num_labels)		)





   def       lowercase_				(self					:						Any		)						->    Dict:
         """simple docstring"""



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





@require_tf
class 					A							(    UpperCAmelCase_		,    UpperCAmelCase_		,    unittest.TestCase					):
   __UpperCAmelCase			:			Union[str, Any]              =      (
       (
           TFDistilBertModel,
           TFDistilBertForMaskedLM,
           TFDistilBertForQuestionAnswering,
           TFDistilBertForSequenceClassification,
           TFDistilBertForTokenClassification,
           TFDistilBertForMultipleChoice,
       )
       if is_tf_available()
       else None
   )
   __UpperCAmelCase			:			Optional[int]              =      (
       {
           'feature-extraction': TFDistilBertModel,
           'fill-mask': TFDistilBertForMaskedLM,
           'question-answering': TFDistilBertForQuestionAnswering,
           'text-classification': TFDistilBertForSequenceClassification,
           'token-classification': TFDistilBertForTokenClassification,
           'zero-shot': TFDistilBertForSequenceClassification,
       }
       if is_tf_available()
       else {}
   )
   __UpperCAmelCase			:			Tuple              =      False
   __UpperCAmelCase			:			str              =      False


   def       lowercase_				(self					:						Union[str, Any]		)						->    str:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertModelTester(self		)
         UpperCAmelCase__								=  ConfigTester(self	,  config_class=__UpperCAmelCase	,  dim=3_7		)


   def       lowercase_				(self					:						Any		)						->    int:
         """simple docstring"""



         self.config_tester.run_common_tests()


   def       lowercase_				(self					:						int		)						->    int:
         """simple docstring"""



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


   def       lowercase_				(self					:						Optional[int]		)						->    int:
         """simple docstring"""



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


   def       lowercase_				(self					:						Any		)						->    Optional[int]:
         """simple docstring"""



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


   def       lowercase_				(self					:						List[str]		)						->    str:
         """simple docstring"""



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


   def       lowercase_				(self					:						Dict		)						->    List[str]:
         """simple docstring"""



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


   def       lowercase_				(self					:						str		)						->    List[Any]:
         """simple docstring"""



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





   @slow
   def       lowercase_				(self					:						Optional[int]		)						->    Optional[int]:
         """simple docstring"""



         for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]		):
               UpperCAmelCase__								=  TFDistilBertModel.from_pretrained(__UpperCAmelCase		)
               self.assertIsNotNone(__UpperCAmelCase		)





@require_tf
class 					A							(    unittest.TestCase					):


   @slow
   def       lowercase_				(self					:						List[Any]		)						->    Dict:
         """simple docstring"""



         UpperCAmelCase__								=  TFDistilBertModel.from_pretrained("distilbert-base-uncased"		)
         UpperCAmelCase__								=  tf.constant([[0, 1, 2, 3, 4, 5]]		)
         UpperCAmelCase__								=  model(__UpperCAmelCase		)[0]

         UpperCAmelCase__								=  [1, 6, 7_6_8]
         self.assertEqual(output.shape	,  __UpperCAmelCase		)

         UpperCAmelCase__								=  tf.constant(
             [
                 [
                     [0.19261885, -0.13732955, 0.4119799],
                     [0.22150156, -0.07422661, 0.39037204],
                     [0.22756018, -0.0896414, 0.3701467],
                 ]
             ]		)
         tf.debugging.assert_near(output[:, :3, :3]	,  __UpperCAmelCase	,  atol=1E-4		)



