<jupyter_start><jupyter_text>Here we will show an example of how to run deepblast on simulation data.<jupyter_code>import os
import pandas as pd
from deepblast.sim import hmm_alignments
import argparse<jupyter_output><empty_output><jupyter_text>We will first simulate multiple sequences from a single PFam family.The resulting generated alignments will be used to train the model.<jupyter_code>fname = '../data/tm_align_output_10k.tab'
cols = [
    'chain1_name', 'chain2_name', 'tmscore1', 'tmscore2', 'rmsd',
    'chain1', 'chain2', 'alignment'
]
align_df = pd.read_table(fname, header=None, sep='\s+')
align_df.columns = cols
n_alignments = align_df.shape[0]<jupyter_output><empty_output><jupyter_text>The simulated alignments will be split into training / testing and validation.<jupyter_code>parts = n_alignments // 10
train_df = align_df.iloc[:parts * 8]
test_df = align_df.iloc[parts * 8:parts * 9]
valid_df = align_df.iloc[parts * 9:]

# save the files to disk.
train_df.to_csv('data/train.txt', sep='\t', index=None, header=None)
test_df.to_csv('data/test.txt', sep='\t', index=None, header=None)
valid_df.to_csv('data/valid.txt', sep='\t', index=None, header=None)<jupyter_output><empty_output><jupyter_text>We will prepare the environment to make sure that the appropriate output directories exist to store the results.<jupyter_code>from deepblast.trainer import LightningAligner
from pytorch_lightning import Trainer

output_dir = 'struct_results'
if not os.path.exists(output_dir):
    os.mkdir(output_dir)
os.getcwd()<jupyter_output><empty_output><jupyter_text>We will now create the arguments.  Below is the way to create this in a python environment.This can also be recreated on a standard command line interface.<jupyter_code>args = [
    '--train-pairs', f'{os.getcwd()}/data/train.txt',
    '--test-pairs', f'{os.getcwd()}/data/test.txt',
    '--valid-pairs', f'{os.getcwd()}/data/valid.txt',
    '--output-directory', output_dir,
    '--epochs', '128',
    '--batch-size', '20',  
    '--num-workers', '30',
    '--layers', '2',
    '--learning-rate', '5e-5',
    '--visualization-fraction', '1',
    '--loss', 'cross_entropy',
    '--scheduler', 'steplr',
    '--gpus', '1'
]<jupyter_output><empty_output><jupyter_text>We will then initialize the alignment model with the parameters we specified earlier.<jupyter_code>parser = argparse.ArgumentParser(add_help=False)
parser = LightningAligner.add_model_specific_args(parser)
parser.add_argument('--num-workers', type=int)
parser.add_argument('--gpus', type=int)
args = parser.parse_args(args)
model = LightningAligner(args)<jupyter_output><empty_output><jupyter_text>We can now train the model.<jupyter_code>trainer = Trainer(
    max_epochs=args.epochs,
    gpus=args.gpus,
    check_val_every_n_epoch=8,
    # profiler=profiler,
    # fast_dev_run=True,
    # auto_scale_batch_size='power'
)

trainer.fit(model)<jupyter_output>GPU available: True, used: True
TPU available: False, using: 0 TPU cores
CUDA_VISIBLE_DEVICES: [0]
/home/juermieboop/miniconda3/envs/pytorch/lib/python3.8/site-packages/pytorch_lightning/utilities/distributed.py:37: UserWarning: Could not log computational graph since the `model.example_input_array` attribute is not set or `input_array` was not given
  warnings.warn(*args, **kwargs)

  | Name    | Type                   | Params
---------------------------------------------------
0 | aligner | NeedlemanWunschAligner | 34 M<jupyter_text>The model diagnostics can be directly visualized in Tensorboard. Here, we show the losses, the accuracy and the alignment results.<jupyter_code>!ls lightning_logs
%load_ext tensorboard
%tensorboard --logdir lightning_logs<jupyter_output><empty_output><jupyter_text>And we did this with just a few million parameters<jupyter_code>model.aligner<jupyter_output><empty_output><jupyter_text>Evaluate model<jupyter_code>!ls lightning_logs/version_3/checkpoints
from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
checkpoint_dir = 'lightning_logs/version_3/checkpoints'
path = f'{checkpoint_dir}/epoch=49.ckpt'
model = LightningAligner.load_from_checkpoint(path)
from deepblast.dataset.dataset import states2matrix, tmstate_f
import numpy as np
import torch
from torch.nn.utils.rnn import pack_sequence
import matplotlib.pyplot as plt
import seaborn as sns
i = 0
test = train_df
x = str.encode(test.iloc[i]['chain1'])
y = str.encode(test.iloc[i]['chain2'])
states = test.iloc[i]['alignment']
x_ = torch.Tensor(model.tokenizer(x)).long().unsqueeze(0)
y_ = torch.Tensor(model.tokenizer(y)).long().unsqueeze(0)
x_ = pack_sequence(x_, enforce_sorted=False).cuda()
y_ = pack_sequence(y_, enforce_sorted=False).cuda()
states = list(map(tmstate_f, states))
A = states2matrix(states)
model = model.cuda()
aln, theta, gap = model.forward(x_, y_)
fig, ax = plt.subplots(1, 3, figsize=(12, 3))
sns.heatmap(A, ax=ax[0])
sns.heatmap(aln.cpu().detach().numpy().squeeze(), ax=ax[1])
sns.heatmap(theta.cpu().detach().numpy().squeeze(), ax=ax[2])
test.iloc[0]['chain1'], test.iloc[0]['chain2']<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1433932451115939, 'fraction_numerical': 0.012048192771084338, 'mean_word_length': 5.534193548387097, 'pattern_counts': {'":': 0, '<': 40, '<?xml version=': 0, '>': 40, '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': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '570400', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 1822, 'n_tokens_neox': 1777, 'n_words': 379, 'path': 'ipynb/struct-benchmark.ipynb', 'repo_name': 'wouterboomsma/deepblast', 'size': '5063'}