"""simple docstring"""







import copy
import logging
import random

import joblib
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

from transformers import AdamW, GPTaLMHeadModel, get_linear_schedule_with_warmup


lowerCamelCase      : Optional[int]       =logging.getLogger(__name__)




def 						_lowercase				(  _SCREAMING_SNAKE_CASE				:				Dict						)    ->     Any:







     '''simple docstring'''


     random.seed(_SCREAMING_SNAKE_CASE						)
     np.random.seed(_SCREAMING_SNAKE_CASE						)
     torch.manual_seed(_SCREAMING_SNAKE_CASE						)
     torch.cuda.manual_seed_all(_SCREAMING_SNAKE_CASE						)




def 						_lowercase				(  _SCREAMING_SNAKE_CASE				:				Any  ,       _SCREAMING_SNAKE_CASE				:				Any  ,       _SCREAMING_SNAKE_CASE				:				Union[str, Any]						)    ->     Optional[Any]:







     '''simple docstring'''


     model.eval()
     __A					:					List[str]								=		next(model.parameters()						).device
     __A					:					Tuple								=		1
     __A					:					Tuple								=		torch.zeros((eval_batch_size, context_len)  ,       dtype=torch.long  ,       device=_SCREAMING_SNAKE_CASE						)
     __A					:					str								=		DataLoader(_SCREAMING_SNAKE_CASE  ,       shuffle=_SCREAMING_SNAKE_CASE  ,       batch_size=_SCREAMING_SNAKE_CASE						)
     __A					:					List[Any]								=		torch.zeros(1  ,       device=_SCREAMING_SNAKE_CASE						)
     __A					:					List[Any]								=		0
     for batch in eval_dataloader:
          batch.to(_SCREAMING_SNAKE_CASE						)
          # pad
          context.zero_()
          for i in range(_SCREAMING_SNAKE_CASE						):
               __A					:					List[Any]								=		batch[i]
          __A					:					Union[str, Any]								=		model(_SCREAMING_SNAKE_CASE  ,       labels=_SCREAMING_SNAKE_CASE						)
          eval_loss += outputs[0].sum().item()
          nb_eval_examples += batch.size(0						)
     __A					:					List[str]								=		eval_loss / nb_eval_examples
     __A					:					Tuple								=		torch.exp(_SCREAMING_SNAKE_CASE						)
     model.train()
     return perplexity




def 						_lowercase				(  _SCREAMING_SNAKE_CASE				:				Dict="gpt2"						)    ->     Any:







     '''simple docstring'''


     __A					:					Union[str, Any]								=		GPTaLMHeadModel.from_pretrained(_SCREAMING_SNAKE_CASE  ,       output_hidden_states=_SCREAMING_SNAKE_CASE						)
     torch.save(model.state_dict()  ,       model_name + 'local.pt'						)
     return model




def 						_lowercase				(  _SCREAMING_SNAKE_CASE				:				Tuple  ,       _SCREAMING_SNAKE_CASE				:				Tuple  ,       _SCREAMING_SNAKE_CASE				:				List[str]						)    ->     Optional[Any]:







     '''simple docstring'''


     __A					:					List[Any]								=		copy.deepcopy(_SCREAMING_SNAKE_CASE						)
     model.to(_SCREAMING_SNAKE_CASE						)

     __A					:					Optional[int]								=		['bias', 'LayerNorm.weight']
     __A					:					List[str]								=		[
         {
             'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay						)],
             'weight_decay': 0.0,
         },
         {'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay						)], 'weight_decay': 0.0},
     ]
     __A					:					Tuple								=		AdamW(_SCREAMING_SNAKE_CASE  ,       lr=5E-5  ,       eps=1E-8						)
     __A					:					Optional[int]								=		get_linear_schedule_with_warmup(_SCREAMING_SNAKE_CASE  ,       0  ,       _SCREAMING_SNAKE_CASE						)
     torch.cuda.empty_cache()
     return model, lm_optimizer, lm_scheduler




def 						_lowercase				(  _SCREAMING_SNAKE_CASE				:				str  ,       _SCREAMING_SNAKE_CASE				:				Dict  ,       _SCREAMING_SNAKE_CASE				:				Optional[int]  ,       _SCREAMING_SNAKE_CASE				:				str						)    ->     Optional[int]:







     '''simple docstring'''


     __A					:					List[str]								=		np.array(_SCREAMING_SNAKE_CASE						).mean()
     __A					:					Dict								=		np.array(_SCREAMING_SNAKE_CASE						).std()
     __A					:					Any								=		(real_perps - avg) / std
     __A					:					List[str]								=		list(zip(_SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE						)						)
     joblib.dump(_SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE						)




def 						_lowercase				(  _SCREAMING_SNAKE_CASE				:				Union[str, Any]  ,       _SCREAMING_SNAKE_CASE				:				Any  ,       _SCREAMING_SNAKE_CASE				:				Any  ,       _SCREAMING_SNAKE_CASE				:				Tuple  ,       _SCREAMING_SNAKE_CASE				:				str  ,       _SCREAMING_SNAKE_CASE				:				Tuple  ,       _SCREAMING_SNAKE_CASE				:				Optional[Any]  ,       _SCREAMING_SNAKE_CASE				:				Tuple="dev.jbl"  ,       _SCREAMING_SNAKE_CASE				:				Union[str, Any]=recopy_gpta  ,       )    ->     Any:







     '''simple docstring'''


     __A					:					List[str]								=		[]
     __A					:					Any								=		[]
     __A					:					Tuple								=		[]

     # Initialize the transformer model
     __A					:					Tuple								=		copy.deepcopy(_SCREAMING_SNAKE_CASE						)
     orig_model.to(device='cpu'						)
     torch.cuda.empty_cache()

     # Compute perplexity of initial transformer model for comparison
     model.train()
     __A						,		__A						,		__A					:					Any								=		recopy_model(_SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE						)

     for step in tqdm(range(_SCREAMING_SNAKE_CASE						)						):
          __A					:					int								=		torch.zeros((1, context_len)  ,       dtype=torch.long  ,       device=_SCREAMING_SNAKE_CASE						)
          __A					:					Tuple								=		random.choice(_SCREAMING_SNAKE_CASE						)
          __A					:					List[str]								=		random.randint(0  ,       len(story[0]						) - context_len - 1						)
          __A					:					Union[str, Any]								=		story[0][start : start + context_len]
          lm_optimizer.zero_grad()
          __A					:					Dict								=		model(_SCREAMING_SNAKE_CASE  ,       labels=_SCREAMING_SNAKE_CASE						)
          __A					:					Tuple								=		outputs[0]
          __A					:					Optional[Any]								=		compute_perplexity(_SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE						)
          model.train()
          lm_loss.backward()
          # Do LM backprop
          torch.nn.utils.clip_grad_norm_(model.parameters()  ,       3.0						)
          lm_optimizer.step()
          lm_scheduler.step()  # Update learning rate schedule

          # Compute perplexity after back-propagating on the selected context
          __A					:					List[str]								=		compute_perplexity(_SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE						)

          # Periodically save the stored (X, IG(X)) pairs
          if step % 1000 == 0 and step > 1:
               intermittent_save(_SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE						)

          # Reset the pretrained model to the original pretrained GPT-2 weights after each iteration
          __A						,		__A						,		__A					:					Optional[int]								=		recopy_model(_SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE						)

          past_perps.append(past_perp.item()						)
          real_perps.append(orig_perp - real_perp.item()						)
          contexts.append(np.array(context.cpu()						)						)

     intermittent_save(_SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE						)




def 						_lowercase				(  _SCREAMING_SNAKE_CASE				:				List[str]  ,       _SCREAMING_SNAKE_CASE				:				Any="data/tokenized_stories_train_wikitext103.jbl"  ,       _SCREAMING_SNAKE_CASE				:				Tuple=100  ,       _SCREAMING_SNAKE_CASE				:				str=1026  ,       _SCREAMING_SNAKE_CASE				:				Dict=True						)    ->     Optional[Any]:







     '''simple docstring'''


     __A					:					Dict								=		joblib.load(_SCREAMING_SNAKE_CASE						)
     print('data loaded'						)
     __A					:					Any								=		[]
     if trim:
          for i, example in enumerate(_SCREAMING_SNAKE_CASE						):
               if len(example[0]						) > min_len:
                    __A					:					Any								=		random.randint(0  ,       len(example[0]						) - context_len - 1						)
                    objective_set.append(example[0, start : start + context_len]						)
               if len(_SCREAMING_SNAKE_CASE						) >= number:
                    break
          __A					:					str								=		[]
          for j in range(i + 1  ,       len(_SCREAMING_SNAKE_CASE						)						):
               if len(data[j][0]						) > min_len:
                    train_data.append(data[j]						)
     else:
          __A					:					Optional[Any]								=		data[0:number]
          __A					:					List[Any]								=		data[number:]

     joblib.dump(_SCREAMING_SNAKE_CASE  ,       'objective_set.jbl'						)
     print('objective set saved'						)
     return train_data, objective_set




def 						_lowercase				(  _SCREAMING_SNAKE_CASE				:				Optional[Any]  ,       _SCREAMING_SNAKE_CASE				:				Optional[int]  ,       _SCREAMING_SNAKE_CASE				:				List[Any]  ,       _SCREAMING_SNAKE_CASE				:				int  ,       _SCREAMING_SNAKE_CASE				:				List[Any]=50  ,       _SCREAMING_SNAKE_CASE				:				Union[str, Any]="secondary_learner.pt"						)    ->     List[str]:







     '''simple docstring'''


     __A					:					Tuple								=		torch.device('cuda:0' if torch.cuda.is_available() else 'cpu'						)
     # We will use the first 512 pairs from our dataset as a test set for
     # our secondary learner and the rest to train
     __A					:					Optional[int]								=		train_dataset[:512]
     __A					:					Union[str, Any]								=		train_dataset[512:]
     __A					:					List[Any]								=		DataLoader(_SCREAMING_SNAKE_CASE  ,       shuffle=_SCREAMING_SNAKE_CASE  ,       batch_size=_SCREAMING_SNAKE_CASE						)
     __A					:					List[str]								=		DataLoader(_SCREAMING_SNAKE_CASE  ,       shuffle=_SCREAMING_SNAKE_CASE  ,       batch_size=_SCREAMING_SNAKE_CASE						)

     # secondary learner model set up
     __A					:					Optional[Any]								=		nn.MSELoss()
     __A					:					Any								=		nn.MSELoss(reduction='sum'						)
     secondary_learner.to(_SCREAMING_SNAKE_CASE						)
     __A					:					Dict								=		torch.optim.Adam(secondary_learner.parameters()  ,       lr=0.0_00_01						)
     secondary_learner.train()

     # TODO in original code this is written as number of actual batches seen
     # not number of items seen but other places it is number of items instead.
     # improve consistency! changed this to epochs for clarity
     __A					:					Dict								=		float('inf'						)
     # Iterate through batches until we've used max_steps batches
     for epoch in range(int(_SCREAMING_SNAKE_CASE						)						):
          __A					:					Any								=		0.0
          secondary_learner.train()
          for step, batch in enumerate(_SCREAMING_SNAKE_CASE						):
               __A					:					int								=		batch[0].to(_SCREAMING_SNAKE_CASE						)
               __A					:					str								=		batch[1].to(_SCREAMING_SNAKE_CASE						)
               __A					:					List[Any]								=		secondary_learner(_SCREAMING_SNAKE_CASE						)
               q_optimizer.zero_grad()
               __A					:					Optional[Any]								=		loss(_SCREAMING_SNAKE_CASE  ,       real_q.float()						)
               q_loss.backward()
               q_optimizer.step()
               tr_q_loss += q_loss.item()

               # model trains fairly quickly so we won't wait for a full epoch
               # eval is triggered at eval_freq and end of epochs
               if (step % eval_freq == 0 and step > 0) or ((step + 1) == len(_SCREAMING_SNAKE_CASE						)):
                    __A					:					List[str]								=		tr_q_loss / (step + 1)

                    secondary_learner.eval()
                    __A					:					int								=		0.0
                    __A					:					Optional[int]								=		0.0
                    __A					:					str								=		[]
                    __A					:					Tuple								=		[]
                    # Compute performance of the secondary learner after this batch
                    for stepa, batcha in enumerate(_SCREAMING_SNAKE_CASE						):
                         __A					:					Optional[int]								=		batcha[0].to(_SCREAMING_SNAKE_CASE						)
                         __A					:					Optional[Any]								=		batcha[1].to(_SCREAMING_SNAKE_CASE						)
                         __A					:					int								=		secondary_learner(_SCREAMING_SNAKE_CASE						)
                         q_lossa += test_loss(_SCREAMING_SNAKE_CASE  ,       _SCREAMING_SNAKE_CASE						).item()
                         sum_qa += torch.sum(_SCREAMING_SNAKE_CASE						).item()
                         for ei, i in enumerate(predicted_qa.cpu().detach().numpy()						):
                              predicted.append(i.item()						)
                         for ei, i in enumerate(real_qa.cpu().detach().numpy()						):
                              actual.append(i.item()						)

                    q_lossa /= len(_SCREAMING_SNAKE_CASE						)
                    print(
                        'Epoch: '  ,       _SCREAMING_SNAKE_CASE  ,       'step: '  ,       _SCREAMING_SNAKE_CASE  ,       'Avg. q:'  ,       sum_qa / len(_SCREAMING_SNAKE_CASE						)  ,       'Train Loss: '  ,       _SCREAMING_SNAKE_CASE  ,       'Test Loss: '  ,       _SCREAMING_SNAKE_CASE  ,       )
                    if q_lossa < best_test_loss:
                         joblib.dump((predicted, actual)  ,       'pred_vs_actual.jbl'						)
                         torch.save(secondary_learner.state_dict()  ,       _SCREAMING_SNAKE_CASE						)
                         __A					:					Optional[Any]								=		q_lossa

               secondary_learner.train()
     return secondary_learner






class 							__snake_case(					nn.Module	):


    '''simple docstring'''





    def __init__(  self					,       __lowerCamelCase	):




         '''simple docstring'''



         super(__lowerCamelCase					,       self	).__init__()
         __A					:					Optional[Any]								=		model.transformer.wte
         __A					:					Union[str, Any]								=		copy.deepcopy(model.transformer.wte.weight	)
         __A					:					Optional[int]								=		nn.Convad(self.embeddings.weight.size(1	)					,       256					,       3					,       padding=1	)
         __A					:					List[str]								=		nn.Sequential(nn.Linear(256					,       32	)					,       nn.Dropout(p=0.1	)					,       nn.Linear(32					,       32	)					,       nn.Linear(32					,       1	)	)
    def 						_a							(  self					,       __lowerCamelCase	):




         '''simple docstring'''



         __A					:					int								=		torch.max(self.conv(self.embeddings(__lowerCamelCase	).squeeze(1	).transpose(1					,       2	)	)					,       2	)[0]
         __A					:					List[str]								=		self.fc(__lowerCamelCase	)
         return qs.squeeze(1	)

    @classmethod
    def 						_a							(  cls					,       __lowerCamelCase					,       __lowerCamelCase	):




         '''simple docstring'''



         __A					:					Optional[Any]								=		cls(__lowerCamelCase	)  # this calls __init__
         __A					:					int								=		torch.load(__lowerCamelCase	)
         secondary_learner.load_state_dict(__lowerCamelCase	)
         __A					:					Any								=		model.transformer.wte
         __A					:					str								=		copy.deepcopy(model.transformer.wte.weight	)
         return secondary_learner




"""simple docstring"""







import unittest

import numpy as np

from transformers import (
    GPTSanJapaneseConfig,
    GPTSanJapaneseForConditionalGeneration,
    GPTSanJapaneseModel,
    GPTSanJapaneseTokenizer,
    is_torch_available,
)
from transformers.generation import GenerationConfig
from transformers.testing_utils import require_torch, slow, tooslow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin






class 							__snake_case:


    '''simple docstring'''





    def __init__(  self					,       __lowerCamelCase					,       __lowerCamelCase=36000					,       __lowerCamelCase=13					,       __lowerCamelCase=7					,       __lowerCamelCase=True					,       __lowerCamelCase=32					,       __lowerCamelCase=42					,       __lowerCamelCase=5					,       __lowerCamelCase=2					,       __lowerCamelCase=4					,       __lowerCamelCase=2					,       __lowerCamelCase=32					,       __lowerCamelCase=80					,       __lowerCamelCase=33					,       __lowerCamelCase=0.0					,       __lowerCamelCase=1e-6					,       __lowerCamelCase=100					,       __lowerCamelCase=0.0					,       ):




         '''simple docstring'''



         __A					:					Tuple								=		vocab_size
         __A					:					Optional[Any]								=		parent
         __A					:					Union[str, Any]								=		batch_size
         __A					:					Optional[int]								=		num_contexts
         # For common tests
         __A					:					List[str]								=		self.num_contexts
         __A					:					Dict								=		is_training
         __A					:					Tuple								=		hidden_size
         __A					:					Dict								=		num_ext_layers
         __A					:					Union[str, Any]								=		ext_size
         __A					:					List[str]								=		num_hidden_layers
         __A					:					Tuple								=		num_attention_heads
         __A					:					Any								=		num_experts
         __A					:					str								=		d_ff
         __A					:					Union[str, Any]								=		d_ext
         __A					:					Optional[int]								=		d_spout
         __A					:					int								=		dropout_rate
         __A					:					Optional[Any]								=		layer_norm_epsilon
         __A					:					Optional[Any]								=		expert_capacity
         __A					:					int								=		router_jitter_noise
    def 						_a							(  self	):




         '''simple docstring'''



         return GPTSanJapaneseConfig.from_pretrained('Tanrei/GPTSAN-japanese'	)
    def 						_a							(  self	):




         '''simple docstring'''



         __A					:					str								=		ids_tensor([self.batch_size, self.seq_length]					,       self.vocab_size	)

         __A					:					List[Any]								=		self.get_config()

         return (config, input_ids)
    def 						_a							(  self	):




         '''simple docstring'''



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

         __A					:					Optional[Any]								=		self.get_config()

         return (config, {"input_ids": input_ids})
    def 						_a							(  self	):




         '''simple docstring'''



         return GPTSanJapaneseConfig(
             vocab_size=self.vocab_size					,       num_contexts=self.seq_length					,       d_model=self.hidden_size					,       d_ff=self.d_ff					,       d_ext=self.d_ext					,       d_spout=self.d_spout					,       num_switch_layers=self.num_hidden_layers - self.num_ext_layers					,       num_ext_layers=self.num_ext_layers					,       num_heads=self.num_attention_heads					,       num_experts=self.num_experts					,       expert_capacity=self.expert_capacity					,       dropout_rate=self.dropout_rate					,       layer_norm_epsilon=self.layer_norm_epsilon					,       router_jitter_noise=self.router_jitter_noise					,       )

    def 						_a							(  self					,       __lowerCamelCase					,       __lowerCamelCase					,       ):




         '''simple docstring'''



         __A					:					Dict								=		GPTSanJapaneseForConditionalGeneration(config=__lowerCamelCase	)
         model.to(__lowerCamelCase	)
         model.eval()
         __A					:					Union[str, Any]								=		model(
             input_ids=__lowerCamelCase					,       )
         self.parent.assertIsNotNone(__lowerCamelCase	)







@require_torch
class 							__snake_case(					A_				,	A_				,	unittest.TestCase	):


    '''simple docstring'''





    _UpperCAmelCase       =			(GPTSanJapaneseModel,) if is_torch_available() else ()
    _UpperCAmelCase       =			(
        {
            "conversational": GPTSanJapaneseForConditionalGeneration,
            "feature-extraction": GPTSanJapaneseForConditionalGeneration,
            "summarization": GPTSanJapaneseForConditionalGeneration,
            "text2text-generation": GPTSanJapaneseForConditionalGeneration,
            "translation": GPTSanJapaneseForConditionalGeneration,
        }
        if is_torch_available()
        else {}
    )
    _UpperCAmelCase       =			False
    _UpperCAmelCase       =			False
    _UpperCAmelCase       =			False
    _UpperCAmelCase       =			False
    _UpperCAmelCase       =			False
    _UpperCAmelCase       =			False
    _UpperCAmelCase       =			False
    _UpperCAmelCase       =			False
    # The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
    _UpperCAmelCase       =			[0.8, 0.9]
    def 						_a							(  self					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase	):




         '''simple docstring'''



         if pipeline_test_casse_name == "SummarizationPipelineTests":
              # TODO: fix `_reorder_cache` is not implemented for this model
              return True
         elif pipeline_test_casse_name == "Text2TextGenerationPipelineTests":
              # TODO: check this.
              return True

         return False
    def 						_a							(  self	):




         '''simple docstring'''



         __A					:					Optional[Any]								=		GPTSanJapaneseTester(self	)
         __A					:					Optional[Any]								=		ConfigTester(self					,       config_class=__lowerCamelCase					,       d_model=37	)
    def 						_a							(  self	):




         '''simple docstring'''



         GPTSanJapaneseConfig()
    def 						_a							(  self	):




         '''simple docstring'''



         __A					:					Tuple								=		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__lowerCamelCase	)
    @unittest.skip(
        reason='skip for now as the computed `max_memory` by `model_split_percents` in the test method will be changed inside `from_pretrained`'	)
    def 						_a							(  self	):




         '''simple docstring'''



         super().test_model_parallelism()

    @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'	)
    def 						_a							(  self	):




         '''simple docstring'''



         pass







@require_torch
class 							__snake_case(					A_				,	A_				,	unittest.TestCase	):


    '''simple docstring'''





    _UpperCAmelCase       =			(GPTSanJapaneseForConditionalGeneration,) if is_torch_available() else ()
    _UpperCAmelCase       =			False
    _UpperCAmelCase       =			False
    _UpperCAmelCase       =			False
    _UpperCAmelCase       =			False
    _UpperCAmelCase       =			False
    _UpperCAmelCase       =			False
    # The small GPTSAN_JAPANESE model needs higher percentages for CPU/MP tests
    _UpperCAmelCase       =			[0.8, 0.9]
    def 						_a							(  self	):




         '''simple docstring'''



         __A					:					List[Any]								=		GPTSanJapaneseTester(self	)
         __A					:					Dict								=		ConfigTester(self					,       config_class=__lowerCamelCase					,       d_model=37	)
    def 						_a							(  self	):




         '''simple docstring'''



         GPTSanJapaneseConfig()
    def 						_a							(  self	):




         '''simple docstring'''



         __A					:					Optional[Any]								=		self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*__lowerCamelCase	)
    @unittest.skip(
        reason='skip for now as the computed `max_memory` by `model_split_percents` in the test method will be changed inside `from_pretrained`'	)
    def 						_a							(  self	):




         '''simple docstring'''



         super().test_model_parallelism()
    @unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'	)
    def 						_a							(  self	):




         '''simple docstring'''



         pass
    @slow
    def 						_a							(  self	):




         '''simple docstring'''



         __A					:					Union[str, Any]								=		GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'	)
         __A					:					List[str]								=		GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'	)
         __A					:					int								=		tokenizer.encode('武田信玄は'					,       return_tensors='pt'	)
         __A					:					Any								=		model(__lowerCamelCase	)
         __A					:					Optional[int]								=		outputs.logits.detach().cpu().numpy()
         # Output of original model created with mesh-tensoflow
         __A					:					Optional[int]								=		[
             # fmt: off
             [-12.0_37_83_98_89_52_63_67, -12.4_33_06_15_99_73_14_45, -14.3_33_84_03_70_17_82_23, -12.4_50_34_59_93_04_19_92, -11.1_66_13_76_95_31_25,
             -11.9_30_13_76_34_27_73_44, -10.6_59_74_04_47_99_80_47, -12.9_09_57_45_08_66_69_92, -13.2_41_04_30_90_82_03_12, -13.3_98_57_95_97_47_31_45,
             -11.1_07_52_48_71_82_61_72, -12.3_68_59_41_69_61_67, -22.97_94_31_15_23_43_75, -10.4_81_06_76_57_47_07_03, -12.4_84_03_07_23_57_17_77,
             -12.8_07_36_06_49_10_88_87, -14.7_69_70_00_50_35_40_04, -12.2_33_57_96_35_62_01_17, -13.4_28_14_54_08_63_03_71, -22.6_24_17_79_32_73_92_58],
             [-7.5_11_14_98_83_27_02_64, -8.2_81_85_17_68_49_36_52, -7.9_43_12_71_55_30_39_55, -7.55_02_13_33_69_44_58, -6.49_86_99_66_50_69_58,
             -7.5_86_79_62_83_72_19_24, -6.9_78_08_59_94_72_04_59, -7.8_39_14_51_83_56_32_32, -8.21_96_41_68_54_85_84, -8.6_95_09_12_47_55_85_94,
             -6.7_06_91_06_10_19_89_75, -6.65_85_79_82_63_54_98_05, -19.5_65_69_86_23_65_72_27, -5.3_53_84_27_35_29_05_27, -8.3_50_68_60_73_30_32_23,
             -8.0_39_38_86_56_61_62_11, -10.8_56_56_92_90_16_11_33, -7.75_15_44_47_55_55_42, -8.8_19_02_21_78_64_99_02, -19.51_53_27_45_36_13_28],
             [-9.73_06_62_34_58_86_23, -10.2_23_92_27_29_49_21_88, -9.9_32_98_14_91_08_88_67, -11.8_57_83_67_23_32_76_37, -7.6_62_62_62_66_47_94_92,
             -11.13_52_91_09_95_48_34, -7.7_65_09_76_18_10_30_27, -11.4_72_92_32_78_80_85_94, -9.5_43_14_99_48_12_01_17, -11.9_05_63_39_26_39_16_02,
             -9.3_66_16_42_07_45_84_96, -11.5_73_42_81_53_99_17, -23.6_99_00_32_19_60_44_92, -9.4_29_59_02_25_21_97_27, -10.42_83_92_41_02_78_32,
             -10.5_85_24_03_64_07_47_07, -10.94_77_19_57_39_74_61, -11.0_95_41_60_69_03_07_62, -10.3_90_24_06_69_25_04_88, -23.7_69_37_29_40_06_34_77],
             [-9.7_28_26_57_62_32_91_02, -9.8_59_71_26_00_70_80_08, -10.09_72_92_90_00_85_45, -9.6_78_52_21_09_98_53_52, -6.8_79_51_99_39_42_26_07,
             -9.68_48_75_48_82_81_25, -4.28_03_42_57_88_87_93_95, -10.0_18_91_42_22_71_72_85, -9.3_08_44_59_30_48_09_57, -10.63_39_45_46_50_87_89,
             -8.0_83_64_67_74_29_19_92, -9.06_30_14_98_41_30_86, -21.9_04_26_63_57_42_18_75, -8.90_16_08_46_71_02_05, -8.8_41_87_60_29_96_82_62,
             -11.8_56_71_99_70_70_31_25, -12.0_79_39_81_55_21_24_02, -11.2_33_75_32_04_34_57_03, -10.1_77_33_86_00_15_86_91, -21.87_25_66_22_31_44_53],
             [-9.6_69_76_45_18_73_77_93, -9.6_14_19_86_84_69_23_83, -9.8_14_51_03_45_45_89_84, -9.9_96_50_19_22_60_74_22, -11.3_75_69_04_60_20_50_78,
             -10.1_13_40_52_27_66_11_33, -10.5_46_86_73_70_60_54_69, -10.04_36_90_68_14_57_52, -10.9_07_80_92_57_50_73_24, -10.5_04_21_61_94_15_28_32,
             -11.1_29_19_90_28_01_51_37, -10.1_51_12_40_00_54_93_16, -21.96_58_67_99_62_15_82, -9.0_86_34_94_87_30_46_88, -11.7_30_33_90_50_29_29_69,
             -10.4_60_66_76_10_16_84_57, -10.2_98_04_99_26_75_78_12, -10.7_84_14_82_16_24_75_59, -10.8_40_69_34_73_81_59_18, -22.03_15_28_47_29_00_39],
             # fmt: on
         ]
         __A					:					Dict								=		np.array(__lowerCamelCase	).flatten()
         __A					:					Tuple								=		output_logits[0, :, :20].flatten()

         def check(__lowerCamelCase					,       __lowerCamelCase					,       __lowerCamelCase=5e-4	):
              return abs(a - b	) < epsilon * max(abs(__lowerCamelCase	)					,       abs(__lowerCamelCase	)	)

         self.assertTrue(np.all([check(target[i]					,       predict[i]	) for i in range(len(__lowerCamelCase	)	)]	)	)
    @slow
    def 						_a							(  self	):




         '''simple docstring'''



         __A					:					Optional[Any]								=		GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'	)
         __A					:					Optional[int]								=		GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'	)
         model.to(__lowerCamelCase	)

         # set deterministically
         __A					:					Any								=		GenerationConfig.from_pretrained('Tanrei/GPTSAN-japanese'	)
         __A					:					str								=		1

         # use different length sentences to test batching
         __A					:					str								=		[
             '甲斐なら武田と言うほど',
             '織田信長は、',
         ]

         __A					:					Union[str, Any]								=		'left'
         __A					:					Union[str, Any]								=		tokenizer(__lowerCamelCase					,       return_tensors='pt'					,       padding=__lowerCamelCase	)
         __A					:					int								=		inputs['input_ids'].to(__lowerCamelCase	)

         self.assertNotEqual(inputs['attention_mask'][0].numpy().tolist()					,       inputs['attention_mask'][1].numpy().tolist()	)

         __A					:					str								=		model.generate(
             input_ids=__lowerCamelCase					,       attention_mask=inputs['attention_mask'].to(__lowerCamelCase	)					,       max_new_tokens=3					,       generation_config=__lowerCamelCase					,       )

         __A					:					Dict								=		tokenizer(sentences[0]					,       return_tensors='pt'	).input_ids.to(__lowerCamelCase	)
         __A					:					List[str]								=		model.generate(
             input_ids=__lowerCamelCase					,       max_new_tokens=3					,       generation_config=__lowerCamelCase	)

         __A					:					Tuple								=		tokenizer(sentences[1]					,       return_tensors='pt'	).input_ids.to(__lowerCamelCase	)
         __A					:					List[Any]								=		model.generate(input_ids=__lowerCamelCase					,       max_new_tokens=3					,       generation_config=__lowerCamelCase	)

         self.assertNotEqual(inputs_non_padded.shape					,       inputs_padded.shape	)

         __A					:					Tuple								=		tokenizer.batch_decode(__lowerCamelCase					,       skip_special_tokens=__lowerCamelCase	)
         __A					:					str								=		tokenizer.decode(output_non_padded[0]					,       skip_special_tokens=__lowerCamelCase	)
         __A					:					Any								=		tokenizer.decode(output_padded[0]					,       skip_special_tokens=__lowerCamelCase	)

         __A					:					Optional[Any]								=		[
             '甲斐なら武田と言うほど甲斐の武田',
             '織田信長は、このような',
         ]
         self.assertListEqual(__lowerCamelCase					,       __lowerCamelCase	)
         self.assertListEqual(__lowerCamelCase					,       [non_padded_sentence, padded_sentence]	)
    @tooslow
    def 						_a							(  self	):




         '''simple docstring'''



         __A					:					Optional[Any]								=		GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'	)
         __A					:					List[str]								=		GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'	)
         # Output of original model created with mesh-tensoflow
         __A					:					List[Any]								=		[
             ('武田信玄は', 35675),
             ('武田信玄は、', 45),
             ('武田信玄は、この', 29),
             ('武田信玄は、このよう', 30642),
             ('武田信玄は、このような', 35680),
             ('武田信玄は、このような「', 8640),
             ('武田信玄は、このような「武田', 31617),
             ('武田信玄は、このような「武田家', 30646),
             ('武田信玄は、このような「武田家の', 31617),
             ('武田信玄は、このような「武田家の家', 31381),
         ]
         for input, output in target:
              __A					:					Tuple								=		tokenizer.encode(__lowerCamelCase					,       return_tensors='pt'	)
              __A					:					List[str]								=		model(__lowerCamelCase	)
              __A					:					Dict								=		outputs.logits.detach().cpu().numpy()[0]
              __A					:					Optional[Any]								=		np.argmax(output_logits[-1]	)
              self.assertEqual(__lowerCamelCase					,       __lowerCamelCase	)
    @slow
    def 						_a							(  self	):




         '''simple docstring'''



         __A					:					str								=		GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'	)
         __A					:					Optional[Any]								=		GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'	)
         model.to(__lowerCamelCase	)

         # set deterministically
         __A					:					List[str]								=		GenerationConfig.from_pretrained('Tanrei/GPTSAN-japanese'	)
         __A					:					List[str]								=		1

         __A					:					str								=		'武田信玄は、'
         __A					:					Optional[Any]								=		tokenizer(__lowerCamelCase					,       return_tensors='pt'	).input_ids.to(__lowerCamelCase	)
         __A					:					Optional[Any]								=		tokenizer([input_text, input_text]					,       return_tensors='pt'	).input_ids.to(__lowerCamelCase	)

         # spout from uniform and one-hot
         __A					:					int								=		[
             # fmt: off
             [0.87_88_22_08, 0.38_42_63_96, 0.33_22_02_48, 0.43_89_04_06, 0.16_56_22_52,
             0.04_80_39_85, 0.21_15_72  , 0.23_18_84_73, 0.37_15_30_68, 0.7_83_63_77 ,
             0.02_16_01_72, 0.38_76_17_19, 0.75_29_07_72, 0.90_19_88_57, 0.34_36_57_77,
             0.64_16_81_69, 0.44_31_84_71, 0.14_57_57_46, 0.92_56_28_81, 0.40_81_21_48,
             0.29_01_91_22, 0.88_86_15_99, 0.65_52_48_46, 0.43_56_34_56, 0.38_17_71_87,
             0.70_83_29_65, 0.81_52_78_92, 0.68_83_28_12, 0.38_83_31_92, 0.4_56_15_22 ,
             0.14_82_88_17, 0.47_24_82_13, 0.54_35_73_35, 0.82_00_95_66, 0.1_33_88_84 ,
             0.02_75_54_17, 0.19_76_46_77, 0.2_42_20_84 , 0.04_75_76_74, 0.65_40_96_06,
             0.0_82_45_89 , 0.03_30_43_83, 0.94_38_76_89, 0.98_76_45_09, 0.82_43_39_01,
             0.27_64_67_41, 0.64_90_74_93, 0.76_00_94_06, 0.30_08_79_15, 0.17_90_46_89,
             0.41_60_17_14, 0.67_04_63_98, 0.10_42_28_22, 0.08_44_73_74, 0.07_35_43_44,
             0.61_42_35_65, 0.70_28_48_66, 0.7_53_23_33 , 0.1_97_20_38 , 0.29_57_56_59,
             0.90_58_38_86, 0.29_26_53_07, 0.50_00_01_75, 0.70_40_76_55, 0.88_93_63  ,
             0.81_90_44_18, 0.66_82_91_28, 0.64_46_88_15, 0.56_56_37_23, 0.85_60_18_75,
             0.94_92_46_72, 0.00_16_67_62, 0.25_22_06_43, 0.74_54_02_19, 0.67_99_32_47,
             0.1_54_96_75 , 0.39_38_53_52, 0.92_15_36_07, 0.63_74_59_31, 0.27_75_90_43,
             0.84_70_22_95, 0.65_90_42_71, 0.58_67_66_14, 0.8_66_69_36 , 0.39_60_74_38,
             0.79_95_49_83, 0.42_22_06_97, 0.39_65_03_81, 0.7_84_98_64 , 0.56_15_02_01,
             0.15_67_89_25, 0.14_74_60_32, 0.34_54_21_14, 0.47_02_67_83, 0.11_95_64_89,
             0.25_42_14_35, 0.33_78_89_01, 0.68_93_48_42, 0.36_42_46_85, 0.71_73_78_98,
             0.38_98_34_49, 0.94_39_37_79, 0.39_57_55_88, 0.36_61_65_53, 0.87_10_46_65,
             0.64_63_02_03, 0.22_51_69_05, 0.88_27_08_04, 0.15_03_13_38, 0.75_14_43_45,
             0.46_45_90_25, 0.85_39_64_54, 0.86_35_56_43, 0.65_13_98_51, 0.70_26_60_61,
             0.30_24_13_89, 0.81_05_64_97, 0.88_86_59_69, 0.38_77_38_07, 0.70_63_58_49,
             0.90_71_84_59, 0.43_24_57_89, 0.28_00_06_54, 0.45_93_55_62, 0.08_77_35_19,
             0.9_55_21_51 , 0.93_90_15_11, 0.22_48_92_88], # uniform
             [1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
              0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
              0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
              0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
              0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
              0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
              0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
              0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
              0., 0., 0., 0., 0., 0., 0., 0.],
             # fmt: on
         ]

         __A					:					Optional[int]								=		model.generate(
             input_ids=__lowerCamelCase					,       spout=spouts[0]					,       max_new_tokens=20					,       generation_config=__lowerCamelCase					,       )

         __A					:					Optional[int]								=		model.generate(
             input_ids=__lowerCamelCase					,       spout=spouts[1]					,       max_new_tokens=20					,       generation_config=__lowerCamelCase					,       )

         __A					:					Union[str, Any]								=		model.generate(
             input_ids=__lowerCamelCase					,       spout=__lowerCamelCase					,       max_new_tokens=20					,       generation_config=__lowerCamelCase					,       )

         __A					:					Union[str, Any]								=		tokenizer.decode(outputa[0]	)
         __A					:					Any								=		tokenizer.decode(outputa[0]	)
         __A					:					int								=		tokenizer.batch_decode(__lowerCamelCase	)

         __A					:					Optional[int]								=		[
             '武田信玄は、武田氏の滅亡後、武田氏の居城であった甲斐武田氏の居城である',
             '武田信玄は、武田家の滅亡を防ぐため、武田家の家臣である武田信虎を討',
         ]
         self.assertListEqual(__lowerCamelCase					,       __lowerCamelCase	)
         self.assertListEqual(__lowerCamelCase					,       [outa_sentence, outa_sentence]	)

    @slow
    def 						_a							(  self	):




         '''simple docstring'''



         __A					:					List[str]								=		GPTSanJapaneseForConditionalGeneration.from_pretrained('Tanrei/GPTSAN-japanese'	)
         __A					:					List[str]								=		GPTSanJapaneseTokenizer.from_pretrained('Tanrei/GPTSAN-japanese'	)
         model.to(__lowerCamelCase	)

         # set deterministically
         __A					:					Optional[Any]								=		GenerationConfig.from_pretrained('Tanrei/GPTSAN-japanese'	)
         __A					:					Union[str, Any]								=		1

         __A					:					Any								=		'武田信玄'
         __A					:					List[str]								=		'織田信長'
         __A					:					List[Any]								=		'は、'
         __A					:					Any								=		'が、'
         __A					:					Any								=		tokenizer(__lowerCamelCase					,       prefix_text=__lowerCamelCase					,       return_tensors='pt'	)
         __A					:					str								=		tokenizer(__lowerCamelCase					,       prefix_text=__lowerCamelCase					,       return_tensors='pt'	)
         __A					:					List[str]								=		tokenizer([[prefix_text_a, input_text_a], [prefix_text_a, input_text_a]]					,       return_tensors='pt'	)

         __A					:					Union[str, Any]								=		model.generate(
             input_ids=input_tok_a.input_ids.to(__lowerCamelCase	)					,       token_type_ids=input_tok_a.token_type_ids.to(__lowerCamelCase	)					,       max_new_tokens=20					,       generation_config=__lowerCamelCase					,       )

         __A					:					Any								=		model.generate(
             input_ids=input_tok_a.input_ids.to(__lowerCamelCase	)					,       token_type_ids=input_tok_a.token_type_ids.to(__lowerCamelCase	)					,       max_new_tokens=20					,       generation_config=__lowerCamelCase					,       )

         __A					:					Dict								=		model.generate(
             input_ids=input_tok_a.input_ids.to(__lowerCamelCase	)					,       token_type_ids=input_tok_a.token_type_ids.to(__lowerCamelCase	)					,       attention_mask=input_tok_a.attention_mask.to(__lowerCamelCase	)					,       max_new_tokens=20					,       generation_config=__lowerCamelCase					,       )

         __A					:					Union[str, Any]								=		tokenizer.decode(outputa[0]	)
         __A					:					Any								=		tokenizer.decode(outputa[0]	)
         __A					:					Any								=		tokenizer.batch_decode(__lowerCamelCase	)

         __A					:					str								=		[
             '武田信玄は、武田氏の祖である武田信虎を、その子・武田信友を擁して',
             '織田信長が、織田信長の妻・お市の方を妻として迎えたという逸話が残',
         ]
         self.assertListEqual(__lowerCamelCase					,       __lowerCamelCase	)
         self.assertListEqual(__lowerCamelCase					,       [outa_sentence, outa_sentence]	)



