<jupyter_start><jupyter_text>Config<jupyter_code>config = {
    "lr": 1e-5,
    "epochs_num": 3000,
    "batch_size": 64,
    "log_each": 1,
    "save_each": 2,
    "device": "cuda:2",
    "x_dim": 1024,
    "z_dim": 8,
    "disc_coef": 5,
    "lambda": 5
}<jupyter_output><empty_output><jupyter_text>Data<jupyter_code>from generation.dataset.shapes_dataset import ShapesDataset
            
dataset = ShapesDataset(4, signal_dim=config['x_dim'])
idx = np.random.choice(range(len(dataset)))
signal = dataset[idx].numpy()
print("Signal size:", signal.shape)
plt.plot(signal)
plt.show()<jupyter_output>Signal size: (1024,)<jupyter_text>Models<jupyter_code>from generation.nets.shapes import Generator, Discriminator
discriminator = Discriminator(config)
test_tensor = dataset[0].unsqueeze(0)
discriminator(test_tensor, debug=True)
generator = Generator(config)
test_z = torch.rand(1, config['z_dim'])
output = generator(test_z, debug=True)
assert(output.shape == test_tensor.shape)<jupyter_output>torch.Size([1, 1, 1024])
torch.Size([1, 8, 1024])
torch.Size([1, 8, 340])
torch.Size([1, 32, 340])
torch.Size([1, 32, 112])
torch.Size([1, 8, 112])
torch.Size([1, 8, 36])
torch.Size([1, 288])
torch.Size([1, 1])
torch.Size([1, 1024])
torch.Size([1, 1, 1024])
torch.Size([1, 8, 1024])
torch.Size([1, 32, 1024])
torch.Size([1, 16, 1024])
torch.Size([1, 8, 1024])
torch.Size([1, 1, 1024])<jupyter_text>Training<jupyter_code>from generation.training.wgan_trainer import WganTrainer

g_optimizer = torch.optim.Adam(generator.parameters(), lr=config['lr'])
d_optimizer = torch.optim.Adam(discriminator.parameters(), lr=config['lr'])

trainer = WganTrainer(generator, discriminator, g_optimizer, \
                      d_optimizer, config)
trainer.run_train(dataset)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1726457399103139, 'fraction_numerical': 0.0633408071748879, 'mean_word_length': 6.083333333333333, 'pattern_counts': {'":': 10, '<': 15, '<?xml version=': 0, '>': 15, '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': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '354267', 'license': 'MIT', 'n_tokens_mistral': 775, 'n_tokens_neox': 691, 'n_words': 119, 'path': 'notebooks/train/shapes.ipynb', 'repo_name': 'whiteRa2bit/hep_generative_models', 'size': '1784'}