前文

本文首要展现了怎么完成微型的 GPT 模型完结文本生成任务，该模型只由 1 个 Transformer 块组成。

Data

这部分代码首要用于预备文本数据集进行言语模型练习，这儿需求事先下载好 aclImdb 数据，而且解压到当前目录。

首先，界说了一个批量大小 batch_size 和一个存储文件名的列表 filenames，其中包含了要处理的文本文件的路径。接下来，经过随机打乱 filenames 列表的顺序，来添加数据集的随机性。接着运用 TensorFlow 的 TextLineDataset 创建一个文本数据集 text_ds，并经过 shuffle 办法对数据集进行洗牌。

界说了一个自界说的标准化函数 custom_standardization，用于对输入字符串进行标准化处理。函数将字符串转换为小写，并运用正则表达式去除 HTML 标签和标点符号。然后创建了一个 TextVectorization 层，经过自界说函数对文本数据进行处理，对 text_ds 中的文本进行矢量化处理，也就是将文本转换为整数序列。并从数据会集自动构建词汇表。

经过 map 办法将处理后的文本转换为模型的输入和标签，即将每个序列中的一句话去掉最终一个字作为输入，然后将对应的同样一个序列从第二个字开端到最终的序列作为标签。从局部来看也就是前一个字是输入，猜测输出后一个字。

最终，运用 prefetch 办法对数据集进行预取操作，以便在模型练习过程中能够高效地加载数据。

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.layers import TextVectorization
import numpy as np
import os
import string
import random
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))
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 = tf.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)
text_ds = text_ds.prefetch(tf.data.AUTOTUNE)

Miniature GPT

Transformer Block

这儿首要是界说了一个 TransformerBlock 类，该类完成了 Transformer 模型中的一个 Transformer Block 。整个 TransformerBlock 类的作用是将输入序列经过自注意力核算和前馈网络变换，得到一个更丰富的表明。

TransformerBlock 类的结构函数 __init__ 承受四个参数：embed_dim 表明嵌入维度，num_heads 表明注意力头数，ff_dim 表明前馈网络的维度，rate 表明 Dropout 的比例。

在 call 办法中，首先获取输入的形状信息，包括批大小和序列长度。然后调用 causal_attention_mask 函数生成一个注意力掩码，用于遮蔽 Transformer 中的未来信息，确保模型只能看到当前方位以及之前的输入信息。这个掩码是一个二维矩阵，维度为 (seq_len, seq_len)。

接下来，运用 MultiHeadAttention 层 self.att 对输入进行自注意力核算，并传入注意力掩码。然后应用榜首个 Dropout 层 self.dropout1 对注意力输出进行随机失活。将输入和注意力输出相加，并经过 LayerNormalization 层 self.layernorm1 进行归一化处理，得到榜首个子层的输出 out1。

接着，将榜首个子层的输出 out1 传入前馈神经网络 self.ffn 进行非线性变换。再次应用 Dropout 层 self.dropout2 对前馈网络的输出进行随机失活。将榜首个子层的输出 out1 和前馈网络的输出相加，并经过 LayerNormalization 层 self.layernorm2 进行归一化处理，得到 Transformer Block 的最终输出。

def causal_attention_mask(batch_size, n_dest, n_src, dtype):
    i = tf.range(n_dest)[:, None]
    j = tf.range(n_src)
    m = i >= j - n_src + n_dest
    mask = tf.cast(m, dtype)
    mask = tf.reshape(mask, [1, n_dest, n_src])
    mult = tf.concat( [tf.expand_dims(batch_size, -1), tf.constant([1, 1], dtype=tf.int32)], 0 )
    return tf.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 = tf.shape(inputs)
        batch_size = input_shape[0]
        seq_len = input_shape[1]
        causal_mask = causal_attention_mask(batch_size, seq_len, seq_len, tf.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)

Token And Position Embedding

这儿界说了一个 TokenAndPositionEmbedding 类，用于得到输入序列的 token 和方位信息的嵌入。

TokenAndPositionEmbedding 类的结构函数 __init__ 承受三个参数：maxlen 表明序列的最大长度，vocab_size 表明词汇表的大小，embed_dim 表明嵌入维度。

在 call 办法中，首先获取输入序列 x 的长度 maxlen。然后运用 tf.range 函数生成一个从 0 到 maxlen-1 的方位向量 positions。接着将方位向量 positions 传入方位嵌入层 self.pos_emb 进行嵌入，得到方位嵌入张量。一起，将输入序列 x 传入符号嵌入层 self.token_emb 进行嵌入，得到 token 的嵌入张量。最终，将 token 嵌入张量和方位嵌入张量相加，得到融合了符号和方位信息的嵌入张量，并将其作为输出返回。

整个 TokenAndPositionEmbedding 类的作用是将输入序列的 token 和方位信息进行嵌入核算，为后续的 Transformer 模型供给丰富的输入表明。在 Transformer 模型中，token 嵌入用于表明每个输入 token 的语义信息，而方位嵌入用于表明每个输入 token 在序列中的方位信息。

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 = tf.shape(x)[-1]
        positions = tf.range(start=0, limit=maxlen, delta=1)
        positions = self.pos_emb(positions)
        x = self.token_emb(x)
        return x + positions

Create Model

这儿界说了一个 create_model 函数，用于创建一个 Transformer 模型。这个模型运用 Transformer 架构来处理输入序列，并在最终经过全连接层进行分类猜测。它能够学习输入序列中的语义和上下文联系，用于生成猜测的单词概率散布。

函数中首先创建了一个输入层 inputs，其形状为 (maxlen,)，数据类型为 tf.int32，用于接收输入序列。接下来界说了一个 TokenAndPositionEmbedding 层，传入参数 maxlen、vocab_size 和 embed_dim，用于将输入序列的符号和方位信息进行嵌入。将输入层 inputs 作为输入传递给嵌入层，得到嵌入后的输出张量 x。

然后创建了一个 TransformerBlock 层，传入参数 embed_dim、num_heads 和 feed_forward_dim，用于对嵌入后的序列进行 Transformer 操作。将嵌入后的张量 x 传递给 TransformerBlock 层，得到处理后的输出张量 x。

接下来经过一个全连接层 layers.Dense 对输出张量 x 进行猜测，输出一个形状为 (vocab_size,) 的张量 outputs，也就是核算出来下一个猜测的单词的概率散布。

最终，界说了丢失函数 loss_fn 为稀疏分类穿插熵丢失函数，并运用 "adam" 优化器进行模型的编译。

vocab_size = 20000
maxlen = 80  
embed_dim = 256  
num_heads = 2  
feed_forward_dim = 256 
def create_model():
    inputs = layers.Input(shape=(maxlen,), dtype=tf.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 = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
    model.compile( "adam", loss=[loss_fn, None], )
    return model

Text Generator

这儿界说了一个 TextGenerator 类，用于在练习过程作为回调函数来生成文本，展现在不同练习 epoch 下面的文本生成作用。

结构函数 __init__ 接收参数 max_tokens、start_tokens、index_to_word、top_k 和 print_every，用于装备文本生成的相关参数。

sample_from 办法用于从给定的 logits（对数概率）中进行采样，依据概率散布挑选下一个猜测的单词。它首先运用 tf.math.top_k 挑选概率最高的前 k 个单词，然后进行 softmax 归一化，得到概率散布。最终，运用 np.random.choice 办法依据概率散布进行采样，挑选下一个猜测的单词。

on_epoch_end 办法在每个练习周期完毕时调用，用于生成文本。它经过循环生成文本的过程，从给定的开始文本开端，逐渐生成下一个单词，直到达到指定的最大生成单词数。在每次生成单词后，将其添加到已生成的列表中，并更新开始文本。最终，将生成的文本转换为字符串，并打印输出。

接下来，界说了一个开始提示文本 start_prompt 为 this movie is very good，并依据词汇表和开始提示文本生成了开始符号 start_tokens。然后，创建了一个 TextGenerator 目标 text_gen_callback，传入生成文本所需的参数。

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 = tf.math.top_k(logits, k=self.k, sorted=True)
        indices = np.asarray(indices).astype("int32")
        preds = keras.activations.softmax(tf.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)
            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 very good"
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)

Train

该部分就是创建了一个文本生成模型，练习 30 个 epoch ，而且调用 text_gen_callback 目标，在每次 epoch 完毕的时分进行文本的生成。

model = create_model()
model.fit(text_ds, epochs=30, callbacks=[text_gen_callback])

下面对部分结果进行打印展现。可以看出来生成的文本作用一般，可能和数据集质量以及模型的复杂度有关。

Epoch 1/30
0s 169ms/stepse_5_loss: 5.59
generated text:
this movie is very good movie . the worst movie is about the plot . the story of course of course the story line is a great story about it . it is so well . the plot is a great plot of the
Epoch 2/30
0s 17ms/step- loss: 4.7109 - dense_5_loss: 4.71
generated text:
this movie is a great movie . a wonderful movie about it , it was just the characters that they were not a movie but the way the acting was not a bad script that is bad . but the script was bad ,
...
Epoch 12/30
0s 18ms/step- loss: 3.6976 - dense_5_loss: 3.69
generated text:
this movie is one of the best movies i have ever seen and i have seen it on vhs uncut and i 've seen the first time . i watched this film for the first and was all of it . it was great
Epoch 13/30
0s 19ms/step- loss: 3.6531 - dense_5_loss: 3.65
generated text:
this movie is one of the worst actors i have ever seen . it is the worst bollywood movie i have ever seen . i have no idea it . the acting was terrible and the directing is bad , but it was bad
...
Epoch 29/30
0s 18ms/step- loss: 3.2507 - dense_5_loss: 3.25
generated text:
this movie is so [UNK] and the acting is awful , but the script is poor . the plot is laughable and the ending is terrible . there isn 't anything about this movie that was so bad it doesn 't make any sense
Epoch 30/30
0s 17ms/step- loss: 3.2359 - dense_5_loss: 3.23
generated text:
this movie is not a great time , but this movie is one of those actors that are [UNK] and that you can not take up to the screen . the plot is simple . it doesn 't matter what 's going on and

mini GPT 实现文本生成

前文

Data

Miniature GPT

Transformer Block

Token And Position Embedding

Create Model

Text Generator

Train

作者信息

mini GPT 实现文本生成

前文

Data

Miniature GPT

Transformer Block

Token And Position Embedding

Create Model

Text Generator

Train

相关文章

【云原生】Redis on k8s 编排部署讲解与实战操作

【青训营 x 字节前端训练营】✍🏻 笔记创作活动

Bert+LSTM+CRF命名实体识别

协程(16) | 优雅地实现一个挂起函数

作者信息