enable flexible tf version for tf.keras.mix_precision global_policy feature

[kerrychu]

according to this post https://stackoverflow.com/questions/67037067/attributeerror-module-tensorflow-keras-mixed-precision-has-no-attribute-set

global_policy is no longer experimental but a feature after tensorflow 2.4

This PR would provide users with flexibility of TensorFlow versions, otherwise, the following error would occur:

AttributeError                            Traceback (most recent call last)
/tmp/ipykernel_1331/2018827728.py in <module>
      4 
      5 # Initializing a model from the original configuration
----> 6 model = RobertaModel.from_config(configuration)

/opt/conda/lib/python3.8/site-packages/tf_transformers/models/roberta/roberta_model.py in from_config(cls, config, return_layer, use_mlm_layer, **kwargs)
    155         # Just create a model and return it with random_weights
    156         # (Distribute strategy fails)
--> 157         model_layer = Encoder(config_dict, **kwargs_copy)
    158         if use_mlm_layer:
    159             model_layer = MaskedLMModel(model_layer, config_dict["embedding_size"], config_dict["layer_norm_epsilon"])

/opt/conda/lib/python3.8/site-packages/tf_transformers/models/roberta/roberta.py in __init__(self, config, mask_mode, name, use_dropout, is_training, use_auto_regressive, use_decoder, batch_size, sequence_length, return_all_layer_outputs, **kwargs)
    147         self.call_fn = self.get_call_method(self._config_dict)
    148         # Initialize model
--> 149         self.model_inputs, self.model_outputs = self.get_model(initialize_only=True)
    150 
    151     def get_model(self: LegacyLayer, initialize_only: bool = False):

/opt/conda/lib/python3.8/site-packages/tf_transformers/models/roberta/roberta.py in get_model(self, initialize_only)
    242                 del inputs["past_length"]
    243 
--> 244         layer_outputs = self(inputs)
    245         if initialize_only:
    246             return inputs, layer_outputs

/opt/conda/lib/python3.8/site-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
     65     except Exception as e:  # pylint: disable=broad-except
     66       filtered_tb = _process_traceback_frames(e.__traceback__)
---> 67       raise e.with_traceback(filtered_tb) from None
     68     finally:
     69       del filtered_tb

/opt/conda/lib/python3.8/site-packages/tf_transformers/models/roberta/roberta.py in tf__call(self, inputs)
      8                 do_return = False
      9                 retval_ = ag__.UndefinedReturnValue()
---> 10                 outputs = ag__.converted_call(ag__.ld(self).call_fn, (ag__.ld(inputs),), None, fscope)
     11                 try:
     12                     do_return = True

/opt/conda/lib/python3.8/site-packages/tf_transformers/models/roberta/roberta.py in tf__call_encoder(self, inputs)
    121                 i = ag__.Undefined('i')
    122                 layer = ag__.Undefined('layer')
--> 123                 ag__.for_stmt(ag__.converted_call(ag__.ld(range), (ag__.ld(self)._config_dict['num_hidden_layers'],), None, fscope), None, loop_body, get_state_5, set_state_5, ('embeddings',), {'iterate_names': 'i'})
    124                 cls_token_tensor = ag__.converted_call(ag__.converted_call(ag__.ld(tf).keras.layers.Lambda, (ag__.autograph_artifact((lambda x: ag__.converted_call(ag__.ld(tf).squeeze, (ag__.ld(x)[:, 0:1, :],), dict(axis=1), fscope))),), None, fscope), (ag__.ld(encoder_outputs)[(- 1)],), None, fscope)
    125                 cls_output = ag__.converted_call(ag__.ld(self)._pooler_layer, (ag__.ld(cls_token_tensor),), None, fscope)

/opt/conda/lib/python3.8/site-packages/tf_transformers/models/roberta/roberta.py in loop_body(itr)
    116                     i = itr
    117                     layer = ag__.ld(self)._transformer_layers[ag__.ld(i)]
--> 118                     (embeddings, _, _) = ag__.converted_call(ag__.ld(layer), ([ag__.ld(embeddings), ag__.ld(attention_mask)],), None, fscope)
    119                     ag__.converted_call(ag__.ld(encoder_outputs).append, (ag__.ld(embeddings),), None, fscope)
    120                 _ = ag__.Undefined('_')

/opt/conda/lib/python3.8/site-packages/tf_transformers/layers/transformer/bert_transformer.py in tf__call(self, inputs, mode, cache_key, cache_value)
     26                     outputs = ag__.converted_call(ag__.ld(self).call_encoder, (ag__.ld(inputs),), dict(cache_key=ag__.ld(cache_key), cache_value=ag__.ld(cache_value)), fscope)
     27                 outputs = ag__.Undefined('outputs')
---> 28                 ag__.if_stmt(ag__.ld(self)._use_decoder, if_body, else_body, get_state, set_state, ('outputs',), 1)
     29                 try:
     30                     do_return = True

/opt/conda/lib/python3.8/site-packages/tf_transformers/layers/transformer/bert_transformer.py in else_body()
     24                 def else_body():
     25                     nonlocal outputs
---> 26                     outputs = ag__.converted_call(ag__.ld(self).call_encoder, (ag__.ld(inputs),), dict(cache_key=ag__.ld(cache_key), cache_value=ag__.ld(cache_value)), fscope)
     27                 outputs = ag__.Undefined('outputs')
     28                 ag__.if_stmt(ag__.ld(self)._use_decoder, if_body, else_body, get_state, set_state, ('outputs',), 1)

/opt/conda/lib/python3.8/site-packages/tf_transformers/layers/transformer/bert_transformer.py in tf__call_encoder(self, inputs, cache_key, cache_value)
     29                 attention_output = ag__.converted_call(ag__.ld(self)._attention_dropout, (ag__.ld(attention_output),), dict(training=ag__.ld(self)._use_dropout), fscope)
     30                 attention_output = ag__.converted_call(ag__.ld(self)._attention_layer_norm, ((ag__.ld(input_tensor) + ag__.ld(attention_output)),), None, fscope)
---> 31                 attention_output = ag__.converted_call(ag__.ld(tf).cast, (ag__.ld(attention_output),), dict(dtype=ag__.converted_call(ag__.ld(tf_utils).get_dtype, (), None, fscope)), fscope)
     32                 intermediate_output = ag__.converted_call(ag__.ld(self)._intermediate_dense, (ag__.ld(attention_output),), None, fscope)
     33                 layer_output = ag__.converted_call(ag__.ld(self)._output_dense, (ag__.ld(intermediate_output),), None, fscope)

/opt/conda/lib/python3.8/site-packages/tf_transformers/utils/tf_utils.py in tf__get_dtype()
     10                 retval_ = ag__.UndefinedReturnValue()
     11                 dtype = ag__.ld(tf).float32
---> 12                 policy = ag__.converted_call(ag__.ld(tf).keras.mixed_precision.experimental.global_policy, (), None, fscope)
     13 
     14                 def get_state():

AttributeError: Exception encountered when calling layer "tf_transformers/roberta" (type RobertaEncoder).

in user code:

    File "/opt/conda/lib/python3.8/site-packages/tf_transformers/models/roberta/roberta.py", line 718, in call  *
        outputs = self.call_fn(inputs)
    File "/opt/conda/lib/python3.8/site-packages/tf_transformers/models/roberta/roberta.py", line 290, in call_encoder  *
        embeddings, _, _ = layer([embeddings, attention_mask])
    File "/opt/conda/lib/python3.8/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler  **
        raise e.with_traceback(filtered_tb) from None
    File "/tmp/__autograph_generated_fileflkmddlu.py", line 28, in tf__call
        ag__.if_stmt(ag__.ld(self)._use_decoder, if_body, else_body, get_state, set_state, ('outputs',), 1)
    File "/tmp/__autograph_generated_fileflkmddlu.py", line 26, in else_body
        outputs = ag__.converted_call(ag__.ld(self).call_encoder, (ag__.ld(inputs),), dict(cache_key=ag__.ld(cache_key), cache_value=ag__.ld(cache_value)), fscope)
    File "/tmp/__autograph_generated_filetd9bb7wo.py", line 31, in tf__call_encoder
        attention_output = ag__.converted_call(ag__.ld(tf).cast, (ag__.ld(attention_output),), dict(dtype=ag__.converted_call(ag__.ld(tf_utils).get_dtype, (), None, fscope)), fscope)
    File "/tmp/__autograph_generated_file9o5z35_o.py", line 12, in tf__get_dtype
        policy = ag__.converted_call(ag__.ld(tf).keras.mixed_precision.experimental.global_policy, (), None, fscope)

    AttributeError: Exception encountered when calling layer "transformer/layer_0" (type TransformerBERT).
    
    in user code:
    
        File "/opt/conda/lib/python3.8/site-packages/tf_transformers/layers/transformer/bert_transformer.py", line 331, in call  *
            outputs = self.call_encoder(inputs, cache_key=cache_key, cache_value=cache_value)
        File "/opt/conda/lib/python3.8/site-packages/tf_transformers/layers/transformer/bert_transformer.py", line 256, in call_encoder  *
            attention_output = tf.cast(attention_output, dtype=tf_utils.get_dtype())
        File "/opt/conda/lib/python3.8/site-packages/tf_transformers/utils/tf_utils.py", line 178, in get_dtype  *
            policy = tf.keras.mixed_precision.experimental.global_policy()
    
        AttributeError: module 'keras.api._v2.keras.mixed_precision' has no attribute 'experimental'
    
    
    Call arguments received by layer "transformer/layer_0" (type TransformerBERT):
      • inputs=['tf.Tensor(shape=(None, None, 768), dtype=float32)', 'tf.Tensor(shape=(None, None, None), dtype=float32)']
      • mode=encoder
      • cache_key=None
      • cache_value=None


Call arguments received by layer "tf_transformers/roberta" (type RobertaEncoder):
  • inputs={'input_ids': 'tf.Tensor(shape=(None, None), dtype=int32)', 'input_mask': 'tf.Tensor(shape=(None, None), dtype=int32)', 'input_type_ids': 'tf.Tensor(shape=(None, None), dtype=int32)'}

This is shown after following the official example

from tf_transformers.models import RobertaConfig, RobertaModel
# Initializing an bert-base-uncased style configuration
configuration = RobertaConfig()

# Initializing a model from the original configuration
model = RobertaModel.from_config(configuration)

settings:

tf_transformers version: 2.0.0
tensorflow text version: 2.9.0
sentencepiece version: 0.1.97
tensorflow version: 2.9.1

[legacyai]

This is done now . Sorry for late .