README.zh-CN.md

Keras Transformer-XL

[中文|English]

Transformer-XL的非官方实现.

安装

pip install keras-transformer-xl

使用

加载预训练权重

在信息文件夹中有一些预训练模型的配置信息，可以根据实际情况进行选择：

import os
from keras_transformer_xl import load_trained_model_from_checkpoint

checkpoint_path = 'foo/bar/sota/enwiki8'
model = load_trained_model_from_checkpoint(
    config_path=os.path.join(checkpoint_path, 'config.json'),
    checkpoint_path=os.path.join(checkpoint_path, 'model.ckpt')
)
model.summary()

关于输入输出

模型有两个输入，第一个是词的下标，第二个是已经记忆的长度，第二个输入要求所有数值相同。MemorySequence对输入进行的切割并加入了记忆长度的输入：

import keras
import numpy as np
from keras_transformer_xl import MemorySequence, build_transformer_xl


class DummySequence(keras.utils.Sequence):

    def __init__(self):
        pass

    def __len__(self):
        return 10

    def __getitem__(self, index):
        return np.ones((3, 5 * (index + 1))), np.ones((3, 5 * (index + 1), 3))


model = build_transformer_xl(
    units=4,
    embed_dim=4,
    hidden_dim=4,
    num_token=3,
    num_block=3,
    num_head=2,
    batch_size=3,
    memory_len=20,
    target_len=10,
)
seq = MemorySequence(
    model=model,
    sequence=DummySequence(),
    target_len=10,
)

model.predict(model, seq, verbose=True)

使用tf.keras

在环境变量里加入TF_KERAS=1启用tensorflow.python.keras。