





'''simple docstring'''





from typing import Optional

import numpy as np
import torch
from torch import nn
from transformers import GPTaConfig, GPTaLMHeadModel
from transformers.modeling_utils import ModuleUtilsMixin

from ...configuration_utils import ConfigMixin, register_to_config
from ...models import ModelMixin






class       __A					(					__UpperCAmelCase    ,       __UpperCAmelCase    ,       __UpperCAmelCase				):

			'''simple docstring'''







			lowerCAmelCase_							  =		[R"""h\.\d+\.attn\.bias""", R"""h\.\d+\.attn\.masked_bias"""]





			@register_to_config
			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = None ,						__lowerCAmelCase = 5_0_2_5_7 ,						__lowerCAmelCase = 1_0_2_4 ,						__lowerCAmelCase = 7_6_8 ,						__lowerCAmelCase = 1_2 ,						__lowerCAmelCase = 1_2 ,						__lowerCAmelCase = None ,						__lowerCAmelCase = "gelu_new" ,						__lowerCAmelCase = 0.1 ,						__lowerCAmelCase = 0.1 ,						__lowerCAmelCase = 0.1 ,						__lowerCAmelCase = 1E-5 ,						__lowerCAmelCase = 0.02 ,						__lowerCAmelCase = True ,						__lowerCAmelCase = True ,						__lowerCAmelCase = False ,						__lowerCAmelCase = False ,						):



				'''simple docstring'''







				super().__init__()

				lowerCamelCase__              =	prefix_length

				if prefix_inner_dim != n_embd and prefix_hidden_dim is None:
					raise ValueError(
					    F'`prefix_hidden_dim` cannot be `None` when `prefix_inner_dim`: {prefix_hidden_dim} and'
					    F' `n_embd`: {n_embd} are not equal.'							)

				lowerCamelCase__              =	prefix_inner_dim
				lowerCamelCase__              =	prefix_hidden_dim

				lowerCamelCase__              =	(
				    nn.Linear(self.prefix_inner_dim ,						self.prefix_hidden_dim							)
				    if self.prefix_hidden_dim is not None
				    else nn.Identity()
				)
				lowerCamelCase__              =	(
				    nn.Linear(self.prefix_hidden_dim ,						lowerCAmelCase_							) if self.prefix_hidden_dim is not None else nn.Identity()
				)

				lowerCamelCase__              =	GPTaConfig(
				    vocab_size=lowerCAmelCase_ ,						n_positions=lowerCAmelCase_ ,						n_embd=lowerCAmelCase_ ,						n_layer=lowerCAmelCase_ ,						n_head=lowerCAmelCase_ ,						n_inner=lowerCAmelCase_ ,						activation_function=lowerCAmelCase_ ,						resid_pdrop=lowerCAmelCase_ ,						embd_pdrop=lowerCAmelCase_ ,						attn_pdrop=lowerCAmelCase_ ,						layer_norm_epsilon=lowerCAmelCase_ ,						initializer_range=lowerCAmelCase_ ,						scale_attn_weights=lowerCAmelCase_ ,						use_cache=lowerCAmelCase_ ,						scale_attn_by_inverse_layer_idx=lowerCAmelCase_ ,						reorder_and_upcast_attn=lowerCAmelCase_ ,						)
				lowerCamelCase__              =	GPTaLMHeadModel(lowerCAmelCase_							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None ,						):



				'''simple docstring'''







				lowerCamelCase__              =	self.transformer.transformer.wte(lowerCAmelCase_							)
				lowerCamelCase__              =	self.encode_prefix(lowerCAmelCase_							)
				lowerCamelCase__              =	self.decode_prefix(lowerCAmelCase_							)
				lowerCamelCase__              =	torch.cat((prefix_embeds, embedding_text) ,						dim=1							)

				if labels is not None:
					lowerCamelCase__              =	self.get_dummy_token(input_ids.shape[0] ,						input_ids.device							)
					lowerCamelCase__              =	torch.cat((dummy_token, input_ids) ,						dim=1							)
				lowerCamelCase__              =	self.transformer(inputs_embeds=lowerCAmelCase_ ,						labels=lowerCAmelCase_ ,						attention_mask=lowerCAmelCase_							)
				if self.prefix_hidden_dim is not None:
					return out, hidden
				else:
					return out





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase							):



				'''simple docstring'''







				return torch.zeros(lowerCAmelCase_ ,						self.prefix_length ,						dtype=torch.intaa ,						device=lowerCAmelCase_							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				return self.encode_prefix(lowerCAmelCase_							)





			@torch.no_grad()
			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	torch.split(lowerCAmelCase_ ,						1 ,						dim=0							)
				lowerCamelCase__              =	[]
				lowerCamelCase__              =	[]
				for feature in features:
					lowerCamelCase__              =	self.decode_prefix(feature.to(lowerCAmelCase_							)							)  # back to the clip feature
					# Only support beam search for now
					lowerCamelCase__       ,			lowerCamelCase__              =	self.generate_beam(
					    input_embeds=lowerCAmelCase_ ,						device=lowerCAmelCase_ ,						eos_token_id=lowerCAmelCase_							)
					generated_tokens.append(output_tokens[0]							)
					generated_seq_lengths.append(seq_lengths[0]							)
				lowerCamelCase__              =	torch.stack(lowerCAmelCase_							)
				lowerCamelCase__              =	torch.stack(lowerCAmelCase_							)
				return generated_tokens, generated_seq_lengths





			@torch.no_grad()
			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase=None ,						__lowerCAmelCase=None ,						__lowerCAmelCase=None ,						__lowerCAmelCase = 5 ,						__lowerCAmelCase = 6_7 ,						__lowerCAmelCase = 1.0 ,						__lowerCAmelCase = None ,						):



				'''simple docstring'''







				lowerCamelCase__              =	eos_token_id
				lowerCamelCase__              =	None
				lowerCamelCase__              =	None
				lowerCamelCase__              =	torch.ones(lowerCAmelCase_ ,						device=lowerCAmelCase_ ,						dtype=torch.int							)
				lowerCamelCase__              =	torch.zeros(lowerCAmelCase_ ,						device=lowerCAmelCase_ ,						dtype=torch.bool							)

				if input_embeds is not None:
					lowerCamelCase__              =	input_embeds
				else:
					lowerCamelCase__              =	self.transformer.transformer.wte(lowerCAmelCase_							)

				for i in range(lowerCAmelCase_							):
					lowerCamelCase__              =	self.transformer(inputs_embeds=lowerCAmelCase_							)
					lowerCamelCase__              =	outputs.logits
					lowerCamelCase__              =	logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
					lowerCamelCase__              =	logits.softmax(-1							).log()

					if scores is None:
						lowerCamelCase__       ,			lowerCamelCase__              =	logits.topk(lowerCAmelCase_ ,						-1							)
						lowerCamelCase__              =	generated.expand(lowerCAmelCase_ ,						*generated.shape[1:]							)
						lowerCamelCase__       ,			lowerCamelCase__              =	next_tokens.permute(1 ,						0							), scores.squeeze(0							)
						if tokens is None:
							lowerCamelCase__              =	next_tokens
						else:
							lowerCamelCase__              =	tokens.expand(lowerCAmelCase_ ,						*tokens.shape[1:]							)
							lowerCamelCase__              =	torch.cat((tokens, next_tokens) ,						dim=1							)
					else:
						lowerCamelCase__              =	-float(np.inf							)
						lowerCamelCase__              =	0
						lowerCamelCase__              =	scores[:, None] + logits
						seq_lengths[~is_stopped] += 1
						lowerCamelCase__              =	scores_sum / seq_lengths[:, None]
						lowerCamelCase__       ,			lowerCamelCase__              =	scores_sum_average.view(-1							).topk(lowerCAmelCase_ ,						-1							)
						lowerCamelCase__              =	next_tokens // scores_sum.shape[1]
						lowerCamelCase__              =	seq_lengths[next_tokens_source]
						lowerCamelCase__              =	next_tokens % scores_sum.shape[1]
						lowerCamelCase__              =	next_tokens.unsqueeze(1							)
						lowerCamelCase__              =	tokens[next_tokens_source]
						lowerCamelCase__              =	torch.cat((tokens, next_tokens) ,						dim=1							)
						lowerCamelCase__              =	generated[next_tokens_source]
						lowerCamelCase__              =	scores_sum_average * seq_lengths
						lowerCamelCase__              =	is_stopped[next_tokens_source]

					lowerCamelCase__              =	self.transformer.transformer.wte(next_tokens.squeeze()							).view(generated.shape[0] ,						1 ,						-1							)
					lowerCamelCase__              =	torch.cat((generated, next_token_embed) ,						dim=1							)
					lowerCamelCase__              =	is_stopped + next_tokens.eq(lowerCAmelCase_							).squeeze()
					if is_stopped.all():
						break

				lowerCamelCase__              =	scores / seq_lengths
				lowerCamelCase__              =	scores.argsort(descending=lowerCAmelCase_							)
				# tokens tensors are already padded to max_seq_length
				lowerCamelCase__              =	[tokens[i] for i in order]
				lowerCamelCase__              =	torch.stack(lowerCAmelCase_ ,						dim=0							)
				lowerCamelCase__              =	torch.tensor([seq_lengths[i] for i in order] ,						dtype=seq_lengths.dtype							)
				return output_texts, seq_lengths










import os
import time
import warnings
from dataclasses import dataclass, field
from enum import Enum
from typing import List, Optional, Union

import torch
from filelock import FileLock
from torch.utils.data import Dataset

from ...tokenization_utils_base import PreTrainedTokenizerBase
from ...utils import logging
from ..processors.glue import glue_convert_examples_to_features, glue_output_modes, glue_processors
from ..processors.utils import InputFeatures


_a						=   logging.get_logger(__name__)






@dataclass
class       __A					:

			'''simple docstring'''







			lowerCAmelCase_							  =		field(metadata={"""help""": """The name of the task to train on: """ + """, """.join(glue_processors.keys()				)}				)
			lowerCAmelCase_							  =		field(
			    metadata={"""help""": """The input data dir. Should contain the .tsv files (or other data files) for the task."""}				)
			lowerCAmelCase_							  =		field(
			    default=128    ,       metadata={
			        """help""": (
			            """The maximum total input sequence length after tokenization. Sequences longer """
			            """than this will be truncated, sequences shorter will be padded."""
			        )
			    }    ,       )
			lowerCAmelCase_							  =		field(
			    default=lowerCAmelCase    ,       metadata={"""help""": """Overwrite the cached training and evaluation sets"""}				)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.task_name.lower()






class       __A					(					lowerCAmelCase				):

			'''simple docstring'''







			lowerCAmelCase_							  =		"""train"""
			lowerCAmelCase_							  =		"""dev"""
			lowerCAmelCase_							  =		"""test"""






class       __A					(					lowerCAmelCase				):

			'''simple docstring'''







			lowerCAmelCase_							  =		42
			lowerCAmelCase_							  =		42
			lowerCAmelCase_							  =		42





			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = None ,						__lowerCAmelCase = Split.train ,						__lowerCAmelCase = None ,						):



				'''simple docstring'''







				warnings.warn(
				    '''This dataset will be removed from the library soon, preprocessing should be handled with the 🤗 Datasets '''
				    '''library. You can have a look at this example script for pointers: '''
				    '''https://github.com/huggingface/transformers/blob/main/examples/pytorch/text-classification/run_glue.py''' ,						__lowerCAmelCase ,						)
				lowerCamelCase__              =	args
				lowerCamelCase__              =	glue_processors[args.task_name]()
				lowerCamelCase__              =	glue_output_modes[args.task_name]
				if isinstance(__lowerCAmelCase ,						__lowerCAmelCase							):
					try:
						lowerCamelCase__              =	Split[mode]
					except KeyError:
						raise KeyError('''mode is not a valid split name'''							)
        # Load data features from cache or dataset file
				lowerCamelCase__              =	os.path.join(
				    cache_dir if cache_dir is not None else args.data_dir ,						F'cached_{mode.value}_{tokenizer.__class__.__name__}_{args.max_seq_length}_{args.task_name}' ,						)
				lowerCamelCase__              =	self.processor.get_labels()
				if args.task_name in ["mnli", "mnli-mm"] and tokenizer.__class__.__name__ in (
				    "RobertaTokenizer",
				    "RobertaTokenizerFast",
				    "XLMRobertaTokenizer",
				    "BartTokenizer",
				    "BartTokenizerFast",
				):
					# HACK(label indices are swapped in RoBERTa pretrained model)
					lowerCamelCase__       ,			lowerCamelCase__              =	label_list[2], label_list[1]
				lowerCamelCase__              =	label_list

				# Make sure only the first process in distributed training processes the dataset,
				# and the others will use the cache.
				lowerCamelCase__              =	cached_features_file + '''.lock'''
				with FileLock(__lowerCAmelCase							):
					if os.path.exists(__lowerCAmelCase							) and not args.overwrite_cache:
						lowerCamelCase__              =	time.time()
						lowerCamelCase__              =	torch.load(__lowerCAmelCase							)
						logger.info(
						    F'Loading features from cached file {cached_features_file} [took %.3f s]' ,						time.time() - start							)
					else:
						logger.info(F'Creating features from dataset file at {args.data_dir}'							)

						if mode == Split.dev:
							lowerCamelCase__              =	self.processor.get_dev_examples(args.data_dir							)
						elif mode == Split.test:
							lowerCamelCase__              =	self.processor.get_test_examples(args.data_dir							)
						else:
							lowerCamelCase__              =	self.processor.get_train_examples(args.data_dir							)
						if limit_length is not None:
							lowerCamelCase__              =	examples[:limit_length]
						lowerCamelCase__              =	glue_convert_examples_to_features(
						    __lowerCAmelCase ,						__lowerCAmelCase ,						max_length=args.max_seq_length ,						label_list=__lowerCAmelCase ,						output_mode=self.output_mode ,						)
						lowerCamelCase__              =	time.time()
						torch.save(self.features ,						__lowerCAmelCase							)
						# ^ This seems to take a lot of time so I want to investigate why and how we can improve.
						logger.info(
						    F'Saving features into cached file {cached_features_file} [took {time.time() - start:.3f} s]'							)





			def __len__(					self							):



				'''simple docstring'''







				return len(self.features							)





			def __getitem__(					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				return self.features[i]





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				return self.label_list



