





import unittest

from parameterized import parameterized

from transformers import OpenLlamaConfig, is_torch_available, set_seed
from transformers.testing_utils import require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
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 OpenLlamaForCausalLM, OpenLlamaForSequenceClassification, OpenLlamaModel

class 		SCREAMING_SNAKE_CASE__    :


      def __init__(self	:				Optional[int]    ,      a__	:				List[Any]    ,      a__	:				Dict=13    ,      a__	:				Tuple=7    ,      a__	:				Optional[Any]=True    ,      a__	:				List[Any]=True    ,      a__	:				List[str]=False    ,      a__	:				Optional[int]=True    ,      a__	:				Dict=99    ,      a__	:				List[Any]=32    ,      a__	:				Tuple=5    ,      a__	:				Union[str, Any]=4    ,      a__	:				List[Any]=37    ,      a__	:				Tuple="gelu"    ,      a__	:				Optional[Any]=0.1    ,      a__	:				Dict=0.1    ,      a__	:				List[Any]=512    ,      a__	:				List[Any]=16    ,      a__	:				Dict=2    ,      a__	:				Union[str, Any]=0.0_2    ,      a__	:				Optional[Any]=3    ,      a__	:				str=4    ,      a__	:				Dict=None    ,      ):





         """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_token_type_ids
         __snake_case								=       use_labels
         __snake_case								=       vocab_size
         __snake_case								=       hidden_size
         __snake_case								=       num_hidden_layers
         __snake_case								=       num_attention_heads
         __snake_case								=       intermediate_size
         __snake_case								=       hidden_act
         __snake_case								=       hidden_dropout_prob
         __snake_case								=       attention_probs_dropout_prob
         __snake_case								=       max_position_embeddings
         __snake_case								=       type_vocab_size
         __snake_case								=       type_sequence_label_size
         __snake_case								=       initializer_range
         __snake_case								=       num_labels
         __snake_case								=       num_choices
         __snake_case								=       scope


      def        a			(self	:				List[Any]				):





         """simple docstring"""



         __snake_case								=       ids_tensor([self.batch_size, self.seq_length]    ,      self.vocab_size				)

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

         __snake_case								=       None
         if self.use_token_type_ids:
            __snake_case								=       ids_tensor([self.batch_size, self.seq_length]    ,      self.type_vocab_size				)

         __snake_case								=       None
         __snake_case								=       None
         __snake_case								=       None
         if self.use_labels:
            __snake_case								=       ids_tensor([self.batch_size]    ,      self.type_sequence_label_size				)
            __snake_case								=       ids_tensor([self.batch_size, self.seq_length]    ,      self.num_labels				)
            __snake_case								=       ids_tensor([self.batch_size]    ,      self.num_choices				)

         __snake_case								=       self.get_config()

         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels


      def        a			(self	:				List[str]				):





         """simple docstring"""



         return OpenLlamaConfig(
             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_size=self.intermediate_size    ,      hidden_act=self.hidden_act    ,      hidden_dropout_prob=self.hidden_dropout_prob    ,      attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,      max_position_embeddings=self.max_position_embeddings    ,      type_vocab_size=self.type_vocab_size    ,      is_decoder=a__    ,      initializer_range=self.initializer_range    ,      use_stable_embedding=a__    ,      )


      def        a			(self	:				List[str]    ,      a__	:				str    ,      a__	:				Union[str, Any]    ,      a__	:				List[Any]    ,      a__	:				int    ,      a__	:				List[str]    ,      a__	:				int    ,      a__	:				List[Any]				):





         """simple docstring"""



         __snake_case								=       OpenLlamaModel(config=a__				)
         model.to(a__				)
         model.eval()
         __snake_case								=       model(a__    ,      attention_mask=a__				)
         __snake_case								=       model(a__				)
         self.parent.assertEqual(result.last_hidden_state.shape    ,      (self.batch_size, self.seq_length, self.hidden_size)				)


      def        a			(self	:				Tuple    ,      a__	:				Optional[Any]    ,      a__	:				str    ,      a__	:				Optional[int]    ,      a__	:				Dict    ,      a__	:				int    ,      a__	:				Optional[Any]    ,      a__	:				str    ,      a__	:				Tuple    ,      a__	:				Optional[Any]    ,      ):





         """simple docstring"""



         __snake_case								=       True
         __snake_case								=       OpenLlamaModel(a__				)
         model.to(a__				)
         model.eval()
         __snake_case								=       model(
             a__    ,      attention_mask=a__    ,      encoder_hidden_states=a__    ,      encoder_attention_mask=a__    ,      )
         __snake_case								=       model(
             a__    ,      attention_mask=a__    ,      encoder_hidden_states=a__    ,      )
         __snake_case								=       model(a__    ,      attention_mask=a__				)
         self.parent.assertEqual(result.last_hidden_state.shape    ,      (self.batch_size, self.seq_length, self.hidden_size)				)


      def        a			(self	:				Union[str, Any]    ,      a__	:				List[Any]    ,      a__	:				Optional[Any]    ,      a__	:				Dict    ,      a__	:				List[str]    ,      a__	:				Optional[Any]    ,      a__	:				str    ,      a__	:				Optional[Any]    ,      a__	:				List[Any]    ,      a__	:				List[Any]    ,      ):





         """simple docstring"""



         __snake_case								=       OpenLlamaForCausalLM(config=a__				)
         model.to(a__				)
         model.eval()
         __snake_case								=       model(a__    ,      attention_mask=a__    ,      labels=a__				)
         self.parent.assertEqual(result.logits.shape    ,      (self.batch_size, self.seq_length, self.vocab_size)				)


      def        a			(self	:				Optional[Any]    ,      a__	:				Tuple    ,      a__	:				Union[str, Any]    ,      a__	:				Optional[int]    ,      a__	:				Optional[int]    ,      a__	:				Optional[int]    ,      a__	:				Dict    ,      a__	:				Any    ,      a__	:				int    ,      a__	:				List[str]    ,      ):





         """simple docstring"""



         __snake_case								=       True
         __snake_case								=       True
         __snake_case								=       OpenLlamaForCausalLM(config=a__				)
         model.to(a__				)
         model.eval()

         # first forward pass
         __snake_case								=       model(
             a__    ,      attention_mask=a__    ,      encoder_hidden_states=a__    ,      encoder_attention_mask=a__    ,      use_cache=a__    ,      )
         __snake_case								=       outputs.past_key_values

         # create hypothetical multiple next token and extent to next_input_ids
         __snake_case								=       ids_tensor((self.batch_size, 3)    ,      config.vocab_size				)
         __snake_case								=       ids_tensor((self.batch_size, 3)    ,      vocab_size=2				)

         # append to next input_ids and
         __snake_case								=       torch.cat([input_ids, next_tokens]    ,      dim=-1				)
         __snake_case								=       torch.cat([input_mask, next_mask]    ,      dim=-1				)

         __snake_case								=       model(
             a__    ,      attention_mask=a__    ,      encoder_hidden_states=a__    ,      encoder_attention_mask=a__    ,      output_hidden_states=a__    ,      )['''hidden_states'''][0]
         __snake_case								=       model(
             a__    ,      attention_mask=a__    ,      encoder_hidden_states=a__    ,      encoder_attention_mask=a__    ,      past_key_values=a__    ,      output_hidden_states=a__    ,      )['''hidden_states'''][0]

         # select random slice
         __snake_case								=       ids_tensor((1,)    ,      output_from_past.shape[-1]				).item()
         __snake_case								=       output_from_no_past[:, -3:, random_slice_idx].detach()
         __snake_case								=       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(a__    ,      a__    ,      atol=1E-3				)				)


      def        a			(self	:				Dict				):





         """simple docstring"""



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






@require_torch
class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase						,				_UpperCAmelCase						,				_UpperCAmelCase						,				unittest.TestCase  ):
      A_						:    int							=   (
          (OpenLlamaModel, OpenLlamaForCausalLM, OpenLlamaForSequenceClassification) if is_torch_available() else ()
      )
      A_						:    Union[str, Any]							=   (OpenLlamaForCausalLM,) if is_torch_available() else ()
      A_						:    Optional[int]							=   (
          {
              'feature-extraction': OpenLlamaModel,
              'text-classification': OpenLlamaForSequenceClassification,
              'text-generation': OpenLlamaForCausalLM,
              'zero-shot': OpenLlamaForSequenceClassification,
          }
          if is_torch_available()
          else {}
      )
      A_						:    Dict							=   False
      A_						:    str							=   False


      def        a			(self	:				Any				):





         """simple docstring"""



         __snake_case								=       OpenLlamaModelTester(self				)
         __snake_case								=       ConfigTester(self    ,      config_class=a__    ,      hidden_size=37				)


      def        a			(self	:				Optional[int]				):





         """simple docstring"""



         self.config_tester.run_common_tests()


      def        a			(self	:				Tuple				):





         """simple docstring"""



         __snake_case								=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a__				)


      def        a			(self	:				Dict				):





         """simple docstring"""



         __snake_case								=       self.model_tester.prepare_config_and_inputs()
         for type in ["absolute", "relative_key", "relative_key_query"]:
            __snake_case								=       type
            self.model_tester.create_and_check_model(*a__				)


      def        a			(self	:				Any				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()
         __snake_case								=       3
         __snake_case								=       input_dict['''input_ids''']
         __snake_case								=       input_ids.ne(1				).to(a__				)
         __snake_case								=       ids_tensor([self.model_tester.batch_size]    ,      self.model_tester.type_sequence_label_size				)
         __snake_case								=       OpenLlamaForSequenceClassification(a__				)
         model.to(a__				)
         model.eval()
         __snake_case								=       model(a__    ,      attention_mask=a__    ,      labels=a__				)
         self.assertEqual(result.logits.shape    ,      (self.model_tester.batch_size, self.model_tester.num_labels)				)


      def        a			(self	:				List[str]				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()
         __snake_case								=       3
         __snake_case								=       '''single_label_classification'''
         __snake_case								=       input_dict['''input_ids''']
         __snake_case								=       input_ids.ne(1				).to(a__				)
         __snake_case								=       ids_tensor([self.model_tester.batch_size]    ,      self.model_tester.type_sequence_label_size				)
         __snake_case								=       OpenLlamaForSequenceClassification(a__				)
         model.to(a__				)
         model.eval()
         __snake_case								=       model(a__    ,      attention_mask=a__    ,      labels=a__				)
         self.assertEqual(result.logits.shape    ,      (self.model_tester.batch_size, self.model_tester.num_labels)				)


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





         """simple docstring"""



         __snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()
         __snake_case								=       3
         __snake_case								=       '''multi_label_classification'''
         __snake_case								=       input_dict['''input_ids''']
         __snake_case								=       input_ids.ne(1				).to(a__				)
         __snake_case								=       ids_tensor(
             [self.model_tester.batch_size, config.num_labels]    ,      self.model_tester.type_sequence_label_size				).to(torch.float				)
         __snake_case								=       OpenLlamaForSequenceClassification(a__				)
         model.to(a__				)
         model.eval()
         __snake_case								=       model(a__    ,      attention_mask=a__    ,      labels=a__				)
         self.assertEqual(result.logits.shape    ,      (self.model_tester.batch_size, self.model_tester.num_labels)				)


      @unittest.skip('''Open-Llama buffers include complex numbers, which breaks this test'''				)
      def        a			(self	:				List[str]				):





         """simple docstring"""



         pass


      @parameterized.expand([('''linear''',), ('''dynamic''',)]				)
      def        a			(self	:				List[str]    ,      a__	:				Union[str, Any]				):





         """simple docstring"""



         __snake_case	,     __snake_case								=       self.model_tester.prepare_config_and_inputs_for_common()
         __snake_case								=       ids_tensor([1, 10]    ,      config.vocab_size				)
         __snake_case								=       ids_tensor([1, int(config.max_position_embeddings * 1.5				)]    ,      config.vocab_size				)

         set_seed(42				)  # Fixed seed at init time so the two models get the same random weights
         __snake_case								=       OpenLlamaModel(a__				)
         original_model.to(a__				)
         original_model.eval()
         __snake_case								=       original_model(a__				).last_hidden_state
         __snake_case								=       original_model(a__				).last_hidden_state

         set_seed(42				)  # Fixed seed at init time so the two models get the same random weights
         __snake_case								=       {'''type''': scaling_type, '''factor''': 1_0.0}
         __snake_case								=       OpenLlamaModel(a__				)
         scaled_model.to(a__				)
         scaled_model.eval()
         __snake_case								=       scaled_model(a__				).last_hidden_state
         __snake_case								=       scaled_model(a__				).last_hidden_state

         # Dynamic scaling does not change the RoPE embeddings until it receives an input longer than the original
         # maximum sequence length, so the outputs for the short input should match.
         if scaling_type == "dynamic":
            self.assertTrue(torch.allclose(a__    ,      a__    ,      atol=1E-5				)				)
         else:
            self.assertFalse(torch.allclose(a__    ,      a__    ,      atol=1E-5				)				)

         # The output should be different for long inputs
         self.assertFalse(torch.allclose(a__    ,      a__    ,      atol=1E-5				)				)








import json
from typing import TYPE_CHECKING, List, Optional, Tuple

from tokenizers import pre_tokenizers

from ...tokenization_utils_base import BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging


if TYPE_CHECKING:
    from transformers.pipelines.conversational import Conversation


snake_case_           =      logging.get_logger(__name__)

snake_case_           =      {'tokenizer_file': 'tokenizer.json'}

snake_case_           =      {
    'tokenizer_file': {
        'bigscience/tokenizer': 'https://huggingface.co/bigscience/tokenizer/blob/main/tokenizer.json',
        'bigscience/bloom-560m': 'https://huggingface.co/bigscience/bloom-560m/blob/main/tokenizer.json',
        'bigscience/bloom-1b1': 'https://huggingface.co/bigscience/bloom-1b1/blob/main/tokenizer.json',
        'bigscience/bloom-1b7': 'https://huggingface.co/bigscience/bloom-1b7/blob/main/tokenizer.json',
        'bigscience/bloom-3b': 'https://huggingface.co/bigscience/bloom-3b/blob/main/tokenizer.json',
        'bigscience/bloom-7b1': 'https://huggingface.co/bigscience/bloom-7b1/blob/main/tokenizer.json',
        'bigscience/bloom': 'https://huggingface.co/bigscience/bloom/blob/main/tokenizer.json',
    },
}

class 		SCREAMING_SNAKE_CASE__    (	_UpperCAmelCase  ):
      A_						:    List[str]							=   VOCAB_FILES_NAMES
      A_						:    Optional[Any]							=   PRETRAINED_VOCAB_FILES_MAP
      A_						:    str							=   ['input_ids', 'attention_mask']
      A_						:    Optional[Any]							=   None


      def __init__(self	:				Optional[int]    ,      a__	:				int=None    ,      a__	:				str=None    ,      a__	:				Any=None    ,      a__	:				List[Any]="<unk>"    ,      a__	:				List[Any]="<s>"    ,      a__	:				Optional[int]="</s>"    ,      a__	:				List[str]="<pad>"    ,      a__	:				Union[str, Any]=False    ,      a__	:				str=False    ,      **a__	:				Optional[Any]    ,      ):





         """simple docstring"""



         super().__init__(
             a__    ,      a__    ,      tokenizer_file=a__    ,      unk_token=a__    ,      bos_token=a__    ,      eos_token=a__    ,      pad_token=a__    ,      add_prefix_space=a__    ,      clean_up_tokenization_spaces=a__    ,      **a__    ,      )
         __snake_case								=       json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()				)
         if pre_tok_state.get('''add_prefix_space'''    ,      a__				) != add_prefix_space:
            __snake_case								=       getattr(a__    ,      pre_tok_state.pop('''type'''				)				)
            __snake_case								=       add_prefix_space
            __snake_case								=       pre_tok_class(**a__				)

         __snake_case								=       add_prefix_space


      def        a			(self	:				int    ,      *a__	:				Tuple    ,      **a__	:				Optional[Any]				):





         """simple docstring"""



         __snake_case								=       kwargs.get('''is_split_into_words'''    ,      a__				)
         if not (self.add_prefix_space or not is_split_into_words):
            raise Exception(
                f"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True to use it with"""
                ''' pretokenized inputs.'''				)

         return super()._batch_encode_plus(*a__    ,      **a__				)


      def        a			(self	:				List[str]    ,      *a__	:				List[str]    ,      **a__	:				List[str]				):





         """simple docstring"""



         __snake_case								=       kwargs.get('''is_split_into_words'''    ,      a__				)

         if not (self.add_prefix_space or not is_split_into_words):
            raise Exception(
                f"""You need to instantiate {self.__class__.__name__} with add_prefix_space=True to use it with"""
                ''' pretokenized inputs.'''				)

         return super()._encode_plus(*a__    ,      **a__				)


      def        a			(self	:				List[Any]    ,      a__	:				str    ,      a__	:				Optional[str] = None				):





         """simple docstring"""



         __snake_case								=       self._tokenizer.model.save(a__    ,      name=a__				)
         return tuple(a__				)


      def        a			(self	:				Tuple    ,      a__	:				"Conversation"				):





         """simple docstring"""



         __snake_case								=       []
         for is_user, text in conversation.iter_texts():
            input_ids.extend(self.encode(a__    ,      add_special_tokens=a__				) + [self.eos_token_id]				)

         if len(a__				) > self.model_max_length:
            __snake_case								=       input_ids[-self.model_max_length :]
         return input_ids

