

import tempfile
import unittest

import numpy as np

import transformers
from transformers import GPTaTokenizer, GPTJConfig, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, tooslow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
					import jax
					import jax.numpy as jnp

					from transformers.modeling_flax_pytorch_utils import (
					    convert_pytorch_state_dict_to_flax,
					    load_flax_weights_in_pytorch_model,
					)
					from transformers.models.gptj.modeling_flax_gptj import FlaxGPTJForCausalLM, FlaxGPTJModel

if is_torch_available():
					import torch

class 				a  :


						def __init__(       self						,					A_						,					A_=14						,					A_=7						,					A_=True						,					A_=True						,					A_=False						,					A_=True						,					A_=99						,					A_=32						,					A_=4						,					A_=4						,					A_=4						,					A_=37						,					A_="gelu"						,					A_=0.1						,					A_=0.1						,					A_=512						,					A_=0.02						,					):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    parent
										_UpperCAmelCase       :		Optional[int] 		=    batch_size
										_UpperCAmelCase       :		Tuple 		=    seq_length
										_UpperCAmelCase       :		Optional[Any] 		=    is_training
										_UpperCAmelCase       :		str 		=    use_input_mask
										_UpperCAmelCase       :		Optional[Any] 		=    use_token_type_ids
										_UpperCAmelCase       :		str 		=    use_labels
										_UpperCAmelCase       :		Dict 		=    vocab_size
										_UpperCAmelCase       :		str 		=    hidden_size
										_UpperCAmelCase       :		List[Any] 		=    rotary_dim
										_UpperCAmelCase       :		Optional[Any] 		=    num_hidden_layers
										_UpperCAmelCase       :		Optional[Any] 		=    num_attention_heads
										_UpperCAmelCase       :		Any 		=    intermediate_size
										_UpperCAmelCase       :		List[str] 		=    hidden_act
										_UpperCAmelCase       :		Optional[int] 		=    hidden_dropout_prob
										_UpperCAmelCase       :		int 		=    attention_probs_dropout_prob
										_UpperCAmelCase       :		Tuple 		=    max_position_embeddings
										_UpperCAmelCase       :		Dict 		=    initializer_range
										_UpperCAmelCase       :		Tuple 		=    None
										_UpperCAmelCase       :		Optional[int] 		=    vocab_size - 1
										_UpperCAmelCase       :		Optional[int] 		=    vocab_size - 1
										_UpperCAmelCase       :		Tuple 		=    vocab_size - 1


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


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

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

										_UpperCAmelCase       :		List[Any] 		=    GPTJConfig(
										    vocab_size=self.vocab_size						,					n_embd=self.hidden_size						,					n_layer=self.num_hidden_layers						,					n_head=self.num_attention_heads						,					n_positions=self.max_position_embeddings						,					use_cache=A_						,					bos_token_id=self.bos_token_id						,					eos_token_id=self.eos_token_id						,					pad_token_id=self.pad_token_id						,					rotary_dim=self.rotary_dim						,					)

										return (config, input_ids, input_mask)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[Any] 		=    self.prepare_config_and_inputs()
										_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase       :		List[str] 		=    config_and_inputs
										_UpperCAmelCase       :		Union[str, Any] 		=    {"input_ids": input_ids, "attention_mask": attention_mask}
										return config, inputs_dict


						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    20
										_UpperCAmelCase       :		Dict 		=    model_class_name(A_	)

										_UpperCAmelCase       :		str 		=    model.init_cache(input_ids.shape[0]						,					A_	)
										_UpperCAmelCase       :		Optional[Any] 		=    jnp.ones((input_ids.shape[0], max_decoder_length)						,					dtype="i4"	)

										_UpperCAmelCase       :		Dict 		=    jnp.broadcast_to(
										    jnp.arange(input_ids.shape[-1] - 1	)[None, :]						,					(input_ids.shape[0], input_ids.shape[-1] - 1)	)
										_UpperCAmelCase       :		Optional[Any] 		=    model(
										    input_ids[:, :-1]						,					attention_mask=A_						,					past_key_values=A_						,					position_ids=A_						,					)

										_UpperCAmelCase       :		int 		=    jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,					dtype="i4"	)
										_UpperCAmelCase       :		Optional[Any] 		=    model(
										    input_ids[:, -1:]						,					attention_mask=A_						,					past_key_values=outputs_cache.past_key_values						,					position_ids=A_						,					)

										_UpperCAmelCase       :		Optional[Any] 		=    model(A_	)

										_UpperCAmelCase       :		Any 		=    np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
										self.parent.assertTrue(diff < 1e-3						,					msg=f'Max diff is {diff}'	)





						def        _UpperCAmelCase				(       self						,					A_						,					A_						,					A_						,					A_	):







										'''simple docstring'''


										_UpperCAmelCase       :		Union[str, Any] 		=    20
										_UpperCAmelCase       :		Dict 		=    model_class_name(A_	)

										_UpperCAmelCase       :		List[Any] 		=    jnp.concatenate(
										    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])	)]						,					axis=-1						,					)

										_UpperCAmelCase       :		Optional[int] 		=    model.init_cache(input_ids.shape[0]						,					A_	)
										_UpperCAmelCase       :		str 		=    jnp.broadcast_to(
										    jnp.arange(input_ids.shape[-1] - 1	)[None, :]						,					(input_ids.shape[0], input_ids.shape[-1] - 1)	)

										_UpperCAmelCase       :		Any 		=    model(
										    input_ids[:, :-1]						,					attention_mask=A_						,					past_key_values=A_						,					position_ids=A_						,					)
										_UpperCAmelCase       :		Tuple 		=    jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,					dtype="i4"	)
										_UpperCAmelCase       :		Dict 		=    model(
										    input_ids[:, -1:]						,					past_key_values=outputs_cache.past_key_values						,					attention_mask=A_						,					position_ids=A_						,					)

										_UpperCAmelCase       :		Any 		=    model(A_						,					attention_mask=A_	)

										_UpperCAmelCase       :		List[Any] 		=    np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
										self.parent.assertTrue(diff < 1e-3						,					msg=f'Max diff is {diff}'	)







@require_flax
class 				a  (	UpperCAmelCase							, UpperCAmelCase							, unittest.TestCase				):
						_lowercase    						=			(FlaxGPTJModel, FlaxGPTJForCausalLM) if is_flax_available() else ()
						_lowercase    						=			(FlaxGPTJForCausalLM,) if is_flax_available() else ()


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[Any] 		=    FlaxGPTJModelTester(self	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										for model_class_name in self.all_model_classes:
														_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase       :		Optional[Any] 		=    self.model_tester.prepare_config_and_inputs()
														self.model_tester.check_use_cache_forward(A_						,					A_						,					A_						,					A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										for model_class_name in self.all_model_classes:
														_UpperCAmelCase		,							_UpperCAmelCase		,							_UpperCAmelCase       :		str 		=    self.model_tester.prepare_config_and_inputs()
														self.model_tester.check_use_cache_forward_with_attn_mask(
														    A_						,					A_						,					A_						,					A_	)


						@tooslow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    GPTaTokenizer.from_pretrained("gpt2"						,					pad_token="<|endoftext|>"						,					padding_side="left"	)
										_UpperCAmelCase       :		Optional[int] 		=    tokenizer(["Hello this is a long string", "Hey"]						,					return_tensors="np"						,					padding=A_						,					truncation=A_	)

										_UpperCAmelCase       :		Optional[Any] 		=    FlaxGPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B"	)
										_UpperCAmelCase       :		Optional[int] 		=    False
										_UpperCAmelCase       :		List[Any] 		=    model.config.eos_token_id

										_UpperCAmelCase       :		Tuple 		=    jax.jit(model.generate	)

										_UpperCAmelCase       :		Any 		=    jit_generate(
										    inputs["input_ids"]						,					attention_mask=inputs["attention_mask"]						,					pad_token_id=tokenizer.pad_token_id	).sequences

										_UpperCAmelCase       :		Tuple 		=    tokenizer.batch_decode(A_						,					skip_special_tokens=A_	)

										_UpperCAmelCase       :		List[Any] 		=    [
										    "Hello this is a long string of text.\n\nI'm trying to get the text of the",
										    "Hey, I'm a little late to the party. I'm going to",
										]

										self.assertListEqual(A_						,					A_	)


						@is_pt_flax_cross_test
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase		,							_UpperCAmelCase       :		List[Any] 		=    self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
														with self.subTest(model_class.__name__	):
																		# prepare inputs
																		_UpperCAmelCase       :		Dict 		=    self._prepare_for_class(A_						,					A_	)
																		_UpperCAmelCase       :		Union[str, Any] 		=    {k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

																		# load corresponding PyTorch class
																		_UpperCAmelCase       :		Optional[Any] 		=    model_class.__name__[4:]  # Skip the "Flax" at the beginning
																		_UpperCAmelCase       :		Optional[int] 		=    getattr(A_						,					A_	)

																		_UpperCAmelCase		,							_UpperCAmelCase       :		List[Any] 		=    pt_inputs["input_ids"].shape
																		_UpperCAmelCase       :		Tuple 		=    np.random.randint(0						,					seq_length - 1						,					size=(batch_size,)	)
																		for batch_idx, start_index in enumerate(A_	):
																						_UpperCAmelCase       :		int 		=    0
																						_UpperCAmelCase       :		Optional[Any] 		=    1
																						_UpperCAmelCase       :		Union[str, Any] 		=    0
																						_UpperCAmelCase       :		Optional[int] 		=    1
																		_UpperCAmelCase       :		Any 		=    pt_model_class(A_	).eval()
																		_UpperCAmelCase       :		Tuple 		=    model_class(A_						,					dtype=jnp.floataa	)

																		_UpperCAmelCase       :		Dict 		=    convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,					A_	)
																		_UpperCAmelCase       :		Any 		=    fx_state

																		with torch.no_grad():
																						_UpperCAmelCase       :		str 		=    pt_model(**A_	).to_tuple()

																		_UpperCAmelCase       :		List[str] 		=    fx_model(**A_	).to_tuple()
																		self.assertEqual(len(A_	)						,					len(A_	)						,					"Output lengths differ between Flax and PyTorch"	)
																		for fx_output, pt_output in zip(A_						,					A_	):
																						self.assert_almost_equals(fx_output[:, -1]						,					pt_output[:, -1].numpy()						,					4e-2	)

																		with tempfile.TemporaryDirectory() as tmpdirname:
																						pt_model.save_pretrained(A_	)
																						_UpperCAmelCase       :		List[str] 		=    model_class.from_pretrained(A_						,					from_pt=A_	)

																		_UpperCAmelCase       :		Union[str, Any] 		=    fx_model_loaded(**A_	).to_tuple()
																		self.assertEqual(
																		    len(A_	)						,					len(A_	)						,					"Output lengths differ between Flax and PyTorch"	)
																		for fx_output_loaded, pt_output in zip(A_						,					A_	):
																						self.assert_almost_equals(fx_output_loaded[:, -1]						,					pt_output[:, -1].numpy()						,					4e-2	)


						@is_pt_flax_cross_test
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase		,							_UpperCAmelCase       :		List[str] 		=    self.model_tester.prepare_config_and_inputs_for_common()
										for model_class in self.all_model_classes:
														with self.subTest(model_class.__name__	):
																		# prepare inputs
																		_UpperCAmelCase       :		Tuple 		=    self._prepare_for_class(A_						,					A_	)
																		_UpperCAmelCase       :		str 		=    {k: torch.tensor(v.tolist()	) for k, v in prepared_inputs_dict.items()}

																		# load corresponding PyTorch class
																		_UpperCAmelCase       :		Optional[Any] 		=    model_class.__name__[4:]  # Skip the "Flax" at the beginning
																		_UpperCAmelCase       :		str 		=    getattr(A_						,					A_	)

																		_UpperCAmelCase       :		Optional[Any] 		=    pt_model_class(A_	).eval()
																		_UpperCAmelCase       :		Dict 		=    model_class(A_						,					dtype=jnp.floataa	)

																		_UpperCAmelCase       :		List[str] 		=    load_flax_weights_in_pytorch_model(A_						,					fx_model.params	)
																		_UpperCAmelCase		,							_UpperCAmelCase       :		int 		=    pt_inputs["input_ids"].shape
																		_UpperCAmelCase       :		Optional[int] 		=    np.random.randint(0						,					seq_length - 1						,					size=(batch_size,)	)
																		for batch_idx, start_index in enumerate(A_	):
																						_UpperCAmelCase       :		Dict 		=    0
																						_UpperCAmelCase       :		Optional[Any] 		=    1
																						_UpperCAmelCase       :		List[Any] 		=    0
																						_UpperCAmelCase       :		Dict 		=    1

																		# make sure weights are tied in PyTorch
																		pt_model.tie_weights()

																		with torch.no_grad():
																						_UpperCAmelCase       :		Tuple 		=    pt_model(**A_	).to_tuple()

																		_UpperCAmelCase       :		List[str] 		=    fx_model(**A_	).to_tuple()
																		self.assertEqual(len(A_	)						,					len(A_	)						,					"Output lengths differ between Flax and PyTorch"	)
																		for fx_output, pt_output in zip(A_						,					A_	):
																						self.assert_almost_equals(fx_output[:, -1]						,					pt_output[:, -1].numpy()						,					4e-2	)

																		with tempfile.TemporaryDirectory() as tmpdirname:
																						fx_model.save_pretrained(A_	)
																						_UpperCAmelCase       :		List[Any] 		=    pt_model_class.from_pretrained(A_						,					from_flax=A_	)

																		with torch.no_grad():
																						_UpperCAmelCase       :		Any 		=    pt_model_loaded(**A_	).to_tuple()

																		self.assertEqual(
																		    len(A_	)						,					len(A_	)						,					"Output lengths differ between Flax and PyTorch"	)
																		for fx_output, pt_output in zip(A_						,					A_	):
																						self.assert_almost_equals(fx_output[:, -1]						,					pt_output[:, -1].numpy()						,					4e-2	)





						@tooslow
						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										for model_class_name in self.all_model_classes:
														_UpperCAmelCase       :		Optional[int] 		=    model_class_name.from_pretrained("EleutherAI/gpt-j-6B"	)
														_UpperCAmelCase       :		Dict 		=    model(np.ones((1, 1)	)	)
														self.assertIsNotNone(A_	)





import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


SCREAMING_SNAKE_CASE_        =      '\nHugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.\n\nIn March 2021, Hugging Face raised $40 million in a Series B funding round.[3]\n\nOn April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]\n'

class 				a  (	unittest.TestCase							, UpperCAmelCase				):


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Tuple 		=    load_tool("text-question-answering"	)
										self.tool.setup()
										_UpperCAmelCase       :		str 		=    load_tool("text-question-answering"						,					remote=A_	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[str] 		=    self.tool(A_						,					"What did Hugging Face do in April 2021?"	)
										self.assertEqual(A_						,					"launched the BigScience Research Workshop"	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		Optional[Any] 		=    self.remote_tool(A_						,					"What did Hugging Face do in April 2021?"	)
										self.assertEqual(A_						,					"launched the BigScience Research Workshop"	)


						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		int 		=    self.tool(text=A_						,					question="What did Hugging Face do in April 2021?"	)
										self.assertEqual(A_						,					"launched the BigScience Research Workshop"	)





						def        _UpperCAmelCase				(       self	):







										'''simple docstring'''


										_UpperCAmelCase       :		List[Any] 		=    self.remote_tool(text=A_						,					question="What did Hugging Face do in April 2021?"	)
										self.assertEqual(A_						,					"launched the BigScience Research Workshop"	)


