
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 			lowerCAmelCase_		:


				"""simple docstring"""



				def __init__(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__=13     , SCREAMING_SNAKE_CASE__=7     , SCREAMING_SNAKE_CASE__=True     , SCREAMING_SNAKE_CASE__=True     , SCREAMING_SNAKE_CASE__=True     , SCREAMING_SNAKE_CASE__=True     , SCREAMING_SNAKE_CASE__=99     , SCREAMING_SNAKE_CASE__=32     , SCREAMING_SNAKE_CASE__=5     , SCREAMING_SNAKE_CASE__=4     , SCREAMING_SNAKE_CASE__=4     , SCREAMING_SNAKE_CASE__="gelu"     , SCREAMING_SNAKE_CASE__=0.0     , SCREAMING_SNAKE_CASE__=0.1     , SCREAMING_SNAKE_CASE__=True     , SCREAMING_SNAKE_CASE__=5_12     , SCREAMING_SNAKE_CASE__=16     , SCREAMING_SNAKE_CASE__=2     , SCREAMING_SNAKE_CASE__=0.02     , SCREAMING_SNAKE_CASE__=3     , SCREAMING_SNAKE_CASE__=4     , SCREAMING_SNAKE_CASE__=None     , )      ->	List[str]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: int     =     parent
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     batch_size
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     seq_length
											SCREAMING_SNAKE_CASE__		: str     =     is_training
											SCREAMING_SNAKE_CASE__		: Tuple     =     use_input_mask
											SCREAMING_SNAKE_CASE__		: Any     =     use_token_type_ids
											SCREAMING_SNAKE_CASE__		: List[str]     =     use_labels
											SCREAMING_SNAKE_CASE__		: List[Any]     =     vocab_size
											SCREAMING_SNAKE_CASE__		: Tuple     =     hidden_size
											SCREAMING_SNAKE_CASE__		: Dict     =     num_hidden_layers
											SCREAMING_SNAKE_CASE__		: Tuple     =     num_attention_heads
											SCREAMING_SNAKE_CASE__		: List[str]     =     intermediate_multiple_size
											SCREAMING_SNAKE_CASE__		: Tuple     =     hidden_act
											SCREAMING_SNAKE_CASE__		: int     =     hidden_dropout
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     attention_dropout
											SCREAMING_SNAKE_CASE__		: int     =     weight_tying
											SCREAMING_SNAKE_CASE__		: str     =     max_position_embeddings
											SCREAMING_SNAKE_CASE__		: List[str]     =     type_vocab_size
											SCREAMING_SNAKE_CASE__		: Dict     =     type_sequence_label_size
											SCREAMING_SNAKE_CASE__		: Tuple     =     initializer_range
											SCREAMING_SNAKE_CASE__		: Tuple     =     num_labels
											SCREAMING_SNAKE_CASE__		: List[str]     =     num_choices
											SCREAMING_SNAKE_CASE__		: int     =     scope



				def 			__magic_name__					(self					)      ->	Union[str, Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: List[Any]     =     ids_tensor([self.batch_size, self.seq_length]     , self.vocab_size					)

											SCREAMING_SNAKE_CASE__		: int     =     None
											if self.use_input_mask:
																		SCREAMING_SNAKE_CASE__		: str     =     random_attention_mask([self.batch_size, self.seq_length]					)

											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     None
											if self.use_labels:
																		SCREAMING_SNAKE_CASE__		: Dict     =     ids_tensor([self.batch_size, self.seq_length]     , self.num_labels					)

											SCREAMING_SNAKE_CASE__		: Any     =     self.get_config()

											return config, input_ids, input_mask, token_labels



				def 			__magic_name__					(self					)      ->	str:



											"""simple docstring"""




											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=SCREAMING_SNAKE_CASE__     , initializer_range=self.initializer_range     , )



				def 			__magic_name__					(self					)      ->	Any:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     self.prepare_config_and_inputs()

											SCREAMING_SNAKE_CASE__		: List[str]     =     True

											return config, input_ids, input_mask, token_labels



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)      ->	Dict:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Tuple     =     GPTNeoXJapaneseModel(config=SCREAMING_SNAKE_CASE__					)
											model.to(SCREAMING_SNAKE_CASE__					)
											model.eval()
											SCREAMING_SNAKE_CASE__		: List[Any]     =     model(SCREAMING_SNAKE_CASE__     , attention_mask=SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: str     =     model(SCREAMING_SNAKE_CASE__					)
											self.parent.assertEqual(result.last_hidden_state.shape     , (self.batch_size, self.seq_length, self.hidden_size)					)



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)      ->	Optional[int]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     True
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     GPTNeoXJapaneseModel(SCREAMING_SNAKE_CASE__					)
											model.to(SCREAMING_SNAKE_CASE__					)
											model.eval()
											SCREAMING_SNAKE_CASE__		: List[Any]     =     model(SCREAMING_SNAKE_CASE__     , attention_mask=SCREAMING_SNAKE_CASE__					)
											self.parent.assertEqual(result.last_hidden_state.shape     , (self.batch_size, self.seq_length, self.hidden_size)					)



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)      ->	Dict:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     GPTNeoXJapaneseForCausalLM(config=SCREAMING_SNAKE_CASE__					)
											model.to(SCREAMING_SNAKE_CASE__					)
											model.eval()
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     model(SCREAMING_SNAKE_CASE__     , attention_mask=SCREAMING_SNAKE_CASE__     , labels=SCREAMING_SNAKE_CASE__					)
											self.parent.assertEqual(result.logits.shape     , (self.batch_size, self.seq_length, self.vocab_size)					)



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)      ->	Optional[int]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Dict     =     True
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     GPTNeoXJapaneseForCausalLM(config=SCREAMING_SNAKE_CASE__					)
											model.to(SCREAMING_SNAKE_CASE__					)
											model.eval()

											# first forward pass
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     model(SCREAMING_SNAKE_CASE__     , attention_mask=SCREAMING_SNAKE_CASE__     , use_cache=SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: List[str]     =     outputs.past_key_values

											# create hypothetical multiple next token and extent to next_input_ids
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     ids_tensor((self.batch_size, 3)     , config.vocab_size					)
											SCREAMING_SNAKE_CASE__		: Any     =     ids_tensor((self.batch_size, 3)     , vocab_size=2					)

											# append to next input_ids and
											SCREAMING_SNAKE_CASE__		: List[Any]     =     torch.cat([input_ids, next_tokens]     , dim=-1					)
											SCREAMING_SNAKE_CASE__		: Dict     =     torch.cat([input_mask, next_mask]     , dim=-1					)

											SCREAMING_SNAKE_CASE__		: Tuple     =     model(SCREAMING_SNAKE_CASE__     , attention_mask=SCREAMING_SNAKE_CASE__     , output_hidden_states=SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: List[str]     =     output_from_no_past["""hidden_states"""][0]
											SCREAMING_SNAKE_CASE__		: Dict     =     model(
											    SCREAMING_SNAKE_CASE__     , attention_mask=SCREAMING_SNAKE_CASE__     , past_key_values=SCREAMING_SNAKE_CASE__     , output_hidden_states=SCREAMING_SNAKE_CASE__     , )["""hidden_states"""][0]

											# select random slice
											SCREAMING_SNAKE_CASE__		: int     =     ids_tensor((1,)     , output_from_past.shape[-1]					).item()
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     output_from_no_past[:, -3:, random_slice_idx].detach()
											SCREAMING_SNAKE_CASE__		: List[Any]     =     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(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , atol=1E-3					)					)



				def 			__magic_name__					(self					)      ->	str:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Optional[int]     =     self.prepare_config_and_inputs()
											SCREAMING_SNAKE_CASE__		: Optional[int]     =     config_and_inputs
											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     {"""input_ids""": input_ids, """attention_mask""": input_mask}
											return config, inputs_dict






@require_torch
class 			lowerCAmelCase_		(a__				,      a__				,      unittest.TestCase						):


				"""simple docstring"""



				__UpperCamelCase   :		Union[str, Any]       =     (GPTNeoXJapaneseModel, GPTNeoXJapaneseForCausalLM) if is_torch_available() else ()
				__UpperCamelCase   :		Optional[Any]       =     (GPTNeoXJapaneseForCausalLM,) if is_torch_available() else ()
				__UpperCamelCase   :		List[Any]       =     (
				    {'''feature-extraction''': GPTNeoXJapaneseModel, '''text-generation''': GPTNeoXJapaneseForCausalLM}
				    if is_torch_available()
				    else {}
				)
				__UpperCamelCase   :		Union[str, Any]       =     False
				__UpperCamelCase   :		List[str]       =     False
				__UpperCamelCase   :		int       =     False
				__UpperCamelCase   :		List[Any]       =     False



				def 			__magic_name__					(self					)      ->	Dict:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: List[Any]     =     GPTNeoXJapaneseModelTester(self					)
											SCREAMING_SNAKE_CASE__		: Dict     =     ConfigTester(self     , config_class=SCREAMING_SNAKE_CASE__     , hidden_size=37					)



				def 			__magic_name__					(self					)      ->	Tuple:



											"""simple docstring"""




											self.config_tester.run_common_tests()



				def 			__magic_name__					(self					)      ->	str:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: str     =     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)



				def 			__magic_name__					(self					)      ->	List[Any]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     self.model_tester.prepare_config_and_inputs_for_decoder()
											self.model_tester.create_and_check_model_as_decoder(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)



				def 			__magic_name__					(self					)      ->	Any:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Tuple     =     self.model_tester.prepare_config_and_inputs_for_decoder()

											SCREAMING_SNAKE_CASE__		: List[Any]     =     None

											self.model_tester.create_and_check_model_as_decoder(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)



				def 			__magic_name__					(self					)      ->	str:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_decoder_model_past_large_inputs(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)



				def 			__magic_name__					(self					)      ->	Optional[int]:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: List[Any]     =     self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_causal_lm(*SCREAMING_SNAKE_CASE__					)



				@slow
				def 			__magic_name__					(self					)      ->	Dict:



											"""simple docstring"""




											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     """abeja/gpt-neox-japanese-2.7b"""

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

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

											SCREAMING_SNAKE_CASE__		: Union[str, Any]     =     GPTNeoXJapaneseTokenizer.from_pretrained(SCREAMING_SNAKE_CASE__					)
											SCREAMING_SNAKE_CASE__		: str     =     GPTNeoXJapaneseForCausalLM.from_pretrained(SCREAMING_SNAKE_CASE__					)

											SCREAMING_SNAKE_CASE__		: Any     =     []
											for prompt in prompts:
																		SCREAMING_SNAKE_CASE__		: List[Any]     =     tokenizer(SCREAMING_SNAKE_CASE__     , return_tensors="""pt"""					).input_ids
																		SCREAMING_SNAKE_CASE__		: List[str]     =     model.generate(SCREAMING_SNAKE_CASE__     , max_length=50					)
																		SCREAMING_SNAKE_CASE__		: Optional[int]     =     tokenizer.batch_decode(SCREAMING_SNAKE_CASE__     , skip_special_tokens=SCREAMING_SNAKE_CASE__					)
																		predicted_outputs += generated_string
											self.assertListEqual(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					)






"""simple docstring"""




import itertools
from dataclasses import dataclass
from typing import Optional

import pandas as pd
import pyarrow as pa

import datasets
from datasets.table import table_cast



@dataclass
class 			lowerCAmelCase_		(datasets.BuilderConfig						):


				"""simple docstring"""



				__UpperCamelCase   :		Optional[datasets.Features]       =     None






class 			lowerCAmelCase_		(datasets.ArrowBasedBuilder						):


				"""simple docstring"""



				__UpperCamelCase   :		Union[str, Any]       =     PandasConfig



				def 			__magic_name__					(self					)      ->	List[str]:



											"""simple docstring"""




											return datasets.DatasetInfo(features=self.config.features					)



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__					)      ->	Optional[int]:



											"""simple docstring"""




											if not self.config.data_files:
																		raise ValueError(F'''At least one data file must be specified, but got data_files={self.config.data_files}'''					)
											SCREAMING_SNAKE_CASE__		: Optional[Any]     =     dl_manager.download_and_extract(self.config.data_files					)
											if isinstance(SCREAMING_SNAKE_CASE__     , (str, list, tuple)					):
																		SCREAMING_SNAKE_CASE__		: Tuple     =     data_files
																		if isinstance(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					):
																									SCREAMING_SNAKE_CASE__		: Optional[Any]     =     [files]
																		# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
																		SCREAMING_SNAKE_CASE__		: str     =     [dl_manager.iter_files(SCREAMING_SNAKE_CASE__					) for file in files]
																		return [datasets.SplitGenerator(name=datasets.Split.TRAIN     , gen_kwargs={"""files""": files}					)]
											SCREAMING_SNAKE_CASE__		: int     =     []
											for split_name, files in data_files.items():
																		if isinstance(SCREAMING_SNAKE_CASE__     , SCREAMING_SNAKE_CASE__					):
																									SCREAMING_SNAKE_CASE__		: List[str]     =     [files]
																		# Use `dl_manager.iter_files` to skip hidden files in an extracted archive
																		SCREAMING_SNAKE_CASE__		: Dict     =     [dl_manager.iter_files(SCREAMING_SNAKE_CASE__					) for file in files]
																		splits.append(datasets.SplitGenerator(name=SCREAMING_SNAKE_CASE__     , gen_kwargs={"""files""": files}					)					)
											return splits



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__					)      ->	pa.Table:



											"""simple docstring"""




											if self.config.features is not None:
																		# more expensive cast to support nested features with keys in a different order
																		# allows str <-> int/float or str to Audio for example
																		SCREAMING_SNAKE_CASE__		: Tuple     =     table_cast(SCREAMING_SNAKE_CASE__     , self.config.features.arrow_schema					)
											return pa_table



				def 			__magic_name__					(self     , SCREAMING_SNAKE_CASE__					)      ->	Optional[Any]:



											"""simple docstring"""




											for i, file in enumerate(itertools.chain.from_iterable(SCREAMING_SNAKE_CASE__					)					):
																		with open(SCREAMING_SNAKE_CASE__     , """rb"""					) as f:
																									SCREAMING_SNAKE_CASE__		: List[str]     =     pa.Table.from_pandas(pd.read_pickle(SCREAMING_SNAKE_CASE__					)					)
																									yield i, self._cast_table(SCREAMING_SNAKE_CASE__					)




