


import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
   from sagemaker import Session, TrainingJobAnalytics
   from sagemaker.huggingface import HuggingFace
@pytest.mark.skipif(
    literal_eval(os.getenv('TEST_SAGEMAKER'   ,'False'				)				) is not True   ,reason='Skipping test because should only be run when releasing minor transformers version'   ,)
@pytest.mark.usefixtures('sm_env'				)
@parameterized_class(
    [
        {
            'framework': 'pytorch',
            'script': 'run_glue_model_parallelism.py',
            'model_name_or_path': 'roberta-large',
            'instance_type': 'ml.p3dn.24xlarge',
            'results': {'train_runtime': 16_00, 'eval_accuracy': 0.3, 'eval_loss': 1.2},
        },
        {
            'framework': 'pytorch',
            'script': 'run_glue.py',
            'model_name_or_path': 'roberta-large',
            'instance_type': 'ml.p3dn.24xlarge',
            'results': {'train_runtime': 16_00, 'eval_accuracy': 0.3, 'eval_loss': 1.2},
        },
    ]				)
class 		__magic_name__						(unittest.TestCase				):





     '''simple docstring'''



     def 	SCREAMING_SNAKE_CASE__						(					self:str   ):
       if self.framework == "pytorch":
         subprocess.run(
             F"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()     ,		encoding='''utf-8'''     ,		check=_a     ,		)
       assert hasattr(self     ,		'''env'''   )

     def 	SCREAMING_SNAKE_CASE__						(					self:List[str]     ,		_a:Tuple   ):
       # configuration for running training on smdistributed Model Parallel
       snake_case__						= {
           '''enabled''': True,
           '''processes_per_host''': 8,
       }
       snake_case__						= {
           '''enabled''': True,
           '''parameters''': {
               '''microbatches''': 4,
               '''placement_strategy''': '''spread''',
               '''pipeline''': '''interleaved''',
               '''optimize''': '''speed''',
               '''partitions''': 4,
               '''ddp''': True,
           },
       }

       snake_case__						= {'''smdistributed''': {'''modelparallel''': smp_options}, '''mpi''': mpi_options}

       snake_case__						= '''trainer''' if self.script == '''run_glue.py''' else '''smtrainer'''
       # creates estimator
       return HuggingFace(
           entry_point=self.script     ,		source_dir=self.env.test_path     ,		role=self.env.role     ,		image_uri=self.env.image_uri     ,		base_job_name=F"""{self.env.base_job_name}-{instance_count}-smp-{name_extension}"""     ,		instance_count=_a     ,		instance_type=self.instance_type     ,		debugger_hook_config=_a     ,		hyperparameters={
               **self.env.hyperparameters,
               '''model_name_or_path''': self.model_name_or_path,
               '''max_steps''': 5_00,
           }     ,		metric_definitions=self.env.metric_definitions     ,		distribution=_a     ,		py_version='''py36'''     ,		)

     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[int]     ,		_a:int   ):
       TrainingJobAnalytics(_a   ).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv"""   )







     @parameterized.expand([(1,)]   )
     def 	SCREAMING_SNAKE_CASE__						(					self:List[str]     ,		_a:Optional[Any]   ):
       # create estimator
       snake_case__						= self.create_estimator(_a   )

       # run training
       estimator.fit()

       # result dataframe
       snake_case__						= TrainingJobAnalytics(estimator.latest_training_job.name   ).dataframe()

       # extract kpis
       snake_case__						= list(result_metrics_df[result_metrics_df.metric_name == '''eval_accuracy''']['''value''']   )
       snake_case__						= list(result_metrics_df[result_metrics_df.metric_name == '''eval_loss''']['''value''']   )
       # get train time from SageMaker job, this includes starting, preprocessing, stopping
       snake_case__						= (
           Session().describe_training_job(estimator.latest_training_job.name   ).get('''TrainingTimeInSeconds'''     ,		99_99_99   )
       )

       # assert kpis
       assert train_runtime <= self.results["train_runtime"]
       assert all(t >= self.results['''eval_accuracy'''] for t in eval_accuracy   )
       assert all(t <= self.results['''eval_loss'''] for t in eval_loss   )

       # dump tests result into json file to share in PR
       with open(F"""{estimator.latest_training_job.name}.json"""     ,		'''w'''   ) as outfile:
         json.dump({'''train_time''': train_runtime, '''eval_accuracy''': eval_accuracy, '''eval_loss''': eval_loss}     ,		_a   )





"""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						( __snake_case					,			__snake_case					,			__snake_case							):
				lowerCamelCase     :       Optional[int]			       =[R'h\.\d+\.attn\.bias', R'h\.\d+\.attn\.masked_bias']







				@register_to_config
				def __init__( self      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase = None      ,	UpperCAmelCase = 5_0257      ,	UpperCAmelCase = 1024      ,	UpperCAmelCase = 768      ,	UpperCAmelCase = 12      ,	UpperCAmelCase = 12      ,	UpperCAmelCase = None      ,	UpperCAmelCase = "gelu_new"      ,	UpperCAmelCase = 0.1      ,	UpperCAmelCase = 0.1      ,	UpperCAmelCase = 0.1      ,	UpperCAmelCase = 1E-5      ,	UpperCAmelCase = 0.0_2      ,	UpperCAmelCase = True      ,	UpperCAmelCase = True      ,	UpperCAmelCase = False      ,	UpperCAmelCase = 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      ,	UpperCAmelCase   ) if self.prefix_hidden_dim is not None else nn.Identity()
					)

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







				def      lowerCamelCase_     ( self      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase = None      ,	UpperCAmelCase = None      ,	):

					'''simple docstring'''






					lowerCAmelCase_      							=			self.transformer.transformer.wte(UpperCAmelCase   )
					lowerCAmelCase_      							=			self.encode_prefix(UpperCAmelCase   )
					lowerCAmelCase_      							=			self.decode_prefix(UpperCAmelCase   )
					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=UpperCAmelCase      ,	labels=UpperCAmelCase      ,	attention_mask=UpperCAmelCase   )
					if self.prefix_hidden_dim is not None:
						return out, hidden
					else:
						return out







				def      lowerCamelCase_     ( self      ,	UpperCAmelCase      ,	UpperCAmelCase   ):

					'''simple docstring'''






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







				def      lowerCamelCase_     ( self      ,	UpperCAmelCase   ):

					'''simple docstring'''






					return self.encode_prefix(UpperCAmelCase   )







				@torch.no_grad()
				def      lowerCamelCase_     ( self      ,	UpperCAmelCase      ,	UpperCAmelCase      ,	UpperCAmelCase   ):

					'''simple docstring'''






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







				@torch.no_grad()
				def      lowerCamelCase_     ( self      ,	UpperCAmelCase=None      ,	UpperCAmelCase=None      ,	UpperCAmelCase=None      ,	UpperCAmelCase = 5      ,	UpperCAmelCase = 67      ,	UpperCAmelCase = 1.0      ,	UpperCAmelCase = None      ,	):

					'''simple docstring'''






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

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

					for i in range(UpperCAmelCase   ):
						lowerCAmelCase_      							=			self.transformer(inputs_embeds=UpperCAmelCase   )
						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(UpperCAmelCase      ,	-1   )
							lowerCAmelCase_      							=			generated.expand(UpperCAmelCase      ,	*generated.shape[1:]   )
							lowerCAmelCase_							,	lowerCAmelCase_      							=			next_tokens.permute(1      ,	0   ), scores.squeeze(0   )
							if tokens is None:
								lowerCAmelCase_      							=			next_tokens
							else:
								lowerCAmelCase_      							=			tokens.expand(UpperCAmelCase      ,	*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(UpperCAmelCase      ,	-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(UpperCAmelCase   ).squeeze()
						if is_stopped.all():
							break

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