

"""simple docstring"""




from __future__ import annotations

import unittest

from transformers import XGLMConfig, XGLMTokenizer, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
    import tensorflow as tf

    from transformers.models.xglm.modeling_tf_xglm import (
        TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFXGLMForCausalLM,
        TFXGLMModel,
    )







@require_tf
class      __lowercase			:
  snake_case_				       =       XGLMConfig
  snake_case_				       =       {}
  snake_case_				       =       """gelu"""

  def __init__(  self		:			Dict					,A		:			List[str]					,A		:			str=14					,A		:			str=7					,A		:			Any=True					,A		:			Tuple=True					,A		:			List[Any]=True					,A		:			Optional[int]=99					,A		:			List[str]=32					,A		:			Optional[int]=2					,A		:			str=4					,A		:			Optional[Any]=37					,A		:			List[Any]="gelu"					,A		:			int=0.1					,A		:			Union[str, Any]=0.1					,A		:			str=512					,A		:			Any=0.0_2					,):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[int]					   =					parent
   UpperCAmelCase__					:  str					   =					batch_size
   UpperCAmelCase__					:  Tuple					   =					seq_length
   UpperCAmelCase__					:  List[str]					   =					is_training
   UpperCAmelCase__					:  Dict					   =					use_input_mask
   UpperCAmelCase__					:  int					   =					use_labels
   UpperCAmelCase__					:  Tuple					   =					vocab_size
   UpperCAmelCase__					:  List[str]					   =					d_model
   UpperCAmelCase__					:  int					   =					num_hidden_layers
   UpperCAmelCase__					:  List[Any]					   =					num_attention_heads
   UpperCAmelCase__					:  List[str]					   =					ffn_dim
   UpperCAmelCase__					:  Any					   =					activation_function
   UpperCAmelCase__					:  str					   =					activation_dropout
   UpperCAmelCase__					:  int					   =					attention_dropout
   UpperCAmelCase__					:  List[Any]					   =					max_position_embeddings
   UpperCAmelCase__					:  Any					   =					initializer_range
   UpperCAmelCase__					:  List[str]					   =					None
   UpperCAmelCase__					:  List[str]					   =					0
   UpperCAmelCase__					:  int					   =					2
   UpperCAmelCase__					:  Tuple					   =					1

  def    __lowercase       (  self		:			List[str]    ):






   '''simple docstring'''
   return XGLMConfig.from_pretrained("""facebook/xglm-564M"""    )

  def    __lowercase       (  self		:			List[str]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					tf.clip_by_value(
       ids_tensor([self.batch_size, self.seq_length]					,self.vocab_size    )					,clip_value_min=0					,clip_value_max=3    )

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

   UpperCAmelCase__					:  Optional[Any]					   =					self.get_config()

   UpperCAmelCase__					:  Tuple					   =					floats_tensor([self.num_hidden_layers, self.num_attention_heads]					,2    )

   return (
       config,
       input_ids,
       input_mask,
       head_mask,
   )

  def    __lowercase       (  self		:			Union[str, Any]    ):






   '''simple docstring'''
   return 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_dict=A					,)




  def    __lowercase       (  self		:			str    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					self.prepare_config_and_inputs()

   (
       (
       UpperCAmelCase__		
   )   ,      (
       UpperCAmelCase__		
   )   ,      (
       UpperCAmelCase__		
   )   ,      (
       UpperCAmelCase__		
   )   ,      		
   )					:  Any					   =					config_and_inputs

   UpperCAmelCase__					:  int					   =					{
       """input_ids""": input_ids,
       """head_mask""": head_mask,
   }

   return config, inputs_dict







@require_tf
class      __lowercase			(    __UpperCAmelCase	,			__UpperCAmelCase	,			unittest.TestCase					):
  snake_case_				       =       (TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
  snake_case_				       =       (TFXGLMForCausalLM,) if is_tf_available() else ()
  snake_case_				       =       (
      {"""feature-extraction""": TFXGLMModel, """text-generation""": TFXGLMForCausalLM} if is_tf_available() else {}
  )
  snake_case_				       =       False
  snake_case_				       =       False
  snake_case_				       =       False

  def    __lowercase       (  self		:			Optional[int]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Optional[int]					   =					TFXGLMModelTester(self    )
   UpperCAmelCase__					:  str					   =					ConfigTester(self					,config_class=A					,n_embd=37    )

  def    __lowercase       (  self		:			Tuple    ):






   '''simple docstring'''
   self.config_tester.run_common_tests()

  @slow
  def    __lowercase       (  self		:			Tuple    ):






   '''simple docstring'''

   for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
    UpperCAmelCase__					:  Any					   =					TFXGLMModel.from_pretrained(A    )
    self.assertIsNotNone(A    )




  @unittest.skip(reason="""Currently, model embeddings are going to undergo a major refactor."""    )
  def    __lowercase       (  self		:			Union[str, Any]    ):






   '''simple docstring'''
   super().test_resize_token_embeddings()







@require_tf
class      __lowercase			(    unittest.TestCase					):

  @slow
  def    __lowercase       (  self		:			Dict					,A		:			Any=True    ):






   '''simple docstring'''
   UpperCAmelCase__					:  Dict					   =					TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""    )
   UpperCAmelCase__					:  Any					   =					tf.convert_to_tensor([[2, 268, 9_865]]					,dtype=tf.intaa    )  # The dog
   # </s> The dog is a very friendly dog. He is very affectionate and loves to play with other
   # fmt: off
   UpperCAmelCase__					:  Dict					   =					[2, 268, 9_865, 67, 11, 1_988, 57_252, 9_865, 5, 984, 67, 1_988, 213_838, 1_658, 53, 70_446, 33, 6_657, 278, 1_581]
   # fmt: on
   UpperCAmelCase__					:  Optional[Any]					   =					model.generate(A					,do_sample=A					,num_beams=1    )
   if verify_outputs:
    self.assertListEqual(output_ids[0].numpy().tolist()					,A    )

  @slow
  def    __lowercase       (  self		:			str    ):






   '''simple docstring'''
   UpperCAmelCase__					:  str					   =					XGLMTokenizer.from_pretrained("""facebook/xglm-564M"""    )
   UpperCAmelCase__					:  Tuple					   =					TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""    )

   tf.random.set_seed(0    )
   UpperCAmelCase__					:  int					   =					tokenizer("""Today is a nice day and"""					,return_tensors="""tf"""    )
   UpperCAmelCase__					:  Union[str, Any]					   =					tokenized.input_ids
   # forces the generation to happen on CPU, to avoid GPU-related quirks (and assure same output regardless of the available devices)
   with tf.device(""":/CPU:0"""    ):
    UpperCAmelCase__					:  List[Any]					   =					model.generate(A					,do_sample=A					,seed=[7, 0]    )
   UpperCAmelCase__					:  int					   =					tokenizer.decode(output_ids[0]					,skip_special_tokens=A    )

   UpperCAmelCase__					:  Optional[int]					   =					(
       """Today is a nice day and warm evening here over Southern Alberta!! Today when they closed schools due"""
   )
   self.assertEqual(A					,A    )




  @slow
  def    __lowercase       (  self		:			Optional[int]    ):






   '''simple docstring'''
   UpperCAmelCase__					:  List[Any]					   =					TFXGLMForCausalLM.from_pretrained("""facebook/xglm-564M"""    )
   UpperCAmelCase__					:  List[Any]					   =					XGLMTokenizer.from_pretrained("""facebook/xglm-564M"""    )

   UpperCAmelCase__					:  Dict					   =					"""left"""

   # use different length sentences to test batching
   UpperCAmelCase__					:  Optional[Any]					   =					[
       """This is an extremelly long sentence that only exists to test the ability of the model to cope with """
       """left-padding, such as in batched generation. The output for the sequence below should be the same """
       """regardless of whether left padding is applied or not. When""",
       """Hello, my dog is a little""",
   ]

   UpperCAmelCase__					:  List[str]					   =					tokenizer(A					,return_tensors="""tf"""					,padding=A    )
   UpperCAmelCase__					:  List[str]					   =					inputs["""input_ids"""]

   UpperCAmelCase__					:  Optional[Any]					   =					model.generate(input_ids=A					,attention_mask=inputs["""attention_mask"""]					,max_new_tokens=12    )

   UpperCAmelCase__					:  Optional[int]					   =					tokenizer(sentences[0]					,return_tensors="""tf"""    ).input_ids
   UpperCAmelCase__					:  Optional[Any]					   =					model.generate(input_ids=A					,max_new_tokens=12    )

   UpperCAmelCase__					:  int					   =					tokenizer(sentences[1]					,return_tensors="""tf"""    ).input_ids
   UpperCAmelCase__					:  List[Any]					   =					model.generate(input_ids=A					,max_new_tokens=12    )

   UpperCAmelCase__					:  Optional[int]					   =					tokenizer.batch_decode(A					,skip_special_tokens=A    )
   UpperCAmelCase__					:  str					   =					tokenizer.decode(output_non_padded[0]					,skip_special_tokens=A    )
   UpperCAmelCase__					:  str					   =					tokenizer.decode(output_padded[0]					,skip_special_tokens=A    )

   UpperCAmelCase__					:  List[str]					   =					[
       """This is an extremelly long sentence that only exists to test the ability of the model to cope with """
       """left-padding, such as in batched generation. The output for the sequence below should be the same """
       """regardless of whether left padding is applied or not. When left padding is applied, the sequence will be """
       """a single""",
       """Hello, my dog is a little bit of a shy one, but he is very friendly""",
   ]
   self.assertListEqual(A					,A    )
   self.assertListEqual(A					,[non_padded_sentence, padded_sentence]    )








"""simple docstring"""




from typing import List, Optional, TypeVar

from .arrow_dataset import Dataset, _concatenate_map_style_datasets, _interleave_map_style_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .info import DatasetInfo
from .iterable_dataset import IterableDataset, _concatenate_iterable_datasets, _interleave_iterable_datasets
from .splits import NamedSplit
from .utils import logging
from .utils.py_utils import Literal


__UpperCAmelCase =				logging.get_logger(__name__)


__UpperCAmelCase =				TypeVar('DatasetType', Dataset, IterableDataset)
def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase = None							,    __UpperCamelCase = None							,    __UpperCamelCase = None							,    __UpperCamelCase = None							,    __UpperCamelCase = "first_exhausted"							,    ):
 '''simple docstring'''





 from .arrow_dataset import Dataset
 from .iterable_dataset import IterableDataset

 if not datasets:
  raise ValueError("""Unable to interleave an empty list of datasets."""      )
 for i, dataset in enumerate(__UpperCamelCase      ):
  if not isinstance(__UpperCamelCase							,    (Dataset, IterableDataset)      ):
   if isinstance(__UpperCamelCase							,    (DatasetDict, IterableDatasetDict)      ):
    if not dataset:
     raise ValueError(
         F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
         """is an empty dataset dictionary."""      )
    raise ValueError(
        F"Dataset at position {i} has at least one split: {list(__UpperCamelCase      )}\n"
        F"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(__UpperCamelCase      )      )}']"      )
   raise ValueError(
       F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(__UpperCamelCase      ).__name__}."      )
  if i == 0:
   UpperCAmelCase__   ,      UpperCAmelCase__					:  List[str]					   =					(
       (Dataset, IterableDataset) if isinstance(__UpperCamelCase							,    __UpperCamelCase      ) else (IterableDataset, Dataset)
   )
  elif not isinstance(__UpperCamelCase							,    __UpperCamelCase      ):
   raise ValueError(
       F"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."      )
 if stopping_strategy not in ["first_exhausted", "all_exhausted"]:
  raise ValueError(F"{stopping_strategy} is not supported. Please enter a valid stopping_strategy."      )
 if dataset_type is Dataset:
  return _interleave_map_style_datasets(
      __UpperCamelCase							,    __UpperCamelCase							,    __UpperCamelCase							,    info=__UpperCamelCase							,    split=__UpperCamelCase							,    stopping_strategy=__UpperCamelCase      )
 else:
  return _interleave_iterable_datasets(
      __UpperCamelCase							,    __UpperCamelCase							,    __UpperCamelCase							,    info=__UpperCamelCase							,    split=__UpperCamelCase							,    stopping_strategy=__UpperCamelCase      )




def   lowerCAmelCase						(					__UpperCamelCase							,    __UpperCamelCase = None							,    __UpperCamelCase = None							,    __UpperCamelCase = 0							,    ):
 '''simple docstring'''





 if not dsets:
  raise ValueError("""Unable to concatenate an empty list of datasets."""      )
 for i, dataset in enumerate(__UpperCamelCase      ):
  if not isinstance(__UpperCamelCase							,    (Dataset, IterableDataset)      ):
   if isinstance(__UpperCamelCase							,    (DatasetDict, IterableDatasetDict)      ):
    if not dataset:
     raise ValueError(
         F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} "
         """is an empty dataset dictionary."""      )
    raise ValueError(
        F"Dataset at position {i} has at least one split: {list(__UpperCamelCase      )}\n"
        F"Please pick one to interleave with the other datasets, for example: dataset['{next(iter(__UpperCamelCase      )      )}']"      )
   raise ValueError(
       F"Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(__UpperCamelCase      ).__name__}."      )
  if i == 0:
   UpperCAmelCase__   ,      UpperCAmelCase__					:  Optional[int]					   =					(
       (Dataset, IterableDataset) if isinstance(__UpperCamelCase							,    __UpperCamelCase      ) else (IterableDataset, Dataset)
   )
  elif not isinstance(__UpperCamelCase							,    __UpperCamelCase      ):
   raise ValueError(
       F"Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."      )
 if dataset_type is Dataset:
  return _concatenate_map_style_datasets(__UpperCamelCase							,    info=__UpperCamelCase							,    split=__UpperCamelCase							,    axis=__UpperCamelCase      )
 else:
  return _concatenate_iterable_datasets(__UpperCamelCase							,    info=__UpperCamelCase							,    split=__UpperCamelCase							,    axis=__UpperCamelCase      )





