




'''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					:     Dict						,		a					:     str						,		a					:     Dict=14						,		a					:     Optional[int]=7						,		a					:     str=True						,		a					:     int=True						,		a					:     str=True						,		a					:     List[str]=99						,		a					:     List[str]=32						,		a					:     int=5						,		a					:     Optional[int]=4						,		a					:     Any=37						,		a					:     Dict="gelu"						,		a					:     Any=0.1						,		a					:     List[Any]=0.1						,		a					:     List[str]=5_12						,		a					:     Optional[int]=0.02						,		a					:     str=None						,		)					->      Union[str, Any]:






    """simple docstring"""


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






 def 	_UpperCamelCase   (self					:     Any						)					->      Optional[int]:






    """simple docstring"""


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

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

    A__         =    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 	_UpperCamelCase   (self					:     List[Any]						)					->      Optional[Any]:






    """simple docstring"""


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






 def 	_UpperCamelCase   (self					:     Optional[int]						)					->      Tuple:






    """simple docstring"""


    A__     , A__     , A__         =    self.prepare_config_and_inputs()

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

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






 def 	_UpperCamelCase   (self					:     List[str]						,		a					:     Optional[int]						,		a					:     int						,		a					:     Optional[Any]						,		a					:     Tuple						)					->      Tuple:






    """simple docstring"""


    A__         =    20
    A__         =    model_class_name(a						)

    A__         =    model.init_cache(input_ids.shape[0]						,		a						)
    A__         =    jnp.ones((input_ids.shape[0], max_decoder_length)						,		dtype='i4'						)

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

    A__         =    jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,		dtype='i4'						)
    A__         =    model(
        input_ids[:, -1:]						,		attention_mask=a						,		past_key_values=outputs_cache.past_key_values						,		position_ids=a						,		)

    A__         =    model(a						)

    A__         =    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					:     Dict						,		a					:     Dict						,		a					:     List[str]						,		a					:     str						,		a					:     Optional[Any]						)					->      Tuple:






    """simple docstring"""


    A__         =    20
    A__         =    model_class_name(a						)

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

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

    A__         =    model(
        input_ids[:, :-1]						,		attention_mask=a						,		past_key_values=a						,		position_ids=a						,		)
    A__         =    jnp.array(input_ids.shape[0] * [[input_ids.shape[-1] - 1]]						,		dtype='i4'						)
    A__         =    model(
        input_ids[:, -1:]						,		past_key_values=outputs_cache.past_key_values						,		attention_mask=a						,		position_ids=a						,		)

    A__         =    model(a						,		attention_mask=a						)
    A__         =    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( UpperCAmelCase					,   UpperCAmelCase					,   unittest.TestCase			):
 __snake_case:       Any									=				(FlaxXGLMModel, FlaxXGLMForCausalLM) if is_flax_available() else ()
 __snake_case:       List[str]									=				(FlaxXGLMForCausalLM,) if is_flax_available() else ()






 def 	_UpperCamelCase   (self					:     str						)					->      int:






    """simple docstring"""


    A__         =    FlaxXGLMModelTester(self						)






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






    """simple docstring"""


    for model_class_name in self.all_model_classes:
       A__     , A__     , A__         =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.check_use_cache_forward(a						,		a						,		a						,		a						)






 def 	_UpperCamelCase   (self					:     List[str]						)					->      int:






    """simple docstring"""


    for model_class_name in self.all_model_classes:
       A__     , A__     , A__         =    self.model_tester.prepare_config_and_inputs()
       self.model_tester.check_use_cache_forward_with_attn_mask(
           a						,		a						,		a						,		a						)






 @slow
 def 	_UpperCamelCase   (self					:     Union[str, Any]						)					->      Optional[int]:






    """simple docstring"""


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

    A__         =    FlaxXGLMForCausalLM.from_pretrained('facebook/xglm-564M'						)
    A__         =    1
    A__         =    False

    A__         =    jax.jit(model.generate						)

    A__         =    jit_generate(inputs['input_ids']						,		attention_mask=inputs['attention_mask']						).sequences

    A__         =    tokenizer.batch_decode(a						,		skip_special_tokens=a						)

    A__         =    [
        '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 	_UpperCamelCase   (self					:     List[Any]						)					->      List[str]:






    """simple docstring"""


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

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

          A__     , A__         =    pt_inputs['input_ids'].shape
          A__         =    np.random.randint(0						,		seq_length - 1						,		size=(batch_size,)						)
          for batch_idx, start_index in enumerate(a						):
             A__         =    0
             A__         =    1
             A__         =    0
             A__         =    1
          A__         =    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.
          A__         =    False
          A__         =    model_class(a						,		dtype=jnp.floataa						)
          A__         =    convert_pytorch_state_dict_to_flax(pt_model.state_dict()						,		a						)
          A__         =    fx_state

          with torch.no_grad():
             A__         =    pt_model(**a						).to_tuple()

          A__         =    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						)
             A__         =    model_class.from_pretrained(a						,		from_pt=a						)

          A__         =    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					:     Dict						)					->      Any:






    """simple docstring"""


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

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

          A__         =    pt_model_class(a						).eval()
          A__         =    False
          A__         =    model_class(a						,		dtype=jnp.floataa						)

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

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

          with torch.no_grad():
             A__         =    pt_model(**a						).to_tuple()

          A__         =    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						)
             A__         =    pt_model_class.from_pretrained(a						,		from_flax=a						)

          with torch.no_grad():
             A__         =    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 	_UpperCamelCase   (self					:     List[str]						)					->      Union[str, Any]:






    """simple docstring"""


    for model_class_name in self.all_model_classes:
       A__         =    model_class_name.from_pretrained('facebook/xglm-564M'						)
       A__         =    model(np.ones((1, 1)						)						)
       self.assertIsNotNone(a						)








'''simple docstring'''







import itertools
import random
import unittest

import numpy as np

from transformers import EncodecFeatureExtractor
from transformers.testing_utils import require_torch
from transformers.utils.import_utils import is_torch_available

from ...test_sequence_feature_extraction_common import SequenceFeatureExtractionTestMixin


if is_torch_available():
    import torch


lowerCAmelCase_    	=							random.Random()







def 	_A			(   UpperCAmelCase			,UpperCAmelCase=1.0			,UpperCAmelCase=None			,UpperCAmelCase=None	):






   '''simple docstring'''







   if rng is None:
      A__         =    global_rng

   A__         =    []
   for batch_idx in range(shape[0]	):
      values.append([]	)
      for _ in range(shape[1]	):
         values[-1].append(rng.random() * scale	)

   return values

@require_torch
class   _snake_case( unittest.TestCase			):






 def __init__(self					:     int						,		a					:     str						,		a					:     Optional[int]=7						,		a					:     Optional[Any]=4_00						,		a					:     Any=20_00						,		a					:     Optional[Any]=1						,		a					:     Optional[int]=0.0						,		a					:     Dict=2_40_00						,		a					:     Tuple=True						,		)					->      Optional[Any]:






    """simple docstring"""


    A__         =    parent
    A__         =    batch_size
    A__         =    min_seq_length
    A__         =    max_seq_length
    A__         =    (self.max_seq_length - self.min_seq_length) // (self.batch_size - 1)
    A__         =    feature_size
    A__         =    padding_value
    A__         =    sampling_rate
    A__         =    return_attention_mask






 def 	_UpperCamelCase   (self					:     int						)					->      List[str]:






    """simple docstring"""


    return {
        "feature_size": self.feature_size,
        "padding_value": self.padding_value,
        "sampling_rate": self.sampling_rate,
        "return_attention_mask": self.return_attention_mask,
    }






 def 	_UpperCamelCase   (self					:     List[Any]						,		a					:     Any=False						,		a					:     Optional[Any]=False						)					->      Tuple:






    """simple docstring"""


    def _flatten(a					:     Optional[Any]						):
       return list(itertools.chain(*a						)						)

    if equal_length:
       A__         =    floats_list((self.batch_size, self.max_seq_length)						)
    else:
       # make sure that inputs increase in size
       A__         =    [
           _flatten(floats_list((x, self.feature_size)						)						)
           for x in range(self.min_seq_length						,		self.max_seq_length						,		self.seq_length_diff						)
       ]

    if numpify:
       A__         =    [np.asarray(a						) for x in audio_inputs]

    return audio_inputs




@require_torch
class   _snake_case( UpperCAmelCase					,   unittest.TestCase			):
 __snake_case:       Union[str, Any]									=				EncodecFeatureExtractor






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






    """simple docstring"""


    A__         =    EnCodecFeatureExtractionTester(self						)






 def 	_UpperCamelCase   (self					:     Dict						)					->      Optional[Any]:






    """simple docstring"""


    A__         =    self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()						)
    # create three inputs of length 800, 1000, and 1200
    A__         =    [floats_list((1, x)						)[0] for x in range(8_00						,		14_00						,		2_00						)]
    A__         =    [np.asarray(a						) for audio_input in audio_inputs]

    # Test not batched input
    A__         =    feat_extract(audio_inputs[0]						,		return_tensors='np'						).input_values
    A__         =    feat_extract(np_audio_inputs[0]						,		return_tensors='np'						).input_values
    self.assertTrue(np.allclose(a						,		a						,		atol=1e-3						)						)

    # Test batched
    A__         =    feat_extract(a						,		padding=a						,		return_tensors='np'						).input_values
    A__         =    feat_extract(a						,		padding=a						,		return_tensors='np'						).input_values
    for enc_seq_a, enc_seq_a in zip(a						,		a						):
       self.assertTrue(np.allclose(a						,		a						,		atol=1e-3						)						)






 def 	_UpperCamelCase   (self					:     Tuple						)					->      Any:






    """simple docstring"""


    A__         =    self.feature_extraction_class(**self.feat_extract_tester.prepare_feat_extract_dict()						)
    A__         =    np.random.rand(1_00						).astype(np.floataa						)
    A__         =    np_audio_inputs.tolist()

    for inputs in [py_audio_inputs, np_audio_inputs]:
       A__         =    feature_extractor.pad([{'input_values': inputs}]						,		return_tensors='np'						)
       self.assertTrue(np_processed.input_values.dtype == np.floataa						)
       A__         =    feature_extractor.pad([{'input_values': inputs}]						,		return_tensors='pt'						)
       self.assertTrue(pt_processed.input_values.dtype == torch.floataa						)






 def 	_UpperCamelCase   (self					:     Optional[int]						,		a					:     Optional[Any]						)					->      Union[str, Any]:






    """simple docstring"""


    from datasets import load_dataset

    A__         =    load_dataset('hf-internal-testing/librispeech_asr_dummy'						,		'clean'						,		split='validation'						)
    # automatic decoding with librispeech
    A__         =    ds.sort('id'						).select(range(a						)						)[:num_samples]['audio']

    return [x["array"] for x in audio_samples]






 def 	_UpperCamelCase   (self					:     Tuple						)					->      Tuple:






    """simple docstring"""


    A__         =    torch.tensor(
        [2.3804e-03, 2.0752e-03, 1.9836e-03, 2.1057e-03, 1.6174e-03,
         3.0518e-04, 9.1553e-05, 3.3569e-04, 9.7656e-04, 1.8311e-03,
         2.0142e-03, 2.1057e-03, 1.7395e-03, 4.5776e-04, -3.9673e-04,
         4.5776e-04, 1.0071e-03, 9.1553e-05, 4.8828e-04, 1.1597e-03,
         7.3242e-04, 9.4604e-04, 1.8005e-03, 1.8311e-03, 8.8501e-04,
         4.2725e-04, 4.8828e-04, 7.3242e-04, 1.0986e-03, 2.1057e-03]						)
    # fmt: on
    A__         =    self._load_datasamples(1						)
    A__         =    EncodecFeatureExtractor()
    A__         =    feature_extractor(a						,		return_tensors='pt'						).input_values
    self.assertEquals(input_values.shape						,		(1, 1, 9_36_80)						)
    self.assertTrue(torch.allclose(input_values[0, 0, :30]						,		a						,		atol=1e-6						)						)






 def 	_UpperCamelCase   (self					:     int						)					->      Optional[Any]:






    """simple docstring"""


    A__         =    torch.tensor(
        [2.3804e-03, 2.0752e-03, 1.9836e-03, 2.1057e-03, 1.6174e-03,
         3.0518e-04, 9.1553e-05, 3.3569e-04, 9.7656e-04, 1.8311e-03,
         2.0142e-03, 2.1057e-03, 1.7395e-03, 4.5776e-04, -3.9673e-04,
         4.5776e-04, 1.0071e-03, 9.1553e-05, 4.8828e-04, 1.1597e-03,
         7.3242e-04, 9.4604e-04, 1.8005e-03, 1.8311e-03, 8.8501e-04,
         4.2725e-04, 4.8828e-04, 7.3242e-04, 1.0986e-03, 2.1057e-03]						)
    # fmt: on
    A__         =    self._load_datasamples(1						)
    A__         =    [np.tile(input_audio[0][None]						,		reps=(2, 1)						)]
    input_audio[0][1] *= 0.5
    A__         =    EncodecFeatureExtractor(feature_size=2						)
    A__         =    feature_extractor(a						,		return_tensors='pt'						).input_values
    self.assertEquals(input_values.shape						,		(1, 2, 9_36_80)						)
    self.assertTrue(torch.allclose(input_values[0, 0, :30]						,		a						,		atol=1e-6						)						)
    self.assertTrue(torch.allclose(input_values[0, 1, :30]						,		EXPECTED_INPUT_VALUES * 0.5						,		atol=1e-6						)						)






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






    """simple docstring"""


    A__         =    self._load_datasamples(2						)
    # would be easier if the stride was like
    A__         =    EncodecFeatureExtractor(feature_size=1						,		chunk_length_s=1						,		overlap=0.01						)

    # pad and trunc raise an error ?
    with self.assertRaisesRegex(
        a						,		'^Both padding and truncation were set. Make sure you only set one.$'						,		):
       A__         =    feature_extractor(
           a						,		padding='max_length'						,		truncation=a						,		return_tensors='pt'						).input_values

    # truncate to chunk
    A__         =    feature_extractor(a						,		truncation=a						,		return_tensors='pt'						).input_values
    self.assertEquals(truncated_outputs.shape						,		(2, 1, 7_15_20)						)  # 2 chunks

    # force truncate to max_length
    A__         =    feature_extractor(
        a						,		truncation=a						,		max_length=4_80_00						,		return_tensors='pt'						).input_values
    self.assertEquals(truncated_outputs.shape						,		(2, 1, 4_80_00)						)

    # pad to chunk
    A__         =    feature_extractor(a						,		padding=a						,		return_tensors='pt'						).input_values
    self.assertEquals(padded_outputs.shape						,		(2, 1, 9_52_80)						)

    # pad to chunk
    A__         =    feature_extractor(a						,		return_tensors='pt'						).input_values
    self.assertEquals(truncated_outputs.shape						,		(2, 1, 9_52_80)						)

    # force pad to max length
    A__         =    feature_extractor(
        a						,		padding='max_length'						,		max_length=10_00_00						,		return_tensors='pt'						).input_values
    self.assertEquals(truncated_outputs.shape						,		(2, 1, 10_00_00)						)

    # force no pad
    with self.assertRaisesRegex(
        a						,		'^Unable to create tensor, you should probably activate padding with \'padding=True\' to have batched tensors with the same length.$'						,		):
       A__         =    feature_extractor(a						,		padding=a						,		return_tensors='pt'						).input_values

    A__         =    feature_extractor(input_audio[0]						,		padding=a						,		return_tensors='pt'						).input_values
    self.assertEquals(truncated_outputs.shape						,		(1, 1, 9_36_80)						)

    # no pad if no chunk_length_s
    A__         =    None
    with self.assertRaisesRegex(
        a						,		'^Unable to create tensor, you should probably activate padding with \'padding=True\' to have batched tensors with the same length.$'						,		):
       A__         =    feature_extractor(a						,		padding=a						,		return_tensors='pt'						).input_values

    A__         =    feature_extractor(input_audio[0]						,		padding=a						,		return_tensors='pt'						).input_values
    self.assertEquals(truncated_outputs.shape						,		(1, 1, 9_36_80)						)

    # no pad if no overlap
    A__         =    2
    A__         =    None
    with self.assertRaisesRegex(
        a						,		'^Unable to create tensor, you should probably activate padding with \'padding=True\' to have batched tensors with the same length.$'						,		):
       A__         =    feature_extractor(a						,		padding=a						,		return_tensors='pt'						).input_values

    A__         =    feature_extractor(input_audio[0]						,		padding=a						,		return_tensors='pt'						).input_values
    self.assertEquals(truncated_outputs.shape						,		(1, 1, 9_36_80)						)


