mirror of
https://github.com/wassname/Clover-Edition.git
synced 2026-09-19 12:10:28 +08:00
173 lines
6.4 KiB
Python
173 lines
6.4 KiB
Python
from __future__ import division
|
|
from __future__ import print_function
|
|
import sys
|
|
|
|
sys.path.append('../')
|
|
import tensorflow as tf
|
|
import os
|
|
import numpy as np
|
|
|
|
tf.enable_eager_execution()
|
|
import transformer
|
|
import argparse
|
|
import pdb
|
|
import re
|
|
from collections import Counter
|
|
from tensorflow.python import debug as tf_debug
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.ops import embedding_ops
|
|
import fastBPE
|
|
import platform
|
|
|
|
use_py3 = platform.python_version()[0] == '3'
|
|
|
|
parser = argparse.ArgumentParser(description='TensorFlow code for generating from CTRL')
|
|
parser.add_argument('--model_dir', type=str, required=True,
|
|
help='location of model checkpoint')
|
|
parser.add_argument('--seed', type=int, default=1337,
|
|
help='random seed for TensorFlow, numpy and PythonHash')
|
|
parser.add_argument('--sequence_len', type=int, default=256,
|
|
help='sequence len of model being fine-tuned (must match also the TFRecords)')
|
|
parser.add_argument('--iterations', type=int, default=1000,
|
|
help='random seed for TensorFlow, numpy and PythonHash')
|
|
|
|
args = parser.parse_args()
|
|
tf.random.set_random_seed(args.seed)
|
|
os.environ['PYTHONHASHSEED'] = str(args.seed)
|
|
np.random.seed(args.seed)
|
|
|
|
# load the vocabulary from file
|
|
vocab = open('../vocab').read().decode(encoding='utf-8').split('\n') if not use_py3 else open('../vocab',
|
|
encoding='utf-8').read().split(
|
|
'\n')
|
|
vocab = list(map(lambda x: x.split(' ')[0], vocab)) + ['<unk>'] + ['\n']
|
|
print('{} unique words'.format(len(vocab)))
|
|
|
|
# length of the vocabulary
|
|
vocab_size = len(vocab)
|
|
|
|
# define the numericalization map
|
|
# idx2word maps the numericalized ID to the word
|
|
# word2idx maps the word to the numericalized ID
|
|
word2idx = {u: i for i, u in enumerate(vocab)}
|
|
idx2word = np.array(vocab)
|
|
|
|
# sequence length to use for the transformer
|
|
# must match the model being fine-tuned
|
|
seq_length = args.sequence_len
|
|
|
|
|
|
def input_fn(params=None):
|
|
print('READING!', params)
|
|
dataset = tf.data.Dataset.list_files(tf.io.gfile.glob('./*.tfrecords'), shuffle=True)
|
|
|
|
tf_data = tf.data.TFRecordDataset(dataset)
|
|
myfeatures = {
|
|
'input': tf.io.FixedLenFeature([256], tf.int64),
|
|
'output': tf.io.FixedLenFeature([256], tf.int64)
|
|
}
|
|
|
|
def _parse_text_function(example_proto):
|
|
blah = tf.io.parse_single_example(example_proto, myfeatures)
|
|
return blah['input'], blah['output']
|
|
|
|
train_data = tf_data.map(_parse_text_function).batch(params['batch_size'], drop_remainder=True).repeat().shuffle(
|
|
10000) # .prefetch(tf.contrib.data.AUTOTUNE)
|
|
|
|
return train_data
|
|
|
|
|
|
# the dimension of the transformer
|
|
embedding_dim = 1280
|
|
|
|
|
|
# Now, we begin defining the model
|
|
# we defer the transformer definition to transformer.py
|
|
# here, we only define the tied softmax layer
|
|
# this layer ties the softmax weights to the input embeddings
|
|
class TiedEmbeddingSoftmax(tf.keras.layers.Layer):
|
|
|
|
def __init__(self, vocab_size=vocab_size, embedding_size=embedding_dim, **kwargs):
|
|
with tf.device('/cpu:0'):
|
|
super(TiedEmbeddingSoftmax, self).__init__()
|
|
self.w = self.add_weight(name='w', shape=(vocab_size, embedding_size),
|
|
initializer='random_normal',
|
|
trainable=True)
|
|
self.b = self.add_weight(name='b', shape=(vocab_size,),
|
|
initializer='zeros',
|
|
trainable=True)
|
|
|
|
def call(self, inputs, embed=True):
|
|
with tf.device('/cpu:0'):
|
|
if embed:
|
|
dtype = tf.keras.backend.dtype(inputs)
|
|
if dtype != 'int32' and dtype != 'int64':
|
|
inputs = math_ops.cast(inputs, 'int32')
|
|
return embedding_ops.embedding_lookup(self.w, inputs)
|
|
else:
|
|
return tf.tensordot(inputs, tf.transpose(self.w), 1) + self.b
|
|
|
|
|
|
# input for the keras model
|
|
tokens = tf.keras.layers.Input(shape=(seq_length,), dtype='int32')
|
|
|
|
# instantiates a tied softmax class
|
|
tied_embedding_softmax = TiedEmbeddingSoftmax()
|
|
|
|
# embedded tokens, before passing it to the transformer
|
|
embedded = tied_embedding_softmax(tokens, embed=True)
|
|
|
|
# the activations after passing it from the transformer
|
|
# for some odd reason, TPUs don't play well with specifying the arguments of the Encoder() function
|
|
# so you have to leave them at their defaults
|
|
transformed = transformer.Encoder()(embedded, training=False)
|
|
|
|
# pass the activations from our tiedsoftmax class
|
|
# this time with embed=False denoting that we are doing the softmax operation
|
|
# and not a lookup
|
|
logits = tied_embedding_softmax(transformed, embed=False)
|
|
|
|
# finally, define the Keras model with inputs as tokens and outputs as the logits we just computed
|
|
model = tf.keras.Model(inputs=tokens, outputs=logits)
|
|
|
|
|
|
# the loss function is a simple categorical crossentropy between the logits and the labels
|
|
def loss(labels, logits):
|
|
loss = tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)
|
|
mean = tf.math.reduce_mean(loss)
|
|
loss = tf.Print(loss, [mean])
|
|
return loss
|
|
|
|
|
|
# the optimizer is not used since this code only supports inference
|
|
# however, to compile the model, we still define it
|
|
optimizer = tf.contrib.estimator.clip_gradients_by_norm(
|
|
tf.train.AdagradOptimizer(learning_rate=3e-3), 0.25)
|
|
|
|
# compile the model with the optimizer and loss
|
|
model.compile(optimizer=optimizer, loss=loss)
|
|
print(model.summary())
|
|
|
|
# IMPORTANT
|
|
# this is where the saved model is presented to the code
|
|
# the model directory should have the model checkpoint and
|
|
# a checkpoint file
|
|
run_config = tf.contrib.tpu.RunConfig(
|
|
model_dir=args.model_dir)
|
|
|
|
# this converts the Keras model to a TensorFlow estimator
|
|
# this step is critical
|
|
# remember to patch the TF 1.14 file before running the code, else you're going to see errors here
|
|
|
|
run_config = tf.contrib.tpu.RunConfig(
|
|
model_dir=args.model_dir,
|
|
session_config=tf.ConfigProto(allow_soft_placement=True, log_device_placement=True),
|
|
tpu_config=tf.contrib.tpu.TPUConfig(iterations_per_loop=100, num_cores_per_replica=1,
|
|
input_partition_dims=[[1, 1], [1, 1]], per_host_input_for_training=3))
|
|
tf.logging.set_verbosity(tf.logging.INFO)
|
|
|
|
estimator_model = tf.keras.estimator.model_to_estimator(keras_model=model, config=run_config)
|
|
|
|
estimator_model.train(input_fn=input_fn, steps=args.iterations)
|
|
|