<jupyter_start><jupyter_text><jupyter_code>%matplotlib
import tensorflow as tf
from tensorflow.keras.layers import Convolution2D, UpSampling2D, AveragePooling2D, Convolution2DTranspose, Input
from tensorflow.keras.models import Model
import numpy as np
import matplotlib.pyplot as plt
from scipy import signal


AUDIO_CHUNK_SIZE = 1024
RATE = 44100

# build the model
kernel_size = (1,5)
input_audio = Input(shape=(1, AUDIO_CHUNK_SIZE, 1))

net = Convolution2D(filters=64, kernel_size=kernel_size,activation='relu',padding='same')(input_audio)
net = AveragePooling2D((1,2))(net)
net = Convolution2D(filters=32, kernel_size=kernel_size,activation='relu',padding='same')(net)
net = AveragePooling2D((1,2))(net)
net = Convolution2D(filters=4, kernel_size=kernel_size, activation='relu',padding='same')(net)
net = AveragePooling2D((1,2))(net)
net = Convolution2D(filters=2, kernel_size=kernel_size, activation='relu',padding='same')(net)

# audio chunk encoded in 128 data points
encoded = AveragePooling2D((1,2), name='encoded_audio')(net)

net = Convolution2DTranspose(filters=2, kernel_size=kernel_size, activation='relu', padding='same')(encoded)
net = UpSampling2D((1,2))(net)
net = Convolution2DTranspose(filters=4, kernel_size=kernel_size, activation='relu', padding='same')(net)
net = UpSampling2D((1,2))(net)
net = Convolution2DTranspose(filters=32, kernel_size=kernel_size, activation='relu', padding='same')(net)
net = UpSampling2D((1,2))(net)
net = Convolution2DTranspose(filters=64, kernel_size=kernel_size, activation='relu', padding='same')(net)
net = UpSampling2D((1,2))(net)

decoded = Convolution2DTranspose(filters=1, kernel_size=kernel_size, activation='linear', padding='same', name='decoded_audio')(net)

autoencoder = Model(input_audio, decoded, name='simple_autoencoder')
autoencoder.summary()
# lets build an audio dataset of random sinus chunks!

def sinus_dataset_generator(num_examples, fs, samples, frequency_range):
    """Builds a dataset of sinus.
    
    Args:
        num_examples: number of examples to generate (int)
        fs: sample rate of the sinus
        samples: number of samples to generate (int)
        frequency_range: a list of two values defining [lower, upper] frequency range (int)
    Returns:
        A numpy array of sinus examples.
    """
    # first example
    sinus_data = (np.sin((2*np.pi*np.arange(samples)*440.0/fs) + 0.0)).astype(np.float32)
    sinus_data = np.reshape(sinus_data, newshape=(1, 1, samples, 1))
    
    for idx in range(0, num_examples-1): 
        # random frequency
        f = np.random.randint(frequency_range[0], frequency_range[1])
        # random phase shift
        phase = np.random.random() * np.pi
        # random gain
        gain = np.random.uniform(0.5, 1.0)

        sinus = (np.sin((2*np.pi*np.arange(samples)*f/fs) + phase) * gain).astype(np.float32)
        # add some noise, mu = 0, sigma = 0.1
        s = np.random.normal(0, 0.1, samples)
        sinus = sinus + s
        # bring it into shape for the model 
        sinus = np.reshape(sinus, newshape=(1, 1, samples, 1))
        sinus_data = np.append(sinus_data, sinus, axis=0)
    return sinus_data

num_examples = 8000
sinus_data = sinus_dataset_generator(num_examples, RATE, AUDIO_CHUNK_SIZE, [30, 8000])

# split into train and eval dataset, roughly a 70/30 split
split = int(num_examples * 0.7)
train_dataset = tf.data.Dataset.from_tensor_slices((sinus_data[:split], sinus_data[:split]))
eval_dataset = tf.data.Dataset.from_tensor_slices((sinus_data[split:], sinus_data[split:]))

# see if its actually a sinus...
fig, ax = plt.subplots(3, 3)
for ii in range(3):
  for jj in range(3):
    ax[ii, jj].plot(np.squeeze(sinus_data[np.random.randint(0, num_examples-1)])[:200], color="tab:blue")
    ax[ii, jj].set_xticks([])
    ax[ii, jj].set_yticks([])
fig.savefig('/tmp/sinus_examples.pdf')
# compile the model with adam optimizer and mean squeared error as its loss function
autoencoder.compile(optimizer='adam', loss='mse', metrics=['mse'])

# some constants for training
BATCH_SIZE = 32
SHUFFLE_BUFFER_SIZE = 100

# shuffle and batch the examples
train_dataset = train_dataset.shuffle(SHUFFLE_BUFFER_SIZE).batch(BATCH_SIZE)
eval_dataset = eval_dataset.batch(BATCH_SIZE)

# train the model
autoencoder.fit(train_dataset, epochs=100, validation_data=eval_dataset)
# plot some examples of sinus and its reconstruction
fig, ax = plt.subplots(3, 3)
for ii in range(3):
  for jj in range(3):
    sinus = sinus_data[ii * 3 + jj + 1]
    sinus = np.expand_dims(sinus, axis=0)
    prediction = autoencoder.predict(sinus, batch_size=1)
    ax[ii, jj].plot(np.squeeze(sinus)[:200], color="tab:blue")
    ax[ii, jj].plot(np.squeeze(prediction)[:200], color="tab:red")
    ax[ii, jj].set_xticks([])
    ax[ii, jj].set_yticks([])
fig.savefig('/tmp/sinus_reconstruction.pdf')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13623248145094807, 'fraction_numerical': 0.030090684253915912, 'mean_word_length': 5.213828425096031, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '346232', 'license': 'MIT', 'n_tokens_mistral': 1785, 'n_tokens_neox': 1702, 'n_words': 375, 'path': 'raw_audio_autoencoder.ipynb', 'repo_name': 'iCorv/raw_audio_autoencoder', 'size': '4852'}