






from copy import deepcopy
from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import Dataset

from .deepspeed import is_deepspeed_zeroa_enabled
from .generation.configuration_utils import GenerationConfig
from .trainer import Trainer
from .utils import logging


if TYPE_CHECKING:
				from .data.data_collator import DataCollator
				from .modeling_utils import PreTrainedModel
				from .tokenization_utils_base import PreTrainedTokenizerBase
				from .trainer_callback import TrainerCallback
				from .trainer_utils import EvalPrediction, PredictionOutput
				from .training_args import TrainingArguments


__A          = logging.get_logger(__name__)





class   SCREAMING_SNAKE_CASE (       lowercase_    ):





							"""simple docstring"""

							def __init__(     self:      Tuple						,   __A:      Union["PreTrainedModel", nn.Module] = None						,   __A:      "TrainingArguments" = None						,   __A:      Optional["DataCollator"] = None						,   __A:      Optional[Dataset] = None						,   __A:      Optional[Union[Dataset, Dict[str, Dataset]]] = None						,   __A:      Optional["PreTrainedTokenizerBase"] = None						,   __A:      Optional[Callable[[], "PreTrainedModel"]] = None						,   __A:      Optional[Callable[["EvalPrediction"], Dict]] = None						,   __A:      Optional[List["TrainerCallback"]] = None						,   __A:      Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None)						,   __A:      Optional[Callable[[torch.Tensor, torch.Tensor], torch.Tensor]] = None						,   )			-> List[Any]:
												super().__init__(
												    model=_lowercase						,   args=_lowercase						,   data_collator=_lowercase						,   train_dataset=_lowercase						,   eval_dataset=_lowercase						,   tokenizer=_lowercase						,   model_init=_lowercase						,   compute_metrics=_lowercase						,   callbacks=_lowercase						,   optimizers=_lowercase						,   preprocess_logits_for_metrics=_lowercase						,   )

												# Override self.model.generation_config if a GenerationConfig is specified in args.
												# Priority: args.generation_config > model.generation_config > default GenerationConfig.
												if self.args.generation_config is not None:
																	_A									=       self.load_generation_config(self.args.generation_config      )
																	_A									=       gen_config

							@staticmethod
							def 					__A       (     __A:      Union[str, GenerationConfig]      )			-> Dict:





												if isinstance(_lowercase						,   _lowercase      ):
																	return deepcopy(_lowercase      )

												# str or Path
												_A									=       Path(_lowercase      ) if isinstance(_lowercase						,   _lowercase      ) else gen_config_arg
												_A									=       None

												# Figuring if it is path pointing to a file, pointing to a directory or else a model id or URL
												# This step is required in order to determine config_file_name
												if pretrained_model_name.is_file():
																	_A									=       pretrained_model_name.name
																	_A									=       pretrained_model_name.parent
												# dir path
												elif pretrained_model_name.is_dir():
																	pass
												# model id or URL
												else:
																	_A									=       gen_config_arg

												_A									=       GenerationConfig.from_pretrained(_lowercase						,   _lowercase      )
												return gen_config

							def 					__A       (     self:      int						,   __A:      Optional[Dataset] = None						,   __A:      Optional[List[str]] = None						,   __A:      str = "eval"						,   **__A:      Any						,   )			-> Dict:
												_A									=       gen_kwargs.copy()
												if gen_kwargs.get('''max_length'''      ) is None and gen_kwargs.get('''max_new_tokens'''      ) is None:
																	_A									=       self.args.generation_max_length
												_A									=       (
												    gen_kwargs['''num_beams'''] if gen_kwargs.get('''num_beams'''      ) is not None else self.args.generation_num_beams
												)
												_A									=       gen_kwargs

												return super().evaluate(_lowercase						,   ignore_keys=_lowercase						,   metric_key_prefix=_lowercase      )

							def 					__A       (     self:      int						,   __A:      Dataset						,   __A:      Optional[List[str]] = None						,   __A:      str = "test"						,   **__A:      int						,   )			-> Optional[int]:
												_A									=       gen_kwargs.copy()
												if gen_kwargs.get('''max_length'''      ) is None and gen_kwargs.get('''max_new_tokens'''      ) is None:
																	_A									=       self.args.generation_max_length
												_A									=       (
												    gen_kwargs['''num_beams'''] if gen_kwargs.get('''num_beams'''      ) is not None else self.args.generation_num_beams
												)
												_A									=       gen_kwargs

												return super().predict(_lowercase						,   ignore_keys=_lowercase						,   metric_key_prefix=_lowercase      )

							def 					__A       (     self:      Any						,   __A:      nn.Module						,   __A:      Dict[str, Union[torch.Tensor, Any]]						,   __A:      bool						,   __A:      Optional[List[str]] = None						,   **__A:      Tuple						,   )			-> Optional[Any]:





												if not self.args.predict_with_generate or prediction_loss_only:
																	return super().prediction_step(
																	    _lowercase						,   _lowercase						,   prediction_loss_only=_lowercase						,   ignore_keys=_lowercase      )

												_A									=       '''labels''' in inputs
												_A									=       self._prepare_inputs(_lowercase      )

												# XXX: adapt synced_gpus for fairscale as well
												# Priority (handled in generate):
												# gen_kwargs > model.generation_config > default GenerationConfig()

												if len(_lowercase      ) == 0 and hasattr(self						,   '''_gen_kwargs'''      ):
																	_A									=       self._gen_kwargs.copy()

												if gen_kwargs.get('''max_length'''      ) is None and gen_kwargs.get('''max_new_tokens'''      ) is None:
																	_A									=       self.model.config.max_length
												_A									=       (
												    gen_kwargs['''num_beams'''] if gen_kwargs.get('''num_beams'''      ) is not None else self.model.config.num_beams
												)
												_A									=       True if is_deepspeed_zeroa_enabled() else False
												_A									=       (
												    gen_kwargs['''synced_gpus'''] if gen_kwargs.get('''synced_gpus'''      ) is not None else default_synced_gpus
												)

												# If the `decoder_input_ids` was created from `labels`, evict the former, so that the model can freely generate
												# (otherwise, it would continue generating from the padded `decoder_input_ids`)
												if (
												    "labels" in inputs
												    and "decoder_input_ids" in inputs
												    and inputs["labels"].shape == inputs["decoder_input_ids"].shape
												):
																	_A									=       {k: v for k, v in inputs.items() if k != '''decoder_input_ids'''}
												_A									=       self.model.generate(**_lowercase						,   **_lowercase      )

												# Temporary hack to ensure the generation config is not initialized for each iteration of the evaluation loop
												# TODO: remove this hack when the legacy code that initializes generation_config from a model config is
												# removed in https://github.com/huggingface/transformers/blob/98d88b23f54e5a23e741833f1e973fdf600cc2c5/src/transformers/generation/utils.py#L1183
												if self.model.generation_config._from_model_config:
																	_A									=       False

												# Retrieves GenerationConfig from model.generation_config
												_A									=       self.model.generation_config
												# in case the batch is shorter than max length, the output should be padded
												if generated_tokens.shape[-1] < gen_config.max_length:
																	_A									=       self._pad_tensors_to_max_len(_lowercase						,   gen_config.max_length      )
												elif gen_config.max_new_tokens is not None and generated_tokens.shape[-1] < gen_config.max_new_tokens + 1:
																	_A									=       self._pad_tensors_to_max_len(_lowercase						,   gen_config.max_new_tokens + 1      )

												with torch.no_grad():
																	if has_labels:
																						with self.compute_loss_context_manager():
																											_A									=       model(**_lowercase      )
																						if self.label_smoother is not None:
																											_A									=       self.label_smoother(_lowercase						,   inputs['''labels''']      ).mean().detach()
																						else:
																											_A									=       (outputs['''loss'''] if isinstance(_lowercase						,   _lowercase      ) else outputs[0]).mean().detach()
																	else:
																						_A									=       None

												if self.args.prediction_loss_only:
																	return loss, None, None

												if has_labels:
																	_A									=       inputs['''labels''']
																	if labels.shape[-1] < gen_config.max_length:
																						_A									=       self._pad_tensors_to_max_len(_lowercase						,   gen_config.max_length      )
																	elif gen_config.max_new_tokens is not None and labels.shape[-1] < gen_config.max_new_tokens + 1:
																						_A									=       self._pad_tensors_to_max_len(_lowercase						,   gen_config.max_new_tokens + 1      )
												else:
																	_A									=       None

												return loss, generated_tokens, labels



							def 					__A       (     self:      Dict						,   __A:      str						,   __A:      List[str]      )			-> List[Any]:





												if self.tokenizer is not None and hasattr(self.tokenizer						,   '''pad_token_id'''      ):
																	# If PAD token is not defined at least EOS token has to be defined
																	_A									=       (
																	    self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else self.tokenizer.eos_token_id
																	)
												else:
																	if self.model.config.pad_token_id is not None:
																						_A									=       self.model.config.pad_token_id
																	else:
																						raise ValueError('''Pad_token_id must be set in the configuration of the model, in order to pad tensors'''      )

												_A									=       pad_token_id * torch.ones(
												    (tensor.shape[0], max_length)						,   dtype=tensor.dtype						,   device=tensor.device      )
												_A									=       tensor
												return padded_tensor










import argparse
import os
import sys
import urllib.request
import zipfile


_snake_case				= ["""CoLA""", """SST""", """MRPC""", """QQP""", """STS""", """MNLI""", """SNLI""", """QNLI""", """RTE""", """WNLI""", """diagnostic"""]
_snake_case				= {
    """CoLA""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FCoLA.zip?alt=media&token=46d5e637-3411-4188-bc44-5809b5bfb5f4""",
    """SST""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FSST-2.zip?alt=media&token=aabc5f6b-e466-44a2-b9b4-cf6337f84ac8""",
    """MRPC""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2Fmrpc_dev_ids.tsv?alt=media&token=ec5c0836-31d5-48f4-b431-7480817f1adc""",
    """QQP""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FQQP.zip?alt=media&token=700c6acf-160d-4d89-81d1-de4191d02cb5""",
    """STS""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FSTS-B.zip?alt=media&token=bddb94a7-8706-4e0d-a694-1109e12273b5""",
    """MNLI""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FMNLI.zip?alt=media&token=50329ea1-e339-40e2-809c-10c40afff3ce""",
    """SNLI""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FSNLI.zip?alt=media&token=4afcfbb2-ff0c-4b2d-a09a-dbf07926f4df""",
    """QNLI""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FQNLIv2.zip?alt=media&token=6fdcf570-0fc5-4631-8456-9505272d1601""",
    """RTE""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FRTE.zip?alt=media&token=5efa7e85-a0bb-4f19-8ea2-9e1840f077fb""",
    """WNLI""": """https://firebasestorage.googleapis.com/v0/b/mtl-sentence-representations.appspot.com/o/data%2FWNLI.zip?alt=media&token=068ad0a0-ded7-4bd7-99a5-5e00222e0faf""",
    """diagnostic""": """https://storage.googleapis.com/mtl-sentence-representations.appspot.com/tsvsWithoutLabels%2FAX.tsv?GoogleAccessId=firebase-adminsdk-0khhl@mtl-sentence-representations.iam.gserviceaccount.com&Expires=2498860800&Signature=DuQ2CSPt2Yfre0C%2BiISrVYrIFaZH1Lc7hBVZDD4ZyR7fZYOMNOUGpi8QxBmTNOrNPjR3z1cggo7WXFfrgECP6FBJSsURv8Ybrue8Ypt%2FTPxbuJ0Xc2FhDi%2BarnecCBFO77RSbfuz%2Bs95hRrYhTnByqu3U%2FYZPaj3tZt5QdfpH2IUROY8LiBXoXS46LE%2FgOQc%2FKN%2BA9SoscRDYsnxHfG0IjXGwHN%2Bf88q6hOmAxeNPx6moDulUF6XMUAaXCSFU%2BnRO2RDL9CapWxj%2BDl7syNyHhB7987hZ80B%2FwFkQ3MEs8auvt5XW1%2Bd4aCU7ytgM69r8JDCwibfhZxpaa4gd50QXQ%3D%3D""",
}

_snake_case				= """https://dl.fbaipublicfiles.com/senteval/senteval_data/msr_paraphrase_train.txt"""
_snake_case				= """https://dl.fbaipublicfiles.com/senteval/senteval_data/msr_paraphrase_test.txt"""


def _A   (    __magic_name__  ,					__magic_name__       ):
   print(f'''Downloading and extracting {task}...'''       )
   lowercase__    =  f'''{task}.zip'''
   urllib.request.urlretrieve(TASKaPATH[task]  ,					__magic_name__       )
   with zipfile.ZipFile(__magic_name__       ) as zip_ref:
      zip_ref.extractall(__magic_name__       )
   os.remove(__magic_name__       )
   print("\tCompleted!"       )


def _A   (    __magic_name__  ,					__magic_name__       ):
   print("Processing MRPC..."       )
   lowercase__    =  os.path.join(__magic_name__  ,					"MRPC"       )
   if not os.path.isdir(__magic_name__       ):
      os.mkdir(__magic_name__       )
   if path_to_data:
      lowercase__    =  os.path.join(__magic_name__  ,					"msr_paraphrase_train.txt"       )
      lowercase__    =  os.path.join(__magic_name__  ,					"msr_paraphrase_test.txt"       )
   else:
      print("Local MRPC data not specified, downloading data from %s" % MRPC_TRAIN       )
      lowercase__    =  os.path.join(__magic_name__  ,					"msr_paraphrase_train.txt"       )
      lowercase__    =  os.path.join(__magic_name__  ,					"msr_paraphrase_test.txt"       )
      urllib.request.urlretrieve(__magic_name__  ,					__magic_name__       )
      urllib.request.urlretrieve(__magic_name__  ,					__magic_name__       )
   if not os.path.isfile(__magic_name__       ):
      raise ValueError(f'''Train data not found at {mrpc_train_file}'''       )
   if not os.path.isfile(__magic_name__       ):
      raise ValueError(f'''Test data not found at {mrpc_test_file}'''       )
   urllib.request.urlretrieve(TASKaPATH["MRPC"]  ,					os.path.join(__magic_name__  ,					"dev_ids.tsv"       )       )

   lowercase__    =  []
   with open(os.path.join(__magic_name__  ,					"dev_ids.tsv"       )  ,					encoding="utf8"       ) as ids_fh:
      for row in ids_fh:
         dev_ids.append(row.strip().split("\t"       )       )

   with open(__magic_name__  ,					encoding="utf8"       ) as data_fh, open(
       os.path.join(__magic_name__  ,					"train.tsv"       )  ,					"w"  ,					encoding="utf8"       ) as train_fh, open(os.path.join(__magic_name__  ,					"dev.tsv"       )  ,					"w"  ,					encoding="utf8"       ) as dev_fh:
      lowercase__    =  data_fh.readline()
      train_fh.write(__magic_name__       )
      dev_fh.write(__magic_name__       )
      for row in data_fh:
         lowercase__   ,							lowercase__   ,							lowercase__   ,							lowercase__   ,							lowercase__    =  row.strip().split("\t"       )
         if [ida, ida] in dev_ids:
            dev_fh.write("%s\t%s\t%s\t%s\t%s\n" % (label, ida, ida, sa, sa)       )
         else:
            train_fh.write("%s\t%s\t%s\t%s\t%s\n" % (label, ida, ida, sa, sa)       )

   with open(__magic_name__  ,					encoding="utf8"       ) as data_fh, open(
       os.path.join(__magic_name__  ,					"test.tsv"       )  ,					"w"  ,					encoding="utf8"       ) as test_fh:
      lowercase__    =  data_fh.readline()
      test_fh.write("index\t#1 ID\t#2 ID\t#1 String\t#2 String\n"       )
      for idx, row in enumerate(__magic_name__       ):
         lowercase__   ,							lowercase__   ,							lowercase__   ,							lowercase__   ,							lowercase__    =  row.strip().split("\t"       )
         test_fh.write("%d\t%s\t%s\t%s\t%s\n" % (idx, ida, ida, sa, sa)       )
   print("\tCompleted!"       )


def _A   (    __magic_name__       ):
   print("Downloading and extracting diagnostic..."       )
   if not os.path.isdir(os.path.join(__magic_name__  ,					"diagnostic"       )       ):
      os.mkdir(os.path.join(__magic_name__  ,					"diagnostic"       )       )
   lowercase__    =  os.path.join(__magic_name__  ,					"diagnostic"  ,					"diagnostic.tsv"       )
   urllib.request.urlretrieve(TASKaPATH["diagnostic"]  ,					__magic_name__       )
   print("\tCompleted!"       )
   return


def _A   (    __magic_name__       ):
   lowercase__    =  task_names.split(","       )
   if "all" in task_names:
      lowercase__    =  TASKS
   else:
      lowercase__    =  []
      for task_name in task_names:
         if task_name not in TASKS:
            raise ValueError(f'''Task {task_name} not found!'''       )
         tasks.append(__magic_name__       )
   return tasks


def _A   (    __magic_name__       ):
   lowercase__    =  argparse.ArgumentParser()
   parser.add_argument("--data_dir"  ,					help="directory to save data to"  ,					type=__magic_name__  ,					default="glue_data"       )
   parser.add_argument(
       "--tasks"  ,					help="tasks to download data for as a comma separated string"  ,					type=__magic_name__  ,					default="all"       )
   parser.add_argument(
       "--path_to_mrpc"  ,					help="path to directory containing extracted MRPC data, msr_paraphrase_train.txt and msr_paraphrase_text.txt"  ,					type=__magic_name__  ,					default=""  ,					)
   lowercase__    =  parser.parse_args(__magic_name__       )

   if not os.path.isdir(args.data_dir       ):
      os.mkdir(args.data_dir       )
   lowercase__    =  get_tasks(args.tasks       )

   for task in tasks:
      if task == "MRPC":
         format_mrpc(args.data_dir  ,					args.path_to_mrpc       )
      elif task == "diagnostic":
         download_diagnostic(args.data_dir       )
      else:
         download_and_extract(__magic_name__  ,					args.data_dir       )


if __name__ == "__main__":
  sys.exit(main(sys.argv[1:]))






