





import unittest

from transformers import GPTNeoXJapaneseConfig, is_torch_available
from transformers.models.gpt_neox_japanese.tokenization_gpt_neox_japanese import GPTNeoXJapaneseTokenizer
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import GPTNeoXJapaneseForCausalLM, GPTNeoXJapaneseModel



class     _UpperCAmelCase   :

  """simple docstring"""


  def __init__(    self			:		Dict ,  lowerCAmelCase_			:		Tuple ,  lowerCAmelCase_			:		Dict=1_3 ,  lowerCAmelCase_			:		Tuple=7 ,  lowerCAmelCase_			:		Optional[int]=True ,  lowerCAmelCase_			:		List[Any]=True ,  lowerCAmelCase_			:		int=True ,  lowerCAmelCase_			:		Any=True ,  lowerCAmelCase_			:		Union[str, Any]=9_9 ,  lowerCAmelCase_			:		Optional[Any]=3_2 ,  lowerCAmelCase_			:		Optional[Any]=5 ,  lowerCAmelCase_			:		str=4 ,  lowerCAmelCase_			:		Optional[Any]=4 ,  lowerCAmelCase_			:		Tuple="gelu" ,  lowerCAmelCase_			:		Dict=0.0 ,  lowerCAmelCase_			:		Optional[Any]=0.1 ,  lowerCAmelCase_			:		List[Any]=True ,  lowerCAmelCase_			:		Optional[Any]=5_1_2 ,  lowerCAmelCase_			:		Optional[Any]=1_6 ,  lowerCAmelCase_			:		Optional[Any]=2 ,  lowerCAmelCase_			:		Tuple=0.02 ,  lowerCAmelCase_			:		List[str]=3 ,  lowerCAmelCase_			:		Dict=4 ,  lowerCAmelCase_			:		List[Any]=None ,  )							->						List[str]:
     __lowerCAmelCase							     =     parent
     __lowerCAmelCase							     =     batch_size
     __lowerCAmelCase							     =     seq_length
     __lowerCAmelCase							     =     is_training
     __lowerCAmelCase							     =     use_input_mask
     __lowerCAmelCase							     =     use_token_type_ids
     __lowerCAmelCase							     =     use_labels
     __lowerCAmelCase							     =     vocab_size
     __lowerCAmelCase							     =     hidden_size
     __lowerCAmelCase							     =     num_hidden_layers
     __lowerCAmelCase							     =     num_attention_heads
     __lowerCAmelCase							     =     intermediate_multiple_size
     __lowerCAmelCase							     =     hidden_act
     __lowerCAmelCase							     =     hidden_dropout
     __lowerCAmelCase							     =     attention_dropout
     __lowerCAmelCase							     =     weight_tying
     __lowerCAmelCase							     =     max_position_embeddings
     __lowerCAmelCase							     =     type_vocab_size
     __lowerCAmelCase							     =     type_sequence_label_size
     __lowerCAmelCase							     =     initializer_range
     __lowerCAmelCase							     =     num_labels
     __lowerCAmelCase							     =     num_choices
     __lowerCAmelCase							     =     scope






  def     lowercase		(    self			:		Tuple					)							->						str:
     __lowerCAmelCase							     =     ids_tensor([self.batch_size, self.seq_length] ,  self.vocab_size					)

     __lowerCAmelCase							     =     None
     if self.use_input_mask:
        __lowerCAmelCase							     =     random_attention_mask([self.batch_size, self.seq_length]					)

     __lowerCAmelCase							     =     None
     if self.use_labels:
        __lowerCAmelCase							     =     ids_tensor([self.batch_size, self.seq_length] ,  self.num_labels					)

     __lowerCAmelCase							     =     self.get_config()

     return config, input_ids, input_mask, token_labels






  def     lowercase		(    self			:		Optional[int]					)							->						str:
     return GPTNeoXJapaneseConfig(
         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_multiple_size=self.intermediate_multiple_size ,  hidden_act=self.hidden_act ,  hidden_dropout=self.hidden_dropout ,  attention_dropout=self.attention_dropout ,  weight_tying=self.weight_tying ,  max_position_embeddings=self.max_position_embeddings ,  type_vocab_size=self.type_vocab_size ,  is_decoder=lowerCAmelCase_ ,  initializer_range=self.initializer_range ,  )






  def     lowercase		(    self			:		int					)							->						Optional[Any]:
     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase							     =     self.prepare_config_and_inputs()

     __lowerCAmelCase							     =     True

     return config, input_ids, input_mask, token_labels






  def     lowercase		(    self			:		Any ,  lowerCAmelCase_			:		str ,  lowerCAmelCase_			:		Dict ,  lowerCAmelCase_			:		Optional[int]					)							->						int:
     __lowerCAmelCase							     =     GPTNeoXJapaneseModel(config=lowerCAmelCase_					)
     model.to(lowerCAmelCase_					)
     model.eval()
     __lowerCAmelCase							     =     model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_					)
     __lowerCAmelCase							     =     model(lowerCAmelCase_					)
     self.parent.assertEqual(result.last_hidden_state.shape ,  (self.batch_size, self.seq_length, self.hidden_size)					)






  def     lowercase		(    self			:		int ,  lowerCAmelCase_			:		str ,  lowerCAmelCase_			:		Optional[Any] ,  lowerCAmelCase_			:		Tuple					)							->						Dict:
     __lowerCAmelCase							     =     True
     __lowerCAmelCase							     =     GPTNeoXJapaneseModel(lowerCAmelCase_					)
     model.to(lowerCAmelCase_					)
     model.eval()
     __lowerCAmelCase							     =     model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_					)
     self.parent.assertEqual(result.last_hidden_state.shape ,  (self.batch_size, self.seq_length, self.hidden_size)					)






  def     lowercase		(    self			:		Any ,  lowerCAmelCase_			:		Dict ,  lowerCAmelCase_			:		Optional[Any] ,  lowerCAmelCase_			:		int ,  lowerCAmelCase_			:		List[Any]					)							->						Dict:
     __lowerCAmelCase							     =     GPTNeoXJapaneseForCausalLM(config=lowerCAmelCase_					)
     model.to(lowerCAmelCase_					)
     model.eval()
     __lowerCAmelCase							     =     model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  labels=lowerCAmelCase_					)
     self.parent.assertEqual(result.logits.shape ,  (self.batch_size, self.seq_length, self.vocab_size)					)






  def     lowercase		(    self			:		List[str] ,  lowerCAmelCase_			:		Union[str, Any] ,  lowerCAmelCase_			:		str ,  lowerCAmelCase_			:		Optional[Any]					)							->						List[Any]:
     __lowerCAmelCase							     =     True
     __lowerCAmelCase							     =     GPTNeoXJapaneseForCausalLM(config=lowerCAmelCase_					)
     model.to(lowerCAmelCase_					)
     model.eval()

     # first forward pass
     __lowerCAmelCase							     =     model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  use_cache=lowerCAmelCase_					)
     __lowerCAmelCase							     =     outputs.past_key_values

     # create hypothetical multiple next token and extent to next_input_ids
     __lowerCAmelCase							     =     ids_tensor((self.batch_size, 3) ,  config.vocab_size					)
     __lowerCAmelCase							     =     ids_tensor((self.batch_size, 3) ,  vocab_size=2					)

     # append to next input_ids and
     __lowerCAmelCase							     =     torch.cat([input_ids, next_tokens] ,  dim=-1					)
     __lowerCAmelCase							     =     torch.cat([input_mask, next_mask] ,  dim=-1					)

     __lowerCAmelCase							     =     model(lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  output_hidden_states=lowerCAmelCase_					)
     __lowerCAmelCase							     =     output_from_no_past['hidden_states'][0]
     __lowerCAmelCase							     =     model(
         lowerCAmelCase_ ,  attention_mask=lowerCAmelCase_ ,  past_key_values=lowerCAmelCase_ ,  output_hidden_states=lowerCAmelCase_ ,  )['hidden_states'][0]

     # select random slice
     __lowerCAmelCase							     =     ids_tensor((1,) ,  output_from_past.shape[-1]					).item()
     __lowerCAmelCase							     =     output_from_no_past[:, -3:, random_slice_idx].detach()
     __lowerCAmelCase							     =     output_from_past[:, :, random_slice_idx].detach()

     self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]					)

     # test that outputs are equal for slice
     self.parent.assertTrue(torch.allclose(lowerCAmelCase_ ,  lowerCAmelCase_ ,  atol=1e-3					)					)






  def     lowercase		(    self			:		List[str]					)							->						Optional[Any]:
     __lowerCAmelCase							     =     self.prepare_config_and_inputs()
     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase							     =     config_and_inputs
     __lowerCAmelCase							     =     {'input_ids': input_ids, 'attention_mask': input_mask}
     return config, inputs_dict





@require_torch
class     _UpperCAmelCase   (   _UpperCamelCase		,   _UpperCamelCase		,   unittest.TestCase   ):

  """simple docstring"""


  a_      =			(GPTNeoXJapaneseModel, GPTNeoXJapaneseForCausalLM) if is_torch_available() else ()
  a_      =			(GPTNeoXJapaneseForCausalLM,) if is_torch_available() else ()
  a_      =			(
      {"""feature-extraction""": GPTNeoXJapaneseModel, """text-generation""": GPTNeoXJapaneseForCausalLM}
      if is_torch_available()
      else {}
  )
  a_      =			False
  a_      =			False
  a_      =			False
  a_      =			False






  def     lowercase		(    self			:		Optional[Any]					)							->						int:
     __lowerCAmelCase							     =     GPTNeoXJapaneseModelTester(self					)
     __lowerCAmelCase							     =     ConfigTester(self ,  config_class=lowerCAmelCase_ ,  hidden_size=3_7					)






  def     lowercase		(    self			:		List[str]					)							->						str:
     self.config_tester.run_common_tests()






  def     lowercase		(    self			:		Optional[Any]					)							->						Tuple:
     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase							     =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_					)






  def     lowercase		(    self			:		Optional[int]					)							->						Tuple:
     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase							     =     self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_model_as_decoder(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_					)






  def     lowercase		(    self			:		Optional[Any]					)							->						Optional[Any]:
     # This regression test was failing with PyTorch < 1.3
     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase							     =     self.model_tester.prepare_config_and_inputs_for_decoder()

     __lowerCAmelCase							     =     None

     self.model_tester.create_and_check_model_as_decoder(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_					)






  def     lowercase		(    self			:		Optional[Any]					)							->						str:
     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase     ,     __lowerCAmelCase							     =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_decoder_model_past_large_inputs(lowerCAmelCase_ ,  lowerCAmelCase_ ,  lowerCAmelCase_					)






  def     lowercase		(    self			:		Dict					)							->						Union[str, Any]:
     __lowerCAmelCase							     =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_for_causal_lm(*lowerCAmelCase_					)






  @slow
  def     lowercase		(    self			:		Union[str, Any]					)							->						Tuple:
     __lowerCAmelCase							     =     'abeja/gpt-neox-japanese-2.7b'

     __lowerCAmelCase							     =     ['データサイエンティストとは、', '100年後に必要とされる会社は、', 'フルリモートの環境で働くために必要なことは、', '国境の長いトンネルを抜けると', '美味しい日本食といえば、']

     __lowerCAmelCase							     =     [
         'データサイエンティストとは、データを分析し、ビジネスに役立つ知見を導き出す専門家のことです。',
         '100年後に必要とされる会社は、「人」が中心の会社です。',
         'フルリモートの環境で働くために必要なことは、「自分の時間をコントロールする」ことです。',
         '国境の長いトンネルを抜けると、そこは雪国だった。',
         '美味しい日本食といえば、やっぱりお寿司ですよね。',
     ]

     __lowerCAmelCase							     =     GPTNeoXJapaneseTokenizer.from_pretrained(lowerCAmelCase_					)
     __lowerCAmelCase							     =     GPTNeoXJapaneseForCausalLM.from_pretrained(lowerCAmelCase_					)

     __lowerCAmelCase							     =     []
     for prompt in prompts:
        __lowerCAmelCase							     =     tokenizer(lowerCAmelCase_ ,  return_tensors='pt'					).input_ids
        __lowerCAmelCase							     =     model.generate(lowerCAmelCase_ ,  max_length=5_0					)
        __lowerCAmelCase							     =     tokenizer.batch_decode(lowerCAmelCase_ ,  skip_special_tokens=lowerCAmelCase_					)
        predicted_outputs += generated_string
     self.assertListEqual(lowerCAmelCase_ ,  lowerCAmelCase_					)








from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import get_down_block, get_mid_block, get_out_block, get_up_block



@dataclass
class     _UpperCAmelCase   (   _UpperCamelCase   ):

  """simple docstring"""


  a_      =			42





class     _UpperCAmelCase   (   _UpperCamelCase		,   _UpperCamelCase   ):

  """simple docstring"""


  @register_to_config
  def __init__(    self			:		List[Any] ,  lowerCAmelCase_			:		int = 6_5_5_3_6 ,  lowerCAmelCase_			:		Optional[int] = None ,  lowerCAmelCase_			:		int = 2 ,  lowerCAmelCase_			:		int = 2 ,  lowerCAmelCase_			:		int = 0 ,  lowerCAmelCase_			:		str = "fourier" ,  lowerCAmelCase_			:		bool = True ,  lowerCAmelCase_			:		bool = False ,  lowerCAmelCase_			:		float = 0.0 ,  lowerCAmelCase_			:		Tuple[str] = ("DownBlock1DNoSkip", "DownBlock1D", "AttnDownBlock1D") ,  lowerCAmelCase_			:		Tuple[str] = ("AttnUpBlock1D", "UpBlock1D", "UpBlock1DNoSkip") ,  lowerCAmelCase_			:		Tuple[str] = "UNetMidBlock1D" ,  lowerCAmelCase_			:		str = None ,  lowerCAmelCase_			:		Tuple[int] = (3_2, 3_2, 6_4) ,  lowerCAmelCase_			:		str = None ,  lowerCAmelCase_			:		int = 8 ,  lowerCAmelCase_			:		int = 1 ,  lowerCAmelCase_			:		bool = False ,  )							->						Optional[int]:
     super().__init__()
     __lowerCAmelCase							     =     sample_size

     # time
     if time_embedding_type == "fourier":
        __lowerCAmelCase							     =     GaussianFourierProjection(
            embedding_size=8 ,  set_W_to_weight=lowerCAmelCase_ ,  log=lowerCAmelCase_ ,  flip_sin_to_cos=lowerCAmelCase_					)
        __lowerCAmelCase							     =     2 * block_out_channels[0]
     elif time_embedding_type == "positional":
        __lowerCAmelCase							     =     Timesteps(
            block_out_channels[0] ,  flip_sin_to_cos=lowerCAmelCase_ ,  downscale_freq_shift=lowerCAmelCase_					)
        __lowerCAmelCase							     =     block_out_channels[0]

     if use_timestep_embedding:
        __lowerCAmelCase							     =     block_out_channels[0] * 4
        __lowerCAmelCase							     =     TimestepEmbedding(
            in_channels=lowerCAmelCase_ ,  time_embed_dim=lowerCAmelCase_ ,  act_fn=lowerCAmelCase_ ,  out_dim=block_out_channels[0] ,  )

     __lowerCAmelCase							     =     nn.ModuleList([]					)
     __lowerCAmelCase							     =     None
     __lowerCAmelCase							     =     nn.ModuleList([]					)
     __lowerCAmelCase							     =     None

     # down
     __lowerCAmelCase							     =     in_channels
     for i, down_block_type in enumerate(lowerCAmelCase_					):
        __lowerCAmelCase							     =     output_channel
        __lowerCAmelCase							     =     block_out_channels[i]

        if i == 0:
           input_channel += extra_in_channels

        __lowerCAmelCase							     =     i == len(lowerCAmelCase_					) - 1

        __lowerCAmelCase							     =     get_down_block(
            lowerCAmelCase_ ,  num_layers=lowerCAmelCase_ ,  in_channels=lowerCAmelCase_ ,  out_channels=lowerCAmelCase_ ,  temb_channels=block_out_channels[0] ,  add_downsample=not is_final_block or downsample_each_block ,  )
        self.down_blocks.append(lowerCAmelCase_					)

     # mid
     __lowerCAmelCase							     =     get_mid_block(
         lowerCAmelCase_ ,  in_channels=block_out_channels[-1] ,  mid_channels=block_out_channels[-1] ,  out_channels=block_out_channels[-1] ,  embed_dim=block_out_channels[0] ,  num_layers=lowerCAmelCase_ ,  add_downsample=lowerCAmelCase_ ,  )

     # up
     __lowerCAmelCase							     =     list(reversed(lowerCAmelCase_					)					)
     __lowerCAmelCase							     =     reversed_block_out_channels[0]
     if out_block_type is None:
        __lowerCAmelCase							     =     out_channels
     else:
        __lowerCAmelCase							     =     block_out_channels[0]

     for i, up_block_type in enumerate(lowerCAmelCase_					):
        __lowerCAmelCase							     =     output_channel
        __lowerCAmelCase							     =     (
            reversed_block_out_channels[i + 1] if i < len(lowerCAmelCase_					) - 1 else final_upsample_channels
        )

        __lowerCAmelCase							     =     i == len(lowerCAmelCase_					) - 1

        __lowerCAmelCase							     =     get_up_block(
            lowerCAmelCase_ ,  num_layers=lowerCAmelCase_ ,  in_channels=lowerCAmelCase_ ,  out_channels=lowerCAmelCase_ ,  temb_channels=block_out_channels[0] ,  add_upsample=not is_final_block ,  )
        self.up_blocks.append(lowerCAmelCase_					)
        __lowerCAmelCase							     =     output_channel

     # out
     __lowerCAmelCase							     =     norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4 ,  3_2					)
     __lowerCAmelCase							     =     get_out_block(
         out_block_type=lowerCAmelCase_ ,  num_groups_out=lowerCAmelCase_ ,  embed_dim=block_out_channels[0] ,  out_channels=lowerCAmelCase_ ,  act_fn=lowerCAmelCase_ ,  fc_dim=block_out_channels[-1] // 4 ,  )






  def     lowercase		(    self			:		Optional[Any] ,  lowerCAmelCase_			:		torch.FloatTensor ,  lowerCAmelCase_			:		Union[torch.Tensor, float, int] ,  lowerCAmelCase_			:		bool = True ,  )							->						Union[UNetaDOutput, Tuple]:
     __lowerCAmelCase							     =     timestep
     if not torch.is_tensor(lowerCAmelCase_					):
        __lowerCAmelCase							     =     torch.tensor([timesteps] ,  dtype=torch.long ,  device=sample.device					)
     elif torch.is_tensor(lowerCAmelCase_					) and len(timesteps.shape					) == 0:
        __lowerCAmelCase							     =     timesteps[None].to(sample.device					)

     __lowerCAmelCase							     =     self.time_proj(lowerCAmelCase_					)
     if self.config.use_timestep_embedding:
        __lowerCAmelCase							     =     self.time_mlp(lowerCAmelCase_					)
     else:
        __lowerCAmelCase							     =     timestep_embed[..., None]
        __lowerCAmelCase							     =     timestep_embed.repeat([1, 1, sample.shape[2]]					).to(sample.dtype					)
        __lowerCAmelCase							     =     timestep_embed.broadcast_to((sample.shape[:1] + timestep_embed.shape[1:])					)

     # 2. down
     __lowerCAmelCase							     =     ()
     for downsample_block in self.down_blocks:
        __lowerCAmelCase     ,     __lowerCAmelCase							     =     downsample_block(hidden_states=lowerCAmelCase_ ,  temb=lowerCAmelCase_					)
        down_block_res_samples += res_samples

     # 3. mid
     if self.mid_block:
        __lowerCAmelCase							     =     self.mid_block(lowerCAmelCase_ ,  lowerCAmelCase_					)

     # 4. up
     for i, upsample_block in enumerate(self.up_blocks					):
        __lowerCAmelCase							     =     down_block_res_samples[-1:]
        __lowerCAmelCase							     =     down_block_res_samples[:-1]
        __lowerCAmelCase							     =     upsample_block(lowerCAmelCase_ ,  res_hidden_states_tuple=lowerCAmelCase_ ,  temb=lowerCAmelCase_					)

     # 5. post-process
     if self.out_block:
        __lowerCAmelCase							     =     self.out_block(lowerCAmelCase_ ,  lowerCAmelCase_					)

     if not return_dict:
        return (sample,)

     return UNetaDOutput(sample=lowerCAmelCase_					)

