Merge branch 'xlnet' into doc-sphinx

pytorch_transformers/modeling_xlnet.py

@@ -122,14 +122,14 @@ def load_tf_weights_in_xlnet(model, config, tf_path):
         import numpy as np
         import tensorflow as tf
     except ImportError:
-        print("Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
+        logger.error("Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
             "https://www.tensorflow.org/install/ for installation instructions.")
         raise
     # Load weights from TF model
     init_vars = tf.train.list_variables(tf_path)
     tf_weights = {}
     for name, shape in init_vars:
-        print("Loading TF weight {} with shape {}".format(name, shape))
+        logger.info("Loading TF weight {} with shape {}".format(name, shape))
         array = tf.train.load_variable(tf_path, name)
         tf_weights[name] = array
 
@@ -137,15 +137,15 @@ def load_tf_weights_in_xlnet(model, config, tf_path):
     tf_to_pt_map = build_tf_xlnet_to_pytorch_map(model, config, tf_weights)
 
     for name, pointer in tf_to_pt_map.items():
-        print("Importing {}".format(name))
+        logger.info("Importing {}".format(name))
         if name not in tf_weights:
-            print("{} not in tf pre-trained weights, skipping".format(name))
+            logger.info("{} not in tf pre-trained weights, skipping".format(name))
             continue
         array = tf_weights[name]
         # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
         # which are not required for using pretrained model
         if 'kernel' in name and ('ff' in name or 'summary' in name or 'logit' in name):
-            print("Transposing")
+            logger.info("Transposing")
             array = np.transpose(array)
         if isinstance(pointer, list):
             # Here we will split the TF weigths
@@ -157,7 +157,7 @@ def load_tf_weights_in_xlnet(model, config, tf_path):
                 except AssertionError as e:
                     e.args += (p_i.shape, arr_i.shape)
                     raise
-                print("Initialize PyTorch weight {} for layer {}".format(name, i))
+                logger.info("Initialize PyTorch weight {} for layer {}".format(name, i))
                 p_i.data = torch.from_numpy(arr_i)
         else:
             try:
@@ -165,13 +165,13 @@ def load_tf_weights_in_xlnet(model, config, tf_path):
             except AssertionError as e:
                 e.args += (pointer.shape, array.shape)
                 raise
-            print("Initialize PyTorch weight {}".format(name))
+            logger.info("Initialize PyTorch weight {}".format(name))
             pointer.data = torch.from_numpy(array)
         tf_weights.pop(name, None)
         tf_weights.pop(name + '/Adam', None)
         tf_weights.pop(name + '/Adam_1', None)
 
-    print("Weights not copied to PyTorch model: {}".format(', '.join(tf_weights.keys())))
+    logger.info("Weights not copied to PyTorch model: {}".format(', '.join(tf_weights.keys())))
     return model
 
 
@@ -252,10 +252,6 @@ class XLNetConfig(PretrainedConfig):
                  layer_norm_eps=1e-12,
 
                  dropout=0.1,
-                 dropatt=0.1,
-                 init="normal",
-                 init_range=0.1,
-                 init_std=0.02,
                  mem_len=None,
                  reuse_len=None,
                  bi_data=False,
@@ -297,11 +293,7 @@ class XLNetConfig(PretrainedConfig):
             self.initializer_range = initializer_range
             self.layer_norm_eps = layer_norm_eps
 
-            self.init = init
-            self.init_range = init_range
-            self.init_std = init_std
             self.dropout = dropout
-            self.dropatt = dropatt
             self.mem_len = mem_len
             self.reuse_len = reuse_len
             self.bi_data = bi_data
@@ -393,7 +385,7 @@ class XLNetRelativeAttention(nn.Module):
         x = x[1:, ...]
         x = x.reshape(x_size[0], x_size[1] - 1, x_size[2], x_size[3])
         # x = x[:, 0:klen, :, :]
-        x = torch.index_select(x, 1, torch.arange(klen))
+        x = torch.index_select(x, 1, torch.arange(klen, device=x.device, dtype=torch.long))
 
         return x