code for feature request to add sampling and training to transfo_xl models

src/transformers/__init__.py

@@ -176,6 +176,7 @@ if is_torch_available():
         load_tf_weights_in_transfo_xl,
         TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_MAP,
     )
+    from .modeling_transfo_xl_utilities import LogUniformSampler
     from .modeling_gpt2 import (
         GPT2PreTrainedModel,
         GPT2Model,
@@ -344,8 +345,11 @@ if is_tf_available():
         TFTransfoXLModel,
         TFTransfoXLLMHeadModel,
         TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_MAP,
+        TFAdaptiveEmbedding,
     )
 
+    from .modeling_tf_transfo_xl_utilities import TFLogUniformSampler
+
     from .modeling_tf_xlnet import (
         TFXLNetPreTrainedModel,
         TFXLNetMainLayer,

src/transformers/modeling_tf_transfo_xl.py

@@ -23,7 +23,7 @@ import tensorflow as tf
 
 from .configuration_transfo_xl import TransfoXLConfig
 from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
-from .modeling_tf_transfo_xl_utilities import TFAdaptiveSoftmaxMask
+from .modeling_tf_transfo_xl_utilities import TFAdaptiveSoftmaxMask, TFLogUniformSampler, sample_logits
 from .modeling_tf_utils import TFPreTrainedModel, get_initializer, keras_serializable, shape_list
 
 
@@ -733,6 +733,25 @@ class TFTransfoXLModel(TFTransfoXLPreTrainedModel):
         return outputs
 
 
+class TFTransfoXLLMHead(tf.keras.layers.Layer):
+    def __init__(self, config, input_embeddings, **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = config.vocab_size
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.input_embeddings = input_embeddings
+
+    def build(self, input_shape):
+        self.bias = self.add_weight(shape=(self.vocab_size,), initializer="zeros", trainable=True, name="bias")
+        super().build(input_shape)
+
+    def call(self, hidden_states):
+        hidden_states = self.input_embeddings(hidden_states, mode="linear")
+        hidden_states = hidden_states + self.bias
+        return hidden_states
+
+
 @add_start_docstrings(
     """The Transformer-XL Model with a language modeling head on top
     (adaptive softmax with weights tied to the adaptive input embeddings)""",
@@ -746,12 +765,25 @@ class TFTransfoXLLMHeadModel(TFTransfoXLPreTrainedModel):
         # use sampled softmax
         if config.sample_softmax > 0:
             raise NotImplementedError
+            # see PT code for implementation
+            self.out_layer = TFTransfoXLLMHead(config, self.transformer.word_emb.weight, name="out_layer")
+            self.sampler = TFLogUniformSampler(config.vocab_size, config.sample_softmax)
         # use adaptive softmax (including standard softmax)
         else:
             self.crit = TFAdaptiveSoftmaxMask(
                 config.vocab_size, config.d_embed, config.d_model, config.cutoffs, div_val=config.div_val, name="crit"
             )
 
+    def get_output_embeddings(self):
+        """ Double-check if you are using adaptive softmax.
+        """
+        if self.sample_softmax > 0:
+            return self.out_layer
+        else:
+            if len(self.crit.out_layers) > 0:
+                return self.crit.out_layers[-1]
+            return None
+
     def reset_length(self, tgt_len, ext_len, mem_len):
         self.transformer.reset_length(tgt_len, ext_len, mem_len)
 
@@ -820,12 +852,19 @@ class TFTransfoXLLMHeadModel(TFTransfoXLPreTrainedModel):
         last_hidden = transformer_outputs[0]
         pred_hid = last_hidden[:, -tgt_len:]
         outputs = transformer_outputs[1:]
+
         if self.sample_softmax > 0 and training:
             raise NotImplementedError
+            # see PT code for implementation
+            assert self.config.tie_weight
+            logit = sample_logits(self.transformer.word_emb, self.out_layer.bias, labels, pred_hid, self.sampler)
+            softmax_output = -tf.nn.log_softmax(logit, -1)[:, :, 0]
+            outputs = [softmax_output] + outputs
+            if labels is not None:
+                # TODO: This is not implemented
+                raise NotImplementedError
         else:
-            # pred_hid = tf.reshape(pred_hid, (-1, shape_list(pred_hid)[-1]))
             softmax_output = self.crit([pred_hid, labels], training=training)
-            # softmax_output = tf.reshape(softmax_output, (bsz, tgt_len, -1))
             outputs = [softmax_output] + outputs
 
         return outputs  # logits, new_mems, (all hidden states), (all attentions)

src/transformers/modeling_tf_transfo_xl_utilities.py

@@ -176,3 +176,76 @@ class TFAdaptiveSoftmaxMask(tf.keras.layers.Layer):
             self.add_metric(loss, name=self.name, aggregation="mean" if return_mean else "")
 
         return out
+
+
+# TODO: (PVP) add tests to compare to PT!
+class TFLogUniformSampler(object):
+    def __init__(self, range_max, n_sample):
+        """
+        Reference : https://github.com/tensorflow/tensorflow/blob/r1.10/tensorflow/python/ops/candidate_sampling_ops.py
+            `P(class) = (log(class + 2) - log(class + 1)) / log(range_max + 1)`
+
+        expected count can be approximated by 1 - (1 - p)^n
+        and we use a numerically stable version -expm1(num_tries * log1p(-p))
+
+        Our implementation fixes num_tries at 2 * n_sample, and the actual #samples will vary from run to run
+        """
+        self.range_max = range_max
+        log_indices = tf.math.log(tf.range(1.0, range_max + 2.0))
+        self.dist = (log_indices[1:] - log_indices[:-1]) / log_indices[-1]
+
+        self.log_q = tf.cast(
+            tf.math.log(-tf.math.expm1(-tf.math.log1p(tf.cast(self.dist, tf.double)) * 2 * n_sample)), tf.float32
+        )
+
+        self.n_sample = n_sample
+
+    def sample(self, labels):
+        """
+            labels: [b1, b2]
+        Return
+            true_log_probs: [b1, b2]
+            samp_log_probs: [n_sample]
+            neg_samples: [n_sample]
+        """
+
+        n_sample = self.n_sample
+        n_tries = 2 * n_sample
+
+        neg_samples = tf.unique(tf.squeeze(tf.random.categorical(tf.expand_dims(self.dist, 0), n_tries), 0))[0]
+        true_log_probs = tf.gather_nd(self.log_q, labels)
+        samp_log_probs = tf.gather(self.log_q, neg_samples)
+        return true_log_probs, samp_log_probs, neg_samples
+
+
+def sample_logits(embedding, bias, labels, inputs, sampler):
+    """
+        embedding: an nn.Embedding layer
+        bias: [n_vocab]
+        labels: [b1, b2]
+        inputs: [b1, b2, n_emb]
+        sampler: you may use a LogUniformSampler
+    Return
+        logits: [b1, b2, 1 + n_sample]
+    """
+    true_log_probs, samp_log_probs, neg_samples = sampler.sample(labels)
+    n_sample = tf.shape(neg_samples, 0)
+    b1, b2 = labels.size(0), labels.size(1)
+    b1, b2 = tf.shape(labels, 0), tf.shape(labels, 1)
+    all_ids = tf.concat([tf.keras.backend.flatten(labels), neg_samples], 0)
+    all_w = embedding(all_ids)
+    true_w = tf.reshape(all_w[:-n_sample], [b1, b2, -1])
+    sample_w = tf.reshape(all_w[-n_sample:], [n_sample, -1])
+
+    all_b = bias[all_ids]
+    true_b = all_b[:-n_sample].view(b1, b2)
+    sample_b = all_b[-n_sample:]
+
+    hit = (labels[:, :, None] == neg_samples).detach()
+
+    true_logits = tf.einsum("ijk,ijk->ij", [true_w, inputs]) + true_b - true_log_probs
+    sample_logits = tf.einsum("lk,ijk->ijl", [sample_w, inputs]) + sample_b - samp_log_probs
+    sample_logits.masked_fill_(hit, -1e30)
+    logits = tf.cat([true_logits[:, :, None], sample_logits], -1)
+
+    return logits

src/transformers/modeling_transfo_xl.py

@@ -810,7 +810,8 @@ class TransfoXLLMHeadModel(TransfoXLPreTrainedModel):
         self.transformer = TransfoXLModel(config)
         self.sample_softmax = config.sample_softmax
         # use sampled softmax
-        if config.sample_softmax > 0:
+        if self.sample_softmax > 0:
+            raise NotImplementedError
             self.out_layer = nn.Linear(config.d_model, config.vocab_size)
             self.sampler = LogUniformSampler(config.vocab_size, config.sample_softmax)
         # use adaptive softmax (including standard softmax)
@@ -827,6 +828,8 @@ class TransfoXLLMHeadModel(TransfoXLPreTrainedModel):
         # sampled softmax
         if self.sample_softmax > 0:
             if self.config.tie_weight:
+                # Here the self.out_layer.weight variable has to be correctly set -> read paper and discuss for this
+                raise NotImplementedError
                 self.out_layer.weight = self.transformer.word_emb.weight
         # adaptive softmax (including standard softmax)
         else:
@@ -909,12 +912,13 @@ class TransfoXLLMHeadModel(TransfoXLPreTrainedModel):
         pred_hid = last_hidden[:, -tgt_len:]
         outputs = transformer_outputs[1:]
         if self.sample_softmax > 0 and self.training:
+            raise NotImplementedError
+            # TODO: code below should work
             assert self.config.tie_weight
             logit = sample_logits(self.transformer.word_emb, self.out_layer.bias, labels, pred_hid, self.sampler)
             softmax_output = -F.log_softmax(logit, -1)[:, :, 0]
             outputs = [softmax_output] + outputs
             if labels is not None:
-                # TODO: This is not implemented
                 raise NotImplementedError
         else:
             softmax_output = self.crit(pred_hid.view(-1, pred_hid.size(-1)), labels)

tests/test_modeling_tf_common.py

@@ -30,7 +30,7 @@ if is_tf_available():
     import tensorflow as tf
     import numpy as np
 
-    from transformers import tf_top_k_top_p_filtering
+    from transformers import tf_top_k_top_p_filtering, TFAdaptiveEmbedding
 
     if _tf_gpu_memory_limit is not None:
         gpus = tf.config.list_physical_devices("GPU")
@@ -348,7 +348,7 @@ class TFModelTesterMixin:
 
         for model_class in self.all_model_classes:
             model = model_class(config)
-            assert isinstance(model.get_input_embeddings(), tf.keras.layers.Layer)
+            assert isinstance(model.get_input_embeddings(), (tf.keras.layers.Layer, TFAdaptiveEmbedding))
             x = model.get_output_embeddings()
             assert x is None or isinstance(x, tf.keras.layers.Layer)
 

tests/test_modeling_tf_transfo_xl.py

@@ -26,7 +26,7 @@ from .utils import CACHE_DIR, require_tf, slow
 
 if is_tf_available():
     import tensorflow as tf
-    from transformers.modeling_tf_transfo_xl import (
+    from transformers import (
         TFTransfoXLModel,
         TFTransfoXLLMHeadModel,
         TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_MAP,
@@ -364,7 +364,7 @@ class TFTransfoXLModelLanguageGenerationTest(unittest.TestCase):
                     0,
                 ]
             ],
-            dtype=tf.int31,
+            dtype=tf.int32,
         )
         #  In 1991 , the remains of Russian Tsar Nicholas II and his family
         #  ( except for Alexei and Maria ) are discovered .
@@ -570,8 +570,5 @@ class TFTransfoXLModelLanguageGenerationTest(unittest.TestCase):
         # Nicholas II and his family were discovered. The voice of <unk> young son,
         # Tsarevich Alexei Nikolaevich, narrates the remainder of the story.<eos>
 
-        # TODO: add this test when trasnfo-xl-lmhead is implemented
-        with self.assertRaises(NotImplementedError):
-            model.generate(input_ids, max_length=200, do_sample=False)
-            print(expected_output_ids)
-        # self.assertListEqual(output_ids[0].numpy().tolist(), expected_output_ids) TODO: (PVP) to add when transfo-xl is implemented
+        output_ids = model.generate(input_ids, max_length=200, do_sample=False)
+        self.assertListEqual(output_ids[0].numpy().tolist(), expected_output_ids)

tests/test_modeling_transfo_xl.py

@@ -129,10 +129,10 @@ class TransfoXLModelTest(ModelTesterMixin, unittest.TestCase):
 
         def check_transfo_xl_model_output(self, result):
             self.parent.assertListEqual(
-                list(result["hidden_states_1"].size()), [self.batch_size, self.seq_length, self.hidden_size]
+                list(result["hidden_states_1"].size()), [self.batch_size, self.seq_length, self.hidden_size],
             )
             self.parent.assertListEqual(
-                list(result["hidden_states_2"].size()), [self.batch_size, self.seq_length, self.hidden_size]
+                list(result["hidden_states_2"].size()), [self.batch_size, self.seq_length, self.hidden_size],
             )
             self.parent.assertListEqual(
                 list(list(mem.size()) for mem in result["mems_1"]),
@@ -166,7 +166,7 @@ class TransfoXLModelTest(ModelTesterMixin, unittest.TestCase):
         def check_transfo_xl_lm_head_output(self, result):
             self.parent.assertListEqual(list(result["loss_1"].size()), [self.batch_size, self.seq_length])
             self.parent.assertListEqual(
-                list(result["lm_logits_1"].size()), [self.batch_size, self.seq_length, self.vocab_size]
+                list(result["lm_logits_1"].size()), [self.batch_size, self.seq_length, self.vocab_size],
             )
             self.parent.assertListEqual(
                 list(list(mem.size()) for mem in result["mems_1"]),
@@ -175,7 +175,7 @@ class TransfoXLModelTest(ModelTesterMixin, unittest.TestCase):
 
             self.parent.assertListEqual(list(result["loss_2"].size()), [self.batch_size, self.seq_length])
             self.parent.assertListEqual(
-                list(result["lm_logits_2"].size()), [self.batch_size, self.seq_length, self.vocab_size]
+                list(result["lm_logits_2"].size()), [self.batch_size, self.seq_length, self.vocab_size],
             )
             self.parent.assertListEqual(
                 list(list(mem.size()) for mem in result["mems_2"]),