208 lines
6.5 KiB
Python
208 lines
6.5 KiB
Python
import tensorflow as tf
|
|
from tensorflow.keras.layers.experimental import preprocessing
|
|
|
|
import numpy as np
|
|
import os
|
|
import time
|
|
|
|
path_to_file = "/Users/williamjones/Downloads/untitled folder/tensorflow/messages.txt"
|
|
|
|
text = open(path_to_file, 'rb').read().decode('utf-8')
|
|
print(f'Length of text: {len(text)} characters')
|
|
print()
|
|
|
|
print('First 250 characters:')
|
|
print(text[:250])
|
|
|
|
print()
|
|
vocab = sorted(set(text))
|
|
print(f'{len(vocab)} unique characters.')
|
|
|
|
example_texts = ['abcdefg', 'xyz']
|
|
chars = tf.strings.unicode_split(example_texts, input_encoding='UTF-8')
|
|
print()
|
|
print(chars)
|
|
|
|
ids_from_chars = preprocessing.StringLookup(vocabulary=list(vocab))
|
|
ids = ids_from_chars(chars)
|
|
print(ids)
|
|
|
|
chars_from_ids = tf.keras.layers.experimental.preprocessing.StringLookup(vocabulary=ids_from_chars.get_vocabulary(), invert=True)
|
|
|
|
chars = chars_from_ids(ids)
|
|
print(chars)
|
|
|
|
def text_from_ids(ids):
|
|
return tf.strings.reduce_join(chars_from_ids(ids), axis=-1)
|
|
|
|
all_ids = ids_from_chars(tf.strings.unicode_split(text, 'UTF-8'))
|
|
print(all_ids)
|
|
|
|
ids_dataset = tf.data.Dataset.from_tensor_slices(all_ids)
|
|
for ids in ids_dataset.take(10):
|
|
print(chars_from_ids(ids).numpy().decode('utf-8'))
|
|
|
|
seq_length = 100
|
|
examples_per_epoch = len(text)
|
|
|
|
sequences = ids_dataset.batch(seq_length+1, drop_remainder=True)
|
|
|
|
for seq in sequences.take(1):
|
|
print(chars_from_ids(seq))
|
|
|
|
for seq in sequences.take(5):
|
|
print(text_from_ids(seq).numpy())
|
|
|
|
def split_input_target(sequence):
|
|
input_text = sequence[:-1]
|
|
target_text = sequence[1:]
|
|
return input_text, target_text
|
|
|
|
split_input_target(list("Tensorflow"))
|
|
|
|
dataset = sequences.map(split_input_target)
|
|
|
|
for input_example, target_example in dataset.take(1):
|
|
print("Input:", text_from_ids(input_example).numpy())
|
|
print("Target:", text_from_ids(target_example).numpy())
|
|
|
|
# Training batches
|
|
BATCH_SIZE = 64
|
|
BUFFER_SIZE = 10000
|
|
print()
|
|
dataset = (dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE))
|
|
print(dataset)
|
|
|
|
# Model
|
|
vocab_size = len(vocab)
|
|
embedding_dim = 256
|
|
rnn_units = 1024
|
|
|
|
class MyModel(tf.keras.Model):
|
|
def __init__(self, vocab_size, embedding_dim, rnn_units):
|
|
super().__init__(self)
|
|
self.embedding = tf.keras.layers.Embedding(vocab_size, embedding_dim)
|
|
self.gru = tf.keras.layers.GRU(rnn_units, return_sequences=True, return_state=True)
|
|
self.dense = tf.keras.layers.Dense(vocab_size)
|
|
|
|
def call(self, inputs, states=None, return_state=False, training=False):
|
|
x = inputs
|
|
x = self.embedding(x, training=training)
|
|
if states is None:
|
|
states = self.gru.get_initial_state(x)
|
|
x, states = self.gru(x, initial_state=states, training=training)
|
|
x = self.dense(x, training=training)
|
|
if return_state:
|
|
return x, states
|
|
else:
|
|
return x
|
|
|
|
model = MyModel(vocab_size=len(ids_from_chars.get_vocabulary()), embedding_dim=embedding_dim, rnn_units=rnn_units)
|
|
print()
|
|
print()
|
|
print()
|
|
# Run the model once before training
|
|
for input_example_batch, target_example_batch in dataset.take(1):
|
|
example_batch_predictions = model(input_example_batch)
|
|
print(example_batch_predictions.shape, '# (batch_size, sequence_length, vocab_size)')
|
|
|
|
model.summary()
|
|
|
|
sampled_indices = tf.random.categorical(example_batch_predictions[0], num_samples=1)
|
|
sampled_indices = tf.squeeze(sampled_indices, axis=-1).numpy()
|
|
|
|
print(sampled_indices)
|
|
|
|
# See text prediced by untrained model
|
|
print("Input:\n", text_from_ids(input_example_batch[0]).numpy())
|
|
print()
|
|
print("Next char predictions:\n", text_from_ids(sampled_indices).numpy())
|
|
|
|
|
|
# Training time!
|
|
print()
|
|
print()
|
|
print()
|
|
print()
|
|
print()
|
|
print()
|
|
loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True)
|
|
example_batch_loss = loss(target_example_batch, example_batch_predictions)
|
|
mean_loss = example_batch_loss.numpy().mean()
|
|
print("Prediction Shape: ", example_batch_predictions.shape, '# (batch_size, sequence_length, vocab_size)')
|
|
print("Mean loss: ", mean_loss)
|
|
|
|
print(tf.exp(mean_loss).numpy())
|
|
|
|
model.compile(optimizer='adam', loss=loss)
|
|
|
|
checkpoint_dir = './training_checkpoints'
|
|
checkpoint_prefix = os.path.join(checkpoint_dir, 'ckpt_{epoch}')
|
|
|
|
checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(filepath=checkpoint_prefix, save_weights_only=True)
|
|
|
|
EPOCHS = 100
|
|
history = model.fit(dataset, epochs=EPOCHS, callbacks=checkpoint_callback)
|
|
|
|
class OneStep(tf.keras.Model):
|
|
def __init__(self, model, chars_from_ids, ids_from_chars, temperature=1.0):
|
|
super().__init__()
|
|
self.temperature = temperature
|
|
self.model = model
|
|
self.chars_from_ids = chars_from_ids
|
|
self.ids_from_chars = ids_from_chars
|
|
|
|
# Create a mask to prevent "" or "[UNK]" from being generated
|
|
skip_ids = self.ids_from_chars(['', '[UNK]'])[:, None]
|
|
sparse_mask = tf.SparseTensor(values=[-float('inf')]*len(skip_ids), indices=skip_ids, dense_shape=[len(ids_from_chars.get_vocabulary())])
|
|
self.prediction_mask = tf.sparse.to_dense(sparse_mask)
|
|
|
|
@tf.function
|
|
def generate_one_step(self, inputs, states=None):
|
|
input_chars = tf.strings.unicode_split(inputs, 'UTF-8')
|
|
input_ids = self.ids_from_chars(input_chars).to_tensor()
|
|
|
|
predicted_logits, states = self.model(inputs=input_ids, states=states, return_state=True)
|
|
|
|
predicted_logits = predicted_logits[:, -1, :]
|
|
predicted_logits = predicted_logits/self.temperature
|
|
|
|
predicted_logits = predicted_logits + self.prediction_mask
|
|
|
|
predicted_ids = tf.random.categorical(predicted_logits, num_samples=1)
|
|
predicted_ids = tf.squeeze(predicted_ids, axis=-1)
|
|
|
|
predicted_chars = self.chars_from_ids(predicted_ids)
|
|
|
|
return predicted_chars, states
|
|
|
|
one_step_model = OneStep(model, chars_from_ids, ids_from_chars)
|
|
|
|
start = time.time()
|
|
states = None
|
|
next_char = tf.constant(['Hello'])
|
|
result = [next_char]
|
|
|
|
for n in range(1000):
|
|
next_char, states = one_step_model.generate_one_step(next_char, states=states)
|
|
result.append(next_char)
|
|
|
|
result = tf.strings.join(result)
|
|
end = time.time()
|
|
|
|
print(result[0].numpy().decode('utf-8'), '\n\n' + '_'*80)
|
|
print('Run time: ', end - start)
|
|
|
|
tf.saved_model.save(one_step_model, 'one_step')
|
|
one_step_reloaded = tf.saved_model.load('one_step')
|
|
|
|
states = None
|
|
next_char = tf.constant(['Hello'])
|
|
result = [next_char]
|
|
|
|
for n in range(100):
|
|
next_char, states = one_step_reloaded.generate_one_step(next_char, states=states)
|
|
result.append(next_char)
|
|
|
|
print(tf.strings.join(result)[0].numpy().decode("utf-8"))
|