def __init__()

in tensorflow_examples/lite/model_maker/core/task/model_spec/text_spec.py [0:0]

• 38 lines of code
• 2 McCabe index (conditional complexity)

  def __init__(
      self,
      uri='https://tfhub.dev/tensorflow/bert_en_uncased_L-12_H-768_A-12/1',
      model_dir=None,
      seq_len=384,
      query_len=64,
      doc_stride=128,
      dropout_rate=0.1,
      initializer_range=0.02,
      learning_rate=8e-5,
      distribution_strategy='mirrored',
      num_gpus=-1,
      tpu='',
      trainable=True,
      predict_batch_size=8,
      do_lower_case=True,
      is_tf2=True,
      tflite_input_name=None,
      tflite_output_name=None,
      init_from_squad_model=False,
      default_batch_size=16,
      name='Bert'):
    """Initialze an instance with model paramaters.

    Args:
      uri: TF-Hub path/url to Bert module.
      model_dir: The location of the model checkpoint files.
      seq_len: Length of the sequence to feed into the model.
      query_len: Length of the query to feed into the model.
      doc_stride: The stride when we do a sliding window approach to take chunks
        of the documents.
      dropout_rate: The rate for dropout.
      initializer_range: The stdev of the truncated_normal_initializer for
        initializing all weight matrices.
      learning_rate: The initial learning rate for Adam.
      distribution_strategy:  A string specifying which distribution strategy to
        use. Accepted values are 'off', 'one_device', 'mirrored',
        'parameter_server', 'multi_worker_mirrored', and 'tpu' -- case
        insensitive. 'off' means not to use Distribution Strategy; 'tpu' means
        to use TPUStrategy using `tpu_address`.
      num_gpus: How many GPUs to use at each worker with the
        DistributionStrategies API. The default is -1, which means utilize all
        available GPUs.
      tpu: TPU address to connect to.
      trainable: boolean, whether pretrain layer is trainable.
      predict_batch_size: Batch size for prediction.
      do_lower_case: boolean, whether to lower case the input text. Should be
        True for uncased models and False for cased models.
      is_tf2: boolean, whether the hub module is in TensorFlow 2.x format.
      tflite_input_name: Dict, input names for the TFLite model.
      tflite_output_name: Dict, output names for the TFLite model.
      init_from_squad_model: boolean, whether to initialize from the model that
        is already retrained on Squad 1.1.
      default_batch_size: Default batch size for training.
      name: Name of the object.
    """
    super(BertQAModelSpec,
          self).__init__(uri, model_dir, seq_len, dropout_rate,
                         initializer_range, learning_rate,
                         distribution_strategy, num_gpus, tpu, trainable,
                         do_lower_case, is_tf2, name, tflite_input_name,
                         default_batch_size)
    self.query_len = query_len
    self.doc_stride = doc_stride
    self.predict_batch_size = predict_batch_size
    if tflite_output_name is None:
      tflite_output_name = {
          'start_logits': 'StatefulPartitionedCall:1',
          'end_logits': 'StatefulPartitionedCall:0'
      }
    self.tflite_output_name = tflite_output_name
    self.init_from_squad_model = init_from_squad_model