



import argparse
import glob
import logging
import os
import time
from argparse import Namespace

import numpy as np
import torch
from lightning_base import BaseTransformer, add_generic_args, generic_train
from torch.utils.data import DataLoader, TensorDataset

from transformers import glue_compute_metrics as compute_metrics
from transformers import glue_convert_examples_to_features as convert_examples_to_features
from transformers import glue_output_modes, glue_tasks_num_labels
from transformers import glue_processors as processors


a__    =	logging.getLogger(__name__)







class 		snake_case						(	__A  ):



   '''simple docstring'''






   snake_case_					: Union[str, Any]          =		"""sequence-classification"""

   def __init__(					self							:			Any				,		lowerCAmelCase							:			Optional[int])				->				List[Any]:

    """simple docstring"""







    if type(lowerCAmelCase) == dict:
     _snake_case    :  Any					=					Namespace(**lowerCAmelCase)
    _snake_case    :  int					=					glue_output_modes[hparams.task]
    _snake_case    :  Any					=					glue_tasks_num_labels[hparams.task]

    super().__init__(lowerCAmelCase				,		lowerCAmelCase				,		self.mode)

   def      UpperCamelCase_  (					self							:			Dict				,		**lowerCAmelCase							:			Any)				->				str:

    """simple docstring"""







    return self.model(**lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			Optional[int]				,		lowerCAmelCase							:			Union[str, Any]				,		lowerCAmelCase							:			Optional[Any])				->				str:

    """simple docstring"""







    _snake_case    :  Optional[int]					=					{'''input_ids''': batch[0], '''attention_mask''': batch[1], '''labels''': batch[3]}

    if self.config.model_type not in ["distilbert", "bart"]:
     _snake_case    :  str					=					batch[2] if self.config.model_type in ['''bert''', '''xlnet''', '''albert'''] else None

    _snake_case    :  Union[str, Any]					=					self(**lowerCAmelCase)
    _snake_case    :  List[Any]					=					outputs[0]

    _snake_case    :  str					=					self.trainer.lr_schedulers[0]['''scheduler''']
    _snake_case    :  Union[str, Any]					=					{'''loss''': loss, '''rate''': lr_scheduler.get_last_lr()[-1]}
    return {"loss": loss, "log": tensorboard_logs}

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

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					self.hparams
    _snake_case    :  str					=					processors[args.task]()
    _snake_case    :  Optional[int]					=					processor.get_labels()

    for mode in ["train", "dev"]:
     _snake_case    :  List[str]					=					self._feature_file(lowerCAmelCase)
     if os.path.exists(lowerCAmelCase) and not args.overwrite_cache:
      logger.info("""Loading features from cached file %s"""				,		lowerCAmelCase)
     else:
      logger.info("""Creating features from dataset file at %s"""				,		args.data_dir)
      _snake_case    :  Optional[int]					=					(
          processor.get_dev_examples(args.data_dir)
          if mode == '''dev'''
          else processor.get_train_examples(args.data_dir)
      )
      _snake_case    :  Optional[int]					=					convert_examples_to_features(
          lowerCAmelCase				,		self.tokenizer				,		max_length=args.max_seq_length				,		label_list=self.labels				,		output_mode=args.glue_output_mode				,		)
      logger.info("""Saving features into cached file %s"""				,		lowerCAmelCase)
      torch.save(lowerCAmelCase				,		lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			List[str]				,		lowerCAmelCase							:			str				,		lowerCAmelCase							:			int				,		lowerCAmelCase							:			bool = False)				->				DataLoader:

    """simple docstring"""







    _snake_case    :  Tuple					=					'''dev''' if mode == '''test''' else mode

    _snake_case    :  Any					=					self._feature_file(lowerCAmelCase)
    logger.info("""Loading features from cached file %s"""				,		lowerCAmelCase)
    _snake_case    :  Optional[Any]					=					torch.load(lowerCAmelCase)
    _snake_case    :  int					=					torch.tensor([f.input_ids for f in features]				,		dtype=torch.long)
    _snake_case    :  str					=					torch.tensor([f.attention_mask for f in features]				,		dtype=torch.long)
    _snake_case    :  Tuple					=					torch.tensor([f.token_type_ids for f in features]				,		dtype=torch.long)
    if self.hparams.glue_output_mode == "classification":
     _snake_case    :  int					=					torch.tensor([f.label for f in features]				,		dtype=torch.long)
    elif self.hparams.glue_output_mode == "regression":
     _snake_case    :  int					=					torch.tensor([f.label for f in features]				,		dtype=torch.float)

    return DataLoader(
        TensorDataset(lowerCAmelCase				,		lowerCAmelCase				,		lowerCAmelCase				,		lowerCAmelCase)				,		batch_size=lowerCAmelCase				,		shuffle=lowerCAmelCase				,		)

   def      UpperCamelCase_  (					self							:			int				,		lowerCAmelCase							:			List[str]				,		lowerCAmelCase							:			Dict)				->				Tuple:

    """simple docstring"""







    _snake_case    :  Any					=					{'''input_ids''': batch[0], '''attention_mask''': batch[1], '''labels''': batch[3]}

    if self.config.model_type not in ["distilbert", "bart"]:
     _snake_case    :  Any					=					batch[2] if self.config.model_type in ['''bert''', '''xlnet''', '''albert'''] else None

    _snake_case    :  Union[str, Any]					=					self(**lowerCAmelCase)
    _snake_case    :  List[Any]					=					outputs[:2]
    _snake_case    :  Union[str, Any]					=					logits.detach().cpu().numpy()
    _snake_case    :  int					=					inputs['''labels'''].detach().cpu().numpy()

    return {"val_loss": tmp_eval_loss.detach().cpu(), "pred": preds, "target": out_label_ids}

   def      UpperCamelCase_  (					self							:			List[str]				,		lowerCAmelCase							:			Dict)				->				tuple:

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					torch.stack([x["""val_loss"""] for x in outputs]).mean().detach().cpu().item()
    _snake_case    :  List[str]					=					np.concatenate([x["""pred"""] for x in outputs]				,		axis=0)

    if self.hparams.glue_output_mode == "classification":
     _snake_case    :  Optional[Any]					=					np.argmax(lowerCAmelCase				,		axis=1)
    elif self.hparams.glue_output_mode == "regression":
     _snake_case    :  int					=					np.squeeze(lowerCAmelCase)

    _snake_case    :  Any					=					np.concatenate([x["""target"""] for x in outputs]				,		axis=0)
    _snake_case    :  int					=					[[] for _ in range(out_label_ids.shape[0])]
    _snake_case    :  Tuple					=					[[] for _ in range(out_label_ids.shape[0])]

    _snake_case    :  Any					=					{**{'''val_loss''': val_loss_mean}, **compute_metrics(self.hparams.task				,		lowerCAmelCase				,		lowerCAmelCase)}

    _snake_case    :  Optional[Any]					=					dict(results.items())
    _snake_case    :  str					=					results
    return ret, preds_list, out_label_list

   def      UpperCamelCase_  (					self							:			Optional[Any]				,		lowerCAmelCase							:			list)				->				dict:

    """simple docstring"""







    _snake_case    :  Dict					=					self._eval_end(lowerCAmelCase)
    _snake_case    :  Any					=					ret['''log''']
    return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}

   def      UpperCamelCase_  (					self							:			int				,		lowerCAmelCase							:			Tuple)				->				dict:

    """simple docstring"""







    _snake_case    :  Union[str, Any]					=					self._eval_end(lowerCAmelCase)
    _snake_case    :  Optional[Any]					=					ret['''log''']
    # `val_loss` is the key returned by `self._eval_end()` but actually refers to `test_loss`
    return {"avg_test_loss": logs["val_loss"], "log": logs, "progress_bar": logs}


   @staticmethod
   def      UpperCamelCase_  (					lowerCAmelCase							:			str				,		lowerCAmelCase							:			Any)				->				Union[str, Any]:

    """simple docstring"""







    BaseTransformer.add_model_specific_args(lowerCAmelCase				,		lowerCAmelCase)
    parser.add_argument(
        """--max_seq_length"""				,		default=128				,		type=lowerCAmelCase				,		help=(
            """The maximum total input sequence length after tokenization. Sequences longer """
            """than this will be truncated, sequences shorter will be padded."""
        )				,		)

    parser.add_argument(
        """--task"""				,		default=""""""				,		type=lowerCAmelCase				,		required=lowerCAmelCase				,		help="""The GLUE task to run"""				,		)
    parser.add_argument(
        """--gpus"""				,		default=0				,		type=lowerCAmelCase				,		help="""The number of GPUs allocated for this, it is by default 0 meaning none"""				,		)

    parser.add_argument(
        """--overwrite_cache"""				,		action="""store_true"""				,		help="""Overwrite the cached training and evaluation sets""")

    return parser



def       lowercase       (   )  ->     Union[str, Any]:
 _snake_case    :  Tuple					=					argparse.ArgumentParser()
 add_generic_args(__A				,       os.getcwd()				)
 _snake_case    :  Any					=					GLUETransformer.add_model_specific_args(__A				,       os.getcwd()				)
 _snake_case    :  Optional[Any]					=					parser.parse_args()

 # If output_dir not provided, a folder will be generated in pwd
 if args.output_dir is None:
  _snake_case    :  Tuple					=					os.path.join(
      """./results"""				,       F'''{args.task}_{time.strftime('%Y%m%d_%H%M%S'				)}'''				,       )
  os.makedirs(args.output_dir				)

 _snake_case    :  str					=					GLUETransformer(__A				)
 _snake_case    :  Any					=					generic_train(__A				,       __A				)

 # Optionally, predict on dev set and write to output_dir
 if args.do_predict:
  _snake_case    :  Tuple					=					sorted(glob.glob(os.path.join(args.output_dir				,       """checkpoint-epoch=*.ckpt"""				)				,       recursive=__A				)				)
  _snake_case    :  int					=					model.load_from_checkpoint(checkpoints[-1]				)
  return trainer.test(__A				)


if __name__ == "__main__":
  main()





from __future__ import annotations

import unittest

from transformers import 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 numpy
  import tensorflow as tf

  from transformers import (
      TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
      TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST,
      TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST,
      BertConfig,
      DPRConfig,
      TFDPRContextEncoder,
      TFDPRQuestionEncoder,
      TFDPRReader,
  )







class 		snake_case						:



   '''simple docstring'''






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

    """simple docstring"""







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

   def      UpperCamelCase_  (					self							:			Dict)				->				List[str]:

    """simple docstring"""







    _snake_case    :  Optional[int]					=					ids_tensor([self.batch_size, self.seq_length]				,		self.vocab_size)

    _snake_case    :  List[str]					=					None
    if self.use_input_mask:
     # follow test_modeling_tf_ctrl.py
     _snake_case    :  List[Any]					=					random_attention_mask([self.batch_size, self.seq_length])

    _snake_case    :  str					=					None
    if self.use_token_type_ids:
     _snake_case    :  Optional[Any]					=					ids_tensor([self.batch_size, self.seq_length]				,		self.type_vocab_size)

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

    _snake_case    :  Union[str, Any]					=					BertConfig(
        vocab_size=self.vocab_size				,		hidden_size=self.hidden_size				,		num_hidden_layers=self.num_hidden_layers				,		num_attention_heads=self.num_attention_heads				,		intermediate_size=self.intermediate_size				,		hidden_act=self.hidden_act				,		hidden_dropout_prob=self.hidden_dropout_prob				,		attention_probs_dropout_prob=self.attention_probs_dropout_prob				,		max_position_embeddings=self.max_position_embeddings				,		type_vocab_size=self.type_vocab_size				,		is_decoder=lowerCAmelCase				,		initializer_range=self.initializer_range				,		)
    _snake_case    :  Optional[int]					=					DPRConfig(projection_dim=self.projection_dim				,		**config.to_dict())

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

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

    """simple docstring"""







    _snake_case    :  Dict					=					TFDPRContextEncoder(config=lowerCAmelCase)
    _snake_case    :  Union[str, Any]					=					model(lowerCAmelCase				,		attention_mask=lowerCAmelCase				,		token_type_ids=lowerCAmelCase)
    _snake_case    :  Optional[int]					=					model(lowerCAmelCase				,		token_type_ids=lowerCAmelCase)
    _snake_case    :  List[str]					=					model(lowerCAmelCase)
    self.parent.assertEqual(result.pooler_output.shape				,		(self.batch_size, self.projection_dim or self.hidden_size))

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

    """simple docstring"""







    _snake_case    :  Optional[int]					=					TFDPRQuestionEncoder(config=lowerCAmelCase)
    _snake_case    :  List[Any]					=					model(lowerCAmelCase				,		attention_mask=lowerCAmelCase				,		token_type_ids=lowerCAmelCase)
    _snake_case    :  List[Any]					=					model(lowerCAmelCase				,		token_type_ids=lowerCAmelCase)
    _snake_case    :  Tuple					=					model(lowerCAmelCase)
    self.parent.assertEqual(result.pooler_output.shape				,		(self.batch_size, self.projection_dim or self.hidden_size))

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

    """simple docstring"""







    _snake_case    :  Optional[Any]					=					TFDPRReader(config=lowerCAmelCase)
    _snake_case    :  Any					=					model(lowerCAmelCase				,		attention_mask=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))
    self.parent.assertEqual(result.relevance_logits.shape				,		(self.batch_size,))


   def      UpperCamelCase_  (					self							:			str)				->				Any:

    """simple docstring"""







    _snake_case    :  List[Any]					=					self.prepare_config_and_inputs()
    (
        (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    (
        _snake_case
    )	,    
    )    :  Union[str, Any]					=					config_and_inputs
    _snake_case    :  Union[str, Any]					=					{"""input_ids""": input_ids}
    return config, inputs_dict







@require_tf
class 		snake_case						(	SCREAMING_SNAKE_CASE_					,SCREAMING_SNAKE_CASE_					,unittest.TestCase  ):



   '''simple docstring'''






   snake_case_					: Any          =		(
       (
           TFDPRContextEncoder,
           TFDPRQuestionEncoder,
           TFDPRReader,
       )
       if is_tf_available()
       else ()
   )
   snake_case_					: List[str]          =		{"""feature-extraction""": TFDPRQuestionEncoder} if is_tf_available() else {}

   snake_case_					: Tuple          =		False
   snake_case_					: str          =		False
   snake_case_					: Tuple          =		False
   snake_case_					: int          =		False
   snake_case_					: List[Any]          =		False

   def      UpperCamelCase_  (					self							:			Tuple)				->				str:

    """simple docstring"""







    _snake_case    :  List[str]					=					TFDPRModelTester(self)
    _snake_case    :  Dict					=					ConfigTester(self				,		config_class=lowerCAmelCase				,		hidden_size=37)

   def      UpperCamelCase_  (					self							:			List[Any])				->				Optional[Any]:

    """simple docstring"""







    self.config_tester.run_common_tests()

   def      UpperCamelCase_  (					self							:			List[str])				->				Optional[int]:

    """simple docstring"""







    _snake_case    :  Any					=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_dpr_context_encoder(*lowerCAmelCase)

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

    """simple docstring"""







    _snake_case    :  str					=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_dpr_question_encoder(*lowerCAmelCase)

   def      UpperCamelCase_  (					self							:			List[str])				->				Tuple:

    """simple docstring"""







    _snake_case    :  str					=					self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_dpr_reader(*lowerCAmelCase)


   @slow
   def      UpperCamelCase_  (					self							:			List[str])				->				Any:

    """simple docstring"""







    for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  List[Any]					=					TFDPRContextEncoder.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)

    for model_name in TF_DPR_CONTEXT_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  List[str]					=					TFDPRContextEncoder.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)

    for model_name in TF_DPR_QUESTION_ENCODER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  Any					=					TFDPRQuestionEncoder.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)

    for model_name in TF_DPR_READER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
     _snake_case    :  Union[str, Any]					=					TFDPRReader.from_pretrained(lowerCAmelCase)
     self.assertIsNotNone(lowerCAmelCase)







@require_tf
class 		snake_case						(	unittest.TestCase  ):



   '''simple docstring'''






   @slow
   def      UpperCamelCase_  (					self							:			List[str])				->				Optional[int]:

    """simple docstring"""







    _snake_case    :  List[str]					=					TFDPRQuestionEncoder.from_pretrained("""facebook/dpr-question_encoder-single-nq-base""")

    _snake_case    :  str					=					tf.constant(
        [[101, 7592, 1010, 2003, 2026, 3899, 1_0140, 1029, 102]])  # [CLS] hello, is my dog cute? [SEP]
    _snake_case    :  Optional[int]					=					model(lowerCAmelCase)[0]  # embedding shape = (1, 768)
    # compare the actual values for a slice.
    _snake_case    :  Union[str, Any]					=					tf.constant(
        [
            [
                0.03_236_253,
                0.12_753_335,
                0.16_818_509,
                0.00_279_786,
                0.3_896_933,
                0.24_264_945,
                0.2_178_971,
                -0.02_335_227,
                -0.08_481_959,
                -0.14_324_117,
            ]
        ])
    self.assertTrue(numpy.allclose(output[:, :10].numpy()				,		expected_slice.numpy()				,		atol=1E-4))
