




'''simple docstring'''





import tempfile
import unittest

import transformers
from transformers import XGLMConfig, XGLMTokenizer, is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax, require_sentencepiece, slow

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


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

		from transformers.modeling_flax_pytorch_utils import (
		    convert_pytorch_state_dict_to_flax,
		    load_flax_weights_in_pytorch_model,
		)
		from transformers.models.xglm.modeling_flax_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel


if is_torch_available():
		import torch




@require_flax
class 				_snake_case	:




			def __init__(      self       ,							a__       ,							a__=14       ,							a__=7       ,							a__=True       ,							a__=True       ,							a__=True       ,							a__=99       ,							a__=32       ,							a__=5       ,							a__=4       ,							a__=37       ,							a__="gelu"       ,							a__=0.1       ,							a__=0.1       ,							a__=512       ,							a__=0.0_2       ,							a__=None       ,							)       ->       List[Any]:

										'''simple docstring'''




										snake_case_			      =   parent
										snake_case_			      =   batch_size
										snake_case_			      =   seq_length
										snake_case_			      =   is_training
										snake_case_			      =   use_input_mask
										snake_case_			      =   use_labels
										snake_case_			      =   vocab_size
										snake_case_			      =   d_model
										snake_case_			      =   num_hidden_layers
										snake_case_			      =   num_attention_heads
										snake_case_			      =   ffn_dim
										snake_case_			      =   activation_function
										snake_case_			      =   activation_dropout
										snake_case_			      =   attention_dropout
										snake_case_			      =   max_position_embeddings
										snake_case_			      =   initializer_range
										snake_case_			      =   None
										snake_case_			      =   0
										snake_case_			      =   2
										snake_case_			      =   1




			def      lowerCAmelCase__       (      self							)       ->       Any:

										'''simple docstring'''




										snake_case_			      =   np.clip(ids_tensor([self.batch_size, self.seq_length]       ,							self.vocab_size							)       ,							3       ,							self.vocab_size							)

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

										snake_case_			      =   XGLMConfig(
										    vocab_size=self.vocab_size       ,							d_model=self.hidden_size       ,							num_layers=self.num_hidden_layers       ,							attention_heads=self.num_attention_heads       ,							ffn_dim=self.ffn_dim       ,							activation_function=self.activation_function       ,							activation_dropout=self.activation_dropout       ,							attention_dropout=self.attention_dropout       ,							max_position_embeddings=self.max_position_embeddings       ,							initializer_range=self.initializer_range       ,							use_cache=a__       ,							bos_token_id=self.bos_token_id       ,							eos_token_id=self.eos_token_id       ,							pad_token_id=self.pad_token_id       ,							)

										return (config, input_ids, input_mask)




			def      lowerCAmelCase__       (      self							)       ->       int:

										'''simple docstring'''




										snake_case_			      =   self.prepare_config_and_inputs()
										snake_case_      ,   snake_case_      ,   snake_case_			      =   config_and_inputs
										snake_case_			      =   {"input_ids": input_ids, "attention_mask": attention_mask}
										return config, inputs_dict




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

										'''simple docstring'''




										snake_case_      ,   snake_case_      ,   snake_case_			      =   self.prepare_config_and_inputs()

										snake_case_			      =   floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
										snake_case_			      =   ids_tensor([self.batch_size, self.seq_length]       ,							vocab_size=2							)

										return (
										    config,
										    input_ids,
										    attention_mask,
										    encoder_hidden_states,
										    encoder_attention_mask,
										)




			def      lowerCAmelCase__       (      self       ,							a__       ,							a__       ,							a__       ,							a__							)       ->       Optional[int]:

										'''simple docstring'''




										snake_case_			      =   20
										snake_case_			      =   model_class_name(a__							)

										snake_case_			      =   model.init_cache(input_ids.shape[0]       ,							a__							)
										snake_case_			      =   jnp.ones((input_ids.shape[0], max_decoder_length)       ,							dtype="i4"							)

										snake_case_			      =   jnp.broadcast_to(
										    jnp.arange(input_ids.shape[-1] - 1							)[None, :]       ,							(input_ids.shape[0], input_ids.shape[-1] - 1)							)
										snake_case_			      =   model(
										    input_ids[:, :-1]       ,							attention_mask=a__       ,							past_key_values=a__       ,							position_ids=a__       ,							)

										snake_case_			      =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,							dtype="i4"							)
										snake_case_			      =   model(
										    input_ids[:, -1:]       ,							attention_mask=a__       ,							past_key_values=outputs_cache.past_key_values       ,							position_ids=a__       ,							)

										snake_case_			      =   model(a__							)

										snake_case_			      =   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      lowerCAmelCase__       (      self       ,							a__       ,							a__       ,							a__       ,							a__							)       ->       Optional[int]:

										'''simple docstring'''




										snake_case_			      =   20
										snake_case_			      =   model_class_name(a__							)

										snake_case_			      =   jnp.concatenate(
										    [attention_mask, jnp.zeros((attention_mask.shape[0], max_decoder_length - attention_mask.shape[1])							)]       ,							axis=-1       ,							)

										snake_case_			      =   model.init_cache(input_ids.shape[0]       ,							a__							)
										snake_case_			      =   jnp.broadcast_to(
										    jnp.arange(input_ids.shape[-1] - 1							)[None, :]       ,							(input_ids.shape[0], input_ids.shape[-1] - 1)							)

										snake_case_			      =   model(
										    input_ids[:, :-1]       ,							attention_mask=a__       ,							past_key_values=a__       ,							position_ids=a__       ,							)
										snake_case_			      =   jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]       ,							dtype="i4"							)
										snake_case_			      =   model(
										    input_ids[:, -1:]       ,							past_key_values=outputs_cache.past_key_values       ,							attention_mask=a__       ,							position_ids=a__       ,							)

										snake_case_			      =   model(a__       ,							attention_mask=a__							)
										snake_case_			      =   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_sentencepiece
@require_flax
class 				_snake_case	(			lowercase_		,		lowercase_		,		unittest.TestCase							):
			lowerCAmelCase_  :   Union[str, Any]		 =					(FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
			lowerCAmelCase_  :   Dict		 =					(FlaxXGLMForCausalLM,) if is_flax_available() else ()




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

										'''simple docstring'''




										snake_case_			      =   FlaxXGLMModelTester(self							)




			def      lowerCAmelCase__       (      self							)       ->       Tuple:

										'''simple docstring'''




										for model_class_name in self.all_model_classes:
																	snake_case_      ,   snake_case_      ,   snake_case_			      =   self.model_tester.prepare_config_and_inputs()
																	self.model_tester.check_use_cache_forward(a__       ,							a__       ,							a__       ,							a__							)




			def      lowerCAmelCase__       (      self							)       ->       Any:

										'''simple docstring'''




										for model_class_name in self.all_model_classes:
																	snake_case_      ,   snake_case_      ,   snake_case_			      =   self.model_tester.prepare_config_and_inputs()
																	self.model_tester.check_use_cache_forward_with_attn_mask(
																	    a__       ,							a__       ,							a__       ,							a__							)




			@slow
			def      lowerCAmelCase__       (      self							)       ->       Dict:

										'''simple docstring'''




										snake_case_			      =   XGLMTokenizer.from_pretrained("XGLM"       ,							padding_side="left"							)
										snake_case_			      =   tokenizer(["Hello this is a long string", "Hey"]       ,							return_tensors="np"       ,							padding=a__       ,							truncation=a__							)

										snake_case_			      =   FlaxXGLMForCausalLM.from_pretrained("facebook/xglm-564M"							)
										snake_case_			      =   1
										snake_case_			      =   False

										snake_case_			      =   jax.jit(model.generate							)

										snake_case_			      =   jit_generate(inputs["input_ids"]       ,							attention_mask=inputs["attention_mask"]							).sequences

										snake_case_			      =   tokenizer.batch_decode(a__       ,							skip_special_tokens=a__							)

										snake_case_			      =   [
										    "Hello this is a long string of questions, but I'm not sure if I'm",
										    "Hey, I'm a newbie to the forum and I'",
										]

										self.assertListEqual(a__       ,							a__							)




			@is_pt_flax_cross_test
			def      lowerCAmelCase__       (      self							)       ->       Tuple:

										'''simple docstring'''




										snake_case_      ,   snake_case_			      =   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
																								snake_case_			      =   self._prepare_for_class(a__       ,							a__							)
																								snake_case_			      =   {k: torch.tensor(v.tolist()							) for k, v in prepared_inputs_dict.items()}

																								# load corresponding PyTorch class
																								snake_case_			      =   model_class.__name__[4:]  # Skip the "Flax" at the beginning
																								snake_case_			      =   getattr(a__       ,							a__							)

																								snake_case_      ,   snake_case_			      =   pt_inputs["input_ids"].shape
																								snake_case_			      =   np.random.randint(0       ,							seq_length - 1       ,							size=(batch_size,)							)
																								for batch_idx, start_index in enumerate(a__							):
																															snake_case_			      =   0
																															snake_case_			      =   1
																															snake_case_			      =   0
																															snake_case_			      =   1
																								snake_case_			      =   pt_model_class(a__							).eval()
																								# Flax models don't use the `use_cache` option and cache is not returned as a default.
																								# So we disable `use_cache` here for PyTorch model.
																								snake_case_			      =   False
																								snake_case_			      =   model_class(a__       ,							dtype=jnp.floataa							)
																								snake_case_			      =   convert_pytorch_state_dict_to_flax(pt_model.state_dict()       ,							a__							)
																								snake_case_			      =   fx_state

																								with torch.no_grad():
																															snake_case_			      =   pt_model(**a__							).to_tuple()

																								snake_case_			      =   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__							)
																															snake_case_			      =   model_class.from_pretrained(a__       ,							from_pt=a__							)

																								snake_case_			      =   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      lowerCAmelCase__       (      self							)       ->       Any:

										'''simple docstring'''




										snake_case_      ,   snake_case_			      =   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
																								snake_case_			      =   self._prepare_for_class(a__       ,							a__							)
																								snake_case_			      =   {k: torch.tensor(v.tolist()							) for k, v in prepared_inputs_dict.items()}

																								# load corresponding PyTorch class
																								snake_case_			      =   model_class.__name__[4:]  # Skip the "Flax" at the beginning
																								snake_case_			      =   getattr(a__       ,							a__							)

																								snake_case_			      =   pt_model_class(a__							).eval()
																								snake_case_			      =   False
																								snake_case_			      =   model_class(a__       ,							dtype=jnp.floataa							)

																								snake_case_			      =   load_flax_weights_in_pytorch_model(a__       ,							fx_model.params							)
																								snake_case_      ,   snake_case_			      =   pt_inputs["input_ids"].shape
																								snake_case_			      =   np.random.randint(0       ,							seq_length - 1       ,							size=(batch_size,)							)
																								for batch_idx, start_index in enumerate(a__							):
																															snake_case_			      =   0
																															snake_case_			      =   1
																															snake_case_			      =   0
																															snake_case_			      =   1

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

																								with torch.no_grad():
																															snake_case_			      =   pt_model(**a__							).to_tuple()

																								snake_case_			      =   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__							)
																															snake_case_			      =   pt_model_class.from_pretrained(a__       ,							from_flax=a__							)

																								with torch.no_grad():
																															snake_case_			      =   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							)




			@slow
			def      lowerCAmelCase__       (      self							)       ->       Optional[int]:

										'''simple docstring'''




										for model_class_name in self.all_model_classes:
																	snake_case_			      =   model_class_name.from_pretrained("facebook/xglm-564M"							)
																	snake_case_			      =   model(np.ones((1, 1)							)							)
																	self.assertIsNotNone(a__							)










'''simple docstring'''





import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
from unittest.mock import Mock, patch

import diffusers.utils.hub_utils




class 				_snake_case	(			unittest.TestCase							):




			@patch("diffusers.utils.hub_utils.get_full_repo_name"							)
			def      lowerCAmelCase__       (      self       ,							a__							)       ->       None:

										'''simple docstring'''




										snake_case_			      =   "full_repo_name"
										with TemporaryDirectory() as tmpdir:
																	# Dummy args values
																	snake_case_			      =   Mock()
																	snake_case_			      =   tmpdir
																	snake_case_			      =   0
																	snake_case_			      =   "hub_token"
																	snake_case_			      =   "dataset_name"
																	snake_case_			      =   0.0_1
																	snake_case_			      =   100_000
																	snake_case_			      =   10_000
																	snake_case_			      =   0.0_1
																	snake_case_			      =   0.0_2
																	snake_case_			      =   0.0_3
																	snake_case_			      =   0.0_0_0_5
																	snake_case_			      =   0.0_0_0_0_0_1
																	snake_case_			      =   1
																	snake_case_			      =   10
																	snake_case_			      =   0.0_0_1
																	snake_case_			      =   0.1
																	snake_case_			      =   0.2
																	snake_case_			      =   True

																	# Model card mush be rendered and saved
																	diffusers.utils.hub_utils.create_model_card(a__       ,							model_name="model_name"							)
																	self.assertTrue((Path(a__							) / "README.md").is_file()							)




