def _collect_previous_mask(input_tensors):
    """Retrieves the output mask(s) of the previous node.

    # Arguments
        input_tensors: A tensor or list of tensors.

    # Returns
        A mask tensor or list of mask tensors.
    """
    input_tensors = to_list(input_tensors)
    masks = []
    for x in input_tensors:
        if hasattr(x, '_keras_history'):
            inbound_layer, node_index, tensor_index = x._keras_history
            node = inbound_layer._inbound_nodes[node_index]
            mask = node.output_masks[tensor_index]             # I got an error here, but I don't know why
            masks.append(mask)
        else:
            masks.append(None)
    if len(masks) == 1:
        return masks[0]
    return masks

load_trained_model_from_checkpoint

~/anaconda3/lib/python3.6/site-packages/keras_xlnet/loader.py in load_trained_model_from_checkpoint(config_path, checkpoint_path, batch_size, memory_len, target_len, in_train_phase, **kwargs)
    160         target_len=target_len,
    161         in_train_phase=in_train_phase,
--> 162         **kwargs)
    163     load_model_weights_from_checkpoint(
    164         model=model,

~/anaconda3/lib/python3.6/site-packages/keras_xlnet/loader.py in build_model_from_config(config_path, batch_size, memory_len, target_len, in_train_phase, **kwargs)
     49         clamp_len=None,
     50         shared_biases=not config['untie_r'],
---> 51         **kwargs)
     52     return model, config
     53 

~/anaconda3/lib/python3.6/site-packages/keras_xlnet/xlnet.py in build_xlnet(units, training, num_token, num_block, num_head, hidden_dim, batch_size, memory_len, target_len, permute, mask_index, dropout, attention_dropout, attention_type, clamp_len, shared_biases)
    213             dropout_rate=dropout,
    214             activation=gelu,
--> 215             name='FeedForward-{}'.format(i + 1),
    216         )
    217         if 0.0 < dropout < 1.0:

~/anaconda3/lib/python3.6/site-packages/keras_position_wise_feed_forward/feed_forward.py in __init__(self, units, activation, use_bias, kernel_initializer, bias_initializer, kernel_regularizer, bias_regularizer, kernel_constraint, bias_constraint, **kwargs)
     45         self.W1, self.b1 = None, None
     46         self.W2, self.b2 = None, None
---> 47         super(FeedForward, self).__init__(**kwargs)
     48 
     49     def get_config(self):

~/anaconda3/lib/python3.6/site-packages/tensorflow/python/training/checkpointable/base.py in _method_wrapper(self, *args, **kwargs)
    424     self._setattr_tracking = False  # pylint: disable=protected-access
    425     try:
--> 426       method(self, *args, **kwargs)
    427     finally:
    428       self._setattr_tracking = previous_value  # pylint: disable=protected-access

~/anaconda3/lib/python3.6/site-packages/tensorflow/python/keras/engine/base_layer.py in __init__(self, trainable, name, dtype, **kwargs)
    135     for kwarg in kwargs:
    136       if kwarg not in allowed_kwargs:
--> 137         raise TypeError('Keyword argument not understood:', kwarg)
    138 
    139     # Mutable properties

TypeError: ('Keyword argument not understood:', 'dropout_rate')

import keras_xlnet

word_dict = load_dict()
tokenizer = Tokenizer(word_dict) # this works in keras_bert

pass

# I re-load the model with these codes
custom_dict = get_custom_objects()
model = load_model(self.model_path, custom_objects=custom_dict)

def _get_model(base_dir, cfg_=None):
    config_file = os.path.join(base_dir, 'xlnet_config.json')
    checkpoint_file = os.path.join(base_dir, 'xlnet_model.ckpt')
    print(cfg_["bs"], cfg_["mem_len"], cfg_['maxlen'])
    model = load_trained_model_from_checkpoint(
        config_path=config_file,
        checkpoint_path=checkpoint_file,
        batch_size=cfg_["bs"],
        memory_len=cfg_["mem_len"],
        target_len=cfg_['maxlen'],
        in_train_phase=False
    )
    return model

model = _get_model("../../../xlnet_model/xlnet_cased_L-12_H-768_A-12/", cfg_=cfg)

token = np.random.randint(0, 123, size=(5, 256)).astype(np.float32)
segm = np.ones_like(token).astype(dtype=np.float32)
mask = np.random.randint(0, 256, size=(token.shape[0], 1)).astype(np.float32)
print(segm.shape, segm.shape, mask.shape)
print(model.inputs)
model.predict([token, segm, mask], verbose=1, batch_size=32)

24 0 256
(5, 256) (5, 256) (5, 1)
[<tf.Tensor 'Input-Token:0' shape=(?, 256) dtype=float32>, <tf.Tensor 'Input-Segment:0' shape=(?, 256) dtype=float32>, <tf.Tensor 'Input-Memory-Length:0' shape=(?, 1) dtype=float32>]
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-3-3061d04374a5> in <module>()
     20 print(segm.shape, segm.shape, mask.shape)
     21 print(model.inputs)
---> 22 model.predict([token, segm, mask], verbose=1, batch_size=32)

~/anaconda3/lib/python3.6/site-packages/keras/engine/training.py in predict(self, x, batch_size, verbose, steps)
   1167                                             batch_size=batch_size,
   1168                                             verbose=verbose,
-> 1169                                             steps=steps)
   1170 
   1171     def train_on_batch(self, x, y,

~/anaconda3/lib/python3.6/site-packages/keras/engine/training_arrays.py in predict_loop(model, f, ins, batch_size, verbose, steps)
    292                 ins_batch[i] = ins_batch[i].toarray()
    293 
--> 294             batch_outs = f(ins_batch)
    295             batch_outs = to_list(batch_outs)
    296             if batch_index == 0:

~/anaconda3/lib/python3.6/site-packages/keras/backend/tensorflow_backend.py in __call__(self, inputs)
   2713                 return self._legacy_call(inputs)
   2714 
-> 2715             return self._call(inputs)
   2716         else:
   2717             if py_any(is_tensor(x) for x in inputs):

~/anaconda3/lib/python3.6/site-packages/keras/backend/tensorflow_backend.py in _call(self, inputs)
   2673             fetched = self._callable_fn(*array_vals, run_metadata=self.run_metadata)
   2674         else:
-> 2675             fetched = self._callable_fn(*array_vals)
   2676         return fetched[:len(self.outputs)]
   2677 

~/anaconda3/lib/python3.6/site-packages/tensorflow/python/client/session.py in __call__(self, *args, **kwargs)
   1397           ret = tf_session.TF_SessionRunCallable(
   1398               self._session._session, self._handle, args, status,
-> 1399               run_metadata_ptr)
   1400         if run_metadata:
   1401           proto_data = tf_session.TF_GetBuffer(run_metadata_ptr)

~/anaconda3/lib/python3.6/site-packages/tensorflow/python/framework/errors_impl.py in __exit__(self, type_arg, value_arg, traceback_arg)
    524             None, None,
    525             compat.as_text(c_api.TF_Message(self.status.status)),
--> 526             c_api.TF_GetCode(self.status.status))
    527     # Delete the underlying status object from memory otherwise it stays alive
    528     # as there is a reference to status from this from the traceback due to

InvalidArgumentError: Input to reshape is a tensor with 983040 values, but the requested shape has 4718592
	 [[{{node Attention-12/Reshape_8}} = Reshape[T=DT_FLOAT, Tshape=DT_INT32, _device="/job:localhost/replica:0/task:0/device:GPU:0"](Attention-12/MatMul_2, Attention-12/Reshape_8/shape)]]
	 [[{{node FeedForward-Normal-12/add_1/_1311}} = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_6655_FeedForward-Normal-12/add_1", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]]

import os
import sys

import numpy as np

from keras_xlnet import Tokenizer, load_trained_model_from_checkpoint, ATTENTION_TYPE_BI



'''Can be found at https://github.com/ymcui/Chinese-PreTrained-XLNet'''
checkpoint_path =  "/users/hdpsbp/HadoopDir/felix/xlnet"
vocab_path = os.path.join(checkpoint_path, 'spiece.model')
config_path = os.path.join(checkpoint_path, 'xlnet_config.json')
model_path = os.path.join(checkpoint_path, 'xlnet_model.ckpt')

# Tokenize inputs
tokenizer = Tokenizer(vocab_path)
text = "给"
tokens = tokenizer.encode(text)
print(tokens) #[19, 841]

tokenizer.decode([19]) #''
pass

import keras_xlnet

pass

ResourceExhaustedError                    Traceback (most recent call last)
<ipython-input-14-3280d689ffd9> in <module>
----> 1 model3.fit([token_input,seg_input,np.array([0]*320)[:,None]],y_train,batch_size=bsz)

/opt/conda/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_freq, max_queue_size, workers, use_multiprocessing, **kwargs)
    778           validation_steps=validation_steps,
    779           validation_freq=validation_freq,
--> 780           steps_name='steps_per_epoch')
    781 
    782   def evaluate(self,

/opt/conda/lib/python3.6/site-packages/tensorflow/python/keras/engine/training_arrays.py in model_iteration(model, inputs, targets, sample_weights, batch_size, epochs, verbose, callbacks, val_inputs, val_targets, val_sample_weights, shuffle, initial_epoch, steps_per_epoch, validation_steps, validation_freq, mode, validation_in_fit, prepared_feed_values_from_dataset, steps_name, **kwargs)
    361 
    362         # Get outputs.
--> 363         batch_outs = f(ins_batch)
    364         if not isinstance(batch_outs, list):
    365           batch_outs = [batch_outs]

/opt/conda/lib/python3.6/site-packages/tensorflow/python/keras/backend.py in __call__(self, inputs)
   3290 
   3291     fetched = self._callable_fn(*array_vals,
-> 3292                                 run_metadata=self.run_metadata)
   3293     self._call_fetch_callbacks(fetched[-len(self._fetches):])
   3294     output_structure = nest.pack_sequence_as(

/opt/conda/lib/python3.6/site-packages/tensorflow/python/client/session.py in __call__(self, *args, **kwargs)
   1456         ret = tf_session.TF_SessionRunCallable(self._session._session,
   1457                                                self._handle, args,
-> 1458                                                run_metadata_ptr)
   1459         if run_metadata:
   1460           proto_data = tf_session.TF_GetBuffer(run_metadata_ptr)

ResourceExhaustedError: 2 root error(s) found.
  (0) Resource exhausted: OOM when allocating tensor with shape[8,12,48400,64] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
	 [[{{node Attention-9/transpose_8}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

	 [[loss/mul/_2627]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

  (1) Resource exhausted: OOM when allocating tensor with shape[8,12,48400,64] and type float on /job:localhost/replica:0/task:0/device:GPU:0 by allocator GPU_0_bfc
	 [[{{node Attention-9/transpose_8}}]]
Hint: If you want to see a list of allocated tensors when OOM happens, add report_tensor_allocations_upon_oom to RunOptions for current allocation info.

0 successful operations.
0 derived errors ignored.

def create_model():
    model = load_trained_model_from_checkpoint(
        config_path=os.path.join(checkpoint_path, 'xlnet_config.json'),
        checkpoint_path=os.path.join(checkpoint_path, 'xlnet_model.ckpt'),
        batch_size=8,
        memory_len=0,
        target_len=220,
        in_train_phase=False,
    )
    sequence_output  = model.get_layer(index=-1).output
    summary = Extract(index=-1)(sequence_output) ##based on our tokenization method, cls token stays on last
    pool_output = keras.layers.Dense(units=1, activation='sigmoid',name='real_output')(summary)
    model3 = keras.models.Model(inputs=model.input, outputs=pool_output)
    return model3
model3 = create_model()
model3.compile('adam','binary_crossentropy')

from keras_xlnet import Tokenizer
dict_path = '../pre_models/chinese_xlnet_mid_L-24_H-768_A-12/spiece.model'
tokenizer = Tokenizer(dict_path)

from keras_xlnet import Tokenizer
dict_path = '/home/hpdbman/gjz/Competition/pre_models/chinese_xlnet_mid_L-24_H-768_A-12/spiece.model'
tokenizer = Tokenizer(dict_path)

def _get_model(base_dir, cfg_=None):
    config_file = os.path.join(base_dir, 'xlnet_config.json')
    checkpoint_file = os.path.join(base_dir, 'xlnet_model.ckpt')
    print(cfg_["bs"], cfg_["mem_len"], cfg_['maxlen'])
    model = load_trained_model_from_checkpoint(
        config_path=config_file,
        checkpoint_path=checkpoint_file,
        batch_size=cfg_["bs"],
        memory_len=cfg_["mem_len"],
        target_len=cfg_['maxlen'],
        in_train_phase=False
    )
    return model

model = _get_model("../../../xlnet_model/xlnet_cased_L-12_H-768_A-12/", cfg_=cfg)