






from __future__ import annotations




def       _SCREAMING_SNAKE_CASE (							__snake_case   ,      __snake_case )				->						list[str]:
       if partitions <= 0:
              raise ValueError("""partitions must be a positive number!""" )
       if partitions > number_of_bytes:
              raise ValueError("""partitions can not > number_of_bytes!""" )
       _UpperCAmelCase							=    number_of_bytes // partitions
       _UpperCAmelCase							=    []
       for i in range(__snake_case ):
              _UpperCAmelCase							=    i * bytes_per_partition + 1
              _UpperCAmelCase							=    (
                  number_of_bytes if i == partitions - 1 else (i + 1) * bytes_per_partition
              )
              allocation_list.append(f"""{start_bytes}-{end_bytes}""" )
       return allocation_list


if __name__ == "__main__":
  import doctest

  doctest.testmod()
from __future__ import annotations

import unittest

import numpy as np

from transformers import OPTConfig, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, slow

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


if is_tf_available():
				import tensorflow as tf

				from transformers import GPTaTokenizer, TFOPTForCausalLM, TFOPTModel


def 	lowerCamelCase					(	UpperCamelCase     :				Any			,       UpperCamelCase     :				str			,       UpperCamelCase     :				Dict=None			,       UpperCamelCase     :				str=None							)   ->    List[str]:
		if attention_mask is None:
				_lowerCamelCase				   =  tf.cast(tf.math.not_equal(UpperCamelCase			,       config.pad_token_id							)			,       tf.inta							)
		return {"input_ids": input_ids, "attention_mask": attention_mask}





@require_tf
class     lowerCAmelCase__     :


		'''simple docstring'''






		lowerCAmelCase_     =   OPTConfig
		lowerCAmelCase_     =   {}
		lowerCAmelCase_     =   'gelu'
		def __init__(       self     :     Dict     ,  snake_case__     :     List[Any]     ,  snake_case__     :     Any=1_3     ,  snake_case__     :     List[str]=7     ,  snake_case__     :     Optional[Any]=True     ,  snake_case__     :     int=False     ,  snake_case__     :     Any=9_9     ,  snake_case__     :     Optional[int]=1_6     ,  snake_case__     :     str=2     ,  snake_case__     :     List[str]=4     ,  snake_case__     :     Union[str, Any]=4     ,  snake_case__     :     Union[str, Any]="gelu"     ,  snake_case__     :     Optional[Any]=0.1     ,  snake_case__     :     Union[str, Any]=0.1     ,  snake_case__     :     Any=2_0     ,  snake_case__     :     Any=2     ,  snake_case__     :     Optional[int]=1     ,  snake_case__     :     Optional[Any]=0     ,  snake_case__     :     int=1_6     ,  snake_case__     :     List[Any]=1_6     ,  )      ->	List[Any]:
				_lowerCamelCase				   =  parent
				_lowerCamelCase				   =  batch_size
				_lowerCamelCase				   =  seq_length
				_lowerCamelCase				   =  is_training
				_lowerCamelCase				   =  use_labels
				_lowerCamelCase				   =  vocab_size
				_lowerCamelCase				   =  hidden_size
				_lowerCamelCase				   =  num_hidden_layers
				_lowerCamelCase				   =  num_attention_heads
				_lowerCamelCase				   =  intermediate_size
				_lowerCamelCase				   =  hidden_act
				_lowerCamelCase				   =  hidden_dropout_prob
				_lowerCamelCase				   =  attention_probs_dropout_prob
				_lowerCamelCase				   =  max_position_embeddings
				_lowerCamelCase				   =  eos_token_id
				_lowerCamelCase				   =  pad_token_id
				_lowerCamelCase				   =  bos_token_id
				_lowerCamelCase				   =  embed_dim
				_lowerCamelCase				   =  word_embed_proj_dim
				_lowerCamelCase				   =  False
		def   _snake_case       (       self     :     List[Any]		)      ->	Optional[int]:
				_lowerCamelCase				   =  ids_tensor([self.batch_size, self.seq_length - 1]     ,  self.vocab_size		)
				_lowerCamelCase				   =  tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size		)     ,  1		)
				_lowerCamelCase				   =  tf.concat([input_ids, eos_tensor]     ,  axis=1		)

				_lowerCamelCase				   =  self.config_cls(
				    vocab_size=self.vocab_size     ,  hidden_size=self.hidden_size     ,  num_hidden_layers=self.num_hidden_layers     ,  num_attention_heads=self.num_attention_heads     ,  ffn_dim=self.intermediate_size     ,  dropout=self.hidden_dropout_prob     ,  attention_dropout=self.attention_probs_dropout_prob     ,  max_position_embeddings=self.max_position_embeddings     ,  eos_token_id=self.eos_token_id     ,  bos_token_id=self.bos_token_id     ,  pad_token_id=self.pad_token_id     ,  embed_dim=self.embed_dim     ,  word_embed_proj_dim=self.word_embed_proj_dim     ,  is_encoder_decoder=snake_case__     ,  **self.config_updates     ,  )
				_lowerCamelCase				   =  prepare_opt_inputs_dict(snake_case__     ,  snake_case__		)
				return config, inputs_dict



		def   _snake_case       (       self     :     Dict     ,  snake_case__     :     Union[str, Any]     ,  snake_case__     :     List[Any]		)      ->	Optional[int]:
				_lowerCamelCase				   =  TFOPTModel(config=snake_case__		)
				_lowerCamelCase				   =  inputs_dict['input_ids']

				_lowerCamelCase				   =  input_ids[:1, :]
				_lowerCamelCase				   =  inputs_dict['attention_mask'][:1, :]
				_lowerCamelCase				   =  1

				# first forward pass
				_lowerCamelCase				   =  model(snake_case__     ,  attention_mask=snake_case__     ,  use_cache=snake_case__		)

				_lowerCamelCase   ,							_lowerCamelCase				   =  outputs.to_tuple()

				# create hypothetical next token and extent to next_input_ids
				_lowerCamelCase				   =  ids_tensor((self.batch_size, 3)     ,  config.vocab_size		)
				_lowerCamelCase				   =  tf.cast(ids_tensor((self.batch_size, 3)     ,  2		)     ,  tf.inta		)

				# append to next input_ids and
				_lowerCamelCase				   =  tf.concat([input_ids, next_tokens]     ,  axis=-1		)
				_lowerCamelCase				   =  tf.concat([attention_mask, next_attn_mask]     ,  axis=-1		)

				_lowerCamelCase				   =  model(snake_case__     ,  attention_mask=snake_case__		)[0]
				_lowerCamelCase				   =  model(snake_case__     ,  attention_mask=snake_case__     ,  past_key_values=snake_case__		)[0]

				self.parent.assertEqual(next_tokens.shape[1]     ,  output_from_past.shape[1]		)

				# select random slice
				_lowerCamelCase				   =  int(ids_tensor((1,)     ,  output_from_past.shape[-1]		)		)
				_lowerCamelCase				   =  output_from_no_past[:, -3:, random_slice_idx]
				_lowerCamelCase				   =  output_from_past[:, :, random_slice_idx]

				# test that outputs are equal for slice
				tf.debugging.assert_near(snake_case__     ,  snake_case__     ,  rtol=1e-3		)





@require_tf
class     lowerCAmelCase__     (				SCREAMING_SNAKE_CASE_     ,SCREAMING_SNAKE_CASE_     ,unittest.TestCase			):


		'''simple docstring'''






		lowerCAmelCase_     =   (TFOPTModel, TFOPTForCausalLM) if is_tf_available() else ()
		lowerCAmelCase_     =   (TFOPTForCausalLM,) if is_tf_available() else ()
		lowerCAmelCase_     =   (
		    {'feature-extraction': TFOPTModel, 'text-generation': TFOPTForCausalLM} if is_tf_available() else {}
		)
		lowerCAmelCase_     =   False
		lowerCAmelCase_     =   False
		lowerCAmelCase_     =   False
		lowerCAmelCase_     =   10
		def   _snake_case       (       self     :     List[Any]		)      ->	int:
				_lowerCamelCase				   =  TFOPTModelTester(self		)
				_lowerCamelCase				   =  ConfigTester(self     ,  config_class=snake_case__		)
		def   _snake_case       (       self     :     List[str]		)      ->	Union[str, Any]:
				self.config_tester.run_common_tests()
		def   _snake_case       (       self     :     Dict		)      ->	Union[str, Any]:
				_lowerCamelCase				   =  self.model_tester.prepare_config_and_inputs_for_common()
				self.model_tester.check_decoder_model_past_large_inputs(*snake_case__		)



		def   _snake_case       (       self     :     Optional[int]		)      ->	Optional[Any]:
				_lowerCamelCase   ,							_lowerCamelCase				   =  self.model_tester.prepare_config_and_inputs_for_common()

				def _get_word_embedding_weight(snake_case__     :     List[Any]     ,  snake_case__     :     Dict		):
						if hasattr(snake_case__     ,  'weight'		):
								return embedding_layer.weight
						else:
								# Here we build the word embeddings weights if not exists.
								# And then we retry to get the attribute once built.
								model.build()
								if hasattr(snake_case__     ,  'weight'		):
										return embedding_layer.weight
								else:
										return None

				for model_class in self.all_model_classes:
						for size in [config.vocab_size - 1_0, config.vocab_size + 1_0]:
								# build the embeddings
								_lowerCamelCase				   =  model_class(config=snake_case__		)
								_lowerCamelCase				   =  _get_word_embedding_weight(snake_case__     ,  model.get_input_embeddings()		)
								_lowerCamelCase				   =  _get_word_embedding_weight(snake_case__     ,  model.get_output_embeddings()		)

								# reshape the embeddings
								model.resize_token_embeddings(snake_case__		)
								_lowerCamelCase				   =  _get_word_embedding_weight(snake_case__     ,  model.get_input_embeddings()		)
								_lowerCamelCase				   =  _get_word_embedding_weight(snake_case__     ,  model.get_output_embeddings()		)

								# check that the resized embeddings size matches the desired size.
								_lowerCamelCase				   =  size if size is not None else config.vocab_size

								self.assertEqual(new_input_embeddings.shape[0]     ,  snake_case__		)

								# check that weights remain the same after resizing
								_lowerCamelCase				   =  True
								for pa, pa in zip(old_input_embeddings.value()     ,  new_input_embeddings.value()		):
										if tf.math.reduce_sum(tf.math.abs(pa - pa		)		) > 0:
												_lowerCamelCase				   =  False
								self.assertTrue(snake_case__		)

								if old_output_embeddings is not None and new_output_embeddings is not None:
										self.assertEqual(new_output_embeddings.shape[0]     ,  snake_case__		)

										_lowerCamelCase				   =  True
										for pa, pa in zip(old_output_embeddings.value()     ,  new_output_embeddings.value()		):
												if tf.math.reduce_sum(tf.math.abs(pa - pa		)		) > 0:
														_lowerCamelCase				   =  False
										self.assertTrue(snake_case__		)



def 	lowerCamelCase					(	UpperCamelCase     :				str							)   ->    List[str]:
		return tf.constant(UpperCamelCase			,       dtype=tf.intaa							)





@require_tf
class     lowerCAmelCase__     (				unittest.TestCase			):


		'''simple docstring'''






		lowerCAmelCase_     =   99
		def   _snake_case       (       self     :     Any		)      ->	int:
				_lowerCamelCase				   =  tf.ones((4, 1)     ,  dtype=tf.intaa		) * 2
				_lowerCamelCase				   =  tf.concat([ids_tensor((4, 6)     ,  self.vocab_size - 3		) + 3, eos_column_vector]     ,  axis=1		)
				_lowerCamelCase				   =  input_ids.shape[0]
				_lowerCamelCase				   =  OPTConfig(
				    vocab_size=self.vocab_size     ,  hidden_size=2_4     ,  num_hidden_layers=2     ,  num_attention_heads=2     ,  ffn_dim=3_2     ,  max_position_embeddings=4_8     ,  eos_token_id=2     ,  pad_token_id=1     ,  bos_token_id=0     ,  )
				return config, input_ids, batch_size





@require_sentencepiece
@require_tf
class     lowerCAmelCase__     (				unittest.TestCase			):


		'''simple docstring'''






		@slow
		def   _snake_case       (       self     :     List[Any]		)      ->	str:
				_lowerCamelCase				   =  TFOPTModel.from_pretrained('facebook/opt-350m'		)
				_lowerCamelCase				   =  _long_tensor([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]]		)
				_lowerCamelCase				   =  tf.not_equal(snake_case__     ,  model.config.pad_token_id		)
				with tf.GradientTape():
						_lowerCamelCase				   =  model(input_ids=snake_case__     ,  attention_mask=snake_case__		).last_hidden_state
				_lowerCamelCase				   =  (1, 1_1, 5_1_2)
				self.assertEqual(output.shape     ,  snake_case__		)
				_lowerCamelCase				   =  tf.constant(
				    [[-0.2873, -1.9218, -0.3033], [-1.2710, -0.1338, -0.1902], [0.4095, 0.1214, -1.3121]]		)
				self.assertTrue(np.allclose(output[:, :3, :3]     ,  snake_case__     ,  atol=4e-3		)		)

				_lowerCamelCase				   =  tf.function(snake_case__     ,  jit_compile=snake_case__		)
				_lowerCamelCase				   =  xla_generate(snake_case__     ,  snake_case__		)[0]
				self.assertTrue(np.allclose(output[:, :3, :3]     ,  snake_case__     ,  atol=4e-2		)		)





@require_tf
@slow
class     lowerCAmelCase__     (				unittest.TestCase			):


		'''simple docstring'''






		def   _snake_case       (       self     :     Tuple		)      ->	List[str]:
				super().setUp()
				_lowerCamelCase				   =  'facebook/opt-350m'



		def   _snake_case       (       self     :     Union[str, Any]		)      ->	Dict:
				_lowerCamelCase				   =  TFOPTForCausalLM.from_pretrained(self.path_model		)
				_lowerCamelCase				   =  GPTaTokenizer.from_pretrained(self.path_model		)

				_lowerCamelCase				   =  [
				    'Today is a beautiful day and I want to',
				    'In the city of',
				    'Paris is the capital of France and',
				    'Computers and mobile phones have taken',
				]
				# verify that prompt without BOS token is identical to Metaseq -> add_special_tokens=False
				_lowerCamelCase				   =  tokenizer(snake_case__     ,  return_tensors='tf'     ,  padding=snake_case__     ,  add_special_tokens=snake_case__		)
				_lowerCamelCase				   =  tf.math.reduce_mean(model(inputs.input_ids     ,  attention_mask=inputs.attention_mask		)[0]     ,  axis=-1		)
				_lowerCamelCase				   =  tf.constant(
				    [
				        [1.3851, -13.8923, -10.5229, -10.7533, -0.2309, -10.2384, -0.5365, -9.0947, -5.1670],
				        [-4.7073, -10.6276, -3.9415, -21.5242, -0.2822, -0.2822, -0.2822, -0.2822, -0.2822],
				        [0.6247, -3.4229, -8.9179, -1.4297, -14.1650, 1.4146, -9.0218, -0.2703, -0.2703],
				        [6.4783, -1.9913, -10.7926, -2.3336, 1.5092, -0.9974, -6.8213, 1.3477, 1.3477],
				    ]		)
				self.assertTrue(np.allclose(snake_case__     ,  snake_case__     ,  atol=1e-4		)		)

				_lowerCamelCase				   =  tf.function(snake_case__     ,  jit_compile=snake_case__		)
				_lowerCamelCase				   =  tf.math.reduce_mean(xla_generate(inputs.input_ids     ,  attention_mask=inputs.attention_mask		)[0]     ,  axis=-1		)
				self.assertTrue(np.allclose(snake_case__     ,  snake_case__     ,  atol=1e-4		)		)





@require_tf
@slow
class     lowerCAmelCase__     (				unittest.TestCase			):


		'''simple docstring'''






		@property
		def   _snake_case       (       self     :     str		)      ->	List[str]:
				return [
				    "Today is a beautiful day and I want",
				    "In the city of",
				    "Paris is the capital of France and",
				    "Computers and mobile phones have taken",
				]
		def   _snake_case       (       self     :     Dict		)      ->	Optional[int]:
				_lowerCamelCase				   =  'facebook/opt-125m'

				_lowerCamelCase				   =  [
				    'Today is a beautiful day and I want to',
				    'In the city of New York, the city',
				    'Paris is the capital of France and the capital',
				    'Computers and mobile phones have taken over the',
				]

				_lowerCamelCase				   =  []
				_lowerCamelCase				   =  GPTaTokenizer.from_pretrained(snake_case__		)
				_lowerCamelCase				   =  TFOPTForCausalLM.from_pretrained(snake_case__		)

				for prompt in self.prompts:
						_lowerCamelCase				   =  tokenizer(snake_case__     ,  return_tensors='tf'		).input_ids

						_lowerCamelCase				   =  model.generate(snake_case__     ,  max_length=1_0		)

						_lowerCamelCase				   =  tokenizer.batch_decode(snake_case__     ,  skip_special_tokens=snake_case__		)
						predicted_outputs += generated_string

				self.assertListEqual(snake_case__     ,  snake_case__		)
		def   _snake_case       (       self     :     int		)      ->	Dict:
				_lowerCamelCase				   =  'facebook/opt-350m'

				_lowerCamelCase				   =  GPTaTokenizer.from_pretrained(snake_case__		)
				_lowerCamelCase				   =  TFOPTForCausalLM.from_pretrained(snake_case__		)

				_lowerCamelCase				   =  'left'

				# use different length sentences to test batching
				_lowerCamelCase				   =  [
				    'Hello, my dog is a little',
				    'Today, I',
				]

				_lowerCamelCase				   =  tokenizer(snake_case__     ,  return_tensors='tf'     ,  padding=snake_case__		)
				_lowerCamelCase				   =  inputs['input_ids']

				_lowerCamelCase				   =  model.generate(input_ids=snake_case__     ,  attention_mask=inputs['attention_mask']		)

				_lowerCamelCase				   =  tokenizer(sentences[0]     ,  return_tensors='tf'		).input_ids
				_lowerCamelCase				   =  model.generate(input_ids=snake_case__		)

				_lowerCamelCase				   =  inputs_non_padded.shape[-1] - tf.math.reduce_sum(
				    tf.cast(inputs['attention_mask'][-1]     ,  tf.intaa		)		)
				_lowerCamelCase				   =  tokenizer(sentences[1]     ,  return_tensors='tf'		).input_ids
				_lowerCamelCase				   =  model.generate(input_ids=snake_case__     ,  max_length=model.config.max_length - num_paddings		)

				_lowerCamelCase				   =  tokenizer.batch_decode(snake_case__     ,  skip_special_tokens=snake_case__		)
				_lowerCamelCase				   =  tokenizer.decode(output_non_padded[0]     ,  skip_special_tokens=snake_case__		)
				_lowerCamelCase				   =  tokenizer.decode(output_padded[0]     ,  skip_special_tokens=snake_case__		)

				_lowerCamelCase				   =  [
				    'Hello, my dog is a little bit of a dork.\nI\'m a little bit',
				    'Today, I was in the middle of a conversation with a friend about the',
				]
				self.assertListEqual(snake_case__     ,  snake_case__		)
				self.assertListEqual(snake_case__     ,  [non_padded_sentence, padded_sentence]		)



		def   _snake_case       (       self     :     List[Any]		)      ->	Optional[Any]:
				_lowerCamelCase				   =  'facebook/opt-350m'

				_lowerCamelCase				   =  [
				    'Today is a beautiful day and I want to',
				    'In the city of San Francisco, the city',
				    'Paris is the capital of France and the capital',
				    'Computers and mobile phones have taken over the',
				]

				_lowerCamelCase				   =  []
				_lowerCamelCase				   =  GPTaTokenizer.from_pretrained(snake_case__		)
				_lowerCamelCase				   =  TFOPTForCausalLM.from_pretrained(snake_case__		)

				for prompt in self.prompts:
						_lowerCamelCase				   =  tokenizer(snake_case__     ,  return_tensors='tf'		).input_ids

						_lowerCamelCase				   =  model.generate(snake_case__     ,  max_length=1_0		)

						_lowerCamelCase				   =  tokenizer.batch_decode(snake_case__     ,  skip_special_tokens=snake_case__		)
						predicted_outputs += generated_string

				self.assertListEqual(snake_case__     ,  snake_case__		)