<jupyter_start><jupyter_text>TODO: convert to script<jupyter_code>import sys
sys.path.append("../src/")

import nussl
from nussl.datasets import transforms as nussl_tfm
from torch.utils.data import Dataset, DataLoader
import torch.optim as optim
import torch.nn as nn
import torch.nn.functional as F
import torch
import matplotlib.pyplot as plt

import models
# rnn separator
rnn_config = {
    'bidirectional': True,
    'dropout': 0.3,
    'filter_length': 256,
    'hidden_size': 50,
    'hop_length': 64,
    'mask_activation': ['softmax'],
    'mask_complex': False,
    'mix_key': 'mix_audio',
    'normalization_class': 'BatchNorm',
    'num_audio_channels': 1,
    'num_filters': 256,
    'num_layers': 1,
    'num_sources': 2,
    'rnn_type': 'lstm',
    'window_type': 'sqrt_hann',
}

rnn_config = nussl.ml.networks.builders.build_recurrent_end_to_end(**rnn_config)

rnn = models.RnnSeparator(rnn_config)
class PreTrainDataset(Dataset):
    def __init__(self):
        # load the two sources
        source1 = nussl.AudioSignal('../sounds/phone/cellphone_ringing.wav')
        source2 = nussl.AudioSignal('../sounds/siren/siren.wav')
        
        # preprocess
        stft = nussl.STFTParams(window_length=512, hop_length=128)
        source1.stft_params = stft
        source2.stft_params = stft
        
        # truncate and mix
        source2.to_mono(overwrite=True)
        source1.truncate_seconds(source2.signal_duration)
        mix = source1 + source2
        sample = {'mix': mix, 'sources': {'source1': source1, 'source2': source2}}
        
        
        # transform
        self.transform = nussl_tfm.Compose([
            nussl_tfm.GetAudio(),
            nussl_tfm.ToSeparationModel()
        ])
        
        sample = self.transform(sample)
        
        # train on many of the same examples
        self.items = [sample] * 300
    
    def __getitem__(self, index):
        return self.items[index]
    
    def __len__(self):
        return len(self.items)
dataset = PreTrainDataset()
dataloader = DataLoader(dataset, batch_size=10)
loss_fn = nussl.ml.train.loss.SISDRLoss()
optimizer = optim.Adam(rnn.parameters(), lr=.001)
losses = []
for i, data in enumerate(dataloader):    
    output = rnn(data)
    
    # preds vs. ground truth
    loss = loss_fn(output['audio'], data['source_audio'])
    print(loss)
    losses.append(loss)
    
    # backprop
    loss.backward()
    optimizer.step()
plt.plot(losses)
SAVE_PATH = constants.PRETRAIN_PATH
# add save code
torch.save(rnn.rnn_model.state_dict(), SAVE_PATH)<jupyter_output><empty_output><jupyter_text>Evaluate<jupyter_code># get sources and mix
source1 = nussl.AudioSignal('../sounds/phone/cellphone_ringing.wav')
source2 = nussl.AudioSignal('../sounds/siren/siren.wav')
stft = nussl.STFTParams(window_length=512, hop_length=128)
source1.stft_params = stft
source2.stft_params = stft

source2.to_mono(overwrite=True)
source1.truncate_seconds(source2.signal_duration)
mix = source1 + source2

print(mix)
mix.embed_audio(display=False)
def visualize_and_embed(sources):
    plt.figure(figsize=(10, 6))
    plt.subplot(211)
    nussl.utils.visualize_sources_as_masks(sources,
        y_axis='linear', db_cutoff=-80, alpha_amount=2.0)
    plt.subplot(212)
    nussl.utils.visualize_sources_as_waveform(
        sources, show_legend=False)
    plt.show()
    nussl.play_utils.multitrack(sources)
# in case we need to load saved model
# rnn = models.RnnSeparator(rnn_config)
# model_dict = torch.load(save_path)
# rnn.rnn_model = nussl.ml.SeparationModel(model_dict['config'])
# rnn.rnn_model.load_state_dict(model_dict['state_dict'])

separator = nussl.separation.deep.DeepAudioEstimation(mix, model_path=None)
separator.model = rnn.rnn_model

print('\n', separator.model)

separator.transform = nussl_tfm.Compose([
    nussl_tfm.GetAudio(),
    nussl_tfm.ToSeparationModel()
])

separator.device = 'cpu'

print('\n', mix.stft_params)
separator.metadata = {
    'sample_rate': 44100,
    'stft_params': mix.stft_params,
    'num_channels': 1
}

estimates = separator()
print('\n', estimates)

# passing "mix" results in a error
ibm = nussl.separation.benchmark.IdealBinaryMask(mix, estimates)
estimates = ibm()

estimates = {
    f'Source {i}': e for i, e in enumerate(estimates)
}

visualize_and_embed(estimates)
estimates['Source 0'].embed_audio(display=False)
estimates['Source 1'].embed_audio(display=False)<jupyter_output><empty_output><jupyter_text>Checkout model<jupyter_code>SAVE_PATH = '../models/pretrained.pth'
model_dict = torch.load(SAVE_PATH)
#print(model_dict)
rnn = models.RnnSeparator(rnn_config)
rnn.rnn_model.load_state_dict(model_dict)
rnn.rnn_model<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1240144896654592, 'fraction_numerical': 0.01662049861495845, 'mean_word_length': 3.9830148619957537, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, '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': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '132070', 'license': 'MIT', 'n_tokens_mistral': 1766, 'n_tokens_neox': 1692, 'n_words': 288, 'path': 'experiment/pretrain.ipynb', 'repo_name': 'pseeth/otoworld', 'size': '4693'}