<jupyter_start><jupyter_text>Load config<jupyter_code>config = json.load(open('../configs/config.json'))<jupyter_output><empty_output><jupyter_text>Load model + checkpoint<jupyter_code>model = create_model(config)
state_dict = torch.load('../ckpts/ckpt.pt')
model.load_state_dict(state_dict)
model.eval()
model.cuda();<jupyter_output><empty_output><jupyter_text>Load dataset<jupyter_code>dset = AllenDataset(config, mode='all')
feat, neigh = dset.__getsingleitem__(0)
feat.shape, len(neigh)
adj_matrix = neighbors_to_adjacency(neigh, list(neigh.keys()))

lapl = compute_eig_lapl(adj_matrix, pos_enc_dim=config['model']['pos_dim'])<jupyter_output><empty_output><jupyter_text>Plot neuron<jupyter_code>plot_neuron(neigh, feat)<jupyter_output><empty_output><jupyter_text>Run inference<jupyter_code>latents = np.zeros((dset.num_samples, config['model']['dim']))

for i in tqdm(range(dset.num_samples)):
    feat, neigh = dset.__getsingleitem__(i)
    adj = torch.from_numpy(neighbors_to_adjacency(neigh, list(neigh.keys()))).float().cuda()[None, ]
    lapl = compute_eig_lapl(adj_matrix, pos_enc_dim=config['model']['pos_dim']).float().cuda()[None, ]
    feat = torch.from_numpy(feat).float().cuda()[None, ]
    
    latents[i] = model.student_encoder.forward(feat, adj, lapl)[0].cpu().detach()

plt.scatter(np.arange(config['model']['dim']), latents.mean(axis=0))

plt.errorbar(np.arange(config['model']['dim']), latents.mean(axis=0), yerr = latents.std(axis=0), fmt='none')
colors = [list(sns.color_palette("mako", n_colors=10)[3])] + [list(sns.color_palette("mako", n_colors=10)[-2])]
z = TSNE(n_components=2, perplexity=30).fit_transform(latents)
plot_tsne(z, np.concatenate([np.zeros(200, dtype=int), np.ones(230, dtype=int)]), targets=['aspiny', 'spiny'], colors=[colors[1], colors[0]])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18932305998899285, 'fraction_numerical': 0.012107870115575124, 'mean_word_length': 9.569767441860465, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, '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': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '606853', 'license': 'MIT', 'n_tokens_mistral': 728, 'n_tokens_neox': 710, 'n_words': 73, 'path': 'ssl_neuron/demos/model_inference.ipynb', 'repo_name': 'marissaweis/ssl_neuron', 'size': '1817'}