


"""simple docstring"""


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     :   Union[str, Any]			    =			logging.getLogger(__name__)
class   _UpperCamelCase  (   lowerCAmelCase__		):
							'''simple docstring'''







							__UpperCAmelCase						:    Optional[Any]				      ="""sequence-classification"""

							def __init__(     self    ,   __a	):
										if type(__a	) == dict:
													__lowerCAmelCase				    =			Namespace(**__a	)
										__lowerCAmelCase				    =			glue_output_modes[hparams.task]
										__lowerCAmelCase				    =			glue_tasks_num_labels[hparams.task]

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

							def 			snake_case				(     self    ,   **__a	):
										return self.model(**__a	)

							def 			snake_case				(     self    ,   __a    ,   __a	):
										__lowerCAmelCase				    =			{"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}

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

										__lowerCAmelCase				    =			self(**__a	)
										__lowerCAmelCase				    =			outputs[0]

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

							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			self.hparams
										__lowerCAmelCase				    =			processors[args.task]()
										__lowerCAmelCase				    =			processor.get_labels()

										for mode in ["train", "dev"]:
													__lowerCAmelCase				    =			self._feature_file(__a	)
													if os.path.exists(__a	) and not args.overwrite_cache:
																logger.info("Loading features from cached file %s"    ,   __a	)
													else:
																logger.info("Creating features from dataset file at %s"    ,   args.data_dir	)
																__lowerCAmelCase				    =			(
																    processor.get_dev_examples(args.data_dir	)
																    if mode == "dev"
																    else processor.get_train_examples(args.data_dir	)
																)
																__lowerCAmelCase				    =			convert_examples_to_features(
																    __a    ,   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"    ,   __a	)
																torch.save(__a    ,   __a	)

							def 			snake_case				(     self    ,   __a    ,   __a    ,   __a = False	):
										__lowerCAmelCase				    =			"dev" if mode == "test" else mode

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

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

							def 			snake_case				(     self    ,   __a    ,   __a	):
										__lowerCAmelCase				    =			{"input_ids": batch[0], "attention_mask": batch[1], "labels": batch[3]}

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

										__lowerCAmelCase				    =			self(**__a	)
										__lowerCAmelCase    ,	__lowerCAmelCase				    =			outputs[:2]
										__lowerCAmelCase				    =			logits.detach().cpu().numpy()
										__lowerCAmelCase				    =			inputs["labels"].detach().cpu().numpy()

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

							def 			snake_case				(     self    ,   __a	):
										__lowerCAmelCase				    =			torch.stack([x["val_loss"] for x in outputs]	).mean().detach().cpu().item()
										__lowerCAmelCase				    =			np.concatenate([x["pred"] for x in outputs]    ,   axis=0	)

										if self.hparams.glue_output_mode == "classification":
													__lowerCAmelCase				    =			np.argmax(__a    ,   axis=1	)
										elif self.hparams.glue_output_mode == "regression":
													__lowerCAmelCase				    =			np.squeeze(__a	)

										__lowerCAmelCase				    =			np.concatenate([x["target"] for x in outputs]    ,   axis=0	)
										__lowerCAmelCase				    =			[[] for _ in range(out_label_ids.shape[0]	)]
										__lowerCAmelCase				    =			[[] for _ in range(out_label_ids.shape[0]	)]

										__lowerCAmelCase				    =			{**{"val_loss": val_loss_mean}, **compute_metrics(self.hparams.task    ,   __a    ,   __a	)}

										__lowerCAmelCase				    =			dict(results.items()	)
										__lowerCAmelCase				    =			results
										return ret, preds_list, out_label_list

							def 			snake_case				(     self    ,   __a	):
										__lowerCAmelCase    ,	__lowerCAmelCase    ,	__lowerCAmelCase				    =			self._eval_end(__a	)
										__lowerCAmelCase				    =			ret["log"]
										return {"val_loss": logs["val_loss"], "log": logs, "progress_bar": logs}

							def 			snake_case				(     self    ,   __a	):
										__lowerCAmelCase    ,	__lowerCAmelCase    ,	__lowerCAmelCase				    =			self._eval_end(__a	)
										__lowerCAmelCase				    =			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 			snake_case				(     __a    ,   __a	):
										BaseTransformer.add_model_specific_args(__a    ,   __a	)
										parser.add_argument(
										    "--max_seq_length"    ,   default=1_28    ,   type=__a    ,   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=__a    ,   required=__a    ,   help="The GLUE task to run"    ,   )
										parser.add_argument(
										    "--gpus"    ,   default=0    ,   type=__a    ,   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 _lowerCamelCase	(		):

			'''simple docstring'''




			__lowerCAmelCase				    =			argparse.ArgumentParser()
			add_generic_args(_UpperCamelCase    , os.getcwd()      )
			__lowerCAmelCase				    =			GLUETransformer.add_model_specific_args(_UpperCamelCase    , os.getcwd()      )
			__lowerCAmelCase				    =			parser.parse_args()

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

			__lowerCAmelCase				    =			GLUETransformer(_UpperCamelCase      )
			__lowerCAmelCase				    =			generic_train(_UpperCamelCase    , _UpperCamelCase      )

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


if __name__ == "__main__":
	main()








"""simple docstring"""


import unittest

from transformers import (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TextaTextGenerationPipeline,
    pipeline,
)
from transformers.testing_utils import is_pipeline_test, require_tf, require_torch
from transformers.utils import is_torch_available

from .test_pipelines_common import ANY


if is_torch_available():
	import torch
@is_pipeline_test
class   _UpperCamelCase  (   unittest.TestCase		):
							'''simple docstring'''







							__UpperCAmelCase						:    Union[str, Any]				      =MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
							__UpperCAmelCase						:    Union[str, Any]				      =TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

							def 			snake_case				(     self    ,   __a    ,   __a    ,   __a	):
										__lowerCAmelCase				    =			TextaTextGenerationPipeline(model=__a    ,   tokenizer=__a	)
										return generator, ["Something to write", "Something else"]

							def 			snake_case				(     self    ,   __a    ,   __a	):
										__lowerCAmelCase				    =			generator("Something there"	)
										self.assertEqual(__a    ,   [{"generated_text": ANY(__a	)}]	)
										# These are encoder decoder, they don't just append to incoming string
										self.assertFalse(outputs[0]["generated_text"].startswith("Something there"	)	)

										__lowerCAmelCase				    =			generator(["This is great !", "Something else"]    ,   num_return_sequences=2    ,   do_sample=__a	)
										self.assertEqual(
										    __a    ,   [
										        [{"generated_text": ANY(__a	)}, {"generated_text": ANY(__a	)}],
										        [{"generated_text": ANY(__a	)}, {"generated_text": ANY(__a	)}],
										    ]    ,   )

										__lowerCAmelCase				    =			generator(
										    ["This is great !", "Something else"]    ,   num_return_sequences=2    ,   batch_size=2    ,   do_sample=__a	)
										self.assertEqual(
										    __a    ,   [
										        [{"generated_text": ANY(__a	)}, {"generated_text": ANY(__a	)}],
										        [{"generated_text": ANY(__a	)}, {"generated_text": ANY(__a	)}],
										    ]    ,   )

										with self.assertRaises(__a	):
													generator(4	)

							@require_torch
							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			pipeline("text2text-generation"    ,   model="patrickvonplaten/t5-tiny-random"    ,   framework="pt"	)
										# do_sample=False necessary for reproducibility
										__lowerCAmelCase				    =			generator("Something there"    ,   do_sample=__a	)
										self.assertEqual(__a    ,   [{"generated_text": ""}]	)

										__lowerCAmelCase				    =			3
										__lowerCAmelCase				    =			generator(
										    "Something there"    ,   num_return_sequences=__a    ,   num_beams=__a    ,   )
										__lowerCAmelCase				    =			[
										    {"generated_text": "Beide Beide Beide Beide Beide Beide Beide Beide Beide"},
										    {"generated_text": "Beide Beide Beide Beide Beide Beide Beide Beide"},
										    {"generated_text": ""},
										]
										self.assertEqual(__a    ,   __a	)

										__lowerCAmelCase				    =			generator("This is a test"    ,   do_sample=__a    ,   num_return_sequences=2    ,   return_tensors=__a	)
										self.assertEqual(
										    __a    ,   [
										        {"generated_token_ids": ANY(torch.Tensor	)},
										        {"generated_token_ids": ANY(torch.Tensor	)},
										    ]    ,   )
										__lowerCAmelCase				    =			generator.model.config.eos_token_id
										__lowerCAmelCase				    =			"<pad>"
										__lowerCAmelCase				    =			generator(
										    ["This is a test", "This is a second test"]    ,   do_sample=__a    ,   num_return_sequences=2    ,   batch_size=2    ,   return_tensors=__a    ,   )
										self.assertEqual(
										    __a    ,   [
										        [
										            {"generated_token_ids": ANY(torch.Tensor	)},
										            {"generated_token_ids": ANY(torch.Tensor	)},
										        ],
										        [
										            {"generated_token_ids": ANY(torch.Tensor	)},
										            {"generated_token_ids": ANY(torch.Tensor	)},
										        ],
										    ]    ,   )



							@require_tf
							def 			snake_case				(     self	):
										__lowerCAmelCase				    =			pipeline("text2text-generation"    ,   model="patrickvonplaten/t5-tiny-random"    ,   framework="tf"	)
										# do_sample=False necessary for reproducibility
										__lowerCAmelCase				    =			generator("Something there"    ,   do_sample=__a	)
										self.assertEqual(__a    ,   [{"generated_text": ""}]	)




