code: # Copyright 2022 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Tests for ALBERT transformer-based text encoder network."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from absl.testing import parameterized
  import numpy as np
  import tensorflow as tf
  
  from tensorflow.python.keras import keras_parameterized  # pylint: disable=g-direct-tensorflow-import
  from official.nlp.modeling.networks import albert_encoder
  
  
  # This decorator runs the test in V1, V2-Eager, and V2-Functional mode. It
  # guarantees forward compatibility of this code for the V2 switchover.
  @keras_parameterized.run_all_keras_modes
  class AlbertEncoderTest(keras_parameterized.TestCase):
  
    def tearDown(self):
      super(AlbertEncoderTest, self).tearDown()
      tf.keras.mixed_precision.set_global_policy("float32")
  
    @parameterized.named_parameters(
        dict(testcase_name="default", expected_dtype=tf.float32),
        dict(testcase_name="with_float16_dtype", expected_dtype=tf.float16),
    )
    def test_network_creation(self, expected_dtype):
      hidden_size = 32
      sequence_length = 21
  
      kwargs = dict(
          vocab_size=100,
          hidden_size=hidden_size,
          num_attention_heads=2,
          num_layers=3)
      if expected_dtype == tf.float16:
        tf.keras.mixed_precision.set_global_policy("mixed_float16")
  
      # Create a small TransformerEncoder for testing.
      test_network = albert_encoder.AlbertEncoder(**kwargs)
  
      # Create the inputs (note that the first dimension is implicit).
      word_ids = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
      mask = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
      type_ids = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
      data, pooled = test_network([word_ids, mask, type_ids])
  
      expected_data_shape = [None, sequence_length, hidden_size]
      expected_pooled_shape = [None, hidden_size]
      self.assertAllEqual(expected_data_shape, data.shape.as_list())
      self.assertAllEqual(expected_pooled_shape, pooled.shape.as_list())
  
      # If float_dtype is set to float16, the data output is float32 (from a layer
      # norm) and pool output should be float16.
      self.assertEqual(tf.float32, data.dtype)
      self.assertEqual(expected_dtype, pooled.dtype)
  
      # ALBERT has additonal 'embedding_hidden_mapping_in' weights and
      # it shares transformer weights.
      self.assertNotEmpty(
          [x for x in test_network.weights if "embedding_projection/" in x.name])
      self.assertNotEmpty(
          [x for x in test_network.weights if "transformer/" in x.name])
      self.assertEmpty(
          [x for x in test_network.weights if "transformer/layer" in x.name])
  
    def test_network_invocation(self):
      hidden_size = 32
      sequence_length = 21
      vocab_size = 57
      num_types = 7
      num_layers = 3
      # Create a small TransformerEncoder for testing.
      test_network = albert_encoder.AlbertEncoder(
          vocab_size=vocab_size,
          embedding_width=8,
          hidden_size=hidden_size,
          num_attention_heads=2,
          num_layers=num_layers,
          type_vocab_size=num_types)
      # Create the inputs (note that the first dimension is implicit).
      word_ids = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
      mask = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
      type_ids = tf.keras.Input(shape=(sequence_length,), dtype=tf.int32)
      data, pooled = test_network([word_ids, mask, type_ids])
  
      # Create a model based off of this network:
      model = tf.keras.Model([word_ids, mask, type_ids], [data, pooled])
  
      # Invoke the model. We can't validate the output data here (the model is too
      # complex) but this will catch structural runtime errors.
      batch_size = 3
      word_id_data = np.random.randint(
          vocab_size, size=(batch_size, sequence_length))
      mask_data = np.random.randint(2, size=(batch_size, sequence_length))
      type_id_data = np.random.randint(
          num_types, size=(batch_size, sequence_length))
      list_outputs = model.predict([word_id_data, mask_data, type_id_data])
  
      # Creates a TransformerEncoder with max_sequence_length != sequence_length
      max_sequence_length = 128
      test_network = albert_encoder.AlbertEncoder(
          vocab_size=vocab_size,
          embedding_width=8,
          hidden_size=hidden_size,
          max_sequence_length=max_sequence_length,
          num_attention_heads=2,
          num_layers=num_layers,
          type_vocab_size=num_types)
      model = tf.keras.Model([word_ids, mask, type_ids], [data, pooled])
      _ = model.predict([word_id_data, mask_data, type_id_data])
  
      # Tests dictionary outputs.
      test_network_dict = albert_encoder.AlbertEncoder(
          vocab_size=vocab_size,
          embedding_width=8,
          hidden_size=hidden_size,
          max_sequence_length=max_sequence_length,
          num_attention_heads=2,
          num_layers=num_layers,
          type_vocab_size=num_types,
          dict_outputs=True)
      _ = test_network_dict([word_ids, mask, type_ids])
      test_network_dict.set_weights(test_network.get_weights())
      list_outputs = test_network([word_id_data, mask_data, type_id_data])
      dict_outputs = test_network_dict(
          dict(
              input_word_ids=word_id_data,
              input_mask=mask_data,
              input_type_ids=type_id_data))
      self.assertAllEqual(list_outputs[0], dict_outputs["sequence_output"])
      self.assertAllEqual(list_outputs[1], dict_outputs["pooled_output"])
      self.assertLen(dict_outputs["pooled_output"], num_layers)
  
    def test_serialize_deserialize(self):
      tf.keras.mixed_precision.set_global_policy("mixed_float16")
      # Create a network object that sets all of its config options.
      kwargs = dict(
          vocab_size=100,
          embedding_width=8,
          hidden_size=32,
          num_layers=3,
          num_attention_heads=2,
          max_sequence_length=21,
          type_vocab_size=12,
          intermediate_size=1223,
          activation="relu",
          dropout_rate=0.05,
          attention_dropout_rate=0.22,
          initializer="glorot_uniform")
      network = albert_encoder.AlbertEncoder(**kwargs)
  
      expected_config = dict(kwargs)
      expected_config["activation"] = tf.keras.activations.serialize(
          tf.keras.activations.get(expected_config["activation"]))
      expected_config["initializer"] = tf.keras.initializers.serialize(
          tf.keras.initializers.get(expected_config["initializer"]))
      self.assertEqual(network.get_config(), expected_config)
  
      # Create another network object from the first object's config.
      new_network = (
          albert_encoder.AlbertEncoder.from_config(
              network.get_config()))
  
      # Validate that the config can be forced to JSON.
      _ = new_network.to_json()
  
      # If the serialization was successful, the new config should match the old.
      self.assertAllEqual(network.get_config(), new_network.get_config())
  
  
  if __name__ == "__main__":
    tf.test.main()
   # Copyright 2022 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  r"""Evaluates image classification accuracy using tflite_imagenet_evaluator.
  
  Usage:
  tflite_imagenet_evaluator_run --tflite_model_path=/PATH/TO/MODEL.tflite
  """
  
  from typing import Sequence
  from absl import app
  from absl import flags
  import tensorflow as tf
  
  from official.core import exp_factory
  from official.projects.edgetpu.vision.serving import tflite_imagenet_evaluator
  from official.projects.edgetpu.vision.tasks import image_classification
  
  
  flags.DEFINE_string('tflite_model_path', None,
                      'Path to the tflite file to be evaluated.')
  flags.DEFINE_integer('num_threads', 16, 'Number of local threads.')
  flags.DEFINE_integer('batch_size', 256, 'Batch size per thread.')
  flags.DEFINE_string(
      'model_name', 'mobilenet_edgetpu_v2_xs',
      'Model name to identify a registered data pipeline setup and use as the '
      'validation dataset.')
  
  FLAGS = flags.FLAGS
  
  
  def main(argv: Sequence[str]):
    if len(argv) > 1:
      raise app.UsageError('Too many command-line arguments.')
  
    with tf.io.gfile.GFile(FLAGS.tflite_model_path, 'rb') as f:
      model_content = f.read()
  
    config = exp_factory.get_exp_config(FLAGS.model_name)
    global_batch_size = FLAGS.num_threads * FLAGS.batch_size
    config.task.validation_data.global_batch_size = global_batch_size
    config.task.validation_data.dtype = 'float32'
  
    task = image_classification.EdgeTPUTask(config.task)
    dataset = task.build_inputs(config.task.validation_data)
  
    evaluator = tflite_imagenet_evaluator.AccuracyEvaluator(
        model_content=model_content,
        dataset=dataset,
        num_threads=FLAGS.num_threads)
  
    evals, corrects = evaluator.evaluate_all()
    accuracy = 100.0 * corrects / evals if evals > 0 else 0
    print('Final accuracy: {}, Evaluated: {}, Correct: {} '.format(
        accuracy, evals, corrects))
  
  
  if __name__ == '__main__':
    flags.mark_flag_as_required('tflite_model_path')
    app.run(main)
   # Copyright 2022 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Masked language task."""
  
  import dataclasses
  import tensorflow as tf
  
  from official.core import base_task
  from official.core import config_definitions as cfg
  from official.core import task_factory
  from official.modeling import tf_utils
  from official.nlp.configs import bert
  from official.nlp.configs import encoders
  from official.nlp.data import data_loader_factory
  from official.nlp.modeling import layers
  from official.nlp.modeling import models
  
  
  @dataclasses.dataclass
  class MaskedLMConfig(cfg.TaskConfig):
    """The model config."""
    model: bert.PretrainerConfig = bert.PretrainerConfig(cls_heads=[
        bert.ClsHeadConfig(
            inner_dim=768, num_classes=2, dropout_rate=0.1, name='next_sentence')
    ])
    # TODO(b/154564893): Mathematically, scale_loss should be True.
    # However, it works better with scale_loss being False.
    scale_loss: bool = False
    train_data: cfg.DataConfig = cfg.DataConfig()
    validation_data: cfg.DataConfig = cfg.DataConfig()
  
  
  @task_factory.register_task_cls(MaskedLMConfig)
  class MaskedLMTask(base_task.Task):
    """Task object for Mask language modeling."""
  
    def _build_encoder(self, encoder_cfg):
      return encoders.build_encoder(encoder_cfg)
  
    def build_model(self, params=None):
      config = params or self.task_config.model
      encoder_cfg = config.encoder
      encoder_network = self._build_encoder(encoder_cfg)
      cls_heads = [
          layers.ClassificationHead(**cfg.as_dict()) for cfg in config.cls_heads
      ] if config.cls_heads else []
      return models.BertPretrainerV2(
          mlm_activation=tf_utils.get_activation(config.mlm_activation),
          mlm_initializer=tf.keras.initializers.TruncatedNormal(
              stddev=config.mlm_initializer_range),
          encoder_network=encoder_network,
          classification_heads=cls_heads)
  
    def build_losses(self,
                     labels,
                     model_outputs,
                     metrics,
                     aux_losses=None) -> tf.Tensor:
      with tf.name_scope('MaskedLMTask/losses'):
        metrics = dict([(metric.name, metric) for metric in metrics])
        lm_prediction_losses = tf.keras.losses.sparse_categorical_crossentropy(
            labels['masked_lm_ids'],
            tf.cast(model_outputs['mlm_logits'], tf.float32),
            from_logits=True)
        lm_label_weights = labels['masked_lm_weights']
        lm_numerator_loss = tf.reduce_sum(lm_prediction_losses *
                                          lm_label_weights)
        lm_denominator_loss = tf.reduce_sum(lm_label_weights)
        mlm_loss = tf.math.divide_no_nan(lm_numerator_loss, lm_denominator_loss)
        metrics['lm_example_loss'].update_state(mlm_loss)
        if 'next_sentence_labels' in labels:
          sentence_labels = labels['next_sentence_labels']
          sentence_outputs = tf.cast(
              model_outputs['next_sentence'], dtype=tf.float32)
          sentence_loss = tf.reduce_mean(
              tf.keras.losses.sparse_categorical_crossentropy(
                  sentence_labels, sentence_outputs, from_logits=True))
          metrics['next_sentence_loss'].update_state(sentence_loss)
          total_loss = mlm_loss + sentence_loss
        else:
          total_loss = mlm_loss
  
        if aux_losses:
          total_loss += tf.add_n(aux_losses)
        return total_loss
  
    def build_inputs(self, params, input_context=None):
      """Returns tf.data.Dataset for pretraining."""
      if params.input_path == 'dummy':
  
        def dummy_data(_):
          dummy_ids = tf.zeros((1, params.seq_length), dtype=tf.int32)
          dummy_lm = tf.zeros((1, params.max_predictions_per_seq), dtype=tf.int32)
          return dict(
              input_word_ids=dummy_ids,
              input_mask=dummy_ids,
              input_type_ids=dummy_ids,
              masked_lm_positions=dummy_lm,
              masked_lm_ids=dummy_lm,
              masked_lm_weights=tf.cast(dummy_lm, dtype=tf.float32),
              next_sentence_labels=tf.zeros((1, 1), dtype=tf.int32))
  
        dataset = tf.data.Dataset.range(1)
        dataset = dataset.repeat()
        dataset = dataset.map(
            dummy_data, num_parallel_calls=tf.data.experimental.AUTOTUNE)
        return dataset
  
      return data_loader_factory.get_data_loader(params).load(input_context)
  
    def build_metrics(self, training=None):
      del training
      metrics = [
          tf.keras.metrics.SparseCategoricalAccuracy(name='masked_lm_accuracy'),
          tf.keras.metrics.Mean(name='lm_example_loss')
      ]
      # TODO(hongkuny): rethink how to manage metrics creation with heads.
      if self.task_config.train_data.use_next_sentence_label:
        metrics.append(
            tf.keras.metrics.SparseCategoricalAccuracy(
                name='next_sentence_accuracy'))
        metrics.append(tf.keras.metrics.Mean(name='next_sentence_loss'))
      return metrics
  
    def process_metrics(self, metrics, labels, model_outputs):
      with tf.name_scope('MaskedLMTask/process_metrics'):
        metrics = dict([(metric.name, metric) for metric in metrics])
        if 'masked_lm_accuracy' in metrics:
          metrics['masked_lm_accuracy'].update_state(
              labels['masked_lm_ids'], model_outputs['mlm_logits'],
              labels['masked_lm_weights'])
        if 'next_sentence_accuracy' in metrics:
          metrics['next_sentence_accuracy'].update_state(
              labels['next_sentence_labels'], model_outputs['next_sentence'])
  
    def train_step(self, inputs, model: tf.keras.Model,
                   optimizer: tf.keras.optimizers.Optimizer, metrics):
      """Does forward and backward.
  
      Args:
        inputs: a dictionary of input tensors.
        model: the model, forward pass definition.
        optimizer: the optimizer for this training step.
        metrics: a nested structure of metrics objects.
  
      Returns:
        A dictionary of logs.
      """
      with tf.GradientTape() as tape:
        outputs = model(inputs, training=True)
        # Computes per-replica loss.
        loss = self.build_losses(
            labels=inputs,
            model_outputs=outputs,
            metrics=metrics,
            aux_losses=model.losses)
        if self.task_config.scale_loss:
          # Scales loss as the default gradients allreduce performs sum inside the
          # optimizer.
          scaled_loss = loss / tf.distribute.get_strategy().num_replicas_in_sync
      tvars = model.trainable_variables
      if self.task_config.scale_loss:
        grads = tape.gradient(scaled_loss, tvars)
      else:
        grads = tape.gradient(loss, tvars)
      optimizer.apply_gradients(list(zip(grads, tvars)))
      self.process_metrics(metrics, inputs, outputs)
      return {self.loss: loss}
  
    def validation_step(self, inputs, model: tf.keras.Model, metrics):
      """Validatation step.
  
      Args:
        inputs: a dictionary of input tensors.
        model: the keras.Model.
        metrics: a nested structure of metrics objects.
  
      Returns:
        A dictionary of logs.
      """
      outputs = self.inference_step(inputs, model)
      loss = self.build_losses(
          labels=inputs,
          model_outputs=outputs,
          metrics=metrics,
          aux_losses=model.losses)
      self.process_metrics(metrics, inputs, outputs)
      return {self.loss: loss}
  
apis: tensorflow.keras.mixed_precision.set_global_policy tensorflow.keras.initializers.get tensorflow.keras.Model numpy.random.randint tensorflow.keras.activations.get tensorflow.test.main tensorflow.keras.Input
apis: tensorflow.io.gfile.GFile
apis: tensorflow.zeros tensorflow.keras.metrics.SparseCategoricalAccuracy tensorflow.add_n tensorflow.distribute.get_strategy tensorflow.keras.initializers.TruncatedNormal tensorflow.keras.metrics.Mean tensorflow.cast tensorflow.name_scope tensorflow.GradientTape tensorflow.keras.losses.sparse_categorical_crossentropy tensorflow.data.Dataset.range tensorflow.math.divide_no_nan tensorflow.reduce_sum