title text generation with a miniature gpt apoorv nandanhttps twitter comnandanapoorv date created 20200529 last modified 20200529 description implement a miniature version of gpt and train it to generate text accelerator gpu introduction this example demonstrates how to implement an autoregressive language model using a miniature version of the gpt model the model consists of a single transformer block with causal masking in its attention layer we use the text from the imdb sentiment classification dataset for training and generate new movie reviews for a given prompt when using this script with your own dataset make sure it has at least 1 million words this example should be run with tfnightly2 3 0dev20200531 or with tensorflow 2 3 or higher references gpthttps www semanticscholar orgpaperimprovinglanguageunderstandingbygenerativeradfordcd18800a0fe0b668a1cc19f2ec95b5003d0a5035 gpt2https www semanticscholar orgpaperlanguagemodelsareunsupervisedmultitasklearnersradfordwu9405cc0d6169988371b2755e573cc28650d14dfe gpt3https arxiv orgabs2005 14165 setup we set the backend to tensorflow the code works with both tensorflow and torch it does not work with jax due to the behavior of jax numpy tile in a jit scope used in causalattentionmask tile in jax does not support a dynamic reps argument you can make the code work in jax by wrapping the inside of the causalattentionmask function in a decorator to prevent jit compilation with jax ensurecompiletimeeval implement a transformer block as a layer mask the upper half of the dot product matrix in self attention this prevents flow of information from future tokens to current token 1 s in the lower triangle counting from the lower right corner implement an embedding layer create two separate embedding layers one for tokens and one for token index positions implement the miniature gpt model prepare the data for wordlevel language modelling download the imdb dataset and combine training and validation sets for a text generation task shell curl o https ai stanford eduamaasdatasentimentaclimdbv1 tar gz tar xf aclimdbv1 tar gz the dataset contains each review in a separate text file the text files are present in four different folders create a list all files create a dataset from text files remove html linebreak tags and handle punctuation lowercased tfstrings lowerinputstring strippedhtml tfstrings regexreplacelowercased br return tfstrings regexreplacestrippedhtml fstring punctuation r 1 create a vectorization layer and adapt it to the text vectorizelayer textvectorization standardizecustomstandardization maxtokensvocabsize 1 outputmodeint outputsequencelengthmaxlen 1 vectorizelayer adapttextds vocab vectorizelayer getvocabulary to get words back from token indices def preparelminputslabelstext text tensorflow expanddimstext 1 tokenizedsentences vectorizelayertext x tokenizedsentences 1 y tokenizedsentences 1 return x y textds textds mappreparelminputslabels numparallelcallstfdata autotune textds textds prefetchtfdata autotune implement a keras callback for generating text class textgeneratorkeras callbacks callback def init self maxtokens starttokens indextoword topk10 printevery1 self maxtokens maxtokens self starttokens starttokens self indextoword indextoword self printevery printevery self k topk def samplefromself logits logits indices ops topklogits kself k sortedtrue indices np asarrayindices astypeint32 preds keras activations softmaxops expanddimslogits 00 preds np asarraypreds astypefloat32 return np random choiceindices ppreds def detokenizeself number return self indextowordnumber def onepochendself epoch logsnone starttokens for in self starttokens if epoch 1 self printevery 0 return numtokensgenerated 0 tokensgenerated while numtokensgenerated self maxtokens padlen maxlen lenstarttokens sampleindex lenstarttokens 1 if padlen 0 x starttokens maxlen sampleindex maxlen 1 elif padlen 0 x starttokens 0 padlen else x starttokens x np arrayx y self model predictx verbose0 sampletoken self samplefromy0sampleindex tokensgenerated appendsampletoken starttokens appendsampletoken numtokensgenerated lentokensgenerated txt join self detokenize for in self starttokens tokensgenerated printfgenerated text ntxtn tokenize starting prompt wordtoindex for index word in enumeratevocab wordtoindexword index startprompt this movie is starttokens wordtoindex get 1 for in startprompt split numtokensgenerated 40 textgencallback textgeneratornumtokensgenerated starttokens vocab train the model model createmodel model fittextds verbose2 epochs25 callbackstextgencallback introduction this example demonstrates how to implement an autoregressive language model using a miniature version of the gpt model the model consists of a single transformer block with causal masking in its attention layer we use the text from the imdb sentiment classification dataset for training and generate new movie reviews for a given prompt when using this script with your own dataset make sure it has at least 1 million words this example should be run with tf nightly 2 3 0 dev20200531 or with tensorflow 2 3 or higher references gpt https www semanticscholar org paper improving language understanding by generative radford cd18800a0fe0b668a1cc19f2ec95b5003d0a5035 gpt 2 https www semanticscholar org paper language models are unsupervised multitask learners radford wu 9405cc0d6169988371b2755e573cc28650d14dfe gpt 3 https arxiv org abs 2005 14165 setup we set the backend to tensorflow the code works with both tensorflow and torch it does not work with jax due to the behavior of jax numpy tile in a jit scope used in causal_attention_mask tile in jax does not support a dynamic reps argument you can make the code work in jax by wrapping the inside of the causal_attention_mask function in a decorator to prevent jit compilation with jax ensure_compile_time_eval implement a transformer block as a layer mask the upper half of the dot product matrix in self attention this prevents flow of information from future tokens to current token 1 s in the lower triangle counting from the lower right corner implement an embedding layer create two separate embedding layers one for tokens and one for token index positions implement the miniature gpt model only consider the top 20k words max sequence size embedding size for each token number of attention heads hidden layer size in feed forward network inside transformer no loss and optimization based on word embeddings from transformer block prepare the data for word level language modelling download the imdb dataset and combine training and validation sets for a text generation task shell curl o https ai stanford edu amaas data sentiment aclimdb_v1 tar gz tar xf aclimdb_v1 tar gz the dataset contains each review in a separate text file the text files are present in four different folders create a list all files create a dataset from text files remove html line break tags and handle punctuation create a vectorization layer and adapt it to the text to get words back from token indices shift word sequences by 1 position so that the target for position i is word at position i 1 the model will use all words up till position i to predict the next word implement a keras callback for generating text a callback to generate text from a trained model 1 feed some starting prompt to the model 2 predict probabilities for the next token 3 sample the next token and add it to the next input arguments max_tokens integer the number of tokens to be generated after prompt start_tokens list of integers the token indices for the starting prompt index_to_word list of strings obtained from the textvectorization layer top_k integer sample from the top_k token predictions print_every integer print after this many epochs tokenize starting prompt train the model note this code should preferably be run on gpu
import os
os.environ['KERAS_BACKEND'] = 'tensorflow'

import keras
from keras import layers
from keras import ops
from keras.layers import TextVectorization
import numpy as np
import os
import string
import random
import tensorflow
import tensorflow.data as tf_data
import tensorflow.strings as tf_strings





def causal_attention_mask(batch_size, n_dest, n_src, dtype):
    
    i = ops.arange(n_dest)[:, None]
    j = ops.arange(n_src)
    m = i >= j - n_src + n_dest
    mask = ops.cast(m, dtype)
    mask = ops.reshape(mask, [1, n_dest, n_src])
    mult = ops.concatenate(
        [ops.expand_dims(batch_size, -1), ops.convert_to_tensor([1, 1])], 0
    )
    return ops.tile(mask, mult)


class TransformerBlock(layers.Layer):
    def __init__(self, embed_dim, num_heads, ff_dim, rate=0.1):
        super().__init__()
        self.att = layers.MultiHeadAttention(num_heads, embed_dim)
        self.ffn = keras.Sequential(
            [
                layers.Dense(ff_dim, activation="relu"),
                layers.Dense(embed_dim),
            ]
        )
        self.layernorm1 = layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = layers.LayerNormalization(epsilon=1e-6)
        self.dropout1 = layers.Dropout(rate)
        self.dropout2 = layers.Dropout(rate)

    def call(self, inputs):
        input_shape = ops.shape(inputs)
        batch_size = input_shape[0]
        seq_len = input_shape[1]
        causal_mask = causal_attention_mask(batch_size, seq_len, seq_len, "bool")
        attention_output = self.att(inputs, inputs, attention_mask=causal_mask)
        attention_output = self.dropout1(attention_output)
        out1 = self.layernorm1(inputs + attention_output)
        ffn_output = self.ffn(out1)
        ffn_output = self.dropout2(ffn_output)
        return self.layernorm2(out1 + ffn_output)





class TokenAndPositionEmbedding(layers.Layer):
    def __init__(self, maxlen, vocab_size, embed_dim):
        super().__init__()
        self.token_emb = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)
        self.pos_emb = layers.Embedding(input_dim=maxlen, output_dim=embed_dim)

    def call(self, x):
        maxlen = ops.shape(x)[-1]
        positions = ops.arange(0, maxlen, 1)
        positions = self.pos_emb(positions)
        x = self.token_emb(x)
        return x + positions



vocab_size = 20000  
maxlen = 80  
embed_dim = 256  
num_heads = 2  
feed_forward_dim = 256  


def create_model():
    inputs = layers.Input(shape=(maxlen,), dtype="int32")
    embedding_layer = TokenAndPositionEmbedding(maxlen, vocab_size, embed_dim)
    x = embedding_layer(inputs)
    transformer_block = TransformerBlock(embed_dim, num_heads, feed_forward_dim)
    x = transformer_block(x)
    outputs = layers.Dense(vocab_size)(x)
    model = keras.Model(inputs=inputs, outputs=[outputs, x])
    loss_fn = keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    model.compile(
        "adam",
        loss=[loss_fn, None],
    )  
    return model







batch_size = 128




filenames = []
directories = [
    "aclImdb/train/pos",
    "aclImdb/train/neg",
    "aclImdb/test/pos",
    "aclImdb/test/neg",
]
for dir in directories:
    for f in os.listdir(dir):
        filenames.append(os.path.join(dir, f))

print(f"{len(filenames)} files")


random.shuffle(filenames)
text_ds = tf_data.TextLineDataset(filenames)
text_ds = text_ds.shuffle(buffer_size=256)
text_ds = text_ds.batch(batch_size)


def custom_standardization(input_string):
    
    lowercased = tf_strings.lower(input_string)
    stripped_html = tf_strings.regex_replace(lowercased, "<br />", " ")
    return tf_strings.regex_replace(stripped_html, f"([{string.punctuation}])", r" \1")



vectorize_layer = TextVectorization(
    standardize=custom_standardization,
    max_tokens=vocab_size - 1,
    output_mode="int",
    output_sequence_length=maxlen + 1,
)
vectorize_layer.adapt(text_ds)
vocab = vectorize_layer.get_vocabulary()  


def prepare_lm_inputs_labels(text):
    
    text = tensorflow.expand_dims(text, -1)
    tokenized_sentences = vectorize_layer(text)
    x = tokenized_sentences[:, :-1]
    y = tokenized_sentences[:, 1:]
    return x, y


text_ds = text_ds.map(prepare_lm_inputs_labels, num_parallel_calls=tf_data.AUTOTUNE)
text_ds = text_ds.prefetch(tf_data.AUTOTUNE)





class TextGenerator(keras.callbacks.Callback):
    

    def __init__(
        self, max_tokens, start_tokens, index_to_word, top_k=10, print_every=1
    ):
        self.max_tokens = max_tokens
        self.start_tokens = start_tokens
        self.index_to_word = index_to_word
        self.print_every = print_every
        self.k = top_k

    def sample_from(self, logits):
        logits, indices = ops.top_k(logits, k=self.k, sorted=True)
        indices = np.asarray(indices).astype("int32")
        preds = keras.activations.softmax(ops.expand_dims(logits, 0))[0]
        preds = np.asarray(preds).astype("float32")
        return np.random.choice(indices, p=preds)

    def detokenize(self, number):
        return self.index_to_word[number]

    def on_epoch_end(self, epoch, logs=None):
        start_tokens = [_ for _ in self.start_tokens]
        if (epoch + 1) % self.print_every != 0:
            return
        num_tokens_generated = 0
        tokens_generated = []
        while num_tokens_generated <= self.max_tokens:
            pad_len = maxlen - len(start_tokens)
            sample_index = len(start_tokens) - 1
            if pad_len < 0:
                x = start_tokens[:maxlen]
                sample_index = maxlen - 1
            elif pad_len > 0:
                x = start_tokens + [0] * pad_len
            else:
                x = start_tokens
            x = np.array([x])
            y, _ = self.model.predict(x, verbose=0)
            sample_token = self.sample_from(y[0][sample_index])
            tokens_generated.append(sample_token)
            start_tokens.append(sample_token)
            num_tokens_generated = len(tokens_generated)
        txt = " ".join(
            [self.detokenize(_) for _ in self.start_tokens + tokens_generated]
        )
        print(f"generated text:\n{txt}\n")



word_to_index = {}
for index, word in enumerate(vocab):
    word_to_index[word] = index

start_prompt = "this movie is"
start_tokens = [word_to_index.get(_, 1) for _ in start_prompt.split()]
num_tokens_generated = 40
text_gen_callback = TextGenerator(num_tokens_generated, start_tokens, vocab)




model = create_model()

model.fit(text_ds, verbose=2, epochs=25, callbacks=[text_gen_callback])