Returns model_fn closure for TPUEstimator .

def model_fn_builder(bert_config, num_labels, init_checkpoint, learning_rate,
                     num_train_steps, num_warmup_steps, use_tpu,
                     use_one_hot_embeddings):
  """Returns `model_fn` closure for TPUEstimator."""

  def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
    """The `model_fn` for TPUEstimator."""

    tf.logging.info("*** Features ***")
    for name in sorted(features.keys()):
      tf.logging.info("  name = %s, shape = %s" % (name, features[name].shape))

    input_ids = features["input_ids"]
    input_mask = features["input_mask"]
    segment_ids = features["segment_ids"]
    label_ids = features["label_ids"]
    is_real_example = None
    if "is_real_example" in features:
      is_real_example = tf.cast(features["is_real_example"], dtype=tf.float32)
    else:
      is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32)

    is_training = (mode == tf.estimator.ModeKeys.TRAIN)

    (total_loss, per_example_loss, logits, probabilities) = create_model(
        bert_config, is_training, input_ids, input_mask, segment_ids, label_ids,
        num_labels, use_one_hot_embeddings)

    tvars = tf.trainable_variables()
    initialized_variable_names = {}
    scaffold_fn = None
    if init_checkpoint:
      (assignment_map, initialized_variable_names
       ) = modeling.get_assignment_map_from_checkpoint(tvars, init_checkpoint)
      if use_tpu:

        def tpu_scaffold():
          tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
          return tf.train.Scaffold()

        scaffold_fn = tpu_scaffold
      else:
        tf.train.init_from_checkpoint(init_checkpoint, assignment_map)

    tf.logging.info("**** Trainable Variables ****")
    for var in tvars:
      init_string = ""
      if var.name in initialized_variable_names:
        init_string = ", *INIT_FROM_CKPT*"
      tf.logging.info("  name = %s, shape = %s%s", var.name, var.shape,
                      init_string)

    output_spec = None
    if mode == tf.estimator.ModeKeys.TRAIN:

      train_op = optimization.create_optimizer(
          total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu)

      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          loss=total_loss,
          train_op=train_op,
          scaffold_fn=scaffold_fn)
    elif mode == tf.estimator.ModeKeys.EVAL:

      def metric_fn(per_example_loss, label_ids, logits, is_real_example):
        predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
        accuracy = tf.metrics.accuracy(
            labels=label_ids, predictions=predictions, weights=is_real_example)
        loss = tf.metrics.mean(values=per_example_loss, weights=is_real_example)
        return {
            "eval_accuracy": accuracy,
            "eval_loss": loss,
        }

      eval_metrics = (metric_fn,
                      [per_example_loss, label_ids, logits, is_real_example])
      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          loss=total_loss,
          eval_metrics=eval_metrics,
          scaffold_fn=scaffold_fn)
    else:
      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          predictions={"probabilities": probabilities},
          scaffold_fn=scaffold_fn)
    return output_spec

  return model_fn

def
model_fn_builder
(
bert_config
,
num_labels
,
init_checkpoint
,
learning_rate
,
num_train_steps
,
num_warmup_steps
,
use_tpu
,
use_one_hot_embeddings
)
:
def
model_fn
(
features
,
labels
,
mode
,
params
)
:
# pylint: disable=unused-argument
"""The `model_fn` for TPUEstimator."""
tf
.
logging
.
info
(
"*** Features ***"
)
for
name
in
sorted
(
features
.
keys
(
)
)
:
tf
.
logging
.
info
(
"  name = %s, shape = %s"
%
(
name
,
features
[
name
]
.
shape
)
)
input_ids
=
features
[
"input_ids"
]
input_mask
=
features
[
"input_mask"
]
segment_ids
=
features
[
"segment_ids"
]
label_ids
=
features
[
"label_ids"
]
is_real_example
=
if
"is_real_example"
in
features
:
is_real_example
=
tf
.
cast
(
features
[
"is_real_example"
]
,
dtype
=
tf
.
float32
)
else
:
is_real_example
=
tf
.
ones
(
tf
.
shape
(
label_ids
)
,
dtype
=
tf
.
float32
)
is_training
=
(
mode
==
tf
.
estimator
.
ModeKeys
.
TRAIN
)
(
total_loss
,
per_example_loss
,
logits
,
probabilities
)
=
create_model
(
bert_config
,
is_training
,
input_ids
,
input_mask
,
segment_ids
,
label_ids
,
num_labels
,
use_one_hot_embeddings
)
tvars
=
tf
.
trainable_variables
(
)
initialized_variable_names
=
{
}
scaffold_fn
=
if
init_checkpoint
:
(
assignment_map
,
initialized_variable_names
)
=
modeling
.
get_assignment_map_from_checkpoint
(
tvars
,
init_checkpoint
)
if
use_tpu
:
def
tpu_scaffold
(
)
:
tf
.
train
.
init_from_checkpoint
(
init_checkpoint
,
assignment_map
)
return
tf
.
train
.
Scaffold
(
)
scaffold_fn
=
tpu_scaffold
else
:
tf
.
train
.
init_from_checkpoint
(
init_checkpoint
,
assignment_map
)
tf
.
logging
.
info
(
"**** Trainable Variables ****"
)
for
var
in
tvars
:
init_string
=
""
if
var
.
name
in
initialized_variable_names
:
init_string
=
", *INIT_FROM_CKPT*"
tf
.
logging
.
info
(
"  name = %s, shape = %s%s"
,
var
.
name
,
var
.
shape
,
init_string
)
output_spec
=
if
mode
==
tf
.
estimator
.
ModeKeys
.
TRAIN
:
train_op
=
optimization
.
create_optimizer
(
total_loss
,
learning_rate
,
num_train_steps
,
num_warmup_steps
,
use_tpu
)
output_spec
=
tf
.
contrib
.
tpu
.
TPUEstimatorSpec
(
mode
=
mode
,
loss
=
total_loss
,
train_op
=
train_op
,
scaffold_fn
=
scaffold_fn
)
elif
mode
==
tf
.
estimator
.
ModeKeys
.
EVAL
:
def
metric_fn
(
per_example_loss
,
label_ids
,
logits
,
is_real_example
)
:
predictions
=
tf
.
argmax
(
logits
,
axis
=
1
,
output_type
=
tf
.
int32
)
accuracy
=
tf
.
metrics
.
accuracy
(
labels
=
label_ids
,
predictions
=
predictions
,
weights
=
is_real_example
)
loss
=
tf
.
metrics
.
mean
(
values
=
per_example_loss
,
weights
=
is_real_example
)
return
{
"eval_accuracy"
:
accuracy
,
"eval_loss"
:
loss
,
}
eval_metrics
=
(
metric_fn
,
[
per_example_loss
,
label_ids
,
logits
,
is_real_example
]
)
output_spec
=
tf
.
contrib
.
tpu
.
TPUEstimatorSpec
(
mode
=
mode
,
loss
=
total_loss
,
eval_metrics
=
eval_metrics
,
scaffold_fn
=
scaffold_fn
)
else
:
output_spec
=
tf
.
contrib
.
tpu
.
TPUEstimatorSpec
(
mode
=
mode
,
predictions
=
{
"probabilities"
:
probabilities
}
,
scaffold_fn
=
scaffold_fn
)
return
output_spec
return
model_fn


Creates an input_fn closure to be passed to TPUEstimator .

def input_fn_builder(features, seq_length, is_training, drop_remainder):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  all_input_ids = []
  all_input_mask = []
  all_segment_ids = []
  all_label_ids = []

  for feature in features:
    all_input_ids.append(feature.input_ids)
    all_input_mask.append(feature.input_mask)
    all_segment_ids.append(feature.segment_ids)
    all_label_ids.append(feature.label_id)

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    num_examples = len(features)

    # This is for demo purposes and does NOT scale to large data sets. We do
    # not use Dataset.from_generator() because that uses tf.py_func which is
    # not TPU compatible. The right way to load data is with TFRecordReader.
    d = tf.data.Dataset.from_tensor_slices({
        "input_ids":
            tf.constant(
                all_input_ids, shape=[num_examples, seq_length],
                dtype=tf.int32),
        "input_mask":
            tf.constant(
                all_input_mask,
                shape=[num_examples, seq_length],
                dtype=tf.int32),
        "segment_ids":
            tf.constant(
                all_segment_ids,
                shape=[num_examples, seq_length],
                dtype=tf.int32),
        "label_ids":
            tf.constant(all_label_ids, shape=[num_examples], dtype=tf.int32),
    })

    if is_training:
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.batch(batch_size=batch_size, drop_remainder=drop_remainder)
    return d

  return input_fn


Convert a set of InputExample s to a list of InputFeatures .

def convert_examples_to_features(examples, label_list, max_seq_length,
                                 tokenizer):
  """Convert a set of `InputExample`s to a list of `InputFeatures`."""

  features = []
  for (ex_index, example) in enumerate(examples):
    if ex_index % 10000 == 0:
      tf.logging.info("Writing example %d of %d" % (ex_index, len(examples)))

    feature = convert_single_example(ex_index, example, label_list,
                                     max_seq_length, tokenizer)

    features.append(feature)
  return features

def
convert_examples_to_features
(
examples
,
label_list
,
max_seq_length
,
tokenizer
)
:
features
=
[
]
for
(
ex_index
,
example
)
in
enumerate
(
examples
)
:
if
ex_index
%
10000
==
0
:
tf
.
logging
.
info
(
"Writing example %d of %d"
%
(
ex_index
,
len
(
examples
)
)
)
feature
=
convert_single_example
(
ex_index
,
example
,
label_list
,
max_seq_length
,
tokenizer
)
features
.
append
(
feature
)
return
features


def create_optimizer(loss, init_lr, num_train_steps, num_warmup_steps, use_tpu):
  """Creates an optimizer training op."""
  global_step = tf.train.get_or_create_global_step()

  learning_rate = tf.constant(value=init_lr, shape=[], dtype=tf.float32)

  # Implements linear decay of the learning rate.
  learning_rate = tf.train.polynomial_decay(
      learning_rate,
      global_step,
      num_train_steps,
      end_learning_rate=0.0,
      power=1.0,
      cycle=False)

  # Implements linear warmup. I.e., if global_step < num_warmup_steps, the
  # learning rate will be `global_step/num_warmup_steps * init_lr`.
  if num_warmup_steps:
    global_steps_int = tf.cast(global_step, tf.int32)
    warmup_steps_int = tf.constant(num_warmup_steps, dtype=tf.int32)

    global_steps_float = tf.cast(global_steps_int, tf.float32)
    warmup_steps_float = tf.cast(warmup_steps_int, tf.float32)

    warmup_percent_done = global_steps_float / warmup_steps_float
    warmup_learning_rate = init_lr * warmup_percent_done

    is_warmup = tf.cast(global_steps_int < warmup_steps_int, tf.float32)
    learning_rate = (
        (1.0 - is_warmup) * learning_rate + is_warmup * warmup_learning_rate)

  # It is recommended that you use this optimizer for fine tuning, since this
  # is how the model was trained (note that the Adam m/v variables are NOT
  # loaded from init_checkpoint.)
  optimizer = AdamWeightDecayOptimizer(
      learning_rate=learning_rate,
      weight_decay_rate=0.01,
      beta_1=0.9,
      beta_2=0.999,
      epsilon=1e-6,
      exclude_from_weight_decay=["LayerNorm", "layer_norm", "bias"])

  if use_tpu:
    optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)

  tvars = tf.trainable_variables()
  grads = tf.gradients(loss, tvars)

  # This is how the model was pre-trained.
  (grads, _) = tf.clip_by_global_norm(grads, clip_norm=1.0)

  train_op = optimizer.apply_gradients(
      zip(grads, tvars), global_step=global_step)

  # Normally the global step update is done inside of `apply_gradients`.
  # However, `AdamWeightDecayOptimizer` doesn't do this. But if you use
  # a different optimizer, you should probably take this line out.
  new_global_step = global_step + 1
  train_op = tf.group(train_op, [global_step.assign(new_global_step)])
  return train_op


def __init__(self,
               learning_rate,
               weight_decay_rate=0.0,
               beta_1=0.9,
               beta_2=0.999,
               epsilon=1e-6,
               exclude_from_weight_decay=None,
               name="AdamWeightDecayOptimizer"):
    """Constructs a AdamWeightDecayOptimizer."""
    super(AdamWeightDecayOptimizer, self).__init__(False, name)

    self.learning_rate = learning_rate
    self.weight_decay_rate = weight_decay_rate
    self.beta_1 = beta_1
    self.beta_2 = beta_2
    self.epsilon = epsilon
    self.exclude_from_weight_decay = exclude_from_weight_decay

def __init__ ( self , learning_rate , weight_decay_rate = 0.0 , beta_1 = 0.9 , beta_2 = 0.999 , epsilon = 1e-6 , exclude_from_weight_decay = None , name = "AdamWeightDecayOptimizer" ) : super ( AdamWeightDecayOptimizer , self ) . __init__ ( False , name ) self . learning_rate = learning_rate self . weight_decay_rate = weight_decay_rate self . beta_1 = beta_1 self . beta_2 = beta_2 self . epsilon = epsilon self . exclude_from_weight_decay = exclude_from_weight_decay


def apply_gradients(self, grads_and_vars, global_step=None, name=None):
    """See base class."""
    assignments = []
    for (grad, param) in grads_and_vars:
      if grad is None or param is None:
        continue

      param_name = self._get_variable_name(param.name)

      m = tf.get_variable(
          name=param_name + "/adam_m",
          shape=param.shape.as_list(),
          dtype=tf.float32,
          trainable=False,
          initializer=tf.zeros_initializer())
      v = tf.get_variable(
          name=param_name + "/adam_v",
          shape=param.shape.as_list(),
          dtype=tf.float32,
          trainable=False,
          initializer=tf.zeros_initializer())

      # Standard Adam update.
      next_m = (
          tf.multiply(self.beta_1, m) + tf.multiply(1.0 - self.beta_1, grad))
      next_v = (
          tf.multiply(self.beta_2, v) + tf.multiply(1.0 - self.beta_2,
                                                    tf.square(grad)))

      update = next_m / (tf.sqrt(next_v) + self.epsilon)

      # Just adding the square of the weights to the loss function is *not*
      # the correct way of using L2 regularization/weight decay with Adam,
      # since that will interact with the m and v parameters in strange ways.
      #
      # Instead we want ot decay the weights in a manner that doesn't interact
      # with the m/v parameters. This is equivalent to adding the square
      # of the weights to the loss with plain (non-momentum) SGD.
      if self._do_use_weight_decay(param_name):
        update += self.weight_decay_rate * param

      update_with_lr = self.learning_rate * update

      next_param = param - update_with_lr

      assignments.extend(
          [param.assign(next_param),
           m.assign(next_m),
           v.assign(next_v)])
    return tf.group(*assignments, name=name)


Whether to use L2 weight decay for param_name .

def _do_use_weight_decay(self, param_name):
    """Whether to use L2 weight decay for `param_name`."""
    if not self.weight_decay_rate:
      return False
    if self.exclude_from_weight_decay:
      for r in self.exclude_from_weight_decay:
        if re.search(r, param_name) is not None:
          return False
    return True

def _do_use_weight_decay ( self , param_name ) : if not self . weight_decay_rate : return False if self . exclude_from_weight_decay : for r in self . exclude_from_weight_decay : if re . search ( r , param_name ) is not None : return False return True


Get the variable name from the tensor name .

def _get_variable_name(self, param_name):
    """Get the variable name from the tensor name."""
    m = re.match("^(.*):\\d+$", param_name)
    if m is not None:
      param_name = m.group(1)
    return param_name

def _get_variable_name ( self , param_name ) : m = re . match ( "^(.*):\\d+$" , param_name ) if m is not None : param_name = m . group ( 1 ) return param_name


Create TF example files from TrainingInstance s .

def write_instance_to_example_files(instances, tokenizer, max_seq_length,
                                    max_predictions_per_seq, output_files):
  """Create TF example files from `TrainingInstance`s."""
  writers = []
  for output_file in output_files:
    writers.append(tf.python_io.TFRecordWriter(output_file))

  writer_index = 0

  total_written = 0
  for (inst_index, instance) in enumerate(instances):
    input_ids = tokenizer.convert_tokens_to_ids(instance.tokens)
    input_mask = [1] * len(input_ids)
    segment_ids = list(instance.segment_ids)
    assert len(input_ids) <= max_seq_length

    while len(input_ids) < max_seq_length:
      input_ids.append(0)
      input_mask.append(0)
      segment_ids.append(0)

    assert len(input_ids) == max_seq_length
    assert len(input_mask) == max_seq_length
    assert len(segment_ids) == max_seq_length

    masked_lm_positions = list(instance.masked_lm_positions)
    masked_lm_ids = tokenizer.convert_tokens_to_ids(instance.masked_lm_labels)
    masked_lm_weights = [1.0] * len(masked_lm_ids)

    while len(masked_lm_positions) < max_predictions_per_seq:
      masked_lm_positions.append(0)
      masked_lm_ids.append(0)
      masked_lm_weights.append(0.0)

    next_sentence_label = 1 if instance.is_random_next else 0

    features = collections.OrderedDict()
    features["input_ids"] = create_int_feature(input_ids)
    features["input_mask"] = create_int_feature(input_mask)
    features["segment_ids"] = create_int_feature(segment_ids)
    features["masked_lm_positions"] = create_int_feature(masked_lm_positions)
    features["masked_lm_ids"] = create_int_feature(masked_lm_ids)
    features["masked_lm_weights"] = create_float_feature(masked_lm_weights)
    features["next_sentence_labels"] = create_int_feature([next_sentence_label])

    tf_example = tf.train.Example(features=tf.train.Features(feature=features))

    writers[writer_index].write(tf_example.SerializeToString())
    writer_index = (writer_index + 1) % len(writers)

    total_written += 1

    if inst_index < 20:
      tf.logging.info("*** Example ***")
      tf.logging.info("tokens: %s" % " ".join(
          [tokenization.printable_text(x) for x in instance.tokens]))

      for feature_name in features.keys():
        feature = features[feature_name]
        values = []
        if feature.int64_list.value:
          values = feature.int64_list.value
        elif feature.float_list.value:
          values = feature.float_list.value
        tf.logging.info(
            "%s: %s" % (feature_name, " ".join([str(x) for x in values])))

  for writer in writers:
    writer.close()

  tf.logging.info("Wrote %d total instances", total_written)

def
write_instance_to_example_files
(
instances
,
tokenizer
,
max_seq_length
,
max_predictions_per_seq
,
output_files
)
:
writers
=
[
]
for
output_file
in
output_files
:
writers
.
append
(
tf
.
python_io
.
TFRecordWriter
(
output_file
)
)
writer_index
=
0
total_written
=
0
for
(
inst_index
,
instance
)
in
enumerate
(
instances
)
:
input_ids
=
tokenizer
.
convert_tokens_to_ids
(
instance
.
tokens
)
input_mask
=
[
1
]
*
len
(
input_ids
)
segment_ids
=
list
(
instance
.
segment_ids
)
assert
len
(
input_ids
)
<=
max_seq_length
while
len
(
input_ids
)
<
max_seq_length
:
input_ids
.
append
(
0
)
input_mask
.
append
(
0
)
segment_ids
.
append
(
0
)
assert
len
(
input_ids
)
==
max_seq_length
assert
len
(
input_mask
)
==
max_seq_length
assert
len
(
segment_ids
)
==
max_seq_length
masked_lm_positions
=
list
(
instance
.
masked_lm_positions
)
masked_lm_ids
=
tokenizer
.
convert_tokens_to_ids
(
instance
.
masked_lm_labels
)
masked_lm_weights
=
[
1.0
]
*
len
(
masked_lm_ids
)
while
len
(
masked_lm_positions
)
<
max_predictions_per_seq
:
masked_lm_positions
.
append
(
0
)
masked_lm_ids
.
append
(
0
)
masked_lm_weights
.
append
(
0.0
)
next_sentence_label
=
1
if
instance
.
is_random_next
else
0
features
=
collections
.
OrderedDict
(
)
features
[
"input_ids"
]
=
create_int_feature
(
input_ids
)
features
[
"input_mask"
]
=
create_int_feature
(
input_mask
)
features
[
"segment_ids"
]
=
create_int_feature
(
segment_ids
)
features
[
"masked_lm_positions"
]
=
create_int_feature
(
masked_lm_positions
)
features
[
"masked_lm_ids"
]
=
create_int_feature
(
masked_lm_ids
)
features
[
"masked_lm_weights"
]
=
create_float_feature
(
masked_lm_weights
)
features
[
"next_sentence_labels"
]
=
create_int_feature
(
[
next_sentence_label
]
)
tf_example
=
tf
.
train
.
Example
(
features
=
tf
.
train
.
Features
(
feature
=
features
)
)
writers
[
writer_index
]
.
write
(
tf_example
.
SerializeToString
(
)
)
writer_index
=
(
writer_index
+
1
)
%
len
(
writers
)
total_written
+=
1
if
inst_index
<
20
:
tf
.
logging
.
info
(
"*** Example ***"
)
tf
.
logging
.
info
(
"tokens: %s"
%
" "
.
join
(
[
tokenization
.
printable_text
(
x
)
for
x
in
instance
.
tokens
]
)
)
for
feature_name
in
features
.
keys
(
)
:
feature
=
features
[
feature_name
]
values
=
[
]
if
feature
.
int64_list
.
value
:
values
=
feature
.
int64_list
.
value
elif
feature
.
float_list
.
value
:
values
=
feature
.
float_list
.
value
tf
.
logging
.
info
(
"%s: %s"
%
(
feature_name
,
" "
.
join
(
[
str
(
x
)
for
x
in
values
]
)
)
)
for
writer
in
writers
:
writer
.
close
(
)
tf
.
logging
.
info
(
"Wrote %d total instances"
,
total_written
)


Create TrainingInstance s from raw text .

def create_training_instances(input_files, tokenizer, max_seq_length,
                              dupe_factor, short_seq_prob, masked_lm_prob,
                              max_predictions_per_seq, rng):
  """Create `TrainingInstance`s from raw text."""
  all_documents = [[]]

  # Input file format:
  # (1) One sentence per line. These should ideally be actual sentences, not
  # entire paragraphs or arbitrary spans of text. (Because we use the
  # sentence boundaries for the "next sentence prediction" task).
  # (2) Blank lines between documents. Document boundaries are needed so
  # that the "next sentence prediction" task doesn't span between documents.
  for input_file in input_files:
    with tf.gfile.GFile(input_file, "r") as reader:
      while True:
        line = tokenization.convert_to_unicode(reader.readline())
        if not line:
          break
        line = line.strip()

        # Empty lines are used as document delimiters
        if not line:
          all_documents.append([])
        tokens = tokenizer.tokenize(line)
        if tokens:
          all_documents[-1].append(tokens)

  # Remove empty documents
  all_documents = [x for x in all_documents if x]
  rng.shuffle(all_documents)

  vocab_words = list(tokenizer.vocab.keys())
  instances = []
  for _ in range(dupe_factor):
    for document_index in range(len(all_documents)):
      instances.extend(
          create_instances_from_document(
              all_documents, document_index, max_seq_length, short_seq_prob,
              masked_lm_prob, max_predictions_per_seq, vocab_words, rng))

  rng.shuffle(instances)
  return instances


Creates TrainingInstance s for a single document .

def create_instances_from_document(
    all_documents, document_index, max_seq_length, short_seq_prob,
    masked_lm_prob, max_predictions_per_seq, vocab_words, rng):
  """Creates `TrainingInstance`s for a single document."""
  document = all_documents[document_index]

  # Account for [CLS], [SEP], [SEP]
  max_num_tokens = max_seq_length - 3

  # We *usually* want to fill up the entire sequence since we are padding
  # to `max_seq_length` anyways, so short sequences are generally wasted
  # computation. However, we *sometimes*
  # (i.e., short_seq_prob == 0.1 == 10% of the time) want to use shorter
  # sequences to minimize the mismatch between pre-training and fine-tuning.
  # The `target_seq_length` is just a rough target however, whereas
  # `max_seq_length` is a hard limit.
  target_seq_length = max_num_tokens
  if rng.random() < short_seq_prob:
    target_seq_length = rng.randint(2, max_num_tokens)

  # We DON'T just concatenate all of the tokens from a document into a long
  # sequence and choose an arbitrary split point because this would make the
  # next sentence prediction task too easy. Instead, we split the input into
  # segments "A" and "B" based on the actual "sentences" provided by the user
  # input.
  instances = []
  current_chunk = []
  current_length = 0
  i = 0
  while i < len(document):
    segment = document[i]
    current_chunk.append(segment)
    current_length += len(segment)
    if i == len(document) - 1 or current_length >= target_seq_length:
      if current_chunk:
        # `a_end` is how many segments from `current_chunk` go into the `A`
        # (first) sentence.
        a_end = 1
        if len(current_chunk) >= 2:
          a_end = rng.randint(1, len(current_chunk) - 1)

        tokens_a = []
        for j in range(a_end):
          tokens_a.extend(current_chunk[j])

        tokens_b = []
        # Random next
        is_random_next = False
        if len(current_chunk) == 1 or rng.random() < 0.5:
          is_random_next = True
          target_b_length = target_seq_length - len(tokens_a)

          # This should rarely go for more than one iteration for large
          # corpora. However, just to be careful, we try to make sure that
          # the random document is not the same as the document
          # we're processing.
          for _ in range(10):
            random_document_index = rng.randint(0, len(all_documents) - 1)
            if random_document_index != document_index:
              break

          random_document = all_documents[random_document_index]
          random_start = rng.randint(0, len(random_document) - 1)
          for j in range(random_start, len(random_document)):
            tokens_b.extend(random_document[j])
            if len(tokens_b) >= target_b_length:
              break
          # We didn't actually use these segments so we "put them back" so
          # they don't go to waste.
          num_unused_segments = len(current_chunk) - a_end
          i -= num_unused_segments
        # Actual next
        else:
          is_random_next = False
          for j in range(a_end, len(current_chunk)):
            tokens_b.extend(current_chunk[j])
        truncate_seq_pair(tokens_a, tokens_b, max_num_tokens, rng)

        assert len(tokens_a) >= 1
        assert len(tokens_b) >= 1

        tokens = []
        segment_ids = []
        tokens.append("[CLS]")
        segment_ids.append(0)
        for token in tokens_a:
          tokens.append(token)
          segment_ids.append(0)

        tokens.append("[SEP]")
        segment_ids.append(0)

        for token in tokens_b:
          tokens.append(token)
          segment_ids.append(1)
        tokens.append("[SEP]")
        segment_ids.append(1)

        (tokens, masked_lm_positions,
         masked_lm_labels) = create_masked_lm_predictions(
             tokens, masked_lm_prob, max_predictions_per_seq, vocab_words, rng)
        instance = TrainingInstance(
            tokens=tokens,
            segment_ids=segment_ids,
            is_random_next=is_random_next,
            masked_lm_positions=masked_lm_positions,
            masked_lm_labels=masked_lm_labels)
        instances.append(instance)
      current_chunk = []
      current_length = 0
    i += 1

  return instances


Creates the predictions for the masked LM objective .

def create_masked_lm_predictions(tokens, masked_lm_prob,
                                 max_predictions_per_seq, vocab_words, rng):
  """Creates the predictions for the masked LM objective."""

  cand_indexes = []
  for (i, token) in enumerate(tokens):
    if token == "[CLS]" or token == "[SEP]":
      continue
    # Whole Word Masking means that if we mask all of the wordpieces
    # corresponding to an original word. When a word has been split into
    # WordPieces, the first token does not have any marker and any subsequence
    # tokens are prefixed with ##. So whenever we see the ## token, we
    # append it to the previous set of word indexes.
    #
    # Note that Whole Word Masking does *not* change the training code
    # at all -- we still predict each WordPiece independently, softmaxed
    # over the entire vocabulary.
    if (FLAGS.do_whole_word_mask and len(cand_indexes) >= 1 and
        token.startswith("##")):
      cand_indexes[-1].append(i)
    else:
      cand_indexes.append([i])

  rng.shuffle(cand_indexes)

  output_tokens = list(tokens)

  num_to_predict = min(max_predictions_per_seq,
                       max(1, int(round(len(tokens) * masked_lm_prob))))

  masked_lms = []
  covered_indexes = set()
  for index_set in cand_indexes:
    if len(masked_lms) >= num_to_predict:
      break
    # If adding a whole-word mask would exceed the maximum number of
    # predictions, then just skip this candidate.
    if len(masked_lms) + len(index_set) > num_to_predict:
      continue
    is_any_index_covered = False
    for index in index_set:
      if index in covered_indexes:
        is_any_index_covered = True
        break
    if is_any_index_covered:
      continue
    for index in index_set:
      covered_indexes.add(index)

      masked_token = None
      # 80% of the time, replace with [MASK]
      if rng.random() < 0.8:
        masked_token = "[MASK]"
      else:
        # 10% of the time, keep original
        if rng.random() < 0.5:
          masked_token = tokens[index]
        # 10% of the time, replace with random word
        else:
          masked_token = vocab_words[rng.randint(0, len(vocab_words) - 1)]

      output_tokens[index] = masked_token

      masked_lms.append(MaskedLmInstance(index=index, label=tokens[index]))
  assert len(masked_lms) <= num_to_predict
  masked_lms = sorted(masked_lms, key=lambda x: x.index)

  masked_lm_positions = []
  masked_lm_labels = []
  for p in masked_lms:
    masked_lm_positions.append(p.index)
    masked_lm_labels.append(p.label)

  return (output_tokens, masked_lm_positions, masked_lm_labels)


Truncates a pair of sequences to a maximum sequence length .

def truncate_seq_pair(tokens_a, tokens_b, max_num_tokens, rng):
  """Truncates a pair of sequences to a maximum sequence length."""
  while True:
    total_length = len(tokens_a) + len(tokens_b)
    if total_length <= max_num_tokens:
      break

    trunc_tokens = tokens_a if len(tokens_a) > len(tokens_b) else tokens_b
    assert len(trunc_tokens) >= 1

    # We want to sometimes truncate from the front and sometimes from the
    # back to add more randomness and avoid biases.
    if rng.random() < 0.5:
      del trunc_tokens[0]
    else:
      trunc_tokens.pop()


def __init__(self, guid, text, label=None):
        """Constructs a InputExample.

        Args:
          guid: Unique id for the example.
          text_a: string. The untokenized text of the first sequence. For single
            sequence tasks, only this sequence must be specified.
          label: (Optional) string. The label of the example. This should be
            specified for train and dev examples, but not for test examples.
        """
        self.guid = guid
        self.text = text
        self.label = label

def __init__ ( self , guid , text , label = None ) : self . guid = guid self . text = text self . label = label


Gets a collection of InputExample s for the train set .

def get_train_examples(self, data_dir):
        """Gets a collection of `InputExample`s for the train set."""
        raise NotImplementedError()

def get_train_examples ( self , data_dir ) : raise NotImplementedError ( )


Gets a collection of InputExample s for the dev set .

def get_dev_examples(self, data_dir):
        """Gets a collection of `InputExample`s for the dev set."""
        raise NotImplementedError()

def get_dev_examples ( self , data_dir ) : raise NotImplementedError ( )


Gets the list of labels for this data set .

def get_labels(self):
        """Gets the list of labels for this data set."""
        raise NotImplementedError()

def get_labels ( self ) : raise NotImplementedError ( )


def _read_data(cls, input_file):
        """Reads a BIO data."""
        with open(input_file) as f:
            lines = []
            words = []
            labels = []
            for line in f:
                contends = line.strip()
                word = line.strip().split(' ')[0]
                label = line.strip().split(' ')[-1]
                if contends.startswith("-DOCSTART-"):
                    words.append('')
                    continue
                if len(contends) == 0 and words[-1] == '.':
                    l = ' '.join([label for label in labels if len(label) > 0])
                    w = ' '.join([word for word in words if len(word) > 0])
                    lines.append([l, w])
                    words = []
                    labels = []
                    continue
                if len(contends) == 0:
                    continue
                words.append(word)
                labels.append(label)
            return lines

def
_read_data
(
cls
,
input_file
)
:
with
open
(
input_file
)
as
f
:
lines
=
[
]
words
=
[
]
labels
=
[
]
for
line
in
f
:
contends
=
line
.
strip
(
)
word
=
line
.
strip
(
)
.
split
(
' '
)
[
0
]
label
=
line
.
strip
(
)
.
split
(
' '
)
[
1
]
if
contends
.
startswith
(
"-DOCSTART-"
)
:
words
.
append
(
''
)
continue
if
len
(
contends
)
==
0
and
words
[
1
]
==
'.'
:
l
=
' '
.
join
(
[
label
for
label
in
labels
if
len
(
label
)
>
0
]
)
w
=
' '
.
join
(
[
word
for
word
in
words
if
len
(
word
)
>
0
]
)
lines
.
append
(
[
l
,
w
]
)
words
=
[
]
labels
=
[
]
continue
if
len
(
contends
)
==
0
:
continue
words
.
append
(
word
)
labels
.
append
(
label
)
return
lines


def input_fn_builder(features, seq_length):
  """Creates an `input_fn` closure to be passed to TPUEstimator."""

  all_unique_ids = []
  all_input_ids = []
  all_input_mask = []
  all_input_type_ids = []

  for feature in features:
    all_unique_ids.append(feature.unique_id)
    all_input_ids.append(feature.input_ids)
    all_input_mask.append(feature.input_mask)
    all_input_type_ids.append(feature.input_type_ids)

  def input_fn(params):
    """The actual input function."""
    batch_size = params["batch_size"]

    num_examples = len(features)

    # This is for demo purposes and does NOT scale to large data sets. We do
    # not use Dataset.from_generator() because that uses tf.py_func which is
    # not TPU compatible. The right way to load data is with TFRecordReader.
    d = tf.data.Dataset.from_tensor_slices({
        "unique_ids":
            tf.constant(all_unique_ids, shape=[num_examples], dtype=tf.int32),
        "input_ids":
            tf.constant(
                all_input_ids, shape=[num_examples, seq_length],
                dtype=tf.int32),
        "input_mask":
            tf.constant(
                all_input_mask,
                shape=[num_examples, seq_length],
                dtype=tf.int32),
        "input_type_ids":
            tf.constant(
                all_input_type_ids,
                shape=[num_examples, seq_length],
                dtype=tf.int32),
    })

    d = d.batch(batch_size=batch_size, drop_remainder=False)
    return d

  return input_fn


def model_fn_builder(bert_config, init_checkpoint, layer_indexes, use_tpu,
                     use_one_hot_embeddings):
  """Returns `model_fn` closure for TPUEstimator."""

  def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
    """The `model_fn` for TPUEstimator."""

    unique_ids = features["unique_ids"]
    input_ids = features["input_ids"]
    input_mask = features["input_mask"]
    input_type_ids = features["input_type_ids"]

    model = modeling.BertModel(
        config=bert_config,
        is_training=False,
        input_ids=input_ids,
        input_mask=input_mask,
        token_type_ids=input_type_ids,
        use_one_hot_embeddings=use_one_hot_embeddings)

    if mode != tf.estimator.ModeKeys.PREDICT:
      raise ValueError("Only PREDICT modes are supported: %s" % (mode))

    tvars = tf.trainable_variables()
    scaffold_fn = None
    (assignment_map,
     initialized_variable_names) = modeling.get_assignment_map_from_checkpoint(
         tvars, init_checkpoint)
    if use_tpu:

      def tpu_scaffold():
        tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
        return tf.train.Scaffold()

      scaffold_fn = tpu_scaffold
    else:
      tf.train.init_from_checkpoint(init_checkpoint, assignment_map)

    tf.logging.info("**** Trainable Variables ****")
    for var in tvars:
      init_string = ""
      if var.name in initialized_variable_names:
        init_string = ", *INIT_FROM_CKPT*"
      tf.logging.info("  name = %s, shape = %s%s", var.name, var.shape,
                      init_string)

    all_layers = model.get_all_encoder_layers()

    predictions = {
        "unique_id": unique_ids,
    }

    for (i, layer_index) in enumerate(layer_indexes):
      predictions["layer_output_%d" % i] = all_layers[layer_index]

    output_spec = tf.contrib.tpu.TPUEstimatorSpec(
        mode=mode, predictions=predictions, scaffold_fn=scaffold_fn)
    return output_spec

  return model_fn

def
model_fn_builder
(
bert_config
,
init_checkpoint
,
layer_indexes
,
use_tpu
,
use_one_hot_embeddings
)
:
def
model_fn
(
features
,
labels
,
mode
,
params
)
:
# pylint: disable=unused-argument
"""The `model_fn` for TPUEstimator."""
unique_ids
=
features
[
"unique_ids"
]
input_ids
=
features
[
"input_ids"
]
input_mask
=
features
[
"input_mask"
]
input_type_ids
=
features
[
"input_type_ids"
]
model
=
modeling
.
BertModel
(
config
=
bert_config
,
is_training
=
False
,
input_ids
=
input_ids
,
input_mask
=
input_mask
,
token_type_ids
=
input_type_ids
,
use_one_hot_embeddings
=
use_one_hot_embeddings
)
if
mode
!=
tf
.
estimator
.
ModeKeys
.
PREDICT
:
raise
ValueError
(
"Only PREDICT modes are supported: %s"
%
(
mode
)
)
tvars
=
tf
.
trainable_variables
(
)
scaffold_fn
=
(
assignment_map
,
initialized_variable_names
)
=
modeling
.
get_assignment_map_from_checkpoint
(
tvars
,
init_checkpoint
)
if
use_tpu
:
def
tpu_scaffold
(
)
:
tf
.
train
.
init_from_checkpoint
(
init_checkpoint
,
assignment_map
)
return
tf
.
train
.
Scaffold
(
)
scaffold_fn
=
tpu_scaffold
else
:
tf
.
train
.
init_from_checkpoint
(
init_checkpoint
,
assignment_map
)
tf
.
logging
.
info
(
"**** Trainable Variables ****"
)
for
var
in
tvars
:
init_string
=
""
if
var
.
name
in
initialized_variable_names
:
init_string
=
", *INIT_FROM_CKPT*"
tf
.
logging
.
info
(
"  name = %s, shape = %s%s"
,
var
.
name
,
var
.
shape
,
init_string
)
all_layers
=
model
.
get_all_encoder_layers
(
)
predictions
=
{
"unique_id"
:
unique_ids
,
}
for
(
i
,
layer_index
)
in
enumerate
(
layer_indexes
)
:
predictions
[
"layer_output_%d"
%
i
]
=
all_layers
[
layer_index
]
output_spec
=
tf
.
contrib
.
tpu
.
TPUEstimatorSpec
(
mode
=
mode
,
predictions
=
predictions
,
scaffold_fn
=
scaffold_fn
)
return
output_spec
return
model_fn


Loads a data file into a list of InputBatch s .

def convert_examples_to_features(examples, seq_length, tokenizer):
  """Loads a data file into a list of `InputBatch`s."""

  features = []
  for (ex_index, example) in enumerate(examples):
    tokens_a = tokenizer.tokenize(example.text_a)

    tokens_b = None
    if example.text_b:
      tokens_b = tokenizer.tokenize(example.text_b)

    if tokens_b:
      # Modifies `tokens_a` and `tokens_b` in place so that the total
      # length is less than the specified length.
      # Account for [CLS], [SEP], [SEP] with "- 3"
      _truncate_seq_pair(tokens_a, tokens_b, seq_length - 3)
    else:
      # Account for [CLS] and [SEP] with "- 2"
      if len(tokens_a) > seq_length - 2:
        tokens_a = tokens_a[0:(seq_length - 2)]

    # The convention in BERT is:
    # (a) For sequence pairs:
    #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]
    #  type_ids: 0     0  0    0    0     0       0 0     1  1  1  1   1 1
    # (b) For single sequences:
    #  tokens:   [CLS] the dog is hairy . [SEP]
    #  type_ids: 0     0   0   0  0     0 0
    #
    # Where "type_ids" are used to indicate whether this is the first
    # sequence or the second sequence. The embedding vectors for `type=0` and
    # `type=1` were learned during pre-training and are added to the wordpiece
    # embedding vector (and position vector). This is not *strictly* necessary
    # since the [SEP] token unambiguously separates the sequences, but it makes
    # it easier for the model to learn the concept of sequences.
    #
    # For classification tasks, the first vector (corresponding to [CLS]) is
    # used as as the "sentence vector". Note that this only makes sense because
    # the entire model is fine-tuned.
    tokens = []
    input_type_ids = []
    tokens.append("[CLS]")
    input_type_ids.append(0)
    for token in tokens_a:
      tokens.append(token)
      input_type_ids.append(0)
    tokens.append("[SEP]")
    input_type_ids.append(0)

    if tokens_b:
      for token in tokens_b:
        tokens.append(token)
        input_type_ids.append(1)
      tokens.append("[SEP]")
      input_type_ids.append(1)

    input_ids = tokenizer.convert_tokens_to_ids(tokens)

    # The mask has 1 for real tokens and 0 for padding tokens. Only real
    # tokens are attended to.
    input_mask = [1] * len(input_ids)

    # Zero-pad up to the sequence length.
    while len(input_ids) < seq_length:
      input_ids.append(0)
      input_mask.append(0)
      input_type_ids.append(0)

    assert len(input_ids) == seq_length
    assert len(input_mask) == seq_length
    assert len(input_type_ids) == seq_length

    if ex_index < 5:
      tf.logging.info("*** Example ***")
      tf.logging.info("unique_id: %s" % (example.unique_id))
      tf.logging.info("tokens: %s" % " ".join(
          [tokenization.printable_text(x) for x in tokens]))
      tf.logging.info("input_ids: %s" % " ".join([str(x) for x in input_ids]))
      tf.logging.info("input_mask: %s" % " ".join([str(x) for x in input_mask]))
      tf.logging.info(
          "input_type_ids: %s" % " ".join([str(x) for x in input_type_ids]))

    features.append(
        InputFeatures(
            unique_id=example.unique_id,
            tokens=tokens,
            input_ids=input_ids,
            input_mask=input_mask,
            input_type_ids=input_type_ids))
  return features


Truncates a sequence pair in place to the maximum length .

def _truncate_seq_pair(tokens_a, tokens_b, max_length):
  """Truncates a sequence pair in place to the maximum length."""

  # This is a simple heuristic which will always truncate the longer sequence
  # one token at a time. This makes more sense than truncating an equal percent
  # of tokens from each, since if one sequence is very short then each token
  # that's truncated likely contains more information than a longer sequence.
  while True:
    total_length = len(tokens_a) + len(tokens_b)
    if total_length <= max_length:
      break
    if len(tokens_a) > len(tokens_b):
      tokens_a.pop()
    else:
      tokens_b.pop()


Read a list of InputExample s from an input file .

def read_examples(input_file):
  """Read a list of `InputExample`s from an input file."""
  examples = []
  unique_id = 0
  with tf.gfile.GFile(input_file, "r") as reader:
    while True:
      line = tokenization.convert_to_unicode(reader.readline())
      if not line:
        break
      line = line.strip()
      text_a = None
      text_b = None
      m = re.match(r"^(.*) \|\|\| (.*)$", line)
      if m is None:
        text_a = line
      else:
        text_a = m.group(1)
        text_b = m.group(2)
      examples.append(
          InputExample(unique_id=unique_id, text_a=text_a, text_b=text_b))
      unique_id += 1
  return examples

def
read_examples
(
input_file
)
:
examples
=
[
]
unique_id
=
0
with
tf
.
gfile
.
GFile
(
input_file
,
"r"
)
as
reader
:
while
True
:
line
=
tokenization
.
convert_to_unicode
(
reader
.
readline
(
)
)
if
not
line
:
break
line
=
line
.
strip
(
)
text_a
=
text_b
=
m
=
re
.
match
(
r"^(.*) \|\|\| (.*)$"
,
line
)
if
m
is
:
text_a
=
line
else
:
text_a
=
m
.
group
(
1
)
text_b
=
m
.
group
(
2
)
examples
.
append
(
InputExample
(
unique_id
=
unique_id
,
text_a
=
text_a
,
text_b
=
text_b
)
)
unique_id
+=
1
return
examples


def create_model(is_training, input_ids, input_mask, segment_ids, labels,
                 num_labels, bert_hub_module_handle):
  """Creates a classification model."""
  tags = set()
  if is_training:
    tags.add("train")
  bert_module = hub.Module(bert_hub_module_handle, tags=tags, trainable=True)
  bert_inputs = dict(
      input_ids=input_ids,
      input_mask=input_mask,
      segment_ids=segment_ids)
  bert_outputs = bert_module(
      inputs=bert_inputs,
      signature="tokens",
      as_dict=True)

  # In the demo, we are doing a simple classification task on the entire
  # segment.
  #
  # If you want to use the token-level output, use
  # bert_outputs["sequence_output"] instead.
  output_layer = bert_outputs["pooled_output"]

  hidden_size = output_layer.shape[-1].value

  output_weights = tf.get_variable(
      "output_weights", [num_labels, hidden_size],
      initializer=tf.truncated_normal_initializer(stddev=0.02))

  output_bias = tf.get_variable(
      "output_bias", [num_labels], initializer=tf.zeros_initializer())

  with tf.variable_scope("loss"):
    if is_training:
      # I.e., 0.1 dropout
      output_layer = tf.nn.dropout(output_layer, keep_prob=0.9)

    logits = tf.matmul(output_layer, output_weights, transpose_b=True)
    logits = tf.nn.bias_add(logits, output_bias)
    probabilities = tf.nn.softmax(logits, axis=-1)
    log_probs = tf.nn.log_softmax(logits, axis=-1)

    one_hot_labels = tf.one_hot(labels, depth=num_labels, dtype=tf.float32)

    per_example_loss = -tf.reduce_sum(one_hot_labels * log_probs, axis=-1)
    loss = tf.reduce_mean(per_example_loss)

    return (loss, per_example_loss, logits, probabilities)


def model_fn_builder(num_labels, learning_rate, num_train_steps,
                     num_warmup_steps, use_tpu, bert_hub_module_handle):
  """Returns `model_fn` closure for TPUEstimator."""

  def model_fn(features, labels, mode, params):  # pylint: disable=unused-argument
    """The `model_fn` for TPUEstimator."""

    tf.logging.info("*** Features ***")
    for name in sorted(features.keys()):
      tf.logging.info("  name = %s, shape = %s" % (name, features[name].shape))

    input_ids = features["input_ids"]
    input_mask = features["input_mask"]
    segment_ids = features["segment_ids"]
    label_ids = features["label_ids"]

    is_training = (mode == tf.estimator.ModeKeys.TRAIN)

    (total_loss, per_example_loss, logits, probabilities) = create_model(
        is_training, input_ids, input_mask, segment_ids, label_ids, num_labels,
        bert_hub_module_handle)

    output_spec = None
    if mode == tf.estimator.ModeKeys.TRAIN:
      train_op = optimization.create_optimizer(
          total_loss, learning_rate, num_train_steps, num_warmup_steps, use_tpu)

      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          loss=total_loss,
          train_op=train_op)
    elif mode == tf.estimator.ModeKeys.EVAL:

      def metric_fn(per_example_loss, label_ids, logits):
        predictions = tf.argmax(logits, axis=-1, output_type=tf.int32)
        accuracy = tf.metrics.accuracy(label_ids, predictions)
        loss = tf.metrics.mean(per_example_loss)
        return {
            "eval_accuracy": accuracy,
            "eval_loss": loss,
        }

      eval_metrics = (metric_fn, [per_example_loss, label_ids, logits])
      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode,
          loss=total_loss,
          eval_metrics=eval_metrics)
    elif mode == tf.estimator.ModeKeys.PREDICT:
      output_spec = tf.contrib.tpu.TPUEstimatorSpec(
          mode=mode, predictions={"probabilities": probabilities})
    else:
      raise ValueError(
          "Only TRAIN, EVAL and PREDICT modes are supported: %s" % (mode))

    return output_spec

  return model_fn


Get the vocab file and casing info from the Hub module .

def create_tokenizer_from_hub_module(bert_hub_module_handle):
  """Get the vocab file and casing info from the Hub module."""
  with tf.Graph().as_default():
    bert_module = hub.Module(bert_hub_module_handle)
    tokenization_info = bert_module(signature="tokenization_info", as_dict=True)
    with tf.Session() as sess:
      vocab_file, do_lower_case = sess.run([tokenization_info["vocab_file"],
                                            tokenization_info["do_lower_case"]])
  return tokenization.FullTokenizer(
      vocab_file=vocab_file, do_lower_case=do_lower_case)

def create_tokenizer_from_hub_module ( bert_hub_module_handle ) : with tf . Graph ( ) . as_default ( ) : bert_module = hub . Module ( bert_hub_module_handle ) tokenization_info = bert_module ( signature = "tokenization_info" , as_dict = True ) with tf . Session ( ) as sess : vocab_file , do_lower_case = sess . run ( [ tokenization_info [ "vocab_file" ] , tokenization_info [ "do_lower_case" ] ] ) return tokenization . FullTokenizer ( vocab_file = vocab_file , do_lower_case = do_lower_case )


def gelu(x):
  """Gaussian Error Linear Unit.

  This is a smoother version of the RELU.
  Original paper: https://arxiv.org/abs/1606.08415
  Args:
    x: float Tensor to perform activation.

  Returns:
    `x` with the GELU activation applied.
  """
  cdf = 0.5 * (1.0 + tf.tanh(
      (np.sqrt(2 / np.pi) * (x + 0.044715 * tf.pow(x, 3)))))
  return x * cdf

def gelu ( x ) : cdf = 0.5 * ( 1.0 + tf . tanh ( ( np . sqrt ( 2 / np . pi ) * ( x + 0.044715 * tf . pow ( x , 3 ) ) ) ) ) return x * cdf


def embedding_lookup(x, n_token, d_embed, initializer, use_tpu=True,
                     scope='embedding', reuse=None, dtype=tf.float32):
  """TPU and GPU embedding_lookup function."""
  with tf.variable_scope(scope, reuse=reuse):
    lookup_table = tf.get_variable('lookup_table', [n_token, d_embed],
                                   dtype=dtype, initializer=initializer)
    if use_tpu:
      one_hot_idx = tf.one_hot(x, n_token, dtype=dtype)
      if one_hot_idx.shape.ndims == 2:
        return tf.einsum('in,nd->id', one_hot_idx, lookup_table), lookup_table
      else:
        return tf.einsum('ibn,nd->ibd', one_hot_idx, lookup_table), lookup_table
    else:
      return tf.nn.embedding_lookup(lookup_table, x), lookup_table

def embedding_lookup ( x , n_token , d_embed , initializer , use_tpu = True , scope = 'embedding' , reuse = None , dtype = tf . float32 ) : with tf . variable_scope ( scope , reuse = reuse ) : lookup_table = tf . get_variable ( 'lookup_table' , [ n_token , d_embed ] , dtype = dtype , initializer = initializer ) if use_tpu : one_hot_idx = tf . one_hot ( x , n_token , dtype = dtype ) if one_hot_idx . shape . ndims == 2 : return tf . einsum ( 'in,nd->id' , one_hot_idx , lookup_table ) , lookup_table else : return tf . einsum ( 'ibn,nd->ibd' , one_hot_idx , lookup_table ) , lookup_table else : return tf . nn . embedding_lookup ( lookup_table , x ) , lookup_table


Position - wise Feed - forward Network .

def positionwise_ffn(inp, d_model, d_inner, dropout, kernel_initializer,
                     activation_type='relu', scope='ff', is_training=True,
                     reuse=None):
  """Position-wise Feed-forward Network."""
  if activation_type == 'relu':
    activation = tf.nn.relu
  elif activation_type == 'gelu':
    activation = gelu
  else:
    raise ValueError('Unsupported activation type {}'.format(activation_type))

  output = inp
  with tf.variable_scope(scope, reuse=reuse):
    output = tf.layers.dense(output, d_inner, activation=activation,
                             kernel_initializer=kernel_initializer,
                             name='layer_1')
    output = tf.layers.dropout(output, dropout, training=is_training,
                               name='drop_1')
    output = tf.layers.dense(output, d_model,
                             kernel_initializer=kernel_initializer,
                             name='layer_2')
    output = tf.layers.dropout(output, dropout, training=is_training,
                               name='drop_2')
    output = tf.contrib.layers.layer_norm(output + inp, begin_norm_axis=-1,
                                          scope='LayerNorm')
  return output

def
positionwise_ffn
(
inp
,
d_model
,
d_inner
,
dropout
,
kernel_initializer
,
activation_type
=
'relu'
,
scope
=
'ff'
,
is_training
=
True
,
reuse
=
)
:
if
activation_type
==
'relu'
:
activation
=
tf
.
nn
.
relu
elif
activation_type
==
'gelu'
:
activation
=
gelu
else
:
raise
ValueError
(
'Unsupported activation type {}'
.
format
(
activation_type
)
)
output
=
inp
with
tf
.
variable_scope
(
scope
,
reuse
=
reuse
)
:
output
=
tf
.
layers
.
dense
(
output
,
d_inner
,
activation
=
activation
,
kernel_initializer
=
kernel_initializer
,
name
=
'layer_1'
)
output
=
tf
.
layers
.
dropout
(
output
,
dropout
,
training
=
is_training
,
name
=
'drop_1'
)
output
=
tf
.
layers
.
dense
(
output
,
d_model
,
kernel_initializer
=
kernel_initializer
,
name
=
'layer_2'
)
output
=
tf
.
layers
.
dropout
(
output
,
dropout
,
training
=
is_training
,
name
=
'drop_2'
)
output
=
tf
.
contrib
.
layers
.
layer_norm
(
output
+
inp
,
begin_norm_axis
=
1
,
scope
=
'LayerNorm'
)
return
output


Project hidden states to a specific head with a 4D - shape .

def head_projection(h, d_model, n_head, d_head, kernel_initializer, name):
  """Project hidden states to a specific head with a 4D-shape."""
  proj_weight = tf.get_variable('{}/kernel'.format(name),
                                [d_model, n_head, d_head], dtype=h.dtype,
                                initializer=kernel_initializer)
  head = tf.einsum('ibh,hnd->ibnd', h, proj_weight)

  return head

def head_projection ( h , d_model , n_head , d_head , kernel_initializer , name ) : proj_weight = tf . get_variable ( '{}/kernel' . format ( name ) , [ d_model , n_head , d_head ] , dtype = h . dtype , initializer = kernel_initializer ) head = tf . einsum ( 'ibh,hnd->ibnd' , h , proj_weight ) return head


def post_attention(h, attn_vec, d_model, n_head, d_head, dropout, is_training,
                   kernel_initializer, residual=True):
  """Post-attention processing."""
  # post-attention projection (back to `d_model`)
  proj_o = tf.get_variable('o/kernel', [d_model, n_head, d_head],
                           dtype=h.dtype, initializer=kernel_initializer)
  attn_out = tf.einsum('ibnd,hnd->ibh', attn_vec, proj_o)

  attn_out = tf.layers.dropout(attn_out, dropout, training=is_training)
  if residual:
    output = tf.contrib.layers.layer_norm(attn_out + h, begin_norm_axis=-1,
                                          scope='LayerNorm')
  else:
    output = tf.contrib.layers.layer_norm(attn_out, begin_norm_axis=-1,
                                          scope='LayerNorm')

  return output


Core absolute positional attention operations .

def abs_attn_core(q_head, k_head, v_head, attn_mask, dropatt, is_training,
                  scale):
  """Core absolute positional attention operations."""

  attn_score = tf.einsum('ibnd,jbnd->ijbn', q_head, k_head)
  attn_score *= scale
  if attn_mask is not None:
    attn_score = attn_score - 1e30 * attn_mask

  # attention probability
  attn_prob = tf.nn.softmax(attn_score, 1)
  attn_prob = tf.layers.dropout(attn_prob, dropatt, training=is_training)

  # attention output
  attn_vec = tf.einsum('ijbn,jbnd->ibnd', attn_prob, v_head)

  return attn_vec

def
abs_attn_core
(
q_head
,
k_head
,
v_head
,
attn_mask
,
dropatt
,
is_training
,
scale
)
:
attn_score
=
tf
.
einsum
(
'ibnd,jbnd->ijbn'
,
q_head
,
k_head
)
attn_score
*=
scale
if
attn_mask
is
not
:
attn_score
=
attn_score
1e30
*
attn_mask
# attention probability
attn_prob
=
tf
.
nn
.
softmax
(
attn_score
,
1
)
attn_prob
=
tf
.
layers
.
dropout
(
attn_prob
,
dropatt
,
training
=
is_training
)
# attention output
attn_vec
=
tf
.
einsum
(
'ijbn,jbnd->ibnd'
,
attn_prob
,
v_head
)
return
attn_vec


Core relative positional attention operations .

def rel_attn_core(q_head, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat,
                  r_w_bias, r_r_bias, r_s_bias, attn_mask, dropatt, is_training,
                  scale):
  """Core relative positional attention operations."""

  # content based attention score
  ac = tf.einsum('ibnd,jbnd->ijbn', q_head + r_w_bias, k_head_h)

  # position based attention score
  bd = tf.einsum('ibnd,jbnd->ijbn', q_head + r_r_bias, k_head_r)
  bd = rel_shift(bd, klen=tf.shape(ac)[1])

  # segment based attention score
  if seg_mat is None:
    ef = 0
  else:
    ef = tf.einsum('ibnd,snd->ibns', q_head + r_s_bias, seg_embed)
    ef = tf.einsum('ijbs,ibns->ijbn', seg_mat, ef)

  # merge attention scores and perform masking
  attn_score = (ac + bd + ef) * scale
  if attn_mask is not None:
    # attn_score = attn_score * (1 - attn_mask) - 1e30 * attn_mask
    attn_score = attn_score - 1e30 * attn_mask

  # attention probability
  attn_prob = tf.nn.softmax(attn_score, 1)
  attn_prob = tf.layers.dropout(attn_prob, dropatt, training=is_training)

  # attention output
  attn_vec = tf.einsum('ijbn,jbnd->ibnd', attn_prob, v_head_h)

  return attn_vec


perform relative shift to form the relative attention score .

def rel_shift(x, klen=-1):
  """perform relative shift to form the relative attention score."""
  x_size = tf.shape(x)

  x = tf.reshape(x, [x_size[1], x_size[0], x_size[2], x_size[3]])
  x = tf.slice(x, [1, 0, 0, 0], [-1, -1, -1, -1])
  x = tf.reshape(x, [x_size[0], x_size[1] - 1, x_size[2], x_size[3]])
  x = tf.slice(x, [0, 0, 0, 0], [-1, klen, -1, -1])

  return x


def _create_mask(qlen, mlen, dtype=tf.float32, same_length=False):
  """create causal attention mask."""
  attn_mask = tf.ones([qlen, qlen], dtype=dtype)
  mask_u = tf.matrix_band_part(attn_mask, 0, -1)
  mask_dia = tf.matrix_band_part(attn_mask, 0, 0)
  attn_mask_pad = tf.zeros([qlen, mlen], dtype=dtype)
  ret = tf.concat([attn_mask_pad, mask_u - mask_dia], 1)
  if same_length:
    mask_l = tf.matrix_band_part(attn_mask, -1, 0)
    ret = tf.concat([ret[:, :qlen] + mask_l - mask_dia, ret[:, qlen:]], 1)

  return ret

def _create_mask ( qlen , mlen , dtype = tf . float32 , same_length = False ) : attn_mask = tf . ones ( [ qlen , qlen ] , dtype = dtype ) mask_u = tf . matrix_band_part ( attn_mask , 0 , - 1 ) mask_dia = tf . matrix_band_part ( attn_mask , 0 , 0 ) attn_mask_pad = tf . zeros ( [ qlen , mlen ] , dtype = dtype ) ret = tf . concat ( [ attn_mask_pad , mask_u - mask_dia ] , 1 ) if same_length : mask_l = tf . matrix_band_part ( attn_mask , - 1 , 0 ) ret = tf . concat ( [ ret [ : , : qlen ] + mask_l - mask_dia , ret [ : , qlen : ] ] , 1 ) return ret


def _cache_mem(curr_out, prev_mem, mem_len, reuse_len=None):
  """cache hidden states into memory."""
  if mem_len is None or mem_len == 0:
    return None
  else:
    if reuse_len is not None and reuse_len > 0:
      curr_out = curr_out[:reuse_len]

    if prev_mem is None:
      new_mem = curr_out[-mem_len:]
    else:
      new_mem = tf.concat([prev_mem, curr_out], 0)[-mem_len:]

  return tf.stop_gradient(new_mem)

def _cache_mem ( curr_out , prev_mem , mem_len , reuse_len = None ) : if mem_len is None or mem_len == 0 : return None else : if reuse_len is not None and reuse_len > 0 : curr_out = curr_out [ : reuse_len ] if prev_mem is None : new_mem = curr_out [ - mem_len : ] else : new_mem = tf . concat ( [ prev_mem , curr_out ] , 0 ) [ - mem_len : ] return tf . stop_gradient ( new_mem )


def relative_positional_encoding(qlen, klen, d_model, clamp_len, attn_type,
                                 bi_data, bsz=None, dtype=None):
  """create relative positional encoding."""
  freq_seq = tf.range(0, d_model, 2.0)
  if dtype is not None and dtype != tf.float32:
    freq_seq = tf.cast(freq_seq, dtype=dtype)
  inv_freq = 1 / (10000 ** (freq_seq / d_model))

  if attn_type == 'bi':
    # beg, end = klen - 1, -qlen
    beg, end = klen, -qlen
  elif attn_type == 'uni':
    # beg, end = klen - 1, -1
    beg, end = klen, -1
  else:
    raise ValueError('Unknown `attn_type` {}.'.format(attn_type))

  if bi_data:
    fwd_pos_seq = tf.range(beg, end, -1.0)
    bwd_pos_seq = tf.range(-beg, -end, 1.0)

    if dtype is not None and dtype != tf.float32:
      fwd_pos_seq = tf.cast(fwd_pos_seq, dtype=dtype)
      bwd_pos_seq = tf.cast(bwd_pos_seq, dtype=dtype)

    if clamp_len > 0:
      fwd_pos_seq = tf.clip_by_value(fwd_pos_seq, -clamp_len, clamp_len)
      bwd_pos_seq = tf.clip_by_value(bwd_pos_seq, -clamp_len, clamp_len)

    if bsz is not None:
      # With bi_data, the batch size should be divisible by 2.
      assert bsz%2 == 0
      fwd_pos_emb = positional_embedding(fwd_pos_seq, inv_freq, bsz//2)
      bwd_pos_emb = positional_embedding(bwd_pos_seq, inv_freq, bsz//2)
    else:
      fwd_pos_emb = positional_embedding(fwd_pos_seq, inv_freq)
      bwd_pos_emb = positional_embedding(bwd_pos_seq, inv_freq)

    pos_emb = tf.concat([fwd_pos_emb, bwd_pos_emb], axis=1)
  else:
    fwd_pos_seq = tf.range(beg, end, -1.0)
    if dtype is not None and dtype != tf.float32:
      fwd_pos_seq = tf.cast(fwd_pos_seq, dtype=dtype)
    if clamp_len > 0:
      fwd_pos_seq = tf.clip_by_value(fwd_pos_seq, -clamp_len, clamp_len)
    pos_emb = positional_embedding(fwd_pos_seq, inv_freq, bsz)

  return pos_emb


Standard multi - head attention with absolute positional embedding .

def multihead_attn(q, k, v, attn_mask, d_model, n_head, d_head, dropout,
                   dropatt, is_training, kernel_initializer, residual=True,
                   scope='abs_attn', reuse=None):
  """Standard multi-head attention with absolute positional embedding."""

  scale = 1 / (d_head ** 0.5)
  with tf.variable_scope(scope, reuse=reuse):
    # attention heads
    q_head = head_projection(
        q, d_model, n_head, d_head, kernel_initializer, 'q')
    k_head = head_projection(
        k, d_model, n_head, d_head, kernel_initializer, 'k')
    v_head = head_projection(
        v, d_model, n_head, d_head, kernel_initializer, 'v')

    # attention vector
    attn_vec = abs_attn_core(q_head, k_head, v_head, attn_mask, dropatt,
                             is_training, scale)

    # post processing
    output = post_attention(v, attn_vec, d_model, n_head, d_head, dropout,
                            is_training, kernel_initializer, residual)

  return output

def
multihead_attn
(
q
,
k
,
v
,
attn_mask
,
d_model
,
n_head
,
d_head
,
dropout
,
dropatt
,
is_training
,
kernel_initializer
,
residual
=
True
,
scope
=
'abs_attn'
,
reuse
=
)
:
scale
=
1
/
(
d_head
**
0.5
)
with
tf
.
variable_scope
(
scope
,
reuse
=
reuse
)
:
# attention heads
q_head
=
head_projection
(
q
,
d_model
,
n_head
,
d_head
,
kernel_initializer
,
'q'
)
k_head
=
head_projection
(
k
,
d_model
,
n_head
,
d_head
,
kernel_initializer
,
'k'
)
v_head
=
head_projection
(
v
,
d_model
,
n_head
,
d_head
,
kernel_initializer
,
'v'
)
# attention vector
attn_vec
=
abs_attn_core
(
q_head
,
k_head
,
v_head
,
attn_mask
,
dropatt
,
is_training
,
scale
)
# post processing
output
=
post_attention
(
v
,
attn_vec
,
d_model
,
n_head
,
d_head
,
dropout
,
is_training
,
kernel_initializer
,
residual
)
return
output


Multi - head attention with relative positional encoding .

def rel_multihead_attn(h, r, r_w_bias, r_r_bias, seg_mat, r_s_bias, seg_embed,
                       attn_mask, mems, d_model, n_head, d_head, dropout,
                       dropatt, is_training, kernel_initializer,
                       scope='rel_attn', reuse=None):
  """Multi-head attention with relative positional encoding."""

  scale = 1 / (d_head ** 0.5)
  with tf.variable_scope(scope, reuse=reuse):
    if mems is not None and mems.shape.ndims > 1:
      cat = tf.concat([mems, h], 0)
    else:
      cat = h

    # content heads
    q_head_h = head_projection(
        h, d_model, n_head, d_head, kernel_initializer, 'q')
    k_head_h = head_projection(
        cat, d_model, n_head, d_head, kernel_initializer, 'k')
    v_head_h = head_projection(
        cat, d_model, n_head, d_head, kernel_initializer, 'v')

    # positional heads
    k_head_r = head_projection(
        r, d_model, n_head, d_head, kernel_initializer, 'r')

    # core attention ops
    attn_vec = rel_attn_core(
        q_head_h, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
        r_r_bias, r_s_bias, attn_mask, dropatt, is_training, scale)

    # post processing
    output = post_attention(h, attn_vec, d_model, n_head, d_head, dropout,
                            is_training, kernel_initializer)

  return output

def
rel_multihead_attn
(
h
,
r
,
r_w_bias
,
r_r_bias
,
seg_mat
,
r_s_bias
,
seg_embed
,
attn_mask
,
mems
,
d_model
,
n_head
,
d_head
,
dropout
,
dropatt
,
is_training
,
kernel_initializer
,
scope
=
'rel_attn'
,
reuse
=
)
:
scale
=
1
/
(
d_head
**
0.5
)
with
tf
.
variable_scope
(
scope
,
reuse
=
reuse
)
:
if
mems
is
not
and
mems
.
shape
.
ndims
>
1
:
cat
=
tf
.
concat
(
[
mems
,
h
]
,
0
)
else
:
cat
=
h
# content heads
q_head_h
=
head_projection
(
h
,
d_model
,
n_head
,
d_head
,
kernel_initializer
,
'q'
)
k_head_h
=
head_projection
(
cat
,
d_model
,
n_head
,
d_head
,
kernel_initializer
,
'k'
)
v_head_h
=
head_projection
(
cat
,
d_model
,
n_head
,
d_head
,
kernel_initializer
,
'v'
)
# positional heads
k_head_r
=
head_projection
(
r
,
d_model
,
n_head
,
d_head
,
kernel_initializer
,
'r'
)
# core attention ops
attn_vec
=
rel_attn_core
(
q_head_h
,
k_head_h
,
v_head_h
,
k_head_r
,
seg_embed
,
seg_mat
,
r_w_bias
,
r_r_bias
,
r_s_bias
,
attn_mask
,
dropatt
,
is_training
,
scale
)
# post processing
output
=
post_attention
(
h
,
attn_vec
,
d_model
,
n_head
,
d_head
,
dropout
,
is_training
,
kernel_initializer
)
return
output


Two - stream attention with relative positional encoding .

def two_stream_rel_attn(h, g, r, mems, r_w_bias, r_r_bias, seg_mat, r_s_bias,
                        seg_embed, attn_mask_h, attn_mask_g, target_mapping,
                        d_model, n_head, d_head, dropout, dropatt, is_training,
                        kernel_initializer, scope='rel_attn'):
  """Two-stream attention with relative positional encoding."""

  scale = 1 / (d_head ** 0.5)
  with tf.variable_scope(scope, reuse=False):

    # content based attention score
    if mems is not None and mems.shape.ndims > 1:
      cat = tf.concat([mems, h], 0)
    else:
      cat = h

    # content-based key head
    k_head_h = head_projection(
        cat, d_model, n_head, d_head, kernel_initializer, 'k')

    # content-based value head
    v_head_h = head_projection(
        cat, d_model, n_head, d_head, kernel_initializer, 'v')

    # position-based key head
    k_head_r = head_projection(
        r, d_model, n_head, d_head, kernel_initializer, 'r')

    ##### h-stream
    # content-stream query head
    q_head_h = head_projection(
        h, d_model, n_head, d_head, kernel_initializer, 'q')

    # core attention ops
    attn_vec_h = rel_attn_core(
        q_head_h, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
        r_r_bias, r_s_bias, attn_mask_h, dropatt, is_training, scale)

    # post processing
    output_h = post_attention(h, attn_vec_h, d_model, n_head, d_head, dropout,
                              is_training, kernel_initializer)

  with tf.variable_scope(scope, reuse=True):
    ##### g-stream
    # query-stream query head
    q_head_g = head_projection(
        g, d_model, n_head, d_head, kernel_initializer, 'q')

    # core attention ops
    if target_mapping is not None:
      q_head_g = tf.einsum('mbnd,mlb->lbnd', q_head_g, target_mapping)
      attn_vec_g = rel_attn_core(
          q_head_g, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
          r_r_bias, r_s_bias, attn_mask_g, dropatt, is_training, scale)
      attn_vec_g = tf.einsum('lbnd,mlb->mbnd', attn_vec_g, target_mapping)
    else:
      attn_vec_g = rel_attn_core(
          q_head_g, k_head_h, v_head_h, k_head_r, seg_embed, seg_mat, r_w_bias,
          r_r_bias, r_s_bias, attn_mask_g, dropatt, is_training, scale)

    # post processing
    output_g = post_attention(g, attn_vec_g, d_model, n_head, d_head, dropout,
                              is_training, kernel_initializer)

    return output_h, output_g

def
two_stream_rel_attn
(
h
,
g
,
r
,
mems
,
r_w_bias
,
r_r_bias
,
seg_mat
,
r_s_bias
,
seg_embed
,
attn_mask_h
,
attn_mask_g
,
target_mapping
,
d_model
,
n_head
,
d_head
,
dropout
,
dropatt
,
is_training
,
kernel_initializer
,
scope
=
'rel_attn'
)
:
scale
=
1
/
(
d_head
**
0.5
)
with
tf
.
variable_scope
(
scope
,
reuse
=
False
)
:
# content based attention score
if
mems
is
not
and
mems
.
shape
.
ndims
>
1
:
cat
=
tf
.
concat
(
[
mems
,
h
]
,
0
)
else
:
cat
=
h
# content-based key head
k_head_h
=
head_projection
(
cat
,
d_model
,
n_head
,
d_head
,
kernel_initializer
,
'k'
)
# content-based value head
v_head_h
=
head_projection
(
cat
,
d_model
,
n_head
,
d_head
,
kernel_initializer
,
'v'
)
# position-based key head
k_head_r
=
head_projection
(
r
,
d_model
,
n_head
,
d_head
,
kernel_initializer
,
'r'
)
##### h-stream
# content-stream query head
q_head_h
=
head_projection
(
h
,
d_model
,
n_head
,
d_head
,
kernel_initializer
,
'q'
)
# core attention ops
attn_vec_h
=
rel_attn_core
(
q_head_h
,
k_head_h
,
v_head_h
,
k_head_r
,
seg_embed
,
seg_mat
,
r_w_bias
,
r_r_bias
,
r_s_bias
,
attn_mask_h
,
dropatt
,
is_training
,
scale
)
# post processing
output_h
=
post_attention
(
h
,
attn_vec_h
,
d_model
,
n_head
,
d_head
,
dropout
,
is_training
,
kernel_initializer
)
with
tf
.
variable_scope
(
scope
,
reuse
=
True
)
:
##### g-stream
# query-stream query head
q_head_g
=
head_projection
(
g
,
d_model
,
n_head
,
d_head
,
kernel_initializer
,
'q'
)
# core attention ops
if
target_mapping
is
not
:
q_head_g
=
tf
.
einsum
(
'mbnd,mlb->lbnd'
,
q_head_g
,
target_mapping
)
attn_vec_g
=
rel_attn_core
(
q_head_g
,
k_head_h
,
v_head_h
,
k_head_r
,
seg_embed
,
seg_mat
,
r_w_bias
,
r_r_bias
,
r_s_bias
,
attn_mask_g
,
dropatt
,
is_training
,
scale
)
attn_vec_g
=
tf
.
einsum
(
'lbnd,mlb->mbnd'
,
attn_vec_g
,
target_mapping
)
else
:
attn_vec_g
=
rel_attn_core
(
q_head_g
,
k_head_h
,
v_head_h
,
k_head_r
,
seg_embed
,
seg_mat
,
r_w_bias
,
r_r_bias
,
r_s_bias
,
attn_mask_g
,
dropatt
,
is_training
,
scale
)
# post processing
output_g
=
post_attention
(
g
,
attn_vec_g
,
d_model
,
n_head
,
d_head
,
dropout
,
is_training
,
kernel_initializer
)
return
output_h
,
output_g


Defines a Transformer - XL computation graph with additional support for XLNet .

def transformer_xl(inp_k, n_token, n_layer, d_model, n_head,
                d_head, d_inner, dropout, dropatt, attn_type,
                bi_data, initializer, is_training, mem_len=None,
                inp_q=None, mems=None,
                same_length=False, clamp_len=-1, untie_r=False,
                use_tpu=True, input_mask=None,
                perm_mask=None, seg_id=None, reuse_len=None,
                ff_activation='relu', target_mapping=None,
                use_bfloat16=False, scope='transformer', **kwargs):
  """
    Defines a Transformer-XL computation graph with additional
    support for XLNet.

    Args:

    inp_k: int32 Tensor in shape [len, bsz], the input token IDs.
    seg_id: int32 Tensor in shape [len, bsz], the input segment IDs.
    input_mask: float32 Tensor in shape [len, bsz], the input mask.
      0 for real tokens and 1 for padding.
    mems: a list of float32 Tensors in shape [mem_len, bsz, d_model], memory
      from previous batches. The length of the list equals n_layer.
      If None, no memory is used.
    perm_mask: float32 Tensor in shape [len, len, bsz].
      If perm_mask[i, j, k] = 0, i attend to j in batch k;
      if perm_mask[i, j, k] = 1, i does not attend to j in batch k.
      If None, each position attends to all the others.
    target_mapping: float32 Tensor in shape [num_predict, len, bsz].
      If target_mapping[i, j, k] = 1, the i-th predict in batch k is
      on the j-th token.
      Only used during pretraining for partial prediction.
      Set to None during finetuning.
    inp_q: float32 Tensor in shape [len, bsz].
      1 for tokens with losses and 0 for tokens without losses.
      Only used during pretraining for two-stream attention.
      Set to None during finetuning.

    n_layer: int, the number of layers.
    d_model: int, the hidden size.
    n_head: int, the number of attention heads.
    d_head: int, the dimension size of each attention head.
    d_inner: int, the hidden size in feed-forward layers.
    ff_activation: str, "relu" or "gelu".
    untie_r: bool, whether to untie the biases in attention.
    n_token: int, the vocab size.

    is_training: bool, whether in training mode.
    use_tpu: bool, whether TPUs are used.
    use_bfloat16: bool, use bfloat16 instead of float32.
    dropout: float, dropout rate.
    dropatt: float, dropout rate on attention probabilities.
    init: str, the initialization scheme, either "normal" or "uniform".
    init_range: float, initialize the parameters with a uniform distribution
      in [-init_range, init_range]. Only effective when init="uniform".
    init_std: float, initialize the parameters with a normal distribution
      with mean 0 and stddev init_std. Only effective when init="normal".
    mem_len: int, the number of tokens to cache.
    reuse_len: int, the number of tokens in the currect batch to be cached
      and reused in the future.
    bi_data: bool, whether to use bidirectional input pipeline.
      Usually set to True during pretraining and False during finetuning.
    clamp_len: int, clamp all relative distances larger than clamp_len.
      -1 means no clamping.
    same_length: bool, whether to use the same attention length for each token.
    summary_type: str, "last", "first", "mean", or "attn". The method
      to pool the input to get a vector representation.
    initializer: A tf initializer.
    scope: scope name for the computation graph.
  """
  tf.logging.info('memory input {}'.format(mems))
  tf_float = tf.bfloat16 if use_bfloat16 else tf.float32
  tf.logging.info('Use float type {}'.format(tf_float))

  new_mems = []
  with tf.variable_scope(scope):
    if untie_r:
      r_w_bias = tf.get_variable('r_w_bias', [n_layer, n_head, d_head],
                                 dtype=tf_float, initializer=initializer)
      r_r_bias = tf.get_variable('r_r_bias', [n_layer, n_head, d_head],
                                 dtype=tf_float, initializer=initializer)
    else:
      r_w_bias = tf.get_variable('r_w_bias', [n_head, d_head],
                                 dtype=tf_float, initializer=initializer)
      r_r_bias = tf.get_variable('r_r_bias', [n_head, d_head],
                                 dtype=tf_float, initializer=initializer)

    bsz = tf.shape(inp_k)[1]
    qlen = tf.shape(inp_k)[0]
    mlen = tf.shape(mems[0])[0] if mems is not None else 0
    klen = mlen + qlen

    ##### Attention mask
    # causal attention mask
    if attn_type == 'uni':
      attn_mask = _create_mask(qlen, mlen, tf_float, same_length)
      attn_mask = attn_mask[:, :, None, None]
    elif attn_type == 'bi':
      attn_mask = None
    else:
      raise ValueError('Unsupported attention type: {}'.format(attn_type))

    # data mask: input mask & perm mask
    if input_mask is not None and perm_mask is not None:
      data_mask = input_mask[None] + perm_mask
    elif input_mask is not None and perm_mask is None:
      data_mask = input_mask[None]
    elif input_mask is None and perm_mask is not None:
      data_mask = perm_mask
    else:
      data_mask = None

    if data_mask is not None:
      # all mems can be attended to
      mems_mask = tf.zeros([tf.shape(data_mask)[0], mlen, bsz],
                           dtype=tf_float)
      data_mask = tf.concat([mems_mask, data_mask], 1)
      if attn_mask is None:
        attn_mask = data_mask[:, :, :, None]
      else:
        attn_mask += data_mask[:, :, :, None]

    if attn_mask is not None:
      attn_mask = tf.cast(attn_mask > 0, dtype=tf_float)

    if attn_mask is not None:
      non_tgt_mask = -tf.eye(qlen, dtype=tf_float)
      non_tgt_mask = tf.concat([tf.zeros([qlen, mlen], dtype=tf_float),
                                non_tgt_mask], axis=-1)
      non_tgt_mask = tf.cast((attn_mask + non_tgt_mask[:, :, None, None]) > 0,
                             dtype=tf_float)
    else:
      non_tgt_mask = None

    ##### Word embedding
    word_emb_k, lookup_table = embedding_lookup(
        x=inp_k,
        n_token=n_token,
        d_embed=d_model,
        initializer=initializer,
        use_tpu=use_tpu,
        dtype=tf_float,
        scope='word_embedding')

    if inp_q is not None:
      with tf.variable_scope('mask_emb'):
        mask_emb = tf.get_variable('mask_emb', [1, 1, d_model], dtype=tf_float)
        if target_mapping is not None:
          word_emb_q = tf.tile(mask_emb, [tf.shape(target_mapping)[0], bsz, 1])
        else:
          inp_q_ext = inp_q[:, :, None]
          word_emb_q = inp_q_ext * mask_emb + (1 - inp_q_ext) * word_emb_k
    output_h = tf.layers.dropout(word_emb_k, dropout, training=is_training)
    if inp_q is not None:
      output_g = tf.layers.dropout(word_emb_q, dropout, training=is_training)

    ##### Segment embedding
    if seg_id is not None:
      if untie_r:
        r_s_bias = tf.get_variable('r_s_bias', [n_layer, n_head, d_head],
                                   dtype=tf_float, initializer=initializer)
      else:
        # default case (tie)
        r_s_bias = tf.get_variable('r_s_bias', [n_head, d_head],
                                   dtype=tf_float, initializer=initializer)

      seg_embed = tf.get_variable('seg_embed', [n_layer, 2, n_head, d_head],
                                  dtype=tf_float, initializer=initializer)

      # Convert `seg_id` to one-hot `seg_mat`
      mem_pad = tf.zeros([mlen, bsz], dtype=tf.int32)
      cat_ids = tf.concat([mem_pad, seg_id], 0)

      # `1` indicates not in the same segment [qlen x klen x bsz]
      seg_mat = tf.cast(
          tf.logical_not(tf.equal(seg_id[:, None], cat_ids[None, :])),
          tf.int32)
      seg_mat = tf.one_hot(seg_mat, 2, dtype=tf_float)
    else:
      seg_mat = None

    ##### Positional encoding
    pos_emb = relative_positional_encoding(
        qlen, klen, d_model, clamp_len, attn_type, bi_data,
        bsz=bsz, dtype=tf_float)
    pos_emb = tf.layers.dropout(pos_emb, dropout, training=is_training)

    ##### Attention layers
    if mems is None:
      mems = [None] * n_layer

    for i in range(n_layer):
      # cache new mems
      new_mems.append(_cache_mem(output_h, mems[i], mem_len, reuse_len))

      # segment bias
      if seg_id is None:
        r_s_bias_i = None
        seg_embed_i = None
      else:
        r_s_bias_i = r_s_bias if not untie_r else r_s_bias[i]
        seg_embed_i = seg_embed[i]

      with tf.variable_scope('layer_{}'.format(i)):
        if inp_q is not None:
          output_h, output_g = two_stream_rel_attn(
              h=output_h,
              g=output_g,
              r=pos_emb,
              r_w_bias=r_w_bias if not untie_r else r_w_bias[i],
              r_r_bias=r_r_bias if not untie_r else r_r_bias[i],
              seg_mat=seg_mat,
              r_s_bias=r_s_bias_i,
              seg_embed=seg_embed_i,
              attn_mask_h=non_tgt_mask,
              attn_mask_g=attn_mask,
              mems=mems[i],
              target_mapping=target_mapping,
              d_model=d_model,
              n_head=n_head,
              d_head=d_head,
              dropout=dropout,
              dropatt=dropatt,
              is_training=is_training,
              kernel_initializer=initializer)
          reuse = True
        else:
          reuse = False

          output_h = rel_multihead_attn(
              h=output_h,
              r=pos_emb,
              r_w_bias=r_w_bias if not untie_r else r_w_bias[i],
              r_r_bias=r_r_bias if not untie_r else r_r_bias[i],
              seg_mat=seg_mat,
              r_s_bias=r_s_bias_i,
              seg_embed=seg_embed_i,
              attn_mask=non_tgt_mask,
              mems=mems[i],
              d_model=d_model,
              n_head=n_head,
              d_head=d_head,
              dropout=dropout,
              dropatt=dropatt,
              is_training=is_training,
              kernel_initializer=initializer,
              reuse=reuse)

        if inp_q is not None:
          output_g = positionwise_ffn(
              inp=output_g,
              d_model=d_model,
              d_inner=d_inner,
              dropout=dropout,
              kernel_initializer=initializer,
              activation_type=ff_activation,
              is_training=is_training)

        output_h = positionwise_ffn(
            inp=output_h,
            d_model=d_model,
            d_inner=d_inner,
            dropout=dropout,
            kernel_initializer=initializer,
            activation_type=ff_activation,
            is_training=is_training,
            reuse=reuse)

    if inp_q is not None:
      output = tf.layers.dropout(output_g, dropout, training=is_training)
    else:
      output = tf.layers.dropout(output_h, dropout, training=is_training)

    return output, new_mems, lookup_table


def lm_loss(hidden, target, n_token, d_model, initializer, lookup_table=None,
            tie_weight=False, bi_data=True, use_tpu=False):
  """doc."""

  with tf.variable_scope('lm_loss'):
    if tie_weight:
      assert lookup_table is not None, \
          'lookup_table cannot be None for tie_weight'
      softmax_w = lookup_table
    else:
      softmax_w = tf.get_variable('weight', [n_token, d_model],
                                  dtype=hidden.dtype, initializer=initializer)

    softmax_b = tf.get_variable('bias', [n_token], dtype=hidden.dtype,
                                initializer=tf.zeros_initializer())

    logits = tf.einsum('ibd,nd->ibn', hidden, softmax_w) + softmax_b

    if use_tpu:
      one_hot_target = tf.one_hot(target, n_token, dtype=logits.dtype)
      loss = -tf.reduce_sum(tf.nn.log_softmax(logits) * one_hot_target, -1)
    else:
      loss = tf.nn.sparse_softmax_cross_entropy_with_logits(labels=target,
                                                            logits=logits)

    return loss


Different classification tasks may not may not share the same parameters to summarize the sequence features .

def summarize_sequence(summary_type, hidden, d_model, n_head, d_head, dropout,
                       dropatt, input_mask, is_training, initializer,
                       scope=None, reuse=None, use_proj=True):

  """
      Different classification tasks may not may not share the same parameters
      to summarize the sequence features.

      If shared, one can keep the `scope` to the default value `None`.
      Otherwise, one should specify a different `scope` for each task.
  """

  with tf.variable_scope(scope, 'sequnece_summary', reuse=reuse):
    if summary_type == 'last':
      summary = hidden[-1]
    elif summary_type == 'first':
      summary = hidden[0]
    elif summary_type == 'mean':
      summary = tf.reduce_mean(hidden, axis=0)
    elif summary_type == 'attn':
      bsz = tf.shape(hidden)[1]

      summary_bias = tf.get_variable('summary_bias', [d_model],
                                     dtype=hidden.dtype,
                                     initializer=initializer)
      summary_bias = tf.tile(summary_bias[None, None], [1, bsz, 1])

      if input_mask is not None:
        input_mask = input_mask[None, :, :, None]

      summary = multihead_attn(summary_bias, hidden, hidden, input_mask,
                               d_model, n_head, d_head, dropout, dropatt,
                               is_training, initializer, residual=False)
      summary = summary[0]
    else:
      raise ValueError('Unsupported summary type {}'.format(summary_type))

    # use another projection as in BERT
    if use_proj:
      summary = tf.layers.dense(
          summary,
          d_model,
          activation=tf.tanh,
          kernel_initializer=initializer,
          name='summary')

    # dropout
    summary = tf.layers.dropout(
        summary, dropout, training=is_training,
        name='dropout')

  return summary

def
summarize_sequence
(
summary_type
,
hidden
,
d_model
,
n_head
,
d_head
,
dropout
,
dropatt
,
input_mask
,
is_training
,
initializer
,
scope
=
,
reuse
=
,
use_proj
=
True
)
:
with
tf
.
variable_scope
(
scope
,
'sequnece_summary'
,
reuse
=
reuse
)
:
if
summary_type
==
'last'
:
summary
=
hidden
[
1
]
elif
summary_type
==
'first'
:
summary
=
hidden
[
0
]
elif
summary_type
==
'mean'
:
summary
=
tf
.
reduce_mean
(
hidden
,
axis
=
0
)
elif
summary_type
==
'attn'
:
bsz
=
tf
.
shape
(
hidden
)
[
1
]
summary_bias
=
tf
.
get_variable
(
'summary_bias'
,
[
d_model
]
,
dtype
=
hidden
.
dtype
,
initializer
=
initializer
)
summary_bias
=
tf
.
tile
(
summary_bias
[
,
]
,
[
1
,
bsz
,
1
]
)
if
input_mask
is
not
:
input_mask
=
input_mask
[
,
:
,
:
,
]
summary
=
multihead_attn
(
summary_bias
,
hidden
,
hidden
,
input_mask
,
d_model
,
n_head
,
d_head
,
dropout
,
dropatt
,
is_training
,
initializer
,
residual
=
False
)
summary
=
summary
[
0
]
else
:
raise
ValueError
(
'Unsupported summary type {}'
.
format
(
summary_type
)
)
# use another projection as in BERT
if
use_proj
:
summary
=
tf
.
layers
.
dense
(
summary
,
d_model
,
activation
=
tf
.
tanh
,
kernel_initializer
=
initializer
,
name
=
'summary'
)
# dropout
summary
=
tf
.
layers
.
dropout
(
summary
,
dropout
,
training
=
is_training
,
name
=
'dropout'
)
return
summary


Different classification tasks should use different scope names to ensure different dense layers ( parameters ) are used to produce the logits .

def classification_loss(hidden, labels, n_class, initializer, scope, reuse=None,
                        return_logits=False):
  """
      Different classification tasks should use different scope names to ensure
      different dense layers (parameters) are used to produce the logits.

      An exception will be in transfer learning, where one hopes to transfer
      the classification weights.
  """

  with tf.variable_scope(scope, reuse=reuse):
    logits = tf.layers.dense(
        hidden,
        n_class,
        kernel_initializer=initializer,
        name='logit')

    one_hot_target = tf.one_hot(labels, n_class, dtype=hidden.dtype)
    loss = -tf.reduce_sum(tf.nn.log_softmax(logits) * one_hot_target, -1)

    if return_logits:
      return loss, logits

    return loss

def classification_loss ( hidden , labels , n_class , initializer , scope , reuse = None , return_logits = False ) : with tf . variable_scope ( scope , reuse = reuse ) : logits = tf . layers . dense ( hidden , n_class , kernel_initializer = initializer , name = 'logit' ) one_hot_target = tf . one_hot ( labels , n_class , dtype = hidden . dtype ) loss = - tf . reduce_sum ( tf . nn . log_softmax ( logits ) * one_hot_target , - 1 ) if return_logits : return loss , logits return loss


def _get_initializer(FLAGS):
  """Get variable intializer."""
  if FLAGS.init == "uniform":
    initializer = tf.initializers.random_uniform(
        minval=-FLAGS.init_range,
        maxval=FLAGS.init_range,
        seed=None)
  elif FLAGS.init == "normal":
    initializer = tf.initializers.random_normal(
        stddev=FLAGS.init_std,
        seed=None)
  else:
    raise ValueError("Initializer {} not supported".format(FLAGS.init))
  return initializer

def
_get_initializer
(
FLAGS
)
:
if
FLAGS
.
init
==
"uniform"
:
initializer
=
tf
.
initializers
.
random_uniform
(
minval
=
FLAGS
.
init_range
,
maxval
=
FLAGS
.
init_range
,
seed
=
)
elif
FLAGS
.
init
==
"normal"
:
initializer
=
tf
.
initializers
.
random_normal
(
stddev
=
FLAGS
.
init_std
,
seed
=
)
else
:
raise
ValueError
(
"Initializer {} not supported"
.
format
(
FLAGS
.
init
)
)
return
initializer


Constructing an XLNetConfig . One of FLAGS or json_path should be provided .

def __init__(self, FLAGS=None, json_path=None):
    """Constructing an XLNetConfig.
    One of FLAGS or json_path should be provided."""

    assert FLAGS is not None or json_path is not None

    self.keys = ["n_layer", "d_model", "n_head", "d_head", "d_inner",
                 "ff_activation", "untie_r", "n_token"]

    if FLAGS is not None:
      self.init_from_flags(FLAGS)

    if json_path is not None:
      self.init_from_json(json_path)

def __init__ ( self , FLAGS = None , json_path = None ) : assert FLAGS is not None or json_path is not None self . keys = [ "n_layer" , "d_model" , "n_head" , "d_head" , "d_inner" , "ff_activation" , "untie_r" , "n_token" ] if FLAGS is not None : self . init_from_flags ( FLAGS ) if json_path is not None : self . init_from_json ( json_path )


def to_json(self, json_path):
    """Save XLNetConfig to a json file."""
    json_data = {}
    for key in self.keys:
      json_data[key] = getattr(self, key)

    json_dir = os.path.dirname(json_path)
    if not tf.gfile.Exists(json_dir):
      tf.gfile.MakeDirs(json_dir)
    with tf.gfile.Open(json_path, "w") as f:
      json.dump(json_data, f, indent=4, sort_keys=True)

def to_json ( self , json_path ) : json_data = { } for key in self . keys : json_data [ key ] = getattr ( self , key ) json_dir = os . path . dirname ( json_path ) if not tf . gfile . Exists ( json_dir ) : tf . gfile . MakeDirs ( json_dir ) with tf . gfile . Open ( json_path , "w" ) as f : json . dump ( json_data , f , indent = 4 , sort_keys = True )


Args : is_training : bool whether in training mode . use_tpu : bool whether TPUs are used . use_bfloat16 : bool use bfloat16 instead of float32 . dropout : float dropout rate . dropatt : float dropout rate on attention probabilities . init : str the initialization scheme either normal or uniform . init_range : float initialize the parameters with a uniform distribution in [ - init_range init_range ] . Only effective when init = uniform . init_std : float initialize the parameters with a normal distribution with mean 0 and stddev init_std . Only effective when init = normal . mem_len : int the number of tokens to cache . reuse_len : int the number of tokens in the currect batch to be cached and reused in the future . bi_data : bool whether to use bidirectional input pipeline . Usually set to True during pretraining and False during finetuning . clamp_len : int clamp all relative distances larger than clamp_len . - 1 means no clamping . same_length : bool whether to use the same attention length for each token .

def __init__(self, is_training, use_tpu, use_bfloat16, dropout, dropatt,
               init="normal", init_range=0.1, init_std=0.02, mem_len=None,
               reuse_len=None, bi_data=False, clamp_len=-1, same_length=False):
    """
    Args:
      is_training: bool, whether in training mode.
      use_tpu: bool, whether TPUs are used.
      use_bfloat16: bool, use bfloat16 instead of float32.
      dropout: float, dropout rate.
      dropatt: float, dropout rate on attention probabilities.
      init: str, the initialization scheme, either "normal" or "uniform".
      init_range: float, initialize the parameters with a uniform distribution
        in [-init_range, init_range]. Only effective when init="uniform".
      init_std: float, initialize the parameters with a normal distribution
        with mean 0 and stddev init_std. Only effective when init="normal".
      mem_len: int, the number of tokens to cache.
      reuse_len: int, the number of tokens in the currect batch to be cached
        and reused in the future.
      bi_data: bool, whether to use bidirectional input pipeline.
        Usually set to True during pretraining and False during finetuning.
      clamp_len: int, clamp all relative distances larger than clamp_len.
        -1 means no clamping.
      same_length: bool, whether to use the same attention length for each token.
    """

    self.init = init
    self.init_range = init_range
    self.init_std = init_std
    self.is_training = is_training
    self.dropout = dropout
    self.dropatt = dropatt
    self.use_tpu = use_tpu
    self.use_bfloat16 = use_bfloat16
    self.mem_len = mem_len
    self.reuse_len = reuse_len
    self.bi_data = bi_data
    self.clamp_len = clamp_len
    self.same_length = same_length

def __init__ ( self , is_training , use_tpu , use_bfloat16 , dropout , dropatt , init = "normal" , init_range = 0.1 , init_std = 0.02 , mem_len = None , reuse_len = None , bi_data = False , clamp_len = - 1 , same_length = False ) : self . init = init self . init_range = init_range self . init_std = init_std self . is_training = is_training self . dropout = dropout self . dropatt = dropatt self . use_tpu = use_tpu self . use_bfloat16 = use_bfloat16 self . mem_len = mem_len self . reuse_len = reuse_len self . bi_data = bi_data self . clamp_len = clamp_len self . same_length = same_length


Args : xlnet_config : XLNetConfig run_config : RunConfig input_ids : int32 Tensor in shape [ len bsz ] the input token IDs . seg_ids : int32 Tensor in shape [ len bsz ] the input segment IDs . input_mask : float32 Tensor in shape [ len bsz ] the input mask . 0 for real tokens and 1 for padding . mems : a list of float32 Tensors in shape [ mem_len bsz d_model ] memory from previous batches . The length of the list equals n_layer . If None no memory is used . perm_mask : float32 Tensor in shape [ len len bsz ] . If perm_mask [ i j k ] = 0 i attend to j in batch k ; if perm_mask [ i j k ] = 1 i does not attend to j in batch k . If None each position attends to all the others . target_mapping : float32 Tensor in shape [ num_predict len bsz ] . If target_mapping [ i j k ] = 1 the i - th predict in batch k is on the j - th token . Only used during pretraining for partial prediction . Set to None during finetuning . inp_q : float32 Tensor in shape [ len bsz ] . 1 for tokens with losses and 0 for tokens without losses . Only used during pretraining for two - stream attention . Set to None during finetuning .

def __init__(self, xlnet_config, run_config, input_ids, seg_ids, input_mask,
               mems=None, perm_mask=None, target_mapping=None, inp_q=None,
               **kwargs):
    """
    Args:
      xlnet_config: XLNetConfig,
      run_config: RunConfig,
      input_ids: int32 Tensor in shape [len, bsz], the input token IDs.
      seg_ids: int32 Tensor in shape [len, bsz], the input segment IDs.
      input_mask: float32 Tensor in shape [len, bsz], the input mask.
        0 for real tokens and 1 for padding.
      mems: a list of float32 Tensors in shape [mem_len, bsz, d_model], memory
        from previous batches. The length of the list equals n_layer.
        If None, no memory is used.
      perm_mask: float32 Tensor in shape [len, len, bsz].
        If perm_mask[i, j, k] = 0, i attend to j in batch k;
        if perm_mask[i, j, k] = 1, i does not attend to j in batch k.
        If None, each position attends to all the others.
      target_mapping: float32 Tensor in shape [num_predict, len, bsz].
        If target_mapping[i, j, k] = 1, the i-th predict in batch k is
        on the j-th token.
        Only used during pretraining for partial prediction.
        Set to None during finetuning.
      inp_q: float32 Tensor in shape [len, bsz].
        1 for tokens with losses and 0 for tokens without losses.
        Only used during pretraining for two-stream attention.
        Set to None during finetuning.
    """

    initializer = _get_initializer(run_config)

    tfm_args = dict(
        n_token=xlnet_config.n_token,
        initializer=initializer,
        attn_type="bi",
        n_layer=xlnet_config.n_layer,
        d_model=xlnet_config.d_model,
        n_head=xlnet_config.n_head,
        d_head=xlnet_config.d_head,
        d_inner=xlnet_config.d_inner,
        ff_activation=xlnet_config.ff_activation,
        untie_r=xlnet_config.untie_r,

        is_training=run_config.is_training,
        use_bfloat16=run_config.use_bfloat16,
        use_tpu=run_config.use_tpu,
        dropout=run_config.dropout,
        dropatt=run_config.dropatt,

        mem_len=run_config.mem_len,
        reuse_len=run_config.reuse_len,
        bi_data=run_config.bi_data,
        clamp_len=run_config.clamp_len,
        same_length=run_config.same_length
    )

    input_args = dict(
        inp_k=input_ids,
        seg_id=seg_ids,
        input_mask=input_mask,
        mems=mems,
        perm_mask=perm_mask,
        target_mapping=target_mapping,
        inp_q=inp_q)
    tfm_args.update(input_args)

    with tf.variable_scope("model", reuse=tf.AUTO_REUSE):
      (self.output, self.new_mems, self.lookup_table
          ) = modeling.transformer_xl(**tfm_args)

    self.input_mask = input_mask
    self.initializer = initializer
    self.xlnet_config = xlnet_config
    self.run_config = run_config


Args : summary_type : str last first mean or attn . The method to pool the input to get a vector representation . use_summ_proj : bool whether to use a linear projection during pooling .

def get_pooled_out(self, summary_type, use_summ_proj=True):
    """
    Args:
      summary_type: str, "last", "first", "mean", or "attn". The method
        to pool the input to get a vector representation.
      use_summ_proj: bool, whether to use a linear projection during pooling.

    Returns:
      float32 Tensor in shape [bsz, d_model], the pooled representation.
    """

    xlnet_config = self.xlnet_config
    run_config = self.run_config

    with tf.variable_scope("model", reuse=tf.AUTO_REUSE):
      summary = modeling.summarize_sequence(
          summary_type=summary_type,
          hidden=self.output,
          d_model=xlnet_config.d_model,
          n_head=xlnet_config.n_head,
          d_head=xlnet_config.d_head,
          dropout=run_config.dropout,
          dropatt=run_config.dropatt,
          is_training=run_config.is_training,
          input_mask=self.input_mask,
          initializer=self.initializer,
          use_proj=use_summ_proj)

    return summary

def get_pooled_out ( self , summary_type , use_summ_proj = True ) : xlnet_config = self . xlnet_config run_config = self . run_config with tf . variable_scope ( "model" , reuse = tf . AUTO_REUSE ) : summary = modeling . summarize_sequence ( summary_type = summary_type , hidden = self . output , d_model = xlnet_config . d_model , n_head = xlnet_config . n_head , d_head = xlnet_config . d_head , dropout = run_config . dropout , dropatt = run_config . dropatt , is_training = run_config . is_training , input_mask = self . input_mask , initializer = self . initializer , use_proj = use_summ_proj ) return summary


Returns : float32 Tensor in shape [ len bsz d_model ] . The last layer hidden representation of XLNet .

def get_sequence_output(self):
    """
    Returns:
      float32 Tensor in shape [len, bsz, d_model]. The last layer hidden
      representation of XLNet.
    """

    return self.output

def get_sequence_output ( self ) : return self . output


Returns : list of float32 Tensors in shape [ mem_len bsz d_model ] the new memory that concatenates the previous memory with the current input representations . The length of the list equals n_layer .

def get_new_memory(self):
    """
    Returns:
      list of float32 Tensors in shape [mem_len, bsz, d_model], the new
      memory that concatenates the previous memory with the current input
      representations.
      The length of the list equals n_layer.
    """
    return self.new_mems

def get_new_memory ( self ) : return self . new_mems


Returns : float32 Tensor in shape [ n_token d_model ] . The embedding lookup table . Used for tying embeddings between input and output layers .

def get_embedding_table(self):
    """
    Returns:
      float32 Tensor in shape [n_token, d_model]. The embedding lookup table.
      Used for tying embeddings between input and output layers.
    """
    return self.lookup_table

def get_embedding_table ( self ) : return self . lookup_table


Returns : A tf initializer . Used to initialize variables in layers on top of XLNet .

def get_initializer(self):
    """
    Returns:
      A tf initializer. Used to initialize variables in layers on top of XLNet.
    """
    return self.initializer

def get_initializer ( self ) : return self . initializer