<jupyter_start><jupyter_text>Going beyond simple CBMs<jupyter_code>import patches
import torch
import torch.utils as utils
import lettertask
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
import numpy as np
import plotnine as gg
import lazytools_sflippl as lazytools
import time
import pandas as pd
cbm = lettertask.data.CompositionalBinaryModel(
    width=[50, 10, 20, 40, 100, 40],
    change_probability=[0.05, 0.5, 0.2, 0.5, 0.1, 0.2],
    samples=100000,
    seed=None
)
cts = patches.data.Contrastive1DTimeSeries(cbm.to_array())
cdl = utils.data.DataLoader(cts, batch_size=8, drop_last=True)
ce = patches.networks.LinearScaffold(2, data=cts, bias=False)
criterion = patches.losses.ContrastiveLoss(loss=patches.losses.BilinearLoss())
ideal_1 = np.array([1]*50+[0]*(10+20+40+100+40))/np.sqrt(50)
ideal_2 = np.array([0]*(50+10+20+40)+[1]*100+[0]*40)/np.sqrt(100)
optimizer = optim.Adam(ce.parameters(), lr=1e-2)
angle_1 = []
angle_2 = []
loss_traj = []
running_loss = 0
n_epochs = 1
with tqdm(total=n_epochs) as pbar:
    for epoch in range(n_epochs):
        running_loss = 0
        for i, data in enumerate(cdl):
            if i<len(cdl):
                if i % 10 == 0:
                    est = list(ce.parameters())[0].detach().numpy()
                    angle_1.append(np.matmul(ideal_1, est[0,:].T)/np.sqrt(np.matmul(est[0,:], est[0,:].T)) +
                                   np.matmul(ideal_1, est[1,:].T)/np.sqrt(np.matmul(est[1,:], est[1,:].T)))
                    angle_2.append(np.matmul(ideal_2, est[0,:].T)/np.sqrt(np.matmul(est[0,:], est[0,:].T)) +
                                   np.matmul(ideal_2, est[1,:].T)/np.sqrt(np.matmul(est[1,:], est[1,:].T)))
                optimizer.zero_grad()
                code = ce(data)
                loss = criterion(code)
                loss.backward()
                optimizer.step()
                running_loss += loss
                if i % 500 == 499:
                    loss_traj.append(running_loss.detach().numpy()/500)
                    running_loss = 0
                pbar.update(1/len(cdl))
est.shape
loss_df = lazytools.array_to_dataframe(np.array(loss_traj))\
                   .rename(columns={'dim0': 'steps', 'array': 'loss'})
np.matmul(est, est.T)
(gg.ggplot(loss_df, gg.aes(x='steps', y='loss')) +
    gg.geom_line())
angle_1_df = lazytools.array_to_dataframe(np.array(angle_1))\
                      .rename(columns={'dim0': 'steps', 'array': 'angle'})
angle_1_df['latent_dim'] = np.array('1')
angle_2_df = lazytools.array_to_dataframe(np.array(angle_2))\
                      .rename(columns={'dim0': 'steps', 'array': 'angle'})
angle_2_df['latent_dim'] = np.array('2')
angle_df = pd.concat([angle_1_df, angle_2_df])
(gg.ggplot(angle_df, gg.aes(x='steps', y='angle')) +
    gg.geom_line() + 
    gg.facet_wrap('latent_dim'))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16869565217391305, 'fraction_numerical': 0.04382608695652174, 'mean_word_length': 2.8655913978494625, 'pattern_counts': {'":': 0, '<': 6, '<?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': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '485031', 'license': 'MIT', 'n_tokens_mistral': 1171, 'n_tokens_neox': 1078, 'n_words': 174, 'path': 'notebooks/10-complex-cbms.ipynb', 'repo_name': 'sflippl/patches', 'size': '2875'}