




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 					UpperCAmelCase__     (   lowerCamelCase,	lowerCamelCase,	lowerCamelCase=None,	lowerCamelCase=None      ):
  if attention_mask is None:
    lowercase			:Tuple							=  tf.cast(tf.math.not_equal(lowerCamelCase,	config.pad_token_id      ),	tf.inta      )
  return {"input_ids": input_ids, "attention_mask": attention_mask}




@require_tf
class       __lowerCAmelCase				:
       _a           =  OPTConfig
       _a           =  {}
       _a           =  '''gelu'''






       def __init__(      self:    List[Any]				,    _lowerCAmelCase:    str				,    _lowerCAmelCase:    Any=13				,    _lowerCAmelCase:    Dict=7				,    _lowerCAmelCase:    List[str]=True				,    _lowerCAmelCase:    int=False				,    _lowerCAmelCase:    Union[str, Any]=99				,    _lowerCAmelCase:    Optional[Any]=16				,    _lowerCAmelCase:    Optional[Any]=2				,    _lowerCAmelCase:    List[str]=4				,    _lowerCAmelCase:    str=4				,    _lowerCAmelCase:    List[Any]="gelu"				,    _lowerCAmelCase:    Optional[Any]=0.1				,    _lowerCAmelCase:    Optional[Any]=0.1				,    _lowerCAmelCase:    Optional[Any]=20				,    _lowerCAmelCase:    Any=2				,    _lowerCAmelCase:    Optional[int]=1				,    _lowerCAmelCase:    Dict=0				,    _lowerCAmelCase:    Optional[int]=16				,    _lowerCAmelCase:    Union[str, Any]=16				,    ):
         lowercase			:Tuple							=  parent
         lowercase			:Dict							=  batch_size
         lowercase			:Any							=  seq_length
         lowercase			:Any							=  is_training
         lowercase			:Union[str, Any]							=  use_labels
         lowercase			:Any							=  vocab_size
         lowercase			:Optional[Any]							=  hidden_size
         lowercase			:int							=  num_hidden_layers
         lowercase			:Union[str, Any]							=  num_attention_heads
         lowercase			:Tuple							=  intermediate_size
         lowercase			:Union[str, Any]							=  hidden_act
         lowercase			:List[Any]							=  hidden_dropout_prob
         lowercase			:Union[str, Any]							=  attention_probs_dropout_prob
         lowercase			:Tuple							=  max_position_embeddings
         lowercase			:Any							=  eos_token_id
         lowercase			:str							=  pad_token_id
         lowercase			:str							=  bos_token_id
         lowercase			:Optional[int]							=  embed_dim
         lowercase			:Optional[int]							=  word_embed_proj_dim
         lowercase			:int							=  False






       def     SCREAMING_SNAKE_CASE			(      self:    Dict ):
         lowercase			:Union[str, Any]							=  ids_tensor([self.batch_size, self.seq_length - 1]				,    self.vocab_size )
         lowercase			:Optional[int]							=  tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size )				,    1 )
         lowercase			:Optional[Any]							=  tf.concat([input_ids, eos_tensor]				,    axis=1 )

         lowercase			:str							=  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=_lowerCAmelCase				,    **self.config_updates				,    )
         lowercase			:Any							=  prepare_opt_inputs_dict(_lowerCAmelCase				,    _lowerCAmelCase )
         return config, inputs_dict






       def     SCREAMING_SNAKE_CASE			(      self:    Any				,    _lowerCAmelCase:    Optional[int]				,    _lowerCAmelCase:    str ):
         lowercase			:List[str]							=  TFOPTModel(config=_lowerCAmelCase )
         lowercase			:str							=  inputs_dict["input_ids"]

         lowercase			:str							=  input_ids[:1, :]
         lowercase			:Dict							=  inputs_dict["attention_mask"][:1, :]
         lowercase			:Dict							=  1

         # first forward pass
         lowercase			:List[Any]							=  model(_lowerCAmelCase				,    attention_mask=_lowerCAmelCase				,    use_cache=_lowerCAmelCase )

         lowercase		,			lowercase			:int							=  outputs.to_tuple()

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

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

         lowercase			:List[Any]							=  model(_lowerCAmelCase				,    attention_mask=_lowerCAmelCase )[0]
         lowercase			:List[str]							=  model(_lowerCAmelCase				,    attention_mask=_lowerCAmelCase				,    past_key_values=_lowerCAmelCase )[0]

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

         # select random slice
         lowercase			:Tuple							=  int(ids_tensor((1,)				,    output_from_past.shape[-1] ) )
         lowercase			:Union[str, Any]							=  output_from_no_past[:, -3:, random_slice_idx]
         lowercase			:List[str]							=  output_from_past[:, :, random_slice_idx]

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







@require_tf
class       __lowerCAmelCase				(							lowerCAmelCase					,       lowerCAmelCase					,       unittest.TestCase):
       _a           =  (TFOPTModel, TFOPTForCausalLM) if is_tf_available() else ()
       _a           =  (TFOPTForCausalLM,) if is_tf_available() else ()
       _a           =  (
           {'''feature-extraction''': TFOPTModel, '''text-generation''': TFOPTForCausalLM} if is_tf_available() else {}
       )
       _a           =  False
       _a           =  False
       _a           =  False
       _a           =  10






       def     SCREAMING_SNAKE_CASE			(      self:    List[Any] ):
         lowercase			:int							=  TFOPTModelTester(self )
         lowercase			:Dict							=  ConfigTester(self				,    config_class=_lowerCAmelCase )






       def     SCREAMING_SNAKE_CASE			(      self:    Dict ):
         self.config_tester.run_common_tests()






       def     SCREAMING_SNAKE_CASE			(      self:    Union[str, Any] ):
         lowercase			:List[Any]							=  self.model_tester.prepare_config_and_inputs_for_common()
         self.model_tester.check_decoder_model_past_large_inputs(*_lowerCAmelCase )






       def     SCREAMING_SNAKE_CASE			(      self:    Union[str, Any] ):
         lowercase		,			lowercase			:Optional[Any]							=  self.model_tester.prepare_config_and_inputs_for_common()

         def _get_word_embedding_weight(_lowerCAmelCase:    str				,    _lowerCAmelCase:    Optional[int] ):
           if hasattr(_lowerCAmelCase				,    "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(_lowerCAmelCase				,    "weight" ):
               return embedding_layer.weight
             else:
               return None

         for model_class in self.all_model_classes:
           for size in [config.vocab_size - 10, config.vocab_size + 10]:
             # build the embeddings
             lowercase			:int							=  model_class(config=_lowerCAmelCase )
             lowercase			:str							=  _get_word_embedding_weight(_lowerCAmelCase				,    model.get_input_embeddings() )
             lowercase			:Dict							=  _get_word_embedding_weight(_lowerCAmelCase				,    model.get_output_embeddings() )

             # reshape the embeddings
             model.resize_token_embeddings(_lowerCAmelCase )
             lowercase			:Dict							=  _get_word_embedding_weight(_lowerCAmelCase				,    model.get_input_embeddings() )
             lowercase			:Dict							=  _get_word_embedding_weight(_lowerCAmelCase				,    model.get_output_embeddings() )

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

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

             # check that weights remain the same after resizing
             lowercase			:List[Any]							=  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:
                 lowercase			:int							=  False
             self.assertTrue(_lowerCAmelCase )

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

               lowercase			:Tuple							=  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:
                   lowercase			:List[Any]							=  False
               self.assertTrue(_lowerCAmelCase )







def 					UpperCAmelCase__     (   lowerCamelCase      ):
  return tf.constant(lowerCamelCase,	dtype=tf.intaa      )




@require_tf
class       __lowerCAmelCase				(							unittest.TestCase):
       _a           =  99






       def     SCREAMING_SNAKE_CASE			(      self:    List[Any] ):
         lowercase			:Optional[int]							=  tf.ones((4, 1)				,    dtype=tf.intaa ) * 2
         lowercase			:int							=  tf.concat([ids_tensor((4, 6)				,    self.vocab_size - 3 ) + 3, eos_column_vector]				,    axis=1 )
         lowercase			:List[str]							=  input_ids.shape[0]
         lowercase			:int							=  OPTConfig(
             vocab_size=self.vocab_size				,    hidden_size=24				,    num_hidden_layers=2				,    num_attention_heads=2				,    ffn_dim=32				,    max_position_embeddings=48				,    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):






       @slow
       def     SCREAMING_SNAKE_CASE			(      self:    int ):
         lowercase			:Dict							=  TFOPTModel.from_pretrained("facebook/opt-350m" )
         lowercase			:Optional[int]							=  _long_tensor([[0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69, 4_60_78, 15_88, 2]] )
         lowercase			:Union[str, Any]							=  tf.not_equal(_lowerCAmelCase				,    model.config.pad_token_id )
         with tf.GradientTape():
           lowercase			:Dict							=  model(input_ids=_lowerCAmelCase				,    attention_mask=_lowerCAmelCase ).last_hidden_state
         lowercase			:Optional[int]							=  (1, 11, 5_12)
         self.assertEqual(output.shape				,    _lowerCAmelCase )
         lowercase			:Optional[int]							=  tf.constant(
             [[-0.28_73, -1.92_18, -0.30_33], [-1.27_10, -0.13_38, -0.19_02], [0.40_95, 0.12_14, -1.31_21]] )
         self.assertTrue(np.allclose(output[:, :3, :3]				,    _lowerCAmelCase				,    atol=4e-3 ) )

         lowercase			:List[Any]							=  tf.function(_lowerCAmelCase				,    jit_compile=_lowerCAmelCase )
         lowercase			:Optional[int]							=  xla_generate(_lowerCAmelCase				,    _lowerCAmelCase )[0]
         self.assertTrue(np.allclose(output[:, :3, :3]				,    _lowerCAmelCase				,    atol=4e-2 ) )







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






       def     SCREAMING_SNAKE_CASE			(      self:    List[str] ):
         super().setUp()
         lowercase			:List[Any]							=  "facebook/opt-350m"






       def     SCREAMING_SNAKE_CASE			(      self:    List[str] ):
         lowercase			:Optional[Any]							=  TFOPTForCausalLM.from_pretrained(self.path_model )
         lowercase			:Optional[Any]							=  GPTaTokenizer.from_pretrained(self.path_model )

         lowercase			:Optional[Any]							=  [
             "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
         lowercase			:List[Any]							=  tokenizer(_lowerCAmelCase				,    return_tensors="tf"				,    padding=_lowerCAmelCase				,    add_special_tokens=_lowerCAmelCase )
         lowercase			:List[str]							=  tf.math.reduce_mean(model(inputs.input_ids				,    attention_mask=inputs.attention_mask )[0]				,    axis=-1 )
         lowercase			:Any							=  tf.constant(
             [
                 [1.38_51, -13.89_23, -10.52_29, -10.75_33, -0.23_09, -10.23_84, -0.53_65, -9.09_47, -5.16_70],
                 [-4.70_73, -10.62_76, -3.94_15, -21.52_42, -0.28_22, -0.28_22, -0.28_22, -0.28_22, -0.28_22],
                 [0.62_47, -3.42_29, -8.91_79, -1.42_97, -14.16_50, 1.41_46, -9.02_18, -0.27_03, -0.27_03],
                 [6.47_83, -1.99_13, -10.79_26, -2.33_36, 1.50_92, -0.99_74, -6.82_13, 1.34_77, 1.34_77],
             ] )
         self.assertTrue(np.allclose(_lowerCAmelCase				,    _lowerCAmelCase				,    atol=1e-4 ) )

         lowercase			:Dict							=  tf.function(_lowerCAmelCase				,    jit_compile=_lowerCAmelCase )
         lowercase			:int							=  tf.math.reduce_mean(xla_generate(inputs.input_ids				,    attention_mask=inputs.attention_mask )[0]				,    axis=-1 )
         self.assertTrue(np.allclose(_lowerCAmelCase				,    _lowerCAmelCase				,    atol=1e-4 ) )







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






       @property
       def     SCREAMING_SNAKE_CASE			(      self:    Optional[int] ):
         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     SCREAMING_SNAKE_CASE			(      self:    Tuple ):
         lowercase			:Optional[int]							=  "facebook/opt-125m"

         lowercase			:int							=  [
             "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",
         ]

         lowercase			:List[str]							=  []
         lowercase			:Optional[Any]							=  GPTaTokenizer.from_pretrained(_lowerCAmelCase )
         lowercase			:Union[str, Any]							=  TFOPTForCausalLM.from_pretrained(_lowerCAmelCase )

         for prompt in self.prompts:
           lowercase			:Union[str, Any]							=  tokenizer(_lowerCAmelCase				,    return_tensors="tf" ).input_ids

           lowercase			:Optional[Any]							=  model.generate(_lowerCAmelCase				,    max_length=10 )

           lowercase			:Optional[Any]							=  tokenizer.batch_decode(_lowerCAmelCase				,    skip_special_tokens=_lowerCAmelCase )
           predicted_outputs += generated_string

         self.assertListEqual(_lowerCAmelCase				,    _lowerCAmelCase )






       def     SCREAMING_SNAKE_CASE			(      self:    Union[str, Any] ):
         lowercase			:Any							=  "facebook/opt-350m"

         lowercase			:int							=  GPTaTokenizer.from_pretrained(_lowerCAmelCase )
         lowercase			:Dict							=  TFOPTForCausalLM.from_pretrained(_lowerCAmelCase )

         lowercase			:Union[str, Any]							=  "left"

         # use different length sentences to test batching
         lowercase			:Union[str, Any]							=  [
             "Hello, my dog is a little",
             "Today, I",
         ]

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

         lowercase			:Tuple							=  model.generate(input_ids=_lowerCAmelCase				,    attention_mask=inputs["attention_mask"] )

         lowercase			:Dict							=  tokenizer(sentences[0]				,    return_tensors="tf" ).input_ids
         lowercase			:str							=  model.generate(input_ids=_lowerCAmelCase )

         lowercase			:Optional[int]							=  inputs_non_padded.shape[-1] - tf.math.reduce_sum(
             tf.cast(inputs["attention_mask"][-1]				,    tf.intaa ) )
         lowercase			:int							=  tokenizer(sentences[1]				,    return_tensors="tf" ).input_ids
         lowercase			:Dict							=  model.generate(input_ids=_lowerCAmelCase				,    max_length=model.config.max_length - num_paddings )

         lowercase			:Optional[Any]							=  tokenizer.batch_decode(_lowerCAmelCase				,    skip_special_tokens=_lowerCAmelCase )
         lowercase			:str							=  tokenizer.decode(output_non_padded[0]				,    skip_special_tokens=_lowerCAmelCase )
         lowercase			:Optional[int]							=  tokenizer.decode(output_padded[0]				,    skip_special_tokens=_lowerCAmelCase )

         lowercase			:Union[str, Any]							=  [
             "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(_lowerCAmelCase				,    _lowerCAmelCase )
         self.assertListEqual(_lowerCAmelCase				,    [non_padded_sentence, padded_sentence] )






       def     SCREAMING_SNAKE_CASE			(      self:    Tuple ):
         lowercase			:Optional[Any]							=  "facebook/opt-350m"

         lowercase			:str							=  [
             "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",
         ]

         lowercase			:int							=  []
         lowercase			:Any							=  GPTaTokenizer.from_pretrained(_lowerCAmelCase )
         lowercase			:Optional[int]							=  TFOPTForCausalLM.from_pretrained(_lowerCAmelCase )

         for prompt in self.prompts:
           lowercase			:List[str]							=  tokenizer(_lowerCAmelCase				,    return_tensors="tf" ).input_ids

           lowercase			:Union[str, Any]							=  model.generate(_lowerCAmelCase				,    max_length=10 )

           lowercase			:int							=  tokenizer.batch_decode(_lowerCAmelCase				,    skip_special_tokens=_lowerCAmelCase )
           predicted_outputs += generated_string

         self.assertListEqual(_lowerCAmelCase				,    _lowerCAmelCase )









from typing import Union

import fire
import torch
from tqdm import tqdm



def 					UpperCAmelCase__     (   lowerCamelCase,	lowerCamelCase = "cpu",	lowerCamelCase = None      ):
  lowercase			:Optional[Any]							=  torch.load(lowerCamelCase,	map_location=lowerCamelCase      )
  for k, v in tqdm(state_dict.items()      ):
    if not isinstance(lowerCamelCase,	torch.Tensor      ):
      raise TypeError("FP16 conversion only works on paths that are saved state dicts, like pytorch_model.bin"      )
    lowercase			:List[Any]							=  v.half()
  if save_path is None:  # overwrite src_path
    lowercase			:Optional[Any]							=  src_path
  torch.save(lowerCamelCase,	lowerCamelCase      )


if __name__ == "__main__":
     fire.Fire(convert)



