lyrics / lyrics_generation_rnn.py

Upload lyrics_generation_rnn.py

e30cc6b 10 months ago

No virus

14.6 kB

	# -- coding: utf-8 --
	"""lyrics_generation_rnn.ipynb

	Automatically generated by Colaboratory.

	Original file is located at
	https://colab.research.google.com/drive/1MkBq8eqZoPqaVDczKmYhSThcV4r23z25
	"""

	!pip install pickle
	import pickle
	!pip install string
	import string

	import tensorflow as tf
	from string import punctuation
	import numpy as np
	import os
	import time
	import pickle
	model_path='/content/drive/MyDrive/Colab Notebooks'
	# create directory to store pickled files in
	if not os.path.exists(f'/content/drive/MyDrive/Colab Notebooks/pkl'):
	os.mkdir(f'/content/drive/MyDrive/Colab Notebooks/pkl')

	# ----------------------------------------------------------------------

	### LIMITING GPU MEMORY GROWTH ###

	# get list of visible GPUs
	gpus = tf.config.experimental.list_physical_devices('GPU')

	if gpus: # if GPU(s) is detected
	try: # try setting memory growth to true for all GPUs
	for gpu in gpus:
	tf.config.experimental.set_memory_growth(gpu, True) # enabling memory growth
	logical_gpus = tf.config.experimental.list_logical_devices('GPU')
	print('\n', len(gpus), 'Physical GPUs,', len(logical_gpus), 'Logical GPU')
	except RuntimeError as e:
	# memory growth must be set before GPUs have been initialized
	print('\n', e)

	# ----------------------------------------------------------------------

	### READ IN AND CLEAN THE LYRICS DATA ###

	# ****TAKE IN USER INPUT FOR LYRICS (ARTIST NAME? FILE NAME?)****

	# read in the lyrics text file
	text = str(open('/content/drake.txt', 'r').read())
	# artist_name = input('\nPlease ')

	# make all letters lowercase and make line breaks into its own "word"
	words = text.lower().replace('\n', ' \n ')

	# remove punctuation
	for punc in punctuation:
	words = words.replace(punc, '')

	# split the entire words string into a Python list of words
	words = words.split(' ')

	# obtain list of unique words across all lyrics
	vocab = sorted(set(words))
	print(f'\nThere are {len(vocab)} unique words in the lyrics file.')

	# pickle the vocab file - will need it for the generation script
	outfile = open(file='/content/drive/MyDrive/Colab Notebooks/pkl/vocab', mode='wb')
	pickle.dump(vocab, outfile)
	outfile.close()

	# ----------------------------------------------------------------------

	### WORD MAPPING ###

	# map unique characters to indices
	word2idx = {u:i for i, u in enumerate(vocab)}

	# pickle this since it is needed in text generation
	outfile = open(file='/content/drive/MyDrive/Colab Notebooks/pkl/word2idx', mode='wb')
	pickle.dump(word2idx, outfile)
	outfile.close()

	# reverse the map - use this to specify an index to obtain a character
	idx2word = np.array(vocab)

	# pickle this since it is needed in text generation
	outfile = open(file='/content/drive/MyDrive/Colab Notebooks/pkl/idx2word', mode='wb')
	pickle.dump(idx2word, outfile)
	outfile.close()

	# entire text document represented in the above character-to-indices mapping
	words_as_int = np.array([word2idx[c] for c in words])

	# ----------------------------------------------------------------------

	### CREATING TRAINING EXAMPLES & TARGETS ###

	# ****TAKE IN USER INPUT FOR SEQUENCE LENGTH?****

	# max sentence length (in number of words) desired for training
	seq_length = 100
	# seq_length = input('\nPlease enter a desired sequence length (in number of words) to train the model on: ')
	examples_per_epoch = len(words) // (seq_length + 1)

	# create training examples/targets
	word_dataset = tf.data.Dataset.from_tensor_slices(words_as_int)

	# data type of train examples/targets
	print('\n', type(word_dataset))

	# create sequence batches from the word_dataset
	sequences = word_dataset.batch(seq_length + 1, drop_remainder=True)
	print('\n', type(sequences))

	# define the shifting (splitting) function
	def split_input_target(chunk):
	input_text = chunk[:-1] # up to but not including the last character
	target_text = chunk[1:] # everything except for the firs tcharacter
	return input_text, target_text

	# apply the shifting to create input texts and target texts that comprise of our dataset
	dataset = sequences.map(split_input_target)

	# ----------------------------------------------------------------------

	### CREATE TRAINING BATCHES ###

	# batch size
	BATCH_SIZE = 64

	# buffer size to shuffle the dataset
	# (TensorFlow data is designed to work with possibly infinite sequences,
	# so it doesn't attempt to shuffle the entire sequence in memory. Instead,
	# it maintains a buffer in which it shuffles elements)
	BUFFER_SIZE = 10000

	# create a dataset that has been shuffled and batched
	dataset_sb = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)

	# display batch dataset shapes and data types
	print('\n', dataset_sb)

	# ----------------------------------------------------------------------

	### BUILDING THE RNN ###

	# vocabulary length (number of unique words in dataset)
	vocab_size = len(vocab)

	# embedding dimension
	embedding_dim = 256

	# number of RNN units
	rnn_units = 1024

	# pickle model parameters - will need in the generation script
	model_params = [vocab_size, embedding_dim, rnn_units]
	outfile = open(file='/content/drive/MyDrive/Colab Notebooks/pkl/model_params', mode='wb')
	pickle.dump(model_params, outfile)
	outfile.close()

	# helper function to quickly build the RNN model based on vocab size, embedding dimension, number of RNN units, and batch size
	def build_model(vocab_size, embedding_dim, rnn_units, batch_size):

	# initialize sequential model architecture
	model = tf.keras.Sequential()

	# add embedding layer
	model.add(tf.keras.layers.Embedding(
	input_dim = vocab_size,
	output_dim = embedding_dim,
	batch_input_shape=[batch_size, None]
	))

	# add recurrent layer
	model.add(tf.keras.layers.GRU(
	units = rnn_units,
	return_sequences = True,
	stateful = True,
	recurrent_initializer = 'glorot_uniform'
	))

	# add dense layer
	model.add(tf.keras.layers.Dense(units=vocab_size))
	model_path= '/content/drive/MyDrive/Colab Notebooks'

	def save_model(self, model_path):
	# Save the model weights
	self.save_weights(model_path)
	print(f"Model saved to {model_path}")
	return model

	# build the model using the above helper function
	rnn = build_model(
	vocab_size = vocab_size,
	embedding_dim = embedding_dim,
	rnn_units = rnn_units,
	batch_size = BATCH_SIZE
	)

	# check the shape of the output
	for input_example_batch, target_example_batch in dataset_sb.take(1):
	example_batch_predictions = rnn(input_example_batch)
	print('\n', example_batch_predictions.shape, '# (batch_size, sequence_length, vocab_size)')

	# model architecture summary
	print('\n', rnn.summary(), '\n')

	# ----------------------------------------------------------------------

	### SET UP METRICS ###

	# helper function to obtain the loss function
	def loss(labels, logits):
	return tf.keras.losses.sparse_categorical_crossentropy(labels, logits, from_logits=True)

	# compile the model
	rnn.compile(
	optimizer = 'adam',
	loss = loss,
	metrics = ['accuracy']
	)

	# create directory where the checkpoints will be saved
	checkpoint_dir = '/content/drive/MyDrive/Colab Notebooks/training_checkpoints'

	# name of the checkpoint files
	checkpoint_prefix = os.path.join(checkpoint_dir, 'checkpoint')

	# create checkpoints-saving object
	checkpoint_callback = tf.keras.callbacks.ModelCheckpoint(
	filepath = checkpoint_prefix,
	monitor = 'loss',
	save_best_only = True,
	mode = 'min',
	save_weights_only = True
	)

	# ----------------------------------------------------------------------

	### MODEL TRAINING ###

	# set number of desired epochs
	EPOCHS = 200

	# training!
	history = rnn.fit(
	x = dataset_sb,
	epochs = EPOCHS,
	callbacks = [checkpoint_callback]
	)

	build_model.save('/content/drive/MyDrive/Colab Notebooks')

	import tensorflow as tf
	from string import punctuation
	import pickle

	# ----------------------------------------------------------------------

	### LIMITING GPU MEMORY GROWTH ###

	# get list of visible GPUs
	gpus = tf.config.experimental.list_physical_devices('GPU')

	if gpus: # if GPU(s) is detected
	try: # try setting memory growth to true for all GPUs
	for gpu in gpus:
	tf.config.experimental.set_memory_growth(gpu, True) # enabling memory growth
	logical_gpus = tf.config.experimental.list_logical_devices('GPU')
	print('\n', len(gpus), 'Physical GPUs,', len(logical_gpus), 'Logical GPU')
	except RuntimeError as e:
	# memory growth must be set before GPUs have been initialized
	print('\n', e)

	# -------------------------------------------------------------------------

	### MODEL BUILDING FUNCTION FROM TRAINING SCRIPT ###

	# helper function to quickly build the RNN model based on vocab size, embedding dimension, number of RNN units, and batch size
	def build_model(vocab_size, embedding_dim, rnn_units, batch_size):
	model = tf.keras.Sequential()

	model.add(tf.keras.layers.Embedding(
	input_dim = vocab_size,
	output_dim = embedding_dim,
	batch_input_shape=[batch_size, None]
	))

	model.add(tf.keras.layers.GRU(
	units = rnn_units,
	return_sequences = True,
	stateful = True,
	recurrent_initializer = 'glorot_uniform'
	))

	model.add(tf.keras.layers.Dense(units=vocab_size))

	model_path= '/content/drive/MyDrive/Colab Notebooks'

	def save_model(self, model_path):
	# Save the model weights
	self.save_weights(model_path)
	print(f"Model saved to {model_path}")

	return model

	# -------------------------------------------------------------------------

	### INITIATE MODEL AND LOAD IN WEIGHTS FROM CHECKPOINT ###

	# unpickle the model parameters from the training script
	infile = open(file='pkl/model_params', mode='rb')
	vocab_size, embedding_dim, rnn_units = pickle.load(infile)
	infile.close()

	# initiate new model instance
	rnn_cp = build_model(vocab_size, embedding_dim, rnn_units, batch_size=1)

	# load saved weights from checkpoint into new model instance
	rnn_cp.load_weights(tf.train.latest_checkpoint('./training_checkpoints'))

	# build the model with a new input shape
	rnn_cp.build(tf.TensorShape([1, None]))

	# -------------------------------------------------------------------------

	### TEXT PREDICTION FUNCTION ###

	# unpickle the index-word files that were pickled from the training script
	infile = open(file='pkl/word2idx', mode='rb')
	word2idx = pickle.load(infile)
	infile.close()
	infile = open(file='pkl/idx2word', mode='rb')
	idx2word = pickle.load(infile)
	infile.close()
	#build_model.is_valid():
	#build_model.save('/content/drive/MyDrive/Colab Notebooks')

	def generate_text(model, start_string, num_generate=500, temperature=1.0):

	# num of chars to generate
	num_generate = num_generate

	# vectorizing the start string to numbers
	input_eval = [word2idx[s] for s in start_string]
	input_eval = tf.expand_dims(input=input_eval, axis=0) # returns a tensor with a length-1 axis inserted at index `axis`

	# empty string to store results
	text_generated = list()

	# "temperature"
	# low temperature results in more predictable text,
	# high temperature results in more surprising text.
	# feel free to experiment with this parameter
	temperature = 1.0

	# the batch size was defined when we loaded model weights from training

	model.reset_states()
	for i in range(num_generate):
	predictions = model(input_eval)

	# remove the batch dimension
	predictions = tf.squeeze(predictions, 0)

	# use a categorical distribution to predict the character returned by the model
	preidctions = predictions / temperature
	predicted_id = tf.random.categorical(predictions, num_samples=1)[-1, 0].numpy()

	# pass the predicted character as the next input to the model along with the previous hidden state
	input_eval = tf.expand_dims([predicted_id], 0)

	text_generated.append(idx2word[predicted_id])

	return(' '.join(start_string + text_generated))

	# -------------------------------------------------------------------------

	### TAKE IN INPUT STRING AND CHECK IF ALL WORDS IN IT ARE IN THE VOCABULARY ###
	# (this is a requirement for text generation)

	# unpickle the vocabulary file that was pickled from the training script
	infile = open(file='pkl/vocab', mode='rb')
	vocab = pickle.load(infile)
	infile.close()

	# initialize the checking loop
	check = True

	while check:

	# take in user input for starting lyrics
	start_string = input('\nPlease input some text to initiate the lyrics generation (caps insensitive):\n')

	# lowercase
	start_string = start_string.lower()

	# remove punctuation
	for punc in punctuation:
	start_string = start_string.replace(punc, '')

	# create a list where each element is one word from the start string
	start_string = start_string.split(' ')

	# store all words that aren't in the vocabulary
	non_vocab = []

	# for every word in the start string
	for word in start_string:

	# if the word is NOT in the vocabulary
	if word not in vocab:

	# add the word to the non_vocab variable
	non_vocab.append(word)

	# if the non-vocab list is empty (i.e. all words in the start string are in the vocab)
	if non_vocab == []:

	# break out of the loop
	check = False

	# if there are words not in the vocabulary
	else:

	# print what those words are
	print(f'\nWords in the input text not present in the vocabulary are: {", ".join(non_vocab)}')
	print('\nAll input words must be in the vocabulary.')

	# -------------------------------------------------------------------------

	### TEXT GENERATION ###

	# text generation!
	print('\n', generate_text(rnn_cp, start_string=start_string, num_generate=250))

	### SAVE TO FILE??? ###

	# -------------------------------------------------------------------------

	# -------------------------------------------------------------------------



	build_model.save('/content/drive/MyDrive/Colab Notebooks')

	model = build_model

	"""import tensorflow as tf
	build_model.state_dict()
	# Assuming you have a trained model named 'model'
	model = ...

	# Define the path to save the model
	model_path = 'path_to_save_model'

	# Save the entire model (architecture, weights, and optimizer state)
	model.save(model_path)

	[link text](https:// [link text](https://))# Alternatively, you can save only the model weights
	model.save_weights('path_to_save_weights')

	# You can also save the model in a format optimized for serving
	tf.saved_model.save(model, 'path_for_serving')

	"""